diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 585de510e8b..42e9680a7fc 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -9,6 +9,10 @@ import httpx import litellm from litellm import verbose_logger +from litellm.integrations.braintrust_mock_client import ( + should_use_braintrust_mock, + create_mock_braintrust_client, +) from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, @@ -34,6 +38,10 @@ class BraintrustLogger(CustomLogger): self, api_key: Optional[str] = None, api_base: Optional[str] = None ) -> None: super().__init__() + self.is_mock_mode = should_use_braintrust_mock() + if self.is_mock_mode: + create_mock_braintrust_client() + verbose_logger.info("[BRAINTRUST MOCK] Braintrust logger initialized in mock mode") self.validate_environment(api_key=api_key) self.api_base = api_base or os.getenv("BRAINTRUST_API_BASE") or API_BASE self.default_project_id = None @@ -254,6 +262,8 @@ class BraintrustLogger(CustomLogger): json={"events": [request_data]}, headers=self.headers, ) + if self.is_mock_mode: + print_verbose("[BRAINTRUST MOCK] Sync event successfully mocked") except httpx.HTTPStatusError as e: raise Exception(e.response.text) except Exception as e: @@ -399,6 +409,8 @@ class BraintrustLogger(CustomLogger): json={"events": [request_data]}, headers=self.headers, ) + if self.is_mock_mode: + print_verbose("[BRAINTRUST MOCK] Async event successfully mocked") except httpx.HTTPStatusError as e: raise Exception(e.response.text) except Exception as e: diff --git a/litellm/integrations/braintrust_mock_client.py b/litellm/integrations/braintrust_mock_client.py new file mode 100644 index 00000000000..5dd2d79211c --- /dev/null +++ b/litellm/integrations/braintrust_mock_client.py @@ -0,0 +1,117 @@ +""" +Mock HTTP client for Braintrust integration testing. + +This module intercepts Braintrust API calls and returns successful mock responses, +allowing full code execution without making actual network calls. + +Usage: + Set BRAINTRUST_MOCK=true in environment variables or config to enable mock mode. +""" + +import os +import time + +from litellm._logging import verbose_logger +from litellm.integrations.mock_client_factory import MockClientConfig, MockResponse, create_mock_client_factory + +# Use factory for should_use_mock and MockResponse +# Braintrust uses both HTTPHandler (sync) and AsyncHTTPHandler (async) +# Braintrust needs endpoint-specific responses, so we use custom HTTPHandler.post patching +_config = MockClientConfig( + name="BRAINTRUST", + env_var="BRAINTRUST_MOCK", + default_latency_ms=100, + default_status_code=200, + default_json_data={"id": "mock-project-id", "status": "success"}, + url_matchers=[ + ".braintrustdata.com", + "braintrustdata.com", + ".braintrust.dev", + "braintrust.dev", + ], + patch_async_handler=True, # Patch AsyncHTTPHandler.post for async calls + patch_sync_client=False, # HTTPHandler uses self.client.send(), not self.client.post() + patch_http_handler=False, # We use custom patching for endpoint-specific responses +) + +# Get should_use_mock and create_mock_client from factory +# We need to call the factory's create_mock_client to patch AsyncHTTPHandler.post +create_mock_braintrust_factory_client, should_use_braintrust_mock = create_mock_client_factory(_config) + +# Store original HTTPHandler.post method (Braintrust-specific for sync calls with custom logic) +_original_http_handler_post = None +_mocks_initialized = False + +# Default mock latency in seconds +_MOCK_LATENCY_SECONDS = float(os.getenv("BRAINTRUST_MOCK_LATENCY_MS", "100")) / 1000.0 + + +def _is_braintrust_url(url: str) -> bool: + """Check if URL is a Braintrust API URL.""" + url_lower = url.lower() + return "braintrustdata.com" in url_lower or "braintrust.dev" in url_lower + + +def _mock_http_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, files=None, content=None, logging_obj=None): + """Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses.""" + # Only mock Braintrust API calls + if isinstance(url, str) and _is_braintrust_url(url): + verbose_logger.info(f"[BRAINTRUST MOCK] POST to {url}") + time.sleep(_MOCK_LATENCY_SECONDS) + # Return appropriate mock response based on endpoint + if "/project" in url: + # Project creation/retrieval/register endpoint + project_name = json.get("name", "litellm") if json else "litellm" + mock_data = {"id": f"mock-project-id-{project_name}", "name": project_name} + elif "/project_logs" in url: + # Log insertion endpoint + mock_data = {"status": "success"} + else: + mock_data = _config.default_json_data + return MockResponse( + status_code=_config.default_status_code, + json_data=mock_data, + url=url, + elapsed_seconds=_MOCK_LATENCY_SECONDS + ) + if _original_http_handler_post is not None: + return _original_http_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, files=files, content=content, logging_obj=logging_obj) + raise RuntimeError("Original HTTPHandler.post not available") + + +def create_mock_braintrust_client(): + """ + Monkey-patch HTTPHandler.post to intercept Braintrust sync calls. + + Braintrust uses HTTPHandler for sync calls and AsyncHTTPHandler for async calls. + HTTPHandler.post uses self.client.send(), not self.client.post(), so we need + custom patching for sync (similar to Helicone). + AsyncHTTPHandler.post is patched by the factory. + + We use custom patching instead of factory's patch_http_handler because we need + endpoint-specific responses (different for /project vs /project_logs). + + This function is idempotent - it only initializes mocks once, even if called multiple times. + """ + global _original_http_handler_post, _mocks_initialized + + if _mocks_initialized: + return + + verbose_logger.debug("[BRAINTRUST MOCK] Initializing Braintrust mock client...") + + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + if _original_http_handler_post is None: + _original_http_handler_post = HTTPHandler.post + HTTPHandler.post = _mock_http_handler_post # type: ignore + verbose_logger.debug("[BRAINTRUST MOCK] Patched HTTPHandler.post") + + # CRITICAL: Call the factory's initialization function to patch AsyncHTTPHandler.post + # This is required for async calls to be mocked + create_mock_braintrust_factory_client() + + verbose_logger.debug(f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms") + verbose_logger.debug("[BRAINTRUST MOCK] Braintrust mock client initialization complete") + + _mocks_initialized = True diff --git a/litellm/integrations/helicone_mock_client.py b/litellm/integrations/helicone_mock_client.py index 152e9c73cbe..0f4670a1d2c 100644 --- a/litellm/integrations/helicone_mock_client.py +++ b/litellm/integrations/helicone_mock_client.py @@ -8,14 +8,10 @@ Usage: Set HELICONE_MOCK=true in environment variables or config to enable mock mode. """ -import os -import time +from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory -from litellm._logging import verbose_logger -from litellm.integrations.mock_client_factory import MockClientConfig, MockResponse, create_mock_client_factory - -# Use factory for should_use_mock and MockResponse -# HTTPHandler uses self.client.send(), not self.client.post(), so we need custom patching +# Create mock client using factory +# Helicone uses HTTPHandler which internally uses httpx.Client.send(), not httpx.Client.post() _config = MockClientConfig( name="HELICONE", env_var="HELICONE_MOCK", @@ -30,67 +26,7 @@ _config = MockClientConfig( ], patch_async_handler=False, patch_sync_client=False, # HTTPHandler uses self.client.send(), not self.client.post() + patch_http_handler=True, # Patch HTTPHandler.post directly ) -# Get should_use_mock from factory (but don't use its patching since HTTPHandler is different) -_, should_use_helicone_mock = create_mock_client_factory(_config) - -# Store original HTTPHandler.post method (Helicone-specific) -_original_http_handler_post = None -_mocks_initialized = False - -# Default mock latency in seconds -_MOCK_LATENCY_SECONDS = float(os.getenv("HELICONE_MOCK_LATENCY_MS", "100")) / 1000.0 - - -def _is_helicone_url(url: str) -> bool: - """Check if URL is a Helicone API URL.""" - url_lower = url.lower() - return "hconeai.com" in url_lower or "helicone.ai" in url_lower - - -def _mock_http_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, files=None, content=None, logging_obj=None): - """Monkey-patched HTTPHandler.post that intercepts Helicone calls.""" - # Only mock Helicone API calls - if isinstance(url, str) and _is_helicone_url(url): - verbose_logger.info(f"[HELICONE MOCK] POST to {url}") - time.sleep(_MOCK_LATENCY_SECONDS) - return MockResponse( - status_code=_config.default_status_code, - json_data=_config.default_json_data, - url=url, - elapsed_seconds=_MOCK_LATENCY_SECONDS - ) - if _original_http_handler_post is not None: - return _original_http_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, files=files, content=content, logging_obj=logging_obj) - raise RuntimeError("Original HTTPHandler.post not available") - - -def create_mock_helicone_client(): - """ - Monkey-patch HTTPHandler.post to intercept Helicone calls. - - Helicone uses litellm.module_level_client which is an HTTPHandler instance. - HTTPHandler.post uses self.client.send(), not self.client.post(), so we need - custom patching (similar to how GCS has custom GET/DELETE handlers). - - This function is idempotent - it only initializes mocks once, even if called multiple times. - """ - global _original_http_handler_post, _mocks_initialized - - if _mocks_initialized: - return - - verbose_logger.debug("[HELICONE MOCK] Initializing Helicone mock client...") - - from litellm.llms.custom_httpx.http_handler import HTTPHandler - - if _original_http_handler_post is None: - _original_http_handler_post = HTTPHandler.post - HTTPHandler.post = _mock_http_handler_post # type: ignore - verbose_logger.debug("[HELICONE MOCK] Patched HTTPHandler.post") - - verbose_logger.debug(f"[HELICONE MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms") - verbose_logger.debug("[HELICONE MOCK] Helicone mock client initialization complete") - - _mocks_initialized = True +create_mock_helicone_client, should_use_helicone_mock = create_mock_client_factory(_config) diff --git a/litellm/integrations/mock_client_factory.py b/litellm/integrations/mock_client_factory.py index d788570b995..c45707c263d 100644 --- a/litellm/integrations/mock_client_factory.py +++ b/litellm/integrations/mock_client_factory.py @@ -27,6 +27,7 @@ class MockClientConfig: url_matchers: Optional[List[str]] = None # List of strings to match in URLs (e.g., ["storage.googleapis.com"]) patch_async_handler: bool = True # Whether to patch AsyncHTTPHandler.post patch_sync_client: bool = False # Whether to patch httpx.Client.post + patch_http_handler: bool = False # Whether to patch HTTPHandler.post (for sync calls that use HTTPHandler) def __post_init__(self): """Ensure url_matchers is a list.""" @@ -105,6 +106,7 @@ def create_mock_client_factory(config: MockClientConfig): # Store original methods for restoration _original_async_handler_post = None _original_sync_client_post = None + _original_http_handler_post = None _mocks_initialized = False # Calculate mock latency @@ -147,10 +149,27 @@ def create_mock_client_factory(config: MockClientConfig): if _original_sync_client_post is not None: return _original_sync_client_post(self, url, **kwargs) + # Create HTTPHandler mock (for sync calls that use HTTPHandler.post) + def _mock_http_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, files=None, content=None, logging_obj=None): + """Monkey-patched HTTPHandler.post that intercepts API calls.""" + if isinstance(url, str) and _is_mock_url(url): + verbose_logger.info(f"[{config.name} MOCK] POST to {url}") + import time + time.sleep(_MOCK_LATENCY_SECONDS) + return MockResponse( + status_code=config.default_status_code, + json_data=config.default_json_data, + url=url, + elapsed_seconds=_MOCK_LATENCY_SECONDS + ) + if _original_http_handler_post is not None: + return _original_http_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, files=files, content=content, logging_obj=logging_obj) + raise RuntimeError("Original HTTPHandler.post not available") + # Create mock client initialization function def create_mock_client(): """Initialize the mock client by patching HTTP handlers.""" - nonlocal _original_async_handler_post, _original_sync_client_post, _mocks_initialized + nonlocal _original_async_handler_post, _original_sync_client_post, _original_http_handler_post, _mocks_initialized if _mocks_initialized: return @@ -168,6 +187,12 @@ def create_mock_client_factory(config: MockClientConfig): httpx.Client.post = _mock_sync_client_post # type: ignore verbose_logger.debug(f"[{config.name} MOCK] Patched httpx.Client.post") + if config.patch_http_handler and _original_http_handler_post is None: + from litellm.llms.custom_httpx.http_handler import HTTPHandler + _original_http_handler_post = HTTPHandler.post + HTTPHandler.post = _mock_http_handler_post # type: ignore + verbose_logger.debug(f"[{config.name} MOCK] Patched HTTPHandler.post") + verbose_logger.debug(f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms") verbose_logger.debug(f"[{config.name} MOCK] {config.name} mock client initialization complete")