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" )