mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
Add mock support for Braintrust integration and extend mock client factory
- Add braintrust_mock_client.py with mock HTTP client for Braintrust integration testing - Integrate mock client into BraintrustLogger with mock mode detection - Refactor Helicone mock client to fully utilize factory's HTTPHandler.post patching - Extend mock_client_factory to support patching HTTPHandler.post for sync calls - Enable endpoint-specific mock responses for Braintrust (/project vs /project_logs) - All mock clients now properly handle both async (AsyncHTTPHandler) and sync (HTTPHandler) calls
This commit is contained in:
Alexsander Hamir
parent
5706ba9fe2
commit
5c2f55bcd5
4 changed files with 160 additions and 70 deletions
|
|
@ -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:
|
||||
|
|
|
|||
117
litellm/integrations/braintrust_mock_client.py
Normal file
117
litellm/integrations/braintrust_mock_client.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue