From 0afb95d56114a55aeefae729f58257c9063b5005 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 27 Aug 2026 11:44:43 -0400 Subject: [PATCH 1/2] fix: Use apiKey in loadPrompt --- py/src/braintrust/api/_transport.py | 13 +- py/src/braintrust/api/policies.py | 9 + py/src/braintrust/logger.py | 293 ++++++++++++++---- py/src/braintrust/prompt_cache/lru_cache.py | 22 +- .../prompt_cache/parameters_cache.py | 6 +- .../braintrust/prompt_cache/prompt_cache.py | 29 +- .../braintrust/prompt_cache/test_lru_cache.py | 12 + .../prompt_cache/test_prompt_cache.py | 21 ++ py/src/braintrust/test_logger.py | 225 ++++++++++++++ 9 files changed, 554 insertions(+), 76 deletions(-) diff --git a/py/src/braintrust/api/_transport.py b/py/src/braintrust/api/_transport.py index a6109b62c..46ba68be8 100644 --- a/py/src/braintrust/api/_transport.py +++ b/py/src/braintrust/api/_transport.py @@ -23,7 +23,7 @@ BraintrustTransportError, BraintrustTransportRetryExhaustedError, ) -from .policies import RetryMode, RetryPolicy +from .policies import RetryMode, RetryPolicy, is_retryable_request_exception logger = logging.getLogger(__name__) @@ -120,6 +120,9 @@ def make_long_lived(self) -> None: ) self._reset() + def close(self) -> None: + self.session.close() + @staticmethod def sanitize_token(token: str) -> str: return token.rstrip("\n") @@ -272,7 +275,7 @@ def request( **kwargs, ) except requests.exceptions.RequestException as exc: - if not _is_retryable_request_exception(exc): + if not is_retryable_request_exception(exc): error = BraintrustTransportError(method=method, url=url, attempts=attempt, retryable=False) raise error from exc if attempt >= max_attempts: @@ -396,12 +399,6 @@ def _request_body_is_replayable(data: Any, files: Any) -> bool: return files is None and (data is None or isinstance(data, (bytes, str))) -def _is_retryable_request_exception(exc: requests.exceptions.RequestException) -> bool: - return isinstance(exc, (requests.exceptions.ConnectionError, requests.exceptions.Timeout)) and not isinstance( - exc, requests.exceptions.SSLError - ) - - def _parse_retry_after(value: str | None, wall_time: float) -> float | None: if value is None: return None diff --git a/py/src/braintrust/api/policies.py b/py/src/braintrust/api/policies.py index 8bc9b57ef..a86b98582 100644 --- a/py/src/braintrust/api/policies.py +++ b/py/src/braintrust/api/policies.py @@ -3,6 +3,8 @@ import enum from dataclasses import dataclass +import requests + DEFAULT_RETRYABLE_STATUSES = frozenset({408, 429, 500, 502, 503, 504}) DEFAULT_MAX_ATTEMPTS = 4 @@ -11,6 +13,13 @@ DEFAULT_MAX_BACKOFF = 10.0 +def is_retryable_request_exception(exc: requests.exceptions.RequestException) -> bool: + """Return whether a requests transport failure is safe to retry.""" + return isinstance(exc, (requests.exceptions.ConnectionError, requests.exceptions.Timeout)) and not isinstance( + exc, requests.exceptions.SSLError + ) + + class RetryMode(enum.Enum): """The replay safety classification for an API operation.""" diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index ee4ba6779..95a099d65 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -7,6 +7,7 @@ import contextvars import dataclasses import datetime +import hashlib import inspect import io import json @@ -38,14 +39,22 @@ import chevron import exceptiongroup from braintrust.functions.stream import BraintrustStream +from requests import exceptions as requests_exceptions from requests.adapters import HTTPAdapter from . import context, id_gen from .api._routing import normalize_proxy_url from .api._transport import HTTPConnection from .api._transport import RetryRequestExceptionsAdapter as RetryRequestExceptionsAdapter +from .api.auth import LoginResult, OrganizationInfo from .api.client import BraintrustClient, BraintrustOpenApiClient -from .api.errors import BraintrustAPIError, BraintrustHTTPError +from .api.errors import ( + BraintrustAPIError, + BraintrustHTTPError, + BraintrustJSONDecodeError, + BraintrustTransportError, +) +from .api.policies import DEFAULT_RETRYABLE_STATUSES, is_retryable_request_exception from .bt_json import bt_dumps, bt_safe_deep_copy from .db_fields import ( AUDIT_METADATA_FIELD, @@ -432,6 +441,19 @@ def __exit__( NOOP_SPAN_PERMALINK = "https://www.braintrust.dev/noop-span" +@dataclasses.dataclass(frozen=True) +class _LoaderRequestState: + app_url: str + org_id: str + _api_conn: HTTPConnection + + def api_conn(self) -> HTTPConnection: + return self._api_conn + + def close(self) -> None: + self._api_conn.close() + + class BraintrustState: def __init__(self): self.id = str(uuid.uuid4()) @@ -516,6 +538,14 @@ def default_get_api_conn(): self._otel_flush_callback: Any | None = None def reset_login_info(self): + if hasattr(self, "_loader_login_cache"): + self._loader_login_cache.clear() + else: + self._loader_login_cache: LRUCache[str, LazyValue[_LoaderRequestState]] = LRUCache( + max_size=16, + on_remove=self._close_loader_request_state, + ) + self.app_url: str | None = None self.app_public_url: str | None = None self.login_token: str | None = None @@ -533,6 +563,12 @@ def reset_login_info(self): self._client: BraintrustClient | None = None self._user_info: Mapping[str, Any] | None = None + @staticmethod + def _close_loader_request_state(_key: str, lazy_state: LazyValue[_LoaderRequestState]) -> None: + has_succeeded, request_state = lazy_state.get_sync() + if has_succeeded and request_state is not None: + request_state.close() + def reset_parent_state(self): # reset possible parent state for tests self.current_experiment = None @@ -591,6 +627,7 @@ async def flush_otel(self) -> None: def copy_state(self, other: "BraintrustState"): """Copy login information from another BraintrustState instance.""" + self._loader_login_cache.clear() self.__dict__.update( { k: v @@ -608,6 +645,7 @@ def copy_state(self, other: "BraintrustState"): "_last_otel_setting", "_context_manager_lock", "_client_lock", + "_loader_login_cache", ) } ) @@ -683,6 +721,37 @@ def user_info(self) -> Mapping[str, Any]: self._user_info = self.api_conn().get_json("ping") return self._user_info + def loader_request_state( + self, + *, + app_url: str, + api_key: str, + org_name: str | None, + cache_namespace: str, + ) -> "BraintrustState | _LoaderRequestState": + if ( + self.logged_in + and self.login_token == api_key + and self.app_url == app_url + and (org_name is None or self.org_name == org_name) + ): + return self + + with self._client_lock: + try: + lazy_state = self._loader_login_cache.get(cache_namespace) + except KeyError: + lazy_state = LazyValue( + lambda: _login_to_loader_request_state( + app_url=app_url, + api_key=api_key, + org_name=org_name, + ), + use_mutex=True, + ) + self._loader_login_cache.set(cache_namespace, lazy_state) + return lazy_state.get() + def global_bg_logger(self) -> "_BackgroundLogger": return getattr(self._override_bg_logger, "logger", None) or self._global_bg_logger.get() @@ -1830,6 +1899,68 @@ def compute_metadata(): return ret +def _resolve_loader_login_options( + *, + app_url: str | None, + api_key: str | None, + org_name: str | None, +) -> tuple[str, str, str | None, str]: + resolved_app_url = app_url or (_state.app_url if _state.logged_in else None) or _get_app_url() + resolved_api_key = api_key or (_state.login_token if _state.logged_in else None) + if resolved_api_key is None: + resolved_api_key = BraintrustEnv.API_KEY.get(None, use_dotenv=True) + if resolved_api_key is None: + raise ValueError( + "Could not login to Braintrust. You may need to set BRAINTRUST_API_KEY in your environment " + "or nearest .env.braintrust file." + ) + resolved_api_key = HTTPConnection.sanitize_token(resolved_api_key) + + uses_active_credential = _state.logged_in and resolved_api_key == _state.login_token + resolved_org_name = org_name + if resolved_org_name is None: + resolved_org_name = _state.org_name if uses_active_credential else _get_org_name() + + namespace_input = json.dumps( + ["loader-credential", resolved_app_url, resolved_org_name, resolved_api_key], + separators=(",", ":"), + ) + cache_namespace = f"loader-credential:{hashlib.sha256(namespace_input.encode('utf-8')).hexdigest()}" + return resolved_app_url, resolved_api_key, resolved_org_name, cache_namespace + + +def _is_loader_cache_fallback_error(error: BaseException) -> bool: + pending: list[BaseException] = [error] + seen: set[int] = set() + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + + if isinstance(current, (json.JSONDecodeError, BraintrustJSONDecodeError)): + return False + if isinstance(current, BraintrustTransportError): + return current.retryable + + status_code = getattr(current, "status_code", None) + if status_code is None: + response = getattr(current, "response", None) + status_code = getattr(response, "status_code", None) + if isinstance(status_code, int): + return status_code in DEFAULT_RETRYABLE_STATUSES + + if isinstance(current, requests_exceptions.RequestException): + return is_retryable_request_exception(current) + + if current.__cause__ is not None: + pending.append(current.__cause__) + if current.__context__ is not None: + pending.append(current.__context__) + + return False + + def load_prompt( project: str | None = None, slug: str | None = None, @@ -1855,8 +1986,7 @@ def load_prompt( :param no_trace: If true, do not include logging metadata for this prompt when build() is called. :param environment: The environment to load the prompt from. If both `version` and `environment` are provided, `version` takes precedence. :param app_url: The URL of the Braintrust App. Defaults to https://www.braintrust.dev. - :param api_key: The API key to use. If the parameter is not specified, will try to use the `BRAINTRUST_API_KEY` environment variable. If no API - key is specified, will prompt the user to login. + :param api_key: The API key to use for this request, independently of any existing global login. If the parameter is not specified, will use an existing login or try the `BRAINTRUST_API_KEY` environment variable. If no API key is specified, will prompt the user to login. :param org_name: (Optional) The name of a specific organization to connect to. This is useful if you belong to multiple. :returns: The prompt object. """ @@ -1871,12 +2001,22 @@ def load_prompt( raise ValueError("Must specify slug") def compute_metadata(): + resolved_app_url, resolved_api_key, resolved_org_name, cache_namespace = _resolve_loader_login_options( + app_url=app_url, + api_key=api_key, + org_name=org_name, + ) try: - login(org_name=org_name, api_key=api_key, app_url=app_url) + request_state = _state.loader_request_state( + app_url=resolved_app_url, + api_key=resolved_api_key, + org_name=resolved_org_name, + cache_namespace=cache_namespace, + ) if id: # Load prompt by ID using the /v1/prompt/{id} endpoint prompt_args = _populate_args({}, version=version, environment=effective_environment) - response = _state.api_conn().get_json(f"/v1/prompt/{id}", prompt_args) + response = request_state.api_conn().get_json(f"/v1/prompt/{id}", prompt_args) # Wrap single prompt response in objects array to match list API format if response is not None: response = {"objects": [response]} @@ -1889,8 +2029,10 @@ def compute_metadata(): version=version, environment=effective_environment, ) - response = _state.api_conn().get_json("/v1/prompt", args) + response = request_state.api_conn().get_json("/v1/prompt", args) except Exception as server_error: + if not _is_loader_cache_fallback_error(server_error): + raise # If environment or version was specified, don't fall back to cache if effective_environment is not None or version is not None: raise ValueError(f"Prompt not found with specified parameters") from server_error @@ -1898,13 +2040,14 @@ def compute_metadata(): eprint(f"Failed to load prompt, attempting to fall back to cache: {server_error}") try: if id: - return _state._prompt_cache.get(id=id) + return _state._prompt_cache.get(id=id, cache_namespace=cache_namespace) else: return _state._prompt_cache.get( slug, version=str(version) if version else "latest", project_id=project_id, project_name=project, + cache_namespace=cache_namespace, ) except Exception as cache_error: if id: @@ -1934,6 +2077,7 @@ def compute_metadata(): _state._prompt_cache.set( prompt, id=id, + cache_namespace=cache_namespace, ) elif slug: _state._prompt_cache.set( @@ -1942,6 +2086,7 @@ def compute_metadata(): version=str(version) if version else "latest", project_id=project_id, project_name=project, + cache_namespace=cache_namespace, ) except Exception as e: eprint(f"Failed to store prompt in cache: {e}") @@ -2028,7 +2173,7 @@ def load_parameters( :param id: The ID of a specific parameters object to load. If specified, this takes precedence over project and slug. :param environment: The environment to load the parameters from. If both `version` and `environment` are provided, `version` takes precedence. :param app_url: The URL of the Braintrust App. Defaults to https://www.braintrust.dev. - :param api_key: The API key to use. If the parameter is not specified, will try to use the `BRAINTRUST_API_KEY` environment variable. + :param api_key: The API key to use for this request, independently of any existing global login. If the parameter is not specified, will use an existing login or try the `BRAINTRUST_API_KEY` environment variable. :param org_name: The name of a specific organization to connect to. :returns: A `RemoteEvalParameters` object. """ @@ -2040,11 +2185,21 @@ def load_parameters( effective_environment = None if version is not None else environment should_fall_back_to_cache = version is None and effective_environment is None query_args = _populate_args({}, version=version, environment=effective_environment) + resolved_app_url, resolved_api_key, resolved_org_name, cache_namespace = _resolve_loader_login_options( + app_url=app_url, + api_key=api_key, + org_name=org_name, + ) try: - login(org_name=org_name, api_key=api_key, app_url=app_url) + request_state = _state.loader_request_state( + app_url=resolved_app_url, + api_key=resolved_api_key, + org_name=resolved_org_name, + cache_namespace=cache_namespace, + ) if id: - response = _state.api_conn().get_json(f"/v1/function/{id}", query_args) + response = request_state.api_conn().get_json(f"/v1/function/{id}", query_args) if response is not None: response = {"objects": [response]} else: @@ -2055,20 +2210,23 @@ def load_parameters( slug=slug, **query_args, ) - response = _state.api_conn().get_json("/v1/function", args) + response = request_state.api_conn().get_json("/v1/function", args) except Exception as server_error: + if not _is_loader_cache_fallback_error(server_error): + raise if not should_fall_back_to_cache: raise eprint(f"Failed to load parameters, attempting to fall back to cache: {server_error}") try: if id: - return _state._parameters_cache.get(id=id) + return _state._parameters_cache.get(id=id, cache_namespace=cache_namespace) return _state._parameters_cache.get( slug=slug, version=str(version) if version is not None else "latest", project_id=project_id, project_name=project, + cache_namespace=cache_namespace, ) except Exception as cache_error: if id: @@ -2093,7 +2251,7 @@ def load_parameters( parameters = RemoteEvalParameters.from_function_row(response["objects"][0]) try: if id: - _state._parameters_cache.set(parameters, id=id) + _state._parameters_cache.set(parameters, id=id, cache_namespace=cache_namespace) elif slug: _state._parameters_cache.set( parameters, @@ -2101,6 +2259,7 @@ def load_parameters( version=str(version) if version is not None else "latest", project_id=project_id, project_name=project, + cache_namespace=cache_namespace, ) except Exception as exc: eprint(f"Failed to store parameters in cache: {exc}") @@ -2156,6 +2315,66 @@ def register_otel_flush(callback: Any) -> None: _state.span_cache.disable() +def _login_with_api_key(*, app_url: str, api_key: str, org_name: str | None) -> tuple[BraintrustClient, LoginResult]: + if api_key == TEST_API_KEY: + api_url = BraintrustEnv.API_URL.get("https://api.braintrust.ai") + proxy_url = BraintrustEnv.PROXY_URL.get("https://proxy.braintrust.ai") + organization = OrganizationInfo( + id="test-org-id", + name=org_name or "test-org-name", + api_url=api_url, + proxy_url=proxy_url, + realtime_url=None, + is_universal_api=False, + git_metadata=None, + raw={}, + ) + client = BraintrustClient( + api_key=api_key, + app_url=app_url, + api_url=api_url, + proxy_url=proxy_url, + adapter=_http_adapter, + ) + return client, LoginResult(organization=organization, api_url=api_url, proxy_url=proxy_url, response={}) + + client = BraintrustClient(api_key=api_key, app_url=app_url, adapter=_http_adapter) + try: + return client, client.auth.login(org_name=org_name) + except BraintrustHTTPError as exc: + client.close() + masked_api_key = mask_api_key(api_key) + raise ValueError(f"Invalid API key {masked_api_key}: [{exc.status_code}] {exc.response_body}") from exc + except Exception: + client.close() + raise + + +def _authenticated_api_conn(api_url: str, api_key: str) -> HTTPConnection: + conn = HTTPConnection(api_url, adapter=_http_adapter) + conn.set_token(api_key) + conn.make_long_lived() + return conn + + +def _login_to_loader_request_state( + *, + app_url: str, + api_key: str, + org_name: str | None, +) -> _LoaderRequestState: + client, login_result = _login_with_api_key(app_url=app_url, api_key=api_key, org_name=org_name) + try: + conn = _authenticated_api_conn(login_result.api_url, api_key) + finally: + client.close() + return _LoaderRequestState( + app_url=app_url, + org_id=login_result.organization.id, + _api_conn=conn, + ) + + def login_to_state( app_url: str | None = None, api_key: str | None = None, @@ -2175,46 +2394,13 @@ def login_to_state( state.app_public_url = app_public_url state.org_name = org_name - if api_key == TEST_API_KEY: - # A small hook for pseudo-logins. It still constructs the facade so - # concurrent lazy access follows the same state lifecycle as real login. - test_org_info = [ - { - "id": "test-org-id", - "name": org_name or "test-org-name", - "api_url": "https://api.braintrust.ai", - "proxy_url": "https://proxy.braintrust.ai", - } - ] - _check_org_info(state, test_org_info, org_name) - state._client = BraintrustClient( - api_key=TEST_API_KEY, - app_url=state.app_url, - api_url=state.api_url, - proxy_url=state.proxy_url, - adapter=_http_adapter, - ) - state.login_token = TEST_API_KEY - state.logged_in = True - return state - if api_key is None: raise ValueError( "Could not login to Braintrust. You may need to set BRAINTRUST_API_KEY in your environment " "or nearest .env.braintrust file." ) - client = BraintrustClient(api_key=api_key, app_url=state.app_url, adapter=_http_adapter) - try: - login_result = client.auth.login(org_name=org_name) - except BraintrustHTTPError as exc: - client.close() - masked_api_key = mask_api_key(api_key) - raise ValueError(f"Invalid API key {masked_api_key}: [{exc.status_code}] {exc.response_body}") from exc - except Exception: - client.close() - raise - + client, login_result = _login_with_api_key(app_url=app_url, api_key=api_key, org_name=org_name) organization = login_result.organization state._client = client state.org_id = organization.id @@ -2225,12 +2411,16 @@ def login_to_state( state.git_metadata_settings = ( GitMetadataSettings(**organization.git_metadata) if organization.git_metadata else None ) + state.login_token = HTTPConnection.sanitize_token(api_key) + state.logged_in = True + + if api_key == TEST_API_KEY: + return state # Keep un-migrated call sites on isolated legacy sessions. Their mutable # adapters and session headers must not affect the policy-aware client. - conn = state.api_conn() - conn.set_token(api_key) - conn.make_long_lived() + conn = _authenticated_api_conn(login_result.api_url, api_key) + state._api_conn = conn app_connection = state.app_conn() app_connection.set_token(api_key) @@ -2241,9 +2431,6 @@ def login_to_state( proxy_connection.set_token(api_key) proxy_connection.make_long_lived() - state.login_token = HTTPConnection.sanitize_token(api_key) - state.logged_in = True - # Replace the global logger's api_conn with this one. state.login_replace_api_conn(conn) diff --git a/py/src/braintrust/prompt_cache/lru_cache.py b/py/src/braintrust/prompt_cache/lru_cache.py index 126fbd275..3c7f4f56e 100644 --- a/py/src/braintrust/prompt_cache/lru_cache.py +++ b/py/src/braintrust/prompt_cache/lru_cache.py @@ -8,6 +8,7 @@ """ from collections import OrderedDict +from collections.abc import Callable from typing import Generic, TypeVar @@ -28,11 +29,17 @@ class LRUCache(Generic[K, V]): Args: max_size: Maximum number of items to store in the cache. If not specified, the cache will grow unbounded. + on_remove: Optional callback invoked when an entry is replaced, evicted, or cleared. """ - def __init__(self, max_size: int | None = None): + def __init__( + self, + max_size: int | None = None, + on_remove: Callable[[K, V], None] | None = None, + ): self._cache: OrderedDict[K, V] = OrderedDict() self._max_size = max_size + self._on_remove = on_remove def get(self, key: K) -> V: """ @@ -66,14 +73,23 @@ def set(self, key: K, value: V) -> None: key: The key to store. value: The value to store. """ + removed: tuple[K, V] | None = None if key in self._cache: - self._cache.pop(key) + previous = self._cache.pop(key) + if previous is not value: + removed = (key, previous) elif self._max_size and len(self._cache) >= self._max_size: # Remove oldest item (first item in ordered dict). - self._cache.popitem(last=False) + removed = self._cache.popitem(last=False) self._cache[key] = value + if removed is not None and self._on_remove is not None: + self._on_remove(*removed) def clear(self) -> None: """Removes all items from the cache.""" + items = list(self._cache.items()) if self._on_remove is not None else [] self._cache.clear() + if self._on_remove is not None: + for item in items: + self._on_remove(*item) diff --git a/py/src/braintrust/prompt_cache/parameters_cache.py b/py/src/braintrust/prompt_cache/parameters_cache.py index 5f86cb293..ce93046f8 100644 --- a/py/src/braintrust/prompt_cache/parameters_cache.py +++ b/py/src/braintrust/prompt_cache/parameters_cache.py @@ -20,8 +20,9 @@ def get( project_id: str | None = None, project_name: str | None = None, id: str | None = None, + cache_namespace: str | None = None, ) -> RemoteEvalParameters: - cache_key = prompt_cache._create_cache_key(project_id, project_name, slug, version, id) + cache_key = prompt_cache._create_cache_key(project_id, project_name, slug, version, id, cache_namespace) try: return self.memory_cache.get(cache_key) @@ -45,8 +46,9 @@ def set( project_id: str | None = None, project_name: str | None = None, id: str | None = None, + cache_namespace: str | None = None, ) -> None: - cache_key = prompt_cache._create_cache_key(project_id, project_name, slug, version, id) + cache_key = prompt_cache._create_cache_key(project_id, project_name, slug, version, id, cache_namespace) self.memory_cache.set(cache_key, value) if self.disk_cache: self.disk_cache.set(cache_key, value) diff --git a/py/src/braintrust/prompt_cache/prompt_cache.py b/py/src/braintrust/prompt_cache/prompt_cache.py index ac6d8a33f..84130564b 100644 --- a/py/src/braintrust/prompt_cache/prompt_cache.py +++ b/py/src/braintrust/prompt_cache/prompt_cache.py @@ -6,7 +6,7 @@ 2. A persistent disk-based cache that serves as a backing store This allows for efficient prompt retrieval while maintaining persistence across sessions. -The cache is keyed by project identifier (ID or name), prompt slug, and version. +The cache is keyed by an optional namespace, project identifier (ID or name), prompt slug, and version. """ from braintrust import prompt @@ -19,18 +19,23 @@ def _create_cache_key( slug: str | None, version: str = "latest", id: str | None = None, + cache_namespace: str | None = None, ) -> str: """Creates a unique cache key from project identifier, slug and version, or from ID.""" if id: # When caching by ID, we don't need project or slug - return f"id:{id}" + cache_key = f"id:{id}" + else: + prefix = project_id or project_name + if not prefix: + raise ValueError("Either project_id or project_name must be provided") + if not slug: + raise ValueError("Slug must be provided when not using ID") + cache_key = f"{prefix}:{slug}:{version}" - prefix = project_id or project_name - if not prefix: - raise ValueError("Either project_id or project_name must be provided") - if not slug: - raise ValueError("Slug must be provided when not using ID") - return f"{prefix}:{slug}:{version}" + if cache_namespace is None: + return cache_key + return f"{len(cache_namespace)}:{cache_namespace}:{cache_key}" class PromptCache: @@ -64,6 +69,7 @@ def get( project_id: str | None = None, project_name: str | None = None, id: str | None = None, + cache_namespace: str | None = None, ) -> prompt.PromptSchema: """ Retrieve a prompt from the cache. @@ -74,6 +80,7 @@ def get( project_id: The ID of the project containing the prompt. project_name: The name of the project containing the prompt. id: The ID of a specific prompt. If provided, slug and project parameters are ignored. + cache_namespace: An optional namespace used to isolate cache entries. Returns: The cached Prompt object. @@ -82,7 +89,7 @@ def get( ValueError: If neither project_id nor project_name is provided (when not using id). KeyError: If the prompt is not found in the cache. """ - cache_key = _create_cache_key(project_id, project_name, slug, version, id) + cache_key = _create_cache_key(project_id, project_name, slug, version, id, cache_namespace) # First check memory cache. try: @@ -110,6 +117,7 @@ def set( project_id: str | None = None, project_name: str | None = None, id: str | None = None, + cache_namespace: str | None = None, ) -> None: """ Store a prompt in the cache. @@ -121,12 +129,13 @@ def set( project_id: The ID of the project containing the prompt. project_name: The name of the project containing the prompt. id: The ID of a specific prompt. If provided, slug and project parameters are ignored. + cache_namespace: An optional namespace used to isolate cache entries. Raises: ValueError: If neither project_id nor project_name is provided (when not using id). RuntimeError: If there is an error writing to the disk cache. """ - cache_key = _create_cache_key(project_id, project_name, slug, version, id) + cache_key = _create_cache_key(project_id, project_name, slug, version, id, cache_namespace) # Update memory cache. self.memory_cache.set(cache_key, value) diff --git a/py/src/braintrust/prompt_cache/test_lru_cache.py b/py/src/braintrust/prompt_cache/test_lru_cache.py index 6fb286e57..58806a2d6 100644 --- a/py/src/braintrust/prompt_cache/test_lru_cache.py +++ b/py/src/braintrust/prompt_cache/test_lru_cache.py @@ -67,6 +67,18 @@ def test_clear_all_items(self): with self.assertRaises(KeyError): cache.get("b") + def test_on_remove_runs_for_replacement_eviction_and_clear(self): + removed = [] + cache = lru_cache.LRUCache[str, int](max_size=2, on_remove=lambda key, value: removed.append((key, value))) + + cache.set("a", 1) + cache.set("a", 2) + cache.set("b", 3) + cache.set("c", 4) + cache.clear() + + self.assertEqual(removed, [("a", 1), ("a", 2), ("b", 3), ("c", 4)]) + if __name__ == "__main__": unittest.main() diff --git a/py/src/braintrust/prompt_cache/test_prompt_cache.py b/py/src/braintrust/prompt_cache/test_prompt_cache.py index 0e0d70c8f..3f44ff4b0 100644 --- a/py/src/braintrust/prompt_cache/test_prompt_cache.py +++ b/py/src/braintrust/prompt_cache/test_prompt_cache.py @@ -66,6 +66,27 @@ def test_store_and_retrieve_from_memory_cache(self): result = self.cache.get(slug="test-prompt", version="789", project_id="123") self.assertEqual(result.as_dict(), self.test_prompt.as_dict()) + def test_cache_namespace_isolates_memory_and_disk_entries(self): + self.cache.set( + self.test_prompt, + slug="test-prompt", + project_id="123", + cache_namespace="first-credential", + ) + + result = self.cache.get( + slug="test-prompt", + project_id="123", + cache_namespace="first-credential", + ) + self.assertEqual(result.as_dict(), self.test_prompt.as_dict()) + with self.assertRaises(KeyError): + self.cache.get( + slug="test-prompt", + project_id="123", + cache_namespace="second-credential", + ) + def test_work_with_project_name(self): self.cache.set(self.test_prompt, slug="test-prompt", version="789", project_name="test-project") result = self.cache.get(slug="test-prompt", version="789", project_name="test-project") diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index a1ee21d8e..da65e5f17 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -40,6 +40,9 @@ stringify_exception, ) from braintrust.prompt import PromptChatBlock, PromptData, PromptMessage, PromptSchema +from braintrust.prompt_cache.lru_cache import LRUCache +from braintrust.prompt_cache.parameters_cache import ParametersCache +from braintrust.prompt_cache.prompt_cache import PromptCache from braintrust.test_helpers import ( assert_dict_matches, assert_logged_out, @@ -51,6 +54,9 @@ with_memory_logger, # noqa: F401 # type: ignore[reportUnusedImport] with_simulate_login, # noqa: F401 # type: ignore[reportUnusedImport] ) +from braintrust.util import AugmentedHTTPError +from requests import HTTPError +from requests.exceptions import SSLError def test_login_to_state_uses_env_braintrust_api_key(tmp_path, monkeypatch): @@ -64,6 +70,38 @@ def test_login_to_state_uses_env_braintrust_api_key(tmp_path, monkeypatch): assert state.logged_in is True +def test_loader_request_state_closes_connections_on_eviction_and_reset(): + state = BraintrustState() + state._loader_login_cache = LRUCache(max_size=1, on_remove=state._close_loader_request_state) + first_conn = MagicMock() + second_conn = MagicMock() + request_states = [ + logger._LoaderRequestState("https://app.example.com", "org-a", first_conn), + logger._LoaderRequestState("https://app.example.com", "org-b", second_conn), + ] + + with patch.object(logger, "_login_to_loader_request_state", side_effect=request_states): + state.loader_request_state( + app_url="https://app.example.com", + api_key="first-api-key", + org_name=None, + cache_namespace="first", + ) + state.loader_request_state( + app_url="https://app.example.com", + api_key="second-api-key", + org_name=None, + cache_namespace="second", + ) + + first_conn.close.assert_called_once() + second_conn.close.assert_not_called() + + state.reset_login_info() + + second_conn.close.assert_called_once() + + class TestInit(TestCase): @staticmethod def _mock_api_client(): @@ -346,6 +384,187 @@ def _prompt_response(slug: str): } +def _parameters_response(slug: str): + return { + "objects": [ + { + "id": f"parameters-{slug}", + "project_id": "project-123", + "name": "Saved parameters", + "slug": slug, + "_xact_id": "v1", + "function_data": { + "type": "parameters", + "data": {"prefix": slug}, + "__schema": {"type": "object"}, + }, + } + ] + } + + +def _http_error(status_code: int) -> AugmentedHTTPError: + error = AugmentedHTTPError(f"HTTP {status_code}") + error.__cause__ = HTTPError(response=MagicMock(status_code=status_code)) + return error + + +def test_load_prompt_uses_explicit_api_key_without_changing_global_login(): + simulate_login() + original_login_token = logger._state.login_token + prompt_cache = PromptCache(memory_cache=LRUCache(max_size=10)) + request_conn = MagicMock() + request_conn.get_json.return_value = _prompt_response("saved-prompt") + request_state = MagicMock() + request_state.api_conn.return_value = request_conn + + with ( + patch.object(logger._state, "_prompt_cache", prompt_cache), + patch.object(logger, "_login_to_loader_request_state", return_value=request_state) as mock_login_to_state, + ): + prompt = braintrust.load_prompt( + project="test-project", + slug="saved-prompt", + api_key="prompt-api-key", + ) + assert prompt.slug == "saved-prompt" + + mock_login_to_state.assert_called_once_with( + app_url=logger._state.app_url, + api_key="prompt-api-key", + org_name=None, + ) + assert logger._state.login_token == original_login_token + + +def test_load_parameters_uses_explicit_api_key_without_changing_global_login(): + simulate_login() + original_login_token = logger._state.login_token + parameters_cache = ParametersCache(memory_cache=LRUCache(max_size=10)) + request_conn = MagicMock() + request_conn.get_json.return_value = _parameters_response("saved-parameters") + request_state = MagicMock() + request_state.api_conn.return_value = request_conn + + with ( + patch.object(logger._state, "_parameters_cache", parameters_cache), + patch.object(logger, "_login_to_loader_request_state", return_value=request_state) as mock_login_to_state, + ): + parameters = braintrust.load_parameters( + project="test-project", + slug="saved-parameters", + api_key="parameters-api-key", + ) + + assert parameters.data == {"prefix": "saved-parameters"} + mock_login_to_state.assert_called_once_with( + app_url=logger._state.app_url, + api_key="parameters-api-key", + org_name=None, + ) + assert logger._state.login_token == original_login_token + + +@pytest.mark.parametrize( + "server_error", + [ + _http_error(401), + _http_error(501), + json.JSONDecodeError("invalid JSON", "", 0), + SSLError("invalid certificate"), + ], +) +def test_load_prompt_does_not_fall_back_to_cache_for_non_transient_errors(server_error): + simulate_login() + prompt_cache = PromptCache(memory_cache=LRUCache(max_size=10)) + request_conn = MagicMock() + request_conn.get_json.side_effect = [_prompt_response("saved-prompt"), server_error] + request_state = MagicMock() + request_state.api_conn.return_value = request_conn + + with ( + patch.object(logger._state, "_prompt_cache", prompt_cache), + patch.object(logger, "_login_to_loader_request_state", return_value=request_state), + ): + first_prompt = braintrust.load_prompt( + project="test-project", + slug="saved-prompt", + api_key="prompt-api-key", + ) + assert first_prompt.slug == "saved-prompt" + + second_prompt = braintrust.load_prompt( + project="test-project", + slug="saved-prompt", + api_key="prompt-api-key", + ) + with pytest.raises(type(server_error)): + _ = second_prompt.slug + + +def test_load_prompt_uses_same_api_keys_cache_for_transient_errors(): + simulate_login() + prompt_cache = PromptCache(memory_cache=LRUCache(max_size=10)) + request_conn = MagicMock() + request_conn.get_json.side_effect = [_prompt_response("saved-prompt"), _http_error(500)] + request_state = MagicMock() + request_state.api_conn.return_value = request_conn + + with ( + patch.object(logger._state, "_prompt_cache", prompt_cache), + patch.object(logger, "_login_to_loader_request_state", return_value=request_state), + ): + first_prompt = braintrust.load_prompt( + project="test-project", + slug="saved-prompt", + api_key="prompt-api-key", + ) + assert first_prompt.slug == "saved-prompt" + + cached_prompt = braintrust.load_prompt( + project="test-project", + slug="saved-prompt", + api_key="prompt-api-key", + ) + assert cached_prompt.slug == "saved-prompt" + + +def test_load_prompt_does_not_use_another_api_keys_transient_fallback_cache(): + simulate_login() + prompt_cache = PromptCache(memory_cache=LRUCache(max_size=10)) + first_conn = MagicMock() + first_conn.get_json.return_value = _prompt_response("saved-prompt") + second_conn = MagicMock() + second_conn.get_json.side_effect = _http_error(500) + request_states = {} + for api_key, request_conn in (("first-api-key", first_conn), ("second-api-key", second_conn)): + request_state = MagicMock() + request_state.api_conn.return_value = request_conn + request_states[api_key] = request_state + + def login_for_api_key(*, api_key, **_kwargs): + return request_states[api_key] + + with ( + patch.object(logger._state, "_prompt_cache", prompt_cache), + patch.object(logger, "_login_to_loader_request_state", side_effect=login_for_api_key), + ): + first_prompt = braintrust.load_prompt( + project="test-project", + slug="saved-prompt", + api_key="first-api-key", + ) + assert first_prompt.slug == "saved-prompt" + + second_prompt = braintrust.load_prompt( + project="test-project", + slug="saved-prompt", + api_key="second-api-key", + ) + with pytest.raises(ValueError, match="not found on server or in local cache"): + _ = second_prompt.slug + + @pytest.mark.asyncio async def test_load_prompt_async_eagerly_fetches_prompt(with_simulate_login): mock_api_conn = MagicMock() @@ -499,11 +718,17 @@ def test_load_parameters_returns_remote_object(self): assert parameters.id == "params-123" assert parameters.version == "v1" assert parameters.data == {"prefix": "hello"} + cache_namespace = logger._resolve_loader_login_options( + app_url=None, + api_key=None, + org_name=None, + )[-1] assert ( logger._state._parameters_cache.get( slug="saved-parameters", version="latest", project_name="test-project", + cache_namespace=cache_namespace, ).id == "params-123" ) From e687caf3c3a6b308969895d388527de0c2b0f909 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 31 Aug 2026 13:57:41 -0400 Subject: [PATCH 2/2] fixes --- py/src/braintrust/api/_transport.py | 23 ++ py/src/braintrust/api/test_transport.py | 62 +++- py/src/braintrust/logger.py | 269 ++++++++++-------- .../braintrust/prompt_cache/test_lru_cache.py | 1 + py/src/braintrust/test_logger.py | 133 +++++---- 5 files changed, 312 insertions(+), 176 deletions(-) diff --git a/py/src/braintrust/api/_transport.py b/py/src/braintrust/api/_transport.py index 46ba68be8..5323d7690 100644 --- a/py/src/braintrust/api/_transport.py +++ b/py/src/braintrust/api/_transport.py @@ -102,6 +102,9 @@ def __init__(self, base_url: str, adapter: HTTPAdapter | None = None): self.base_url = base_url self.token = None self.adapter = adapter + # An adapter handed to us belongs to the caller. `set_http_adapter` installs + # one instance across every connection, so we must not close it. + self._injected_adapter = adapter self._reset(total=0) @@ -121,6 +124,7 @@ def make_long_lived(self) -> None: self._reset() def close(self) -> None: + _unmount_adapter(self.session, self._injected_adapter) self.session.close() @staticmethod @@ -134,6 +138,7 @@ def set_token(self, token: str) -> None: def _set_adapter(self, adapter: HTTPAdapter | None) -> None: self.adapter = adapter + self._injected_adapter = adapter def _reset(self, **retry_kwargs: Any) -> None: self.session = requests.Session() @@ -205,6 +210,7 @@ def __init__( ): custom_transport = session is not None or adapter is not None self._owns_session = session is None + self._injected_adapter = adapter self.session = session if session is not None else requests.Session() if not persist_cookies and self._owns_session: self.session.cookies.set_policy(_RejectCookiesPolicy()) @@ -218,6 +224,7 @@ def __init__( def close(self) -> None: if self._owns_session: + _unmount_adapter(self.session, self._injected_adapter) self.session.close() def __enter__(self) -> "Transport": @@ -395,6 +402,22 @@ def _retry_delay(policy: RetryPolicy, attempt: int, retry_after: float | None) - return min(policy.max_backoff, policy.backoff_factor * (2 ** (attempt - 1))) +def _unmount_adapter(session: requests.Session, adapter: HTTPAdapter | None) -> None: + """Detach a caller-owned adapter so ``Session.close()`` leaves it open. + + ``requests.Session.close()`` closes every mounted adapter. A single adapter + installed via ``set_http_adapter`` is mounted on many sessions at once, so + closing one session would otherwise clear the connection pools that the + other sessions are still using. + """ + + if adapter is None: + return + for prefix, mounted in list(session.adapters.items()): + if mounted is adapter: + del session.adapters[prefix] + + def _request_body_is_replayable(data: Any, files: Any) -> bool: return files is None and (data is None or isinstance(data, (bytes, str))) diff --git a/py/src/braintrust/api/test_transport.py b/py/src/braintrust/api/test_transport.py index 04e566f5d..3ff688997 100644 --- a/py/src/braintrust/api/test_transport.py +++ b/py/src/braintrust/api/test_transport.py @@ -1,6 +1,7 @@ import datetime import io from email.utils import format_datetime +from unittest import mock import pytest import requests @@ -14,7 +15,7 @@ RetryPolicy, ) from braintrust.api._test_server import scripted_server -from braintrust.api._transport import Transport +from braintrust.api._transport import HTTPConnection, Transport from braintrust.util import AugmentedHTTPError from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry @@ -68,13 +69,17 @@ def close(self): super().close() -def test_transport_closes_owned_session(): +def test_transport_closes_owned_session_without_closing_injected_adapter(): adapter = TrackingAdapter() + transport = Transport(adapter=adapter) + session = transport.session - with Transport(adapter=adapter) as transport: - assert transport.session is not None + with mock.patch.object(session, "close", wraps=session.close) as close_spy: + transport.close() - assert adapter.close_count > 0 + close_spy.assert_called_once() + # The adapter belongs to the caller and may be mounted on other sessions. + assert adapter.close_count == 0 def test_transport_does_not_close_injected_session(): @@ -351,3 +356,50 @@ def test_non_retrying_custom_adapter_can_delegate_retries_to_sdk(): assert response.status_code == 200 assert handler.request_count == 2 + + +def test_http_connection_close_does_not_close_shared_adapter(): + adapter = TrackingAdapter() + first = HTTPConnection("http://localhost", adapter=adapter) + second = HTTPConnection("http://localhost", adapter=adapter) + + first.close() + + assert adapter.close_count == 0 + assert second.session.get_adapter("http://localhost") is adapter + + +def test_http_connection_close_closes_self_created_long_lived_adapter(): + conn = HTTPConnection("http://localhost") + conn.make_long_lived() + adapter = conn.adapter + assert adapter is not None + + with mock.patch.object(adapter, "close", wraps=adapter.close) as close_spy: + conn.close() + + # Mounted on both the http:// and https:// prefixes, so closed once per mount. + assert close_spy.call_count > 0 + + +def test_http_connection_close_closes_long_lived_adapter_replaced_by_set_adapter(): + adapter = TrackingAdapter() + conn = HTTPConnection("http://localhost") + conn.make_long_lived() + conn._set_adapter(adapter) + conn._reset() + + conn.close() + + assert adapter.close_count == 0 + + +def test_http_connection_close_does_not_close_adapter_set_after_construction(): + adapter = TrackingAdapter() + conn = HTTPConnection("http://localhost") + conn._set_adapter(adapter) + conn._reset() + + conn.close() + + assert adapter.close_count == 0 diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 95a099d65..70cf41b69 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -442,21 +442,67 @@ def __exit__( @dataclasses.dataclass(frozen=True) -class _LoaderRequestState: +class _LoaderLoginOptions: + """The credential one loader call runs under, and the cache scope it implies.""" + app_url: str - org_id: str - _api_conn: HTTPConnection + api_key: str + org_name: str | None + cache_namespace: str - def api_conn(self) -> HTTPConnection: - return self._api_conn - def close(self) -> None: - self._api_conn.close() +class _LoaderLoginEntry: + """One credential's loader login, tracking who is responsible for closing it. + + The login runs outside the cache lock, so an entry can be evicted while its + login is still in flight. An evicted entry is never handed to a new caller, + so the callers already holding it own it, and the last one to finish closes + it. Deferring to the last release also keeps an eviction from tearing down a + connection that another caller is still issuing a request on. + """ + + def __init__(self, factory: Callable[[], HTTPConnection]): + self._lazy: LazyValue[HTTPConnection] = LazyValue(factory, use_mutex=True) + self._lock = threading.Lock() + self._evicted = False + self._active = 0 + + def acquire(self) -> HTTPConnection: + with self._lock: + self._active += 1 + try: + return self._lazy.get() + except BaseException: + self.release() + raise + + def release(self) -> None: + with self._lock: + self._active -= 1 + should_close = self._evicted and self._active == 0 + if should_close: + self._close() + + def evict(self) -> None: + with self._lock: + self._evicted = True + should_close = self._active == 0 + if should_close: + self._close() + + def _close(self) -> None: + has_succeeded, conn = self._lazy.get_sync() + if has_succeeded and conn is not None: + conn.close() class BraintrustState: def __init__(self): self.id = str(uuid.uuid4()) + self._loader_login_cache: LRUCache[str, _LoaderLoginEntry] = LRUCache( + max_size=16, + on_remove=self._evict_loader_login_entry, + ) self.current_experiment: Experiment | None = None # We use both a ContextVar and a plain attribute for the current logger: # - _cv_logger (ContextVar): Provides async context isolation so different @@ -538,13 +584,7 @@ def default_get_api_conn(): self._otel_flush_callback: Any | None = None def reset_login_info(self): - if hasattr(self, "_loader_login_cache"): - self._loader_login_cache.clear() - else: - self._loader_login_cache: LRUCache[str, LazyValue[_LoaderRequestState]] = LRUCache( - max_size=16, - on_remove=self._close_loader_request_state, - ) + self._loader_login_cache.clear() self.app_url: str | None = None self.app_public_url: str | None = None @@ -564,10 +604,8 @@ def reset_login_info(self): self._user_info: Mapping[str, Any] | None = None @staticmethod - def _close_loader_request_state(_key: str, lazy_state: LazyValue[_LoaderRequestState]) -> None: - has_succeeded, request_state = lazy_state.get_sync() - if has_succeeded and request_state is not None: - request_state.close() + def _evict_loader_login_entry(_key: str, entry: "_LoaderLoginEntry") -> None: + entry.evict() def reset_parent_state(self): # reset possible parent state for tests @@ -721,36 +759,36 @@ def user_info(self) -> Mapping[str, Any]: self._user_info = self.api_conn().get_json("ping") return self._user_info - def loader_request_state( - self, - *, - app_url: str, - api_key: str, - org_name: str | None, - cache_namespace: str, - ) -> "BraintrustState | _LoaderRequestState": + @contextlib.contextmanager + def loader_conn(self, options: "_LoaderLoginOptions") -> "Iterator[HTTPConnection]": + """Yield the API connection for one loader call, releasing it on exit. + + The global login's connection is shared and outlives the call, so it is + yielded as-is. A per-credential connection is owned by its cache entry, + which needs the release to know when an evicted one is safe to close. + """ + if ( self.logged_in - and self.login_token == api_key - and self.app_url == app_url - and (org_name is None or self.org_name == org_name) + and self.login_token == options.api_key + and self.app_url == options.app_url + and (options.org_name is None or self.org_name == options.org_name) ): - return self + yield self.api_conn() + return with self._client_lock: try: - lazy_state = self._loader_login_cache.get(cache_namespace) + entry = self._loader_login_cache.get(options.cache_namespace) except KeyError: - lazy_state = LazyValue( - lambda: _login_to_loader_request_state( - app_url=app_url, - api_key=api_key, - org_name=org_name, - ), - use_mutex=True, - ) - self._loader_login_cache.set(cache_namespace, lazy_state) - return lazy_state.get() + entry = _LoaderLoginEntry(lambda: _login_loader_conn(options)) + self._loader_login_cache.set(options.cache_namespace, entry) + + conn = entry.acquire() + try: + yield conn + finally: + entry.release() def global_bg_logger(self) -> "_BackgroundLogger": return getattr(self._override_bg_logger, "logger", None) or self._global_bg_logger.get() @@ -1904,17 +1942,9 @@ def _resolve_loader_login_options( app_url: str | None, api_key: str | None, org_name: str | None, -) -> tuple[str, str, str | None, str]: +) -> "_LoaderLoginOptions": resolved_app_url = app_url or (_state.app_url if _state.logged_in else None) or _get_app_url() - resolved_api_key = api_key or (_state.login_token if _state.logged_in else None) - if resolved_api_key is None: - resolved_api_key = BraintrustEnv.API_KEY.get(None, use_dotenv=True) - if resolved_api_key is None: - raise ValueError( - "Could not login to Braintrust. You may need to set BRAINTRUST_API_KEY in your environment " - "or nearest .env.braintrust file." - ) - resolved_api_key = HTTPConnection.sanitize_token(resolved_api_key) + resolved_api_key = _require_api_key(api_key or (_state.login_token if _state.logged_in else None)) uses_active_credential = _state.logged_in and resolved_api_key == _state.login_token resolved_org_name = org_name @@ -1925,17 +1955,24 @@ def _resolve_loader_login_options( ["loader-credential", resolved_app_url, resolved_org_name, resolved_api_key], separators=(",", ":"), ) - cache_namespace = f"loader-credential:{hashlib.sha256(namespace_input.encode('utf-8')).hexdigest()}" - return resolved_app_url, resolved_api_key, resolved_org_name, cache_namespace + return _LoaderLoginOptions( + app_url=resolved_app_url, + api_key=resolved_api_key, + org_name=resolved_org_name, + cache_namespace=f"loader-credential:{hashlib.sha256(namespace_input.encode('utf-8')).hexdigest()}", + ) def _is_loader_cache_fallback_error(error: BaseException) -> bool: - pending: list[BaseException] = [error] + """Return whether an error is transient enough to justify serving a cached value. + + The failure that matters is usually wrapped, so walk down the chain until one + link is classifiable. `seen` only guards against a self-referential chain. + """ + + current: BaseException | None = error seen: set[int] = set() - while pending: - current = pending.pop() - if id(current) in seen: - continue + while current is not None and id(current) not in seen: seen.add(id(current)) if isinstance(current, (json.JSONDecodeError, BraintrustJSONDecodeError)): @@ -1945,18 +1982,14 @@ def _is_loader_cache_fallback_error(error: BaseException) -> bool: status_code = getattr(current, "status_code", None) if status_code is None: - response = getattr(current, "response", None) - status_code = getattr(response, "status_code", None) + status_code = getattr(getattr(current, "response", None), "status_code", None) if isinstance(status_code, int): return status_code in DEFAULT_RETRYABLE_STATUSES if isinstance(current, requests_exceptions.RequestException): return is_retryable_request_exception(current) - if current.__cause__ is not None: - pending.append(current.__cause__) - if current.__context__ is not None: - pending.append(current.__context__) + current = current.__cause__ or current.__context__ return False @@ -2001,35 +2034,31 @@ def load_prompt( raise ValueError("Must specify slug") def compute_metadata(): - resolved_app_url, resolved_api_key, resolved_org_name, cache_namespace = _resolve_loader_login_options( + login_options = _resolve_loader_login_options( app_url=app_url, api_key=api_key, org_name=org_name, ) + cache_namespace = login_options.cache_namespace try: - request_state = _state.loader_request_state( - app_url=resolved_app_url, - api_key=resolved_api_key, - org_name=resolved_org_name, - cache_namespace=cache_namespace, - ) - if id: - # Load prompt by ID using the /v1/prompt/{id} endpoint - prompt_args = _populate_args({}, version=version, environment=effective_environment) - response = request_state.api_conn().get_json(f"/v1/prompt/{id}", prompt_args) - # Wrap single prompt response in objects array to match list API format - if response is not None: - response = {"objects": [response]} - else: - args = _populate_args( - {}, - project_name=project, - project_id=project_id, - slug=slug, - version=version, - environment=effective_environment, - ) - response = request_state.api_conn().get_json("/v1/prompt", args) + with _state.loader_conn(login_options) as conn: + if id: + # Load prompt by ID using the /v1/prompt/{id} endpoint + prompt_args = _populate_args({}, version=version, environment=effective_environment) + response = conn.get_json(f"/v1/prompt/{id}", prompt_args) + # Wrap single prompt response in objects array to match list API format + if response is not None: + response = {"objects": [response]} + else: + args = _populate_args( + {}, + project_name=project, + project_id=project_id, + slug=slug, + version=version, + environment=effective_environment, + ) + response = conn.get_json("/v1/prompt", args) except Exception as server_error: if not _is_loader_cache_fallback_error(server_error): raise @@ -2185,32 +2214,28 @@ def load_parameters( effective_environment = None if version is not None else environment should_fall_back_to_cache = version is None and effective_environment is None query_args = _populate_args({}, version=version, environment=effective_environment) - resolved_app_url, resolved_api_key, resolved_org_name, cache_namespace = _resolve_loader_login_options( + login_options = _resolve_loader_login_options( app_url=app_url, api_key=api_key, org_name=org_name, ) + cache_namespace = login_options.cache_namespace try: - request_state = _state.loader_request_state( - app_url=resolved_app_url, - api_key=resolved_api_key, - org_name=resolved_org_name, - cache_namespace=cache_namespace, - ) - if id: - response = request_state.api_conn().get_json(f"/v1/function/{id}", query_args) - if response is not None: - response = {"objects": [response]} - else: - args = _populate_args( - {"function_type": "parameters"}, - project_name=project, - project_id=project_id, - slug=slug, - **query_args, - ) - response = request_state.api_conn().get_json("/v1/function", args) + with _state.loader_conn(login_options) as conn: + if id: + response = conn.get_json(f"/v1/function/{id}", query_args) + if response is not None: + response = {"objects": [response]} + else: + args = _populate_args( + {"function_type": "parameters"}, + project_name=project, + project_id=project_id, + slug=slug, + **query_args, + ) + response = conn.get_json("/v1/function", args) except Exception as server_error: if not _is_loader_cache_fallback_error(server_error): raise @@ -2315,6 +2340,16 @@ def register_otel_flush(callback: Any) -> None: _state.span_cache.disable() +def _require_api_key(api_key: str | None) -> str: + resolved = api_key or BraintrustEnv.API_KEY.get(None, use_dotenv=True) + if resolved is None: + raise ValueError( + "Could not login to Braintrust. You may need to set BRAINTRUST_API_KEY in your environment " + "or nearest .env.braintrust file." + ) + return HTTPConnection.sanitize_token(resolved) + + def _login_with_api_key(*, app_url: str, api_key: str, org_name: str | None) -> tuple[BraintrustClient, LoginResult]: if api_key == TEST_API_KEY: api_url = BraintrustEnv.API_URL.get("https://api.braintrust.ai") @@ -2357,22 +2392,14 @@ def _authenticated_api_conn(api_url: str, api_key: str) -> HTTPConnection: return conn -def _login_to_loader_request_state( - *, - app_url: str, - api_key: str, - org_name: str | None, -) -> _LoaderRequestState: - client, login_result = _login_with_api_key(app_url=app_url, api_key=api_key, org_name=org_name) +def _login_loader_conn(options: _LoaderLoginOptions) -> HTTPConnection: + client, login_result = _login_with_api_key( + app_url=options.app_url, api_key=options.api_key, org_name=options.org_name + ) try: - conn = _authenticated_api_conn(login_result.api_url, api_key) + return _authenticated_api_conn(login_result.api_url, options.api_key) finally: client.close() - return _LoaderRequestState( - app_url=app_url, - org_id=login_result.organization.id, - _api_conn=conn, - ) def login_to_state( diff --git a/py/src/braintrust/prompt_cache/test_lru_cache.py b/py/src/braintrust/prompt_cache/test_lru_cache.py index 58806a2d6..406ef2212 100644 --- a/py/src/braintrust/prompt_cache/test_lru_cache.py +++ b/py/src/braintrust/prompt_cache/test_lru_cache.py @@ -72,6 +72,7 @@ def test_on_remove_runs_for_replacement_eviction_and_clear(self): cache = lru_cache.LRUCache[str, int](max_size=2, on_remove=lambda key, value: removed.append((key, value))) cache.set("a", 1) + cache.set("a", 1) # Re-setting the same value is not a removal. cache.set("a", 2) cache.set("b", 3) cache.set("c", 4) diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index da65e5f17..772bed16f 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -70,29 +70,25 @@ def test_login_to_state_uses_env_braintrust_api_key(tmp_path, monkeypatch): assert state.logged_in is True +def _loader_options(api_key: str, cache_namespace: str) -> logger._LoaderLoginOptions: + return logger._LoaderLoginOptions( + app_url="https://app.example.com", + api_key=api_key, + org_name=None, + cache_namespace=cache_namespace, + ) + + def test_loader_request_state_closes_connections_on_eviction_and_reset(): state = BraintrustState() - state._loader_login_cache = LRUCache(max_size=1, on_remove=state._close_loader_request_state) + state._loader_login_cache = LRUCache(max_size=1, on_remove=state._evict_loader_login_entry) first_conn = MagicMock() second_conn = MagicMock() - request_states = [ - logger._LoaderRequestState("https://app.example.com", "org-a", first_conn), - logger._LoaderRequestState("https://app.example.com", "org-b", second_conn), - ] - - with patch.object(logger, "_login_to_loader_request_state", side_effect=request_states): - state.loader_request_state( - app_url="https://app.example.com", - api_key="first-api-key", - org_name=None, - cache_namespace="first", - ) - state.loader_request_state( - app_url="https://app.example.com", - api_key="second-api-key", - org_name=None, - cache_namespace="second", - ) + with patch.object(logger, "_login_loader_conn", side_effect=[first_conn, second_conn]): + with state.loader_conn(_loader_options("first-api-key", "first")): + pass + with state.loader_conn(_loader_options("second-api-key", "second")): + pass first_conn.close.assert_called_once() second_conn.close.assert_not_called() @@ -102,6 +98,57 @@ def test_loader_request_state_closes_connections_on_eviction_and_reset(): second_conn.close.assert_called_once() +def test_loader_request_state_closes_state_evicted_while_login_is_pending(): + state = BraintrustState() + state._loader_login_cache = LRUCache(max_size=1, on_remove=state._evict_loader_login_entry) + pending_conn = MagicMock() + login_started = threading.Event() + release_login = threading.Event() + + def login(options): + if options.api_key == "slow-api-key": + login_started.set() + assert release_login.wait(5) + return pending_conn + return MagicMock() + + def run_slow_login(): + with state.loader_conn(_loader_options("slow-api-key", "first")): + pass + + with patch.object(logger, "_login_loader_conn", side_effect=login): + thread = threading.Thread(target=run_slow_login) + thread.start() + try: + assert login_started.wait(5) + # Evicts "first" while its login is still in flight, so the cache can no + # longer close whatever that login resolves into. + with state.loader_conn(_loader_options("fast-api-key", "second")): + pass + finally: + release_login.set() + thread.join(5) + + assert not thread.is_alive() + pending_conn.close.assert_called_once() + + +def test_loader_request_state_defers_close_until_last_holder_releases(): + state = BraintrustState() + state._loader_login_cache = LRUCache(max_size=1, on_remove=state._evict_loader_login_entry) + conn = MagicMock() + + with patch.object(logger, "_login_loader_conn", side_effect=[conn, MagicMock()]): + with state.loader_conn(_loader_options("first-api-key", "first")): + # Evicting while the caller is still issuing its request must not close + # the connection out from under it. + with state.loader_conn(_loader_options("second-api-key", "second")): + pass + conn.close.assert_not_called() + + conn.close.assert_called_once() + + class TestInit(TestCase): @staticmethod def _mock_api_client(): @@ -415,12 +462,10 @@ def test_load_prompt_uses_explicit_api_key_without_changing_global_login(): prompt_cache = PromptCache(memory_cache=LRUCache(max_size=10)) request_conn = MagicMock() request_conn.get_json.return_value = _prompt_response("saved-prompt") - request_state = MagicMock() - request_state.api_conn.return_value = request_conn with ( patch.object(logger._state, "_prompt_cache", prompt_cache), - patch.object(logger, "_login_to_loader_request_state", return_value=request_state) as mock_login_to_state, + patch.object(logger, "_login_loader_conn", return_value=request_conn) as mock_login_conn, ): prompt = braintrust.load_prompt( project="test-project", @@ -429,11 +474,10 @@ def test_load_prompt_uses_explicit_api_key_without_changing_global_login(): ) assert prompt.slug == "saved-prompt" - mock_login_to_state.assert_called_once_with( - app_url=logger._state.app_url, - api_key="prompt-api-key", - org_name=None, - ) + (called_options,) = mock_login_conn.call_args.args + assert called_options.app_url == logger._state.app_url + assert called_options.api_key == "prompt-api-key" + assert called_options.org_name is None assert logger._state.login_token == original_login_token @@ -443,12 +487,10 @@ def test_load_parameters_uses_explicit_api_key_without_changing_global_login(): parameters_cache = ParametersCache(memory_cache=LRUCache(max_size=10)) request_conn = MagicMock() request_conn.get_json.return_value = _parameters_response("saved-parameters") - request_state = MagicMock() - request_state.api_conn.return_value = request_conn with ( patch.object(logger._state, "_parameters_cache", parameters_cache), - patch.object(logger, "_login_to_loader_request_state", return_value=request_state) as mock_login_to_state, + patch.object(logger, "_login_loader_conn", return_value=request_conn) as mock_login_conn, ): parameters = braintrust.load_parameters( project="test-project", @@ -457,11 +499,10 @@ def test_load_parameters_uses_explicit_api_key_without_changing_global_login(): ) assert parameters.data == {"prefix": "saved-parameters"} - mock_login_to_state.assert_called_once_with( - app_url=logger._state.app_url, - api_key="parameters-api-key", - org_name=None, - ) + (called_options,) = mock_login_conn.call_args.args + assert called_options.app_url == logger._state.app_url + assert called_options.api_key == "parameters-api-key" + assert called_options.org_name is None assert logger._state.login_token == original_login_token @@ -479,12 +520,10 @@ def test_load_prompt_does_not_fall_back_to_cache_for_non_transient_errors(server prompt_cache = PromptCache(memory_cache=LRUCache(max_size=10)) request_conn = MagicMock() request_conn.get_json.side_effect = [_prompt_response("saved-prompt"), server_error] - request_state = MagicMock() - request_state.api_conn.return_value = request_conn with ( patch.object(logger._state, "_prompt_cache", prompt_cache), - patch.object(logger, "_login_to_loader_request_state", return_value=request_state), + patch.object(logger, "_login_loader_conn", return_value=request_conn), ): first_prompt = braintrust.load_prompt( project="test-project", @@ -507,12 +546,10 @@ def test_load_prompt_uses_same_api_keys_cache_for_transient_errors(): prompt_cache = PromptCache(memory_cache=LRUCache(max_size=10)) request_conn = MagicMock() request_conn.get_json.side_effect = [_prompt_response("saved-prompt"), _http_error(500)] - request_state = MagicMock() - request_state.api_conn.return_value = request_conn with ( patch.object(logger._state, "_prompt_cache", prompt_cache), - patch.object(logger, "_login_to_loader_request_state", return_value=request_state), + patch.object(logger, "_login_loader_conn", return_value=request_conn), ): first_prompt = braintrust.load_prompt( project="test-project", @@ -536,18 +573,14 @@ def test_load_prompt_does_not_use_another_api_keys_transient_fallback_cache(): first_conn.get_json.return_value = _prompt_response("saved-prompt") second_conn = MagicMock() second_conn.get_json.side_effect = _http_error(500) - request_states = {} - for api_key, request_conn in (("first-api-key", first_conn), ("second-api-key", second_conn)): - request_state = MagicMock() - request_state.api_conn.return_value = request_conn - request_states[api_key] = request_state + conns = {"first-api-key": first_conn, "second-api-key": second_conn} - def login_for_api_key(*, api_key, **_kwargs): - return request_states[api_key] + def login_for_api_key(options): + return conns[options.api_key] with ( patch.object(logger._state, "_prompt_cache", prompt_cache), - patch.object(logger, "_login_to_loader_request_state", side_effect=login_for_api_key), + patch.object(logger, "_login_loader_conn", side_effect=login_for_api_key), ): first_prompt = braintrust.load_prompt( project="test-project", @@ -722,7 +755,7 @@ def test_load_parameters_returns_remote_object(self): app_url=None, api_key=None, org_name=None, - )[-1] + ).cache_namespace assert ( logger._state._parameters_cache.get( slug="saved-parameters",