fix: Use apiKey in loadPrompt

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

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