diff --git a/py/src/braintrust/api/__init__.py b/py/src/braintrust/api/__init__.py index fef1d172..b83591de 100644 --- a/py/src/braintrust/api/__init__.py +++ b/py/src/braintrust/api/__init__.py @@ -1,5 +1,11 @@ -"""Braintrust API client package.""" +"""Public Braintrust API client package.""" +from ._routing import EndpointRouter, RequestTarget +from ._service import ClientContext +from .attachments import AttachmentsAPI +from .auth import AuthAPI, LoginResult, OrganizationInfo +from .client import BraintrustClient +from .datasets import DatasetsAPI from .errors import ( BraintrustAPIError, BraintrustHTTPError, @@ -8,16 +14,35 @@ BraintrustTransportError, BraintrustTransportRetryExhaustedError, ) +from .experiments import ExperimentsAPI +from .functions import FunctionsAPI from .policies import RetryMode, RetryPolicy +from .projects import ProjectsAPI +from .prompts import PromptsAPI +from .queries import QueriesAPI __all__ = [ + "AttachmentsAPI", + "AuthAPI", "BraintrustAPIError", + "BraintrustClient", "BraintrustHTTPError", "BraintrustResponseError", "BraintrustRetryExhaustedError", "BraintrustTransportError", "BraintrustTransportRetryExhaustedError", + "ClientContext", + "DatasetsAPI", + "EndpointRouter", + "ExperimentsAPI", + "FunctionsAPI", + "LoginResult", + "OrganizationInfo", + "ProjectsAPI", + "PromptsAPI", + "QueriesAPI", + "RequestTarget", "RetryMode", "RetryPolicy", ] diff --git a/py/src/braintrust/api/_routing.py b/py/src/braintrust/api/_routing.py new file mode 100644 index 00000000..e0f997dc --- /dev/null +++ b/py/src/braintrust/api/_routing.py @@ -0,0 +1,69 @@ +"""Endpoint routing for Braintrust API requests.""" + +import enum +from dataclasses import dataclass + +from ..util import _urljoin + + +_V1_PROXY_SUFFIX = "/v1/proxy" + + +class RequestTarget(enum.Enum): + """A logical Braintrust request destination.""" + + APP = "app" + API = "api" + PROXY = "proxy" + + +def normalize_proxy_url(proxy_url: str) -> str: + """Normalize a Universal Proxy URL to the API host used by SDK routes.""" + + if proxy_url.endswith(_V1_PROXY_SUFFIX): + return proxy_url[: -len(_V1_PROXY_SUFFIX)] + return proxy_url + + +@dataclass +class EndpointRouter: + """Resolve logical Braintrust targets without changing their configured origins.""" + + app_url: str + api_url: str | None = None + proxy_url: str | None = None + is_universal_api: bool = False + + def configure( + self, + *, + api_url: str | None, + proxy_url: str | None, + is_universal_api: bool = False, + ) -> None: + """Apply URLs discovered during authentication.""" + + self.api_url = api_url + self.proxy_url = proxy_url + self.is_universal_api = is_universal_api + + def base_url(self, target: RequestTarget) -> str: + """Return the configured origin for ``target``.""" + + if target is RequestTarget.APP: + return self.app_url + if target is RequestTarget.API: + if not self.api_url: + raise RuntimeError("API URL is unavailable before organization discovery") + return self.api_url + if target is RequestTarget.PROXY: + base_url = self.proxy_url or self.api_url + if not base_url: + raise RuntimeError("Proxy URL is unavailable before organization discovery") + return normalize_proxy_url(base_url) + raise ValueError(f"Unknown request target: {target!r}") + + def resolve(self, target: RequestTarget, path: str) -> str: + """Resolve ``path`` against the origin for ``target``.""" + + return _urljoin(self.base_url(target), path) diff --git a/py/src/braintrust/api/_service.py b/py/src/braintrust/api/_service.py new file mode 100644 index 00000000..f6522c0b --- /dev/null +++ b/py/src/braintrust/api/_service.py @@ -0,0 +1,51 @@ +"""Shared resource service primitives.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from ._routing import EndpointRouter, RequestTarget +from ._transport import Transport + + +@dataclass(frozen=True) +class ClientContext: + """Organization and credential context shared by resource services.""" + + org_id: str + org_name: str + + +class ResourceAPI: + """Base class for synchronous resource services.""" + + def __init__( + self, + transport: Transport, + router: EndpointRouter, + context: ClientContext, + api_key: str, + ): + self._transport = transport + self._router = router + self._context = context + self._api_key = api_key + + def _request_json( + self, + target: RequestTarget, + method: str, + path: str, + *, + headers: Mapping[str, str] | None = None, + **kwargs: Any, + ) -> Any: + request_headers = {"Authorization": f"Bearer {self._api_key}"} + if headers: + request_headers.update(headers) + return self._transport.request_json( + method, + self._router.resolve(target, path), + headers=request_headers, + **kwargs, + ) diff --git a/py/src/braintrust/api/_transport.py b/py/src/braintrust/api/_transport.py index 6dd96c0f..410d57c5 100644 --- a/py/src/braintrust/api/_transport.py +++ b/py/src/braintrust/api/_transport.py @@ -1,6 +1,7 @@ """Legacy and policy-aware HTTP transport primitives for the Braintrust SDK.""" import datetime +import http.cookiejar import logging import sys import time @@ -28,6 +29,11 @@ logger = logging.getLogger(__name__) +class _RejectCookiesPolicy(http.cookiejar.DefaultCookiePolicy): + def set_ok(self, cookie: Any, request: Any) -> bool: + return False + + class RetryRequestExceptionsAdapter(HTTPAdapter): """An HTTP adapter that automatically retries requests on connection exceptions. @@ -189,6 +195,7 @@ def __init__( session: requests.Session | None = None, adapter: HTTPAdapter | None = None, enable_sdk_retries: bool | None = None, + persist_cookies: bool = True, sleep: Callable[[float], None] = time.sleep, monotonic: Callable[[], float] = time.monotonic, wall_clock: Callable[[], float] = time.time, @@ -196,6 +203,8 @@ def __init__( custom_transport = session is not None or adapter is not None self._owns_session = session is None 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()) self._sdk_retries_enabled = not custom_transport if enable_sdk_retries is None else enable_sdk_retries if adapter is not None: self.session.mount("http://", adapter) diff --git a/py/src/braintrust/api/attachments.py b/py/src/braintrust/api/attachments.py new file mode 100644 index 00000000..2c56c5f9 --- /dev/null +++ b/py/src/braintrust/api/attachments.py @@ -0,0 +1,10 @@ +"""Attachment metadata API service.""" + +from ._service import ResourceAPI + + +class AttachmentsAPI(ResourceAPI): + """Synchronous attachment metadata operations. + + Signed object-storage traffic remains outside the routed transport. + """ diff --git a/py/src/braintrust/api/auth.py b/py/src/braintrust/api/auth.py new file mode 100644 index 00000000..81938fda --- /dev/null +++ b/py/src/braintrust/api/auth.py @@ -0,0 +1,122 @@ +"""Authentication and organization discovery service.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from ..env import BraintrustEnv +from ._routing import EndpointRouter, RequestTarget +from ._transport import HTTPConnection, Transport +from .policies import RetryMode + + +@dataclass(frozen=True) +class OrganizationInfo: + """Organization routing information returned by API-key login.""" + + id: str + name: str + api_url: str | None + proxy_url: str | None + realtime_url: str | None + is_universal_api: bool + git_metadata: Mapping[str, Any] | None + raw: Mapping[str, Any] + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "OrganizationInfo": + """Parse an additive login response while retaining unknown fields.""" + + org_id = value.get("id") + name = value.get("name") + if not isinstance(org_id, str) or not isinstance(name, str): + raise ValueError("Organization login data must include string id and name fields") + + def optional_string(field: str) -> str | None: + result = value.get(field) + return result if isinstance(result, str) and result else None + + git_metadata = value.get("git_metadata") + if not isinstance(git_metadata, Mapping): + git_metadata = None + + return cls( + id=org_id, + name=name, + api_url=optional_string("api_url"), + proxy_url=optional_string("proxy_url"), + realtime_url=optional_string("realtime_url"), + is_universal_api=bool(value.get("is_universal_api", False)), + git_metadata=MappingProxyType(dict(git_metadata)) if git_metadata is not None else None, + raw=MappingProxyType(dict(value)), + ) + + +@dataclass(frozen=True) +class LoginResult: + """Selected organization and the complete login response.""" + + organization: OrganizationInfo + response: Mapping[str, Any] + + +class AuthAPI: + """Authenticate an API key and configure an endpoint router.""" + + def __init__(self, transport: Transport, router: EndpointRouter): + self._transport = transport + self._router = router + + def login( + self, + api_key: str, + *, + org_name: str | None = None, + api_url: str | None = None, + proxy_url: str | None = None, + ) -> LoginResult: + """Log in, select an organization, and apply routing override precedence.""" + + api_key = HTTPConnection.sanitize_token(api_key) + response = self._transport.request_json( + "POST", + self._router.resolve(RequestTarget.APP, "/api/apikey/login"), + headers={"Authorization": f"Bearer {api_key}"}, + retry_mode=RetryMode.SAFE_READ, + ) + if not isinstance(response, Mapping): + raise ValueError("API-key login returned a non-object response") + raw_orgs = response.get("org_info") + if not isinstance(raw_orgs, Sequence) or isinstance(raw_orgs, (str, bytes)): + raise ValueError("API-key login response did not include an organization list") + + organizations = [OrganizationInfo.from_dict(org) for org in raw_orgs if isinstance(org, Mapping)] + organization = self._select_organization(organizations, org_name) + + resolved_api_url = api_url or BraintrustEnv.API_URL.get(organization.api_url) + resolved_proxy_url = proxy_url or BraintrustEnv.PROXY_URL.get(organization.proxy_url) + if not resolved_api_url: + if org_name: + raise ValueError( + f"Unable to log into organization '{org_name}'." + " Are you sure this credential is scoped to the organization?" + ) + raise ValueError("Unable to log into any organization with the provided credential.") + + self._router.configure( + api_url=resolved_api_url, + proxy_url=resolved_proxy_url, + is_universal_api=organization.is_universal_api, + ) + return LoginResult(organization=organization, response=MappingProxyType(dict(response))) + + @staticmethod + def _select_organization(organizations: Sequence[OrganizationInfo], org_name: str | None) -> OrganizationInfo: + if not organizations: + raise ValueError("This user is not part of any organizations.") + for organization in organizations: + if org_name is None or organization.name == org_name: + return organization + choices = ", ".join(organization.name for organization in organizations) + raise ValueError(f"Organization {org_name} not found. Must be one of {choices}") diff --git a/py/src/braintrust/api/client.py b/py/src/braintrust/api/client.py new file mode 100644 index 00000000..883592dd --- /dev/null +++ b/py/src/braintrust/api/client.py @@ -0,0 +1,154 @@ +"""Synchronous Braintrust API client facade.""" + +from typing import Any + +import requests +from requests.adapters import HTTPAdapter + +from ..env import BraintrustEnv, resolve_app_url, resolve_org_name +from ._routing import EndpointRouter +from ._service import ClientContext +from ._transport import HTTPConnection, Transport +from .attachments import AttachmentsAPI +from .auth import AuthAPI, LoginResult +from .datasets import DatasetsAPI +from .experiments import ExperimentsAPI +from .functions import FunctionsAPI +from .projects import ProjectsAPI +from .prompts import PromptsAPI +from .queries import QueriesAPI + + +class BraintrustClient: + """Synchronous resource-oriented client for the Braintrust API. + + The convenience constructor authenticates through the app origin, selects an + organization, and configures API and proxy routing on one shared transport. + """ + + def __init__( + self, + *, + api_key: str | None = None, + org_name: str | None = None, + app_url: str | None = None, + api_url: str | None = None, + proxy_url: str | None = None, + session: requests.Session | None = None, + adapter: HTTPAdapter | None = None, + transport: Transport | None = None, + enable_sdk_retries: bool | None = None, + ): + if transport is not None and (session is not None or adapter is not None or enable_sdk_retries is not None): + raise ValueError("transport cannot be combined with session, adapter, or enable_sdk_retries") + + resolved_api_key = api_key or BraintrustEnv.API_KEY.get(None, use_dotenv=True) + if not resolved_api_key: + 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_org_name = resolve_org_name(org_name) + resolved_app_url = resolve_app_url(app_url) + + self._owns_transport = transport is None + self.transport = transport or Transport( + session=session, + adapter=adapter, + enable_sdk_retries=enable_sdk_retries, + persist_cookies=False, + ) + self.router = EndpointRouter(app_url=resolved_app_url) + auth = AuthAPI(self.transport, self.router) + try: + result = auth.login( + resolved_api_key, + org_name=resolved_org_name, + api_url=api_url, + proxy_url=proxy_url, + ) + except Exception: + if self._owns_transport: + self.transport.close() + raise + + self._initialize_services(result, resolved_api_key) + + @classmethod + def from_transport( + cls, + *, + transport: Transport, + router: EndpointRouter, + api_key: str, + org_id: str, + org_name: str, + login_result: LoginResult | None = None, + ) -> "BraintrustClient": + """Build a client around an already-authenticated transport and router.""" + + client = cls.__new__(cls) + client._owns_transport = False + client.transport = transport + client.router = router + client._initialize_services_from_context( + ClientContext(org_id=org_id, org_name=org_name), + HTTPConnection.sanitize_token(api_key), + login_result, + ) + return client + + def _initialize_services(self, result: LoginResult, api_key: str) -> None: + organization = result.organization + self._initialize_services_from_context( + ClientContext(org_id=organization.id, org_name=organization.name), + api_key, + result, + ) + + def _initialize_services_from_context( + self, + context: ClientContext, + api_key: str, + login_result: LoginResult | None, + ) -> None: + self.context = context + self.api_key = api_key + self._login_result = login_result + service_args: tuple[Any, ...] = (self.transport, self.router, context, api_key) + self.projects = ProjectsAPI(*service_args) + self.experiments = ExperimentsAPI(*service_args) + self.datasets = DatasetsAPI(*service_args) + self.prompts = PromptsAPI(*service_args) + self.functions = FunctionsAPI(*service_args) + self.queries = QueriesAPI(*service_args) + self.attachments = AttachmentsAPI(*service_args) + + @property + def login_result(self) -> LoginResult: + """Return organization discovery details for a bootstrapped client.""" + + if self._login_result is None: + raise RuntimeError("Login details are unavailable for a pre-authenticated client") + return self._login_result + + @property + def org_id(self) -> str: + return self.context.org_id + + @property + def org_name(self) -> str: + return self.context.org_name + + def close(self) -> None: + """Close the transport when it was created by this client.""" + + if self._owns_transport: + self.transport.close() + + def __enter__(self) -> "BraintrustClient": + return self + + def __exit__(self, *_: Any) -> None: + self.close() diff --git a/py/src/braintrust/api/datasets.py b/py/src/braintrust/api/datasets.py new file mode 100644 index 00000000..bc565343 --- /dev/null +++ b/py/src/braintrust/api/datasets.py @@ -0,0 +1,10 @@ +"""Dataset API service.""" + +from ._service import ResourceAPI + + +class DatasetsAPI(ResourceAPI): + """Synchronous dataset operations. + + Endpoint methods are added as dataset call sites migrate to the API client. + """ diff --git a/py/src/braintrust/api/experiments.py b/py/src/braintrust/api/experiments.py new file mode 100644 index 00000000..6673fd1c --- /dev/null +++ b/py/src/braintrust/api/experiments.py @@ -0,0 +1,10 @@ +"""Experiment API service.""" + +from ._service import ResourceAPI + + +class ExperimentsAPI(ResourceAPI): + """Synchronous experiment operations. + + Endpoint methods are added as experiment call sites migrate to the API client. + """ diff --git a/py/src/braintrust/api/functions.py b/py/src/braintrust/api/functions.py new file mode 100644 index 00000000..94716b00 --- /dev/null +++ b/py/src/braintrust/api/functions.py @@ -0,0 +1,10 @@ +"""Function metadata API service.""" + +from ._service import ResourceAPI + + +class FunctionsAPI(ResourceAPI): + """Synchronous function metadata operations. + + Invocation remains a specialized client and is not implemented here. + """ diff --git a/py/src/braintrust/api/projects.py b/py/src/braintrust/api/projects.py new file mode 100644 index 00000000..bb00abbc --- /dev/null +++ b/py/src/braintrust/api/projects.py @@ -0,0 +1,10 @@ +"""Project API service.""" + +from ._service import ResourceAPI + + +class ProjectsAPI(ResourceAPI): + """Synchronous project operations. + + Endpoint methods are added as project call sites migrate to the API client. + """ diff --git a/py/src/braintrust/api/prompts.py b/py/src/braintrust/api/prompts.py new file mode 100644 index 00000000..96e21c8b --- /dev/null +++ b/py/src/braintrust/api/prompts.py @@ -0,0 +1,10 @@ +"""Prompt API service.""" + +from ._service import ResourceAPI + + +class PromptsAPI(ResourceAPI): + """Synchronous prompt operations. + + Endpoint methods are added as prompt call sites migrate to the API client. + """ diff --git a/py/src/braintrust/api/queries.py b/py/src/braintrust/api/queries.py new file mode 100644 index 00000000..f84878d0 --- /dev/null +++ b/py/src/braintrust/api/queries.py @@ -0,0 +1,10 @@ +"""Query API service.""" + +from ._service import ResourceAPI + + +class QueriesAPI(ResourceAPI): + """Synchronous query operations. + + Endpoint methods are added as query call sites migrate to the API client. + """ diff --git a/py/src/braintrust/api/test_client.py b/py/src/braintrust/api/test_client.py new file mode 100644 index 00000000..ebc4ae9d --- /dev/null +++ b/py/src/braintrust/api/test_client.py @@ -0,0 +1,222 @@ +import contextlib +import http.server +import json +import socketserver +import threading + +import pytest +import requests +from braintrust import logger +from braintrust.api import BraintrustClient, BraintrustHTTPError, EndpointRouter, RequestTarget +from braintrust.api._transport import RetryRequestExceptionsAdapter +from braintrust.logger import BraintrustState, login_to_state +from requests.adapters import HTTPAdapter + + +@contextlib.contextmanager +def login_server(orgs, *, status=200, response_headers=None): + body = json.dumps({"org_info": orgs}).encode() + + class LoginHandler(http.server.BaseHTTPRequestHandler): + request_count = 0 + authorization = None + + def log_message(self, format, *args): + pass + + def do_POST(self): + type(self).request_count += 1 + type(self).authorization = self.headers.get("Authorization") + self.send_response(status) + for name, value in (response_headers or {}).items(): + self.send_header(name, value) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), LoginHandler) + server.daemon_threads = True + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", LoginHandler + finally: + server.shutdown() + server.server_close() + + +def test_endpoint_router_preserves_origins_and_proxy_fallback(): + router = EndpointRouter( + app_url="https://app.example.com/", + api_url="https://api.example.com/", + ) + + assert router.resolve(RequestTarget.APP, "/api/apikey/login") == "https://app.example.com/api/apikey/login" + assert router.resolve(RequestTarget.API, "ping") == "https://api.example.com/ping" + assert router.resolve(RequestTarget.PROXY, "function/invoke") == "https://api.example.com/function/invoke" + + router.proxy_url = "https://universal.example.com/v1/proxy" + assert router.resolve(RequestTarget.PROXY, "function/invoke") == "https://universal.example.com/function/invoke" + + +def test_client_bootstraps_selected_org_on_one_session(monkeypatch): + monkeypatch.delenv("BRAINTRUST_API_URL", raising=False) + monkeypatch.delenv("BRAINTRUST_PROXY_URL", raising=False) + orgs = [ + { + "id": "org-1", + "name": "first", + "api_url": "https://api-1.example.com", + "proxy_url": None, + }, + { + "id": "org-2", + "name": "selected", + "api_url": "https://api-2.example.com", + "proxy_url": "https://api-2.example.com/v1/proxy", + "is_universal_api": True, + "new_server_field": {"preserved": True}, + }, + ] + session = requests.Session() + session.cookies.set("existing", "yes") + cookie_policy = session.cookies.get_policy() + with login_server(orgs, response_headers={"Set-Cookie": "accepted=yes"}) as (app_url, handler): + client = BraintrustClient(api_key="secret\n", org_name="selected", app_url=app_url, session=session) + + assert handler.request_count == 1 + assert handler.authorization == "Bearer secret" + assert client.org_id == "org-2" + assert client.org_name == "selected" + assert client.router.api_url == "https://api-2.example.com" + assert client.router.is_universal_api is True + assert client.router.resolve(RequestTarget.PROXY, "ping") == "https://api-2.example.com/ping" + assert client.login_result.organization.raw["new_server_field"] == {"preserved": True} + assert not hasattr(client, "auth") + assert "Authorization" not in session.headers + assert session.cookies.get("existing") == "yes" + assert session.cookies.get_policy() is cookie_policy + assert all( + service._transport is client.transport + for service in ( + client.projects, + client.experiments, + client.datasets, + client.prompts, + client.functions, + client.queries, + client.attachments, + ) + ) + + +def test_sdk_owned_session_rejects_response_cookies(): + orgs = [{"id": "org-1", "name": "org", "api_url": "https://api.example.com"}] + with login_server(orgs, response_headers={"Set-Cookie": "ignored=yes"}) as (app_url, _): + client = BraintrustClient(api_key="secret", app_url=app_url) + + assert not client.transport.session.cookies + + +def test_client_url_override_precedence(monkeypatch): + monkeypatch.setenv("BRAINTRUST_API_URL", "https://api-env.example.com") + monkeypatch.setenv("BRAINTRUST_PROXY_URL", "https://proxy-env.example.com") + orgs = [ + { + "id": "org-1", + "name": "org", + "api_url": "https://api-discovered.example.com", + "proxy_url": "https://proxy-discovered.example.com", + } + ] + with login_server(orgs) as (app_url, _): + env_client = BraintrustClient(api_key="secret", app_url=app_url) + explicit_client = BraintrustClient( + api_key="secret", + app_url=app_url, + api_url="https://api-explicit.example.com", + proxy_url="https://proxy-explicit.example.com", + ) + + assert env_client.router.api_url == "https://api-env.example.com" + assert env_client.router.proxy_url == "https://proxy-env.example.com" + assert explicit_client.router.api_url == "https://api-explicit.example.com" + assert explicit_client.router.proxy_url == "https://proxy-explicit.example.com" + + +def test_custom_adapter_disables_bootstrap_retries(): + orgs = [{"id": "org-1", "name": "org", "api_url": "https://api.example.com"}] + with login_server(orgs, status=503) as (app_url, handler): + with pytest.raises(BraintrustHTTPError): + BraintrustClient(api_key="secret", app_url=app_url, adapter=HTTPAdapter()) + + assert handler.request_count == 1 + + +def test_login_to_state_hydrates_isolated_legacy_connections(monkeypatch): + monkeypatch.delenv("BRAINTRUST_API_URL", raising=False) + monkeypatch.delenv("BRAINTRUST_PROXY_URL", raising=False) + with login_server([]) as (app_url, _): + orgs = [ + { + "id": "org-1", + "name": "org", + "api_url": app_url, + "proxy_url": app_url, + "git_metadata": {}, + } + ] + with login_server(orgs) as (login_url, _): + state = login_to_state(api_key="secret", app_url=login_url, org_name="org") + + assert state.api_client().org_id == "org-1" + assert state.git_metadata_settings is None + assert state._client.transport.session is not state.api_conn().session + assert state._client.transport.session is not state.app_conn().session + assert state.api_conn().session.headers["Authorization"] == "Bearer secret" + assert state.app_conn().session.headers["Authorization"] == "Bearer secret" + assert state.proxy_conn().session.headers["Authorization"] == "Bearer secret" + assert isinstance(state.api_conn().adapter, RetryRequestExceptionsAdapter) + assert isinstance(state.app_conn().adapter, RetryRequestExceptionsAdapter) + assert isinstance(state.proxy_conn().adapter, RetryRequestExceptionsAdapter) + + +def test_legacy_adapter_mutation_keeps_characterized_target_scope(monkeypatch): + state = BraintrustState() + state.app_url = "https://app.example.com" + state.api_url = "https://api.example.com" + state.proxy_url = "https://proxy.example.com" + monkeypatch.setattr(logger, "_state", state) + monkeypatch.setattr(logger, "_http_adapter", None) + + app_connection = state.app_conn() + api_connection = state.api_conn() + proxy_connection = state.proxy_conn() + adapter = HTTPAdapter() + logger.set_http_adapter(adapter) + + assert app_connection.adapter is adapter + assert api_connection.adapter is adapter + assert proxy_connection.adapter is None + + +def test_state_concurrent_lazy_access_bootstraps_once(monkeypatch): + monkeypatch.delenv("BRAINTRUST_API_URL", raising=False) + monkeypatch.delenv("BRAINTRUST_PROXY_URL", raising=False) + state = BraintrustState() + with login_server([]) as (api_url, _): + orgs = [{"id": "org-1", "name": "org", "api_url": api_url, "proxy_url": None}] + with login_server(orgs) as (app_url, handler): + monkeypatch.setenv("BRAINTRUST_API_KEY", "secret") + monkeypatch.setenv("BRAINTRUST_APP_URL", app_url) + clients = [] + threads = [threading.Thread(target=lambda: clients.append(state.api_client())) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert handler.request_count == 1 + assert len(clients) == 8 + assert all(client is clients[0] for client in clients) diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index d7dceb8b..98103994 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -17,6 +17,23 @@ _Parser = Callable[[str], EnvValue | None] BRAINTRUST_ENV_FILE = ".env.braintrust" BRAINTRUST_ENV_SEARCH_PARENT_LIMIT = 64 +DEFAULT_APP_URL = "https://www.braintrust.dev" + + +def resolve_app_url(app_url: str | None = None) -> str: + """Resolve an explicit or environment-configured Braintrust app URL.""" + + if app_url: + return app_url + return os.getenv("BRAINTRUST_APP_URL", DEFAULT_APP_URL) + + +def resolve_org_name(org_name: str | None = None) -> str | None: + """Resolve an explicit or environment-configured organization name.""" + + if org_name: + return org_name + return os.getenv("BRAINTRUST_ORG_NAME") def parse_float(value: str) -> float | None: @@ -197,6 +214,8 @@ def __get__(self, instance: object, owner: type | None = None) -> bool: class BraintrustEnv: API_KEY = EnvVar("BRAINTRUST_API_KEY", EnvParser.STRING) + API_URL = EnvVar("BRAINTRUST_API_URL", EnvParser.STRING) + PROXY_URL = EnvVar("BRAINTRUST_PROXY_URL", EnvParser.STRING) HTTP_TIMEOUT = EnvVar("BRAINTRUST_HTTP_TIMEOUT", EnvParser.FLOAT) SYNC_FLUSH = EnvVar("BRAINTRUST_SYNC_FLUSH", EnvParser.BOOL) MAX_REQUEST_SIZE = EnvVar("BRAINTRUST_MAX_REQUEST_SIZE", EnvParser.INT) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 5f592784..12af0a9f 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -41,8 +41,11 @@ from requests.adapters import HTTPAdapter from . import context, id_gen -from .api._transport import HTTPConnection +from .api._routing import EndpointRouter, normalize_proxy_url +from .api._transport import HTTPConnection, Transport from .api._transport import RetryRequestExceptionsAdapter as RetryRequestExceptionsAdapter +from .api.client import BraintrustClient +from .api.errors import BraintrustHTTPError from .bt_json import bt_dumps, bt_safe_deep_copy from .db_fields import ( AUDIT_METADATA_FIELD, @@ -53,7 +56,8 @@ TRANSACTION_ID_FIELD, VALID_SOURCES, ) -from .env import BraintrustEnv +from .env import DEFAULT_APP_URL as _DEFAULT_APP_URL +from .env import BraintrustEnv, resolve_app_url, resolve_org_name from .generated_types import ( AttachmentReference, AttachmentStatus, @@ -168,8 +172,7 @@ class SpanInternalOptions(TypedDict, total=False): TEST_API_KEY = "___TEST_API_KEY__" - -DEFAULT_APP_URL = "https://www.braintrust.dev" +DEFAULT_APP_URL = _DEFAULT_APP_URL def _get_exporter(): @@ -429,18 +432,6 @@ def __exit__( NOOP_SPAN: Span = _NoopSpan() NOOP_SPAN_PERMALINK = "https://www.braintrust.dev/noop-span" -_V1_PROXY_SUFFIX = "/v1/proxy" - - -def _normalize_proxy_conn_url(proxy_url: str) -> str: - # proxy_url may point at the universal proxy (`{api_url}/v1/proxy`) for - # EU/self-hosted orgs, but proxy_conn only targets Braintrust API endpoints - # (e.g. function/invoke, function/sandbox-list) served at the API host root. - # Drop the suffix so these requests resolve on all data planes. - if proxy_url.endswith(_V1_PROXY_SUFFIX): - return proxy_url[: -len(_V1_PROXY_SUFFIX)] - return proxy_url - class BraintrustState: def __init__(self): @@ -467,6 +458,7 @@ def __init__(self): # Context manager is dynamically selected based on current environment self._context_manager = None self._context_manager_lock = threading.Lock() + self._client_lock = threading.RLock() def default_get_api_conn(): self.login() @@ -532,12 +524,14 @@ def reset_login_info(self): self.org_name: str | None = None self.api_url: str | None = None self.proxy_url: str | None = None + self.is_universal_api: bool = False self.logged_in: bool = False self.git_metadata_settings: GitMetadataSettings | None = None self._app_conn: HTTPConnection | None = None self._api_conn: HTTPConnection | None = None self._proxy_conn: HTTPConnection | None = None + self._client: BraintrustClient | None = None self._user_info: Mapping[str, Any] | None = None def reset_parent_state(self): @@ -614,6 +608,7 @@ def copy_state(self, other: "BraintrustState"): "_context_manager", "_last_otel_setting", "_context_manager_lock", + "_client_lock", ) } ) @@ -625,28 +620,40 @@ def login( org_name: str | None = None, force_login: bool = False, ) -> None: - if not force_login and self.logged_in: - # We have already logged in. If any provided login inputs disagree - # with our existing settings, raise an Exception warning the user to - # try again with `force_login=True`. - def check_updated_param(varname, arg, orig): - if arg is not None and orig is not None and arg != orig: - raise Exception( - f"Re-logging in with different {varname} ({arg}) than original ({orig}). To force re-login, pass `force_login=True`" - ) + with self._client_lock: + if not force_login and self.logged_in: + # We have already logged in. If any provided login inputs disagree + # with our existing settings, raise an Exception warning the user to + # try again with `force_login=True`. + def check_updated_param(varname, arg, orig): + if arg is not None and orig is not None and arg != orig: + raise Exception( + f"Re-logging in with different {varname} ({arg}) than original ({orig}). To force re-login, pass `force_login=True`" + ) + + sanitized_api_key = HTTPConnection.sanitize_token(api_key) if api_key else None + check_updated_param("app_url", app_url, self.app_url) + check_updated_param("api_key", sanitized_api_key, self.login_token) + check_updated_param("org_name", org_name, self.org_name) + return - sanitized_api_key = HTTPConnection.sanitize_token(api_key) if api_key else None - check_updated_param("app_url", app_url, self.app_url) - check_updated_param("api_key", sanitized_api_key, self.login_token) - check_updated_param("org_name", org_name, self.org_name) - return + state = login_to_state( + app_url=app_url, + api_key=api_key, + org_name=org_name, + ) + self.copy_state(state) - state = login_to_state( - app_url=app_url, - api_key=api_key, - org_name=org_name, - ) - self.copy_state(state) + def api_client(self) -> BraintrustClient: + """Return the lazily bootstrapped resource client.""" + + if self._client is None: + with self._client_lock: + if self._client is None: + self.login() + if self._client is None: + raise RuntimeError("Braintrust API client was not initialized during login") + return self._client def app_conn(self): if not self._app_conn: @@ -669,7 +676,7 @@ def proxy_conn(self): if not self._proxy_conn: if not self.proxy_url: raise RuntimeError("Must initialize proxy_url before requesting proxy_conn") - self._proxy_conn = HTTPConnection(_normalize_proxy_conn_url(self.proxy_url), adapter=_http_adapter) + self._proxy_conn = HTTPConnection(normalize_proxy_url(self.proxy_url), adapter=_http_adapter) return self._proxy_conn def user_info(self) -> Mapping[str, Any]: @@ -2160,9 +2167,9 @@ def login_to_state( state.app_public_url = app_public_url state.org_name = org_name - conn = None if api_key == TEST_API_KEY: - # a small hook for pseudo-logins + # 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", @@ -2172,52 +2179,57 @@ def login_to_state( } ] _check_org_info(state, test_org_info, org_name) + router = EndpointRouter(app_url=state.app_url, api_url=state.api_url, proxy_url=state.proxy_url) + state._client = BraintrustClient.from_transport( + transport=Transport(adapter=_http_adapter), + router=router, + api_key=TEST_API_KEY, + org_id=cast(str, state.org_id), + org_name=cast(str, state.org_name), + ) state.login_token = TEST_API_KEY state.logged_in = True return state - elif api_key is not None: - app_conn = HTTPConnection(state.app_url, adapter=_http_adapter) - app_conn.set_token(api_key) - resp = app_conn.post("api/apikey/login") - if not resp.ok: - masked_api_key = mask_api_key(api_key) - raise ValueError(f"Invalid API key {masked_api_key}: [{resp.status_code}] {resp.text}") - info = resp.json() - - _check_org_info(state, info["org_info"], org_name) - - if not state.api_url: - if org_name: - raise ValueError( - f"Unable to log into organization '{org_name}'." - " Are you sure this credential is scoped to the organization?" - ) - else: - raise ValueError("Unable to log into any organization with the provided credential.") - - conn = state.api_conn() - conn.set_token(api_key) - if not conn: + 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." ) - # make_long_lived() allows the connection to retry if it breaks, which we're okay with after - # this point because we know the connection _can_ successfully ping. + try: + client = BraintrustClient(api_key=api_key, org_name=org_name, app_url=state.app_url, adapter=_http_adapter) + except BraintrustHTTPError as exc: + 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 + + organization = client.login_result.organization + state._client = client + state.org_id = client.org_id + state.org_name = client.org_name + state.api_url = client.router.api_url + state.proxy_url = client.router.proxy_url + state.is_universal_api = client.router.is_universal_api + state.git_metadata_settings = ( + GitMetadataSettings(**organization.git_metadata) if organization.git_metadata else None + ) + + # 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() - # Same for the app conn, which we know is valid because we have - # successfully logged in. - state.app_conn().make_long_lived() + app_connection = state.app_conn() + app_connection.set_token(api_key) + app_connection.make_long_lived() - # Set the same token in the API - state.app_conn().set_token(conn.token) if state.proxy_url: - state.proxy_conn().set_token(conn.token) - state.proxy_conn().make_long_lived() - state.login_token = conn.token + proxy_connection = state.proxy_conn() + 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. @@ -2861,8 +2873,9 @@ def _check_org_info(state, org_info, org_name): if org_name is None or orgs["name"] == org_name: state.org_id = orgs["id"] state.org_name = orgs["name"] - state.api_url = os.environ.get("BRAINTRUST_API_URL", orgs["api_url"]) - state.proxy_url = os.environ.get("BRAINTRUST_PROXY_URL", orgs["proxy_url"]) + state.api_url = BraintrustEnv.API_URL.get(orgs.get("api_url")) + state.proxy_url = BraintrustEnv.PROXY_URL.get(orgs.get("proxy_url")) + state.is_universal_api = bool(orgs.get("is_universal_api", False)) state.git_metadata_settings = ( GitMetadataSettings(**orgs["git_metadata"]) if orgs.get("git_metadata") else None ) @@ -5988,15 +6001,11 @@ def get_prompt_versions(project_id: str, prompt_id: str) -> list[str]: def _get_app_url(app_url: str | None = None) -> str: - if app_url: - return app_url - return os.getenv("BRAINTRUST_APP_URL", DEFAULT_APP_URL) + return resolve_app_url(app_url) def _get_org_name(org_name: str | None = None) -> str | None: - if org_name: - return org_name - return os.getenv("BRAINTRUST_ORG_NAME") + return resolve_org_name(org_name) def _get_error_link(msg="") -> str: diff --git a/py/src/braintrust/type_tests/test_api_client.py b/py/src/braintrust/type_tests/test_api_client.py new file mode 100644 index 00000000..aec402b0 --- /dev/null +++ b/py/src/braintrust/type_tests/test_api_client.py @@ -0,0 +1,25 @@ +"""Static and runtime checks for the public API client facade.""" + +from typing import TYPE_CHECKING + +from braintrust.api import BraintrustClient, EndpointRouter, RequestTarget + + +if TYPE_CHECKING: + client = BraintrustClient(api_key="key", org_name="org") + client_with_overrides = BraintrustClient( + api_key="key", + app_url="https://app.example.com", + api_url="https://api.example.com", + proxy_url="https://proxy.example.com", + ) + org_id: str = client.org_id + org_name: str = client.org_name + api_url: str | None = client_with_overrides.router.api_url + + +def test_api_client_public_types() -> None: + router = EndpointRouter(app_url="https://app.example.com", api_url="https://api.example.com") + + assert router.resolve(RequestTarget.API, "ping") == "https://api.example.com/ping" + assert BraintrustClient.__name__ == "BraintrustClient"