From 4e35df9d517c89d85f2bd221c06694218fb3c06d Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Fri, 31 Jul 2026 14:31:04 -0400 Subject: [PATCH] ref(api): establish transport boundary for API client Braintrust HTTP mechanics currently live inside logger.py, coupling transport to logging and leaving no isolated foundation for a resource-oriented API client. Establish this boundary first so retries, structured errors, and endpoint services can be centralized in follow-up work without mixing those behavior changes into the initial refactor. --- py/src/braintrust/api/__init__.py | 3 + py/src/braintrust/api/_transport.py | 160 ++++++++++++++++++++++++++++ py/src/braintrust/logger.py | 152 +------------------------- py/src/braintrust/test_http.py | 46 +++++++- 4 files changed, 210 insertions(+), 151 deletions(-) create mode 100644 py/src/braintrust/api/__init__.py create mode 100644 py/src/braintrust/api/_transport.py diff --git a/py/src/braintrust/api/__init__.py b/py/src/braintrust/api/__init__.py new file mode 100644 index 000000000..57c4cb007 --- /dev/null +++ b/py/src/braintrust/api/__init__.py @@ -0,0 +1,3 @@ +"""Braintrust API client package.""" + +__all__: list[str] = [] diff --git a/py/src/braintrust/api/_transport.py b/py/src/braintrust/api/_transport.py new file mode 100644 index 000000000..fb3267420 --- /dev/null +++ b/py/src/braintrust/api/_transport.py @@ -0,0 +1,160 @@ +"""Legacy HTTP transport primitives used by the Braintrust SDK.""" + +import sys +import time +from collections.abc import Mapping +from typing import Any + +import requests +import urllib3 +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry + +from ..env import BraintrustEnv +from ..util import _urljoin, response_raise_for_status + + +class RetryRequestExceptionsAdapter(HTTPAdapter): + """An HTTP adapter that automatically retries requests on connection exceptions. + + This adapter extends requests' HTTPAdapter to add retry logic for common network-related + exceptions including connection errors, timeouts, and other HTTP errors. It implements + an exponential backoff strategy between retries to avoid overwhelming servers during + intermittent connectivity issues. + + Attributes: + base_num_retries: Maximum number of retries before giving up and re-raising the exception. + backoff_factor: A multiplier used to determine the time to wait between retries. + The actual wait time is calculated as: backoff_factor * (2 ** retry_count). + default_timeout_secs: Default timeout in seconds for requests that don't specify one. + Prevents indefinite hangs on stale connections. + """ + + def __init__( + self, + *args: Any, + base_num_retries: int = 0, + backoff_factor: float = 0.5, + default_timeout_secs: float = 60, + **kwargs: Any, + ): + self.base_num_retries = base_num_retries + self.backoff_factor = backoff_factor + self.default_timeout_secs = default_timeout_secs + super().__init__(*args, **kwargs) + + def send(self, *args, **kwargs): + # Apply default timeout if none provided to prevent indefinite hangs + if kwargs.get("timeout") is None: + kwargs["timeout"] = self.default_timeout_secs + + num_prev_retries = 0 + while True: + try: + response = super().send(*args, **kwargs) + # Fully-download the content to ensure we catch any errors from + # downloading. + if not response.is_redirect and response.content: + pass + return response + except (urllib3.exceptions.HTTPError, requests.exceptions.RequestException) as e: + if num_prev_retries < self.base_num_retries: + if isinstance(e, requests.exceptions.ReadTimeout): + # Clear all connection pools to discard stale connections. This + # fixes hangs caused by NAT gateways silently dropping idle TCP + # connections (e.g., Azure's ~4 min timeout). close() calls + # PoolManager.clear() which is thread-safe: in-flight requests + # keep their checked-out connections, and new requests create + # fresh pools on demand. + self.close() + # Emulates the sleeping logic in the backoff_factor of urllib3 Retry + sleep_s = self.backoff_factor * (2**num_prev_retries) + print("Retrying request after error:", e, file=sys.stderr) + print("Sleeping for", sleep_s, "seconds", file=sys.stderr) + time.sleep(sleep_s) + num_prev_retries += 1 + else: + raise e + + +class HTTPConnection: + def __init__(self, base_url: str, adapter: HTTPAdapter | None = None): + self.base_url = base_url + self.token = None + self.adapter = adapter + + self._reset(total=0) + + def ping(self) -> bool: + try: + resp = self.get("ping") + return resp.ok + except requests.exceptions.ConnectionError: + return False + + def make_long_lived(self) -> None: + if not self.adapter: + timeout_secs = BraintrustEnv.HTTP_TIMEOUT.get(60.0) + self.adapter = RetryRequestExceptionsAdapter( + base_num_retries=10, backoff_factor=0.5, default_timeout_secs=timeout_secs + ) + self._reset() + + @staticmethod + def sanitize_token(token: str) -> str: + return token.rstrip("\n") + + def set_token(self, token: str) -> None: + token = HTTPConnection.sanitize_token(token) + self.token = token + self._set_session_token() + + def _set_adapter(self, adapter: HTTPAdapter | None) -> None: + self.adapter = adapter + + def _reset(self, **retry_kwargs: Any) -> None: + self.session = requests.Session() + + adapter = self.adapter + if adapter is None: + retry = Retry(**retry_kwargs) + adapter = HTTPAdapter(max_retries=retry) + + self.session.mount("http://", adapter) + self.session.mount("https://", adapter) + + self._set_session_token() + + def _set_session_token(self) -> None: + if self.token: + self.session.headers.update({"Authorization": f"Bearer {self.token}"}) + + def get(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: + return self.session.get(_urljoin(self.base_url, path), *args, **kwargs) + + def post(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: + return self.session.post(_urljoin(self.base_url, path), *args, **kwargs) + + def patch(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: + return self.session.patch(_urljoin(self.base_url, path), *args, **kwargs) + + def put(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: + return self.session.put(_urljoin(self.base_url, path), *args, **kwargs) + + def delete(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: + return self.session.delete(_urljoin(self.base_url, path), *args, **kwargs) + + def get_json(self, object_type: str, args: Mapping[str, Any] | None = None) -> Mapping[str, Any]: + resp = self.get(f"/{object_type}", params=args) + response_raise_for_status(resp) + return resp.json() + + def post_json(self, object_type: str, args: Mapping[str, Any] | None = None) -> Any: + resp = self.post(f"/{object_type.lstrip('/')}", json=args) + response_raise_for_status(resp) + return resp.json() + + def patch_json(self, object_type: str, args: Mapping[str, Any] | None = None) -> Any: + resp = self.patch(f"/{object_type.lstrip('/')}", json=args) + response_raise_for_status(resp) + return resp.json() diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 550b0d1d5..5f5927849 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -37,13 +37,12 @@ import chevron import exceptiongroup -import requests -import urllib3 from braintrust.functions.stream import BraintrustStream from requests.adapters import HTTPAdapter -from urllib3.util.retry import Retry from . import context, id_gen +from .api._transport import HTTPConnection +from .api._transport import RetryRequestExceptionsAdapter as RetryRequestExceptionsAdapter from .bt_json import bt_dumps, bt_safe_deep_copy from .db_fields import ( AUDIT_METADATA_FIELD, @@ -97,7 +96,6 @@ GLOBAL_PROJECT, AugmentedHTTPError, LazyValue, - _urljoin, add_azure_blob_headers, bt_iscoroutinefunction, coalesce, @@ -728,152 +726,6 @@ def set_http_adapter(adapter: HTTPAdapter) -> None: _state._api_conn._reset() -class RetryRequestExceptionsAdapter(HTTPAdapter): - """An HTTP adapter that automatically retries requests on connection exceptions. - - This adapter extends requests' HTTPAdapter to add retry logic for common network-related - exceptions including connection errors, timeouts, and other HTTP errors. It implements - an exponential backoff strategy between retries to avoid overwhelming servers during - intermittent connectivity issues. - - Attributes: - base_num_retries: Maximum number of retries before giving up and re-raising the exception. - backoff_factor: A multiplier used to determine the time to wait between retries. - The actual wait time is calculated as: backoff_factor * (2 ** retry_count). - default_timeout_secs: Default timeout in seconds for requests that don't specify one. - Prevents indefinite hangs on stale connections. - """ - - def __init__( - self, - *args: Any, - base_num_retries: int = 0, - backoff_factor: float = 0.5, - default_timeout_secs: float = 60, - **kwargs: Any, - ): - self.base_num_retries = base_num_retries - self.backoff_factor = backoff_factor - self.default_timeout_secs = default_timeout_secs - super().__init__(*args, **kwargs) - - def send(self, *args, **kwargs): - # Apply default timeout if none provided to prevent indefinite hangs - if kwargs.get("timeout") is None: - kwargs["timeout"] = self.default_timeout_secs - - num_prev_retries = 0 - while True: - try: - response = super().send(*args, **kwargs) - # Fully-download the content to ensure we catch any errors from - # downloading. - if not response.is_redirect and response.content: - pass - return response - except (urllib3.exceptions.HTTPError, requests.exceptions.RequestException) as e: - if num_prev_retries < self.base_num_retries: - if isinstance(e, requests.exceptions.ReadTimeout): - # Clear all connection pools to discard stale connections. This - # fixes hangs caused by NAT gateways silently dropping idle TCP - # connections (e.g., Azure's ~4 min timeout). close() calls - # PoolManager.clear() which is thread-safe: in-flight requests - # keep their checked-out connections, and new requests create - # fresh pools on demand. - self.close() - # Emulates the sleeping logic in the backoff_factor of urllib3 Retry - sleep_s = self.backoff_factor * (2**num_prev_retries) - print("Retrying request after error:", e, file=sys.stderr) - print("Sleeping for", sleep_s, "seconds", file=sys.stderr) - time.sleep(sleep_s) - num_prev_retries += 1 - else: - raise e - - -class HTTPConnection: - def __init__(self, base_url: str, adapter: HTTPAdapter | None = None): - self.base_url = base_url - self.token = None - self.adapter = adapter - - self._reset(total=0) - - def ping(self) -> bool: - try: - resp = self.get("ping") - return resp.ok - except requests.exceptions.ConnectionError: - return False - - def make_long_lived(self) -> None: - if not self.adapter: - timeout_secs = BraintrustEnv.HTTP_TIMEOUT.get(60.0) - self.adapter = RetryRequestExceptionsAdapter( - base_num_retries=10, backoff_factor=0.5, default_timeout_secs=timeout_secs - ) - self._reset() - - @staticmethod - def sanitize_token(token: str) -> str: - return token.rstrip("\n") - - def set_token(self, token: str) -> None: - token = HTTPConnection.sanitize_token(token) - self.token = token - self._set_session_token() - - def _set_adapter(self, adapter: HTTPAdapter | None) -> None: - self.adapter = adapter - - def _reset(self, **retry_kwargs: Any) -> None: - self.session = requests.Session() - - adapter = self.adapter - if adapter is None: - retry = Retry(**retry_kwargs) - adapter = HTTPAdapter(max_retries=retry) - - self.session.mount("http://", adapter) - self.session.mount("https://", adapter) - - self._set_session_token() - - def _set_session_token(self) -> None: - if self.token: - self.session.headers.update({"Authorization": f"Bearer {self.token}"}) - - def get(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: - return self.session.get(_urljoin(self.base_url, path), *args, **kwargs) - - def post(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: - return self.session.post(_urljoin(self.base_url, path), *args, **kwargs) - - def patch(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: - return self.session.patch(_urljoin(self.base_url, path), *args, **kwargs) - - def put(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: - return self.session.put(_urljoin(self.base_url, path), *args, **kwargs) - - def delete(self, path: str, *args: Any, **kwargs: Any) -> requests.Response: - return self.session.delete(_urljoin(self.base_url, path), *args, **kwargs) - - def get_json(self, object_type: str, args: Mapping[str, Any] | None = None) -> Mapping[str, Any]: - resp = self.get(f"/{object_type}", params=args) - response_raise_for_status(resp) - return resp.json() - - def post_json(self, object_type: str, args: Mapping[str, Any] | None = None) -> Any: - resp = self.post(f"/{object_type.lstrip('/')}", json=args) - response_raise_for_status(resp) - return resp.json() - - def patch_json(self, object_type: str, args: Mapping[str, Any] | None = None) -> Any: - resp = self.patch(f"/{object_type.lstrip('/')}", json=args) - response_raise_for_status(resp) - return resp.json() - - # Sometimes we'd like to launch network requests concurrently. We provide a # thread pool to accomplish this. Use a multiple of number of CPU cores to limit # concurrency. diff --git a/py/src/braintrust/test_http.py b/py/src/braintrust/test_http.py index b9ede8d84..8895e3339 100644 --- a/py/src/braintrust/test_http.py +++ b/py/src/braintrust/test_http.py @@ -8,7 +8,7 @@ import pytest import requests -from braintrust.logger import HTTPConnection, RetryRequestExceptionsAdapter +from braintrust.api._transport import HTTPConnection, RetryRequestExceptionsAdapter class HangingConnectionHandler(http.server.BaseHTTPRequestHandler): @@ -173,6 +173,50 @@ def test_adapter_resets_pool_on_timeout(self, hanging_server): assert elapsed < 10.0, f"Request took too long: {elapsed:.2f}s" assert HangingConnectionHandler.request_count >= 2 + def test_adapter_fully_buffers_streaming_responses(self): + """The legacy adapter eagerly downloads response bodies even with stream=True.""" + import concurrent.futures + + headers_sent = threading.Event() + release_body = threading.Event() + + class DelayedBodyHandler(http.server.BaseHTTPRequestHandler): + def log_message(self, format, *args): + pass + + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", "2") + self.end_headers() + self.wfile.flush() + headers_sent.set() + release_body.wait(timeout=5) + self.wfile.write(b"ok") + + server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), DelayedBodyHandler) + server.daemon_threads = True + port = server.server_address[1] + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + + try: + session = requests.Session() + session.mount("http://", RetryRequestExceptionsAdapter()) + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + response_future = executor.submit(session.get, f"http://127.0.0.1:{port}", stream=True) + assert headers_sent.wait(timeout=2) + assert not response_future.done() + + release_body.set() + response = response_future.result(timeout=2) + + assert response.content == b"ok" + finally: + release_body.set() + server.shutdown() + server.server_close() + class TestHTTPConnection: """Tests for HTTPConnection timeout configuration."""