mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Add Langfuse mock mode for testing without API calls
- Add langfuse_mock_client module with httpx monkey-patching - Intercept Langfuse API calls and return mock responses - Enable via LANGFUSE_MOCK environment variable - Move mock client imports to module level in langfuse.py - Remove emojis from debug messages
This commit is contained in:
parent
fb7f5351ab
commit
1b869f46fe
2 changed files with 138 additions and 7 deletions
|
|
@ -25,6 +25,10 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
reconstruct_model_name,
|
||||
)
|
||||
from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
|
||||
from litellm.integrations.langfuse.langfuse_mock_client import (
|
||||
create_mock_langfuse_client,
|
||||
should_use_langfuse_mock,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.langfuse import *
|
||||
|
|
@ -119,8 +123,17 @@ class LangFuseLogger:
|
|||
self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval(
|
||||
flush_interval
|
||||
)
|
||||
http_client = _get_httpx_client()
|
||||
self.langfuse_client = http_client.client
|
||||
|
||||
# Check if we should use mock mode
|
||||
if should_use_langfuse_mock():
|
||||
# Use mock client - exercises all code without network calls
|
||||
self.langfuse_client = create_mock_langfuse_client()
|
||||
self.is_mock_mode = True
|
||||
else:
|
||||
# Use real httpx client
|
||||
http_client = _get_httpx_client()
|
||||
self.langfuse_client = http_client.client
|
||||
self.is_mock_mode = False
|
||||
|
||||
parameters = {
|
||||
"public_key": self.public_key,
|
||||
|
|
@ -139,11 +152,16 @@ class LangFuseLogger:
|
|||
|
||||
# set the current langfuse project id in the environ
|
||||
# this is used by Alerting to link to the correct project
|
||||
try:
|
||||
project_id = self.Langfuse.client.projects.get().data[0].id
|
||||
os.environ["LANGFUSE_PROJECT_ID"] = project_id
|
||||
except Exception:
|
||||
project_id = None
|
||||
if self.is_mock_mode:
|
||||
# In mock mode, use a fake project ID
|
||||
os.environ["LANGFUSE_PROJECT_ID"] = "mock-project-id"
|
||||
verbose_logger.debug("Langfuse Mock: Using mock project ID")
|
||||
else:
|
||||
try:
|
||||
project_id = self.Langfuse.client.projects.get().data[0].id
|
||||
os.environ["LANGFUSE_PROJECT_ID"] = project_id
|
||||
except Exception:
|
||||
project_id = None
|
||||
|
||||
if os.getenv("UPSTREAM_LANGFUSE_SECRET_KEY") is not None:
|
||||
upstream_langfuse_debug = (
|
||||
|
|
|
|||
113
litellm/integrations/langfuse/langfuse_mock_client.py
Normal file
113
litellm/integrations/langfuse/langfuse_mock_client.py
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
"""
|
||||
Mock httpx client for Langfuse integration testing.
|
||||
|
||||
This module intercepts Langfuse API calls and returns successful mock responses,
|
||||
allowing full code execution without making actual network calls.
|
||||
|
||||
Usage:
|
||||
Set LANGFUSE_MOCK=true in environment variables or config to enable mock mode.
|
||||
"""
|
||||
|
||||
import httpx
|
||||
import json
|
||||
from typing import Dict, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
# Store original post method for restoration
|
||||
_original_httpx_post = None
|
||||
|
||||
|
||||
class MockLangfuseResponse:
|
||||
"""Mock httpx.Response that satisfies Langfuse SDK requirements."""
|
||||
|
||||
def __init__(self, status_code: int = 200, json_data: Optional[Dict] = None, url: Optional[str] = None):
|
||||
self.status_code = status_code
|
||||
self._json_data = json_data or {"status": "success"}
|
||||
self.headers = httpx.Headers({})
|
||||
self.is_success = status_code < 400
|
||||
self.is_error = status_code >= 400
|
||||
self.is_redirect = 300 <= status_code < 400
|
||||
self.url = httpx.URL(url) if url else httpx.URL("")
|
||||
self.elapsed = httpx.Timeout(0.0)
|
||||
self._text = json.dumps(self._json_data)
|
||||
self._content = self._text.encode("utf-8")
|
||||
|
||||
@property
|
||||
def text(self) -> str:
|
||||
"""Return response text."""
|
||||
return self._text
|
||||
|
||||
@property
|
||||
def content(self) -> bytes:
|
||||
"""Return response content."""
|
||||
return self._content
|
||||
|
||||
def json(self) -> Dict:
|
||||
"""Return JSON response data."""
|
||||
return self._json_data
|
||||
|
||||
def read(self) -> bytes:
|
||||
"""Read response content."""
|
||||
return self._content
|
||||
|
||||
def raise_for_status(self):
|
||||
"""Raise exception for error status codes."""
|
||||
if self.status_code >= 400:
|
||||
raise Exception(f"HTTP {self.status_code}")
|
||||
|
||||
|
||||
def _mock_httpx_post(self, url, **kwargs):
|
||||
"""Monkey-patched httpx.Client.post that intercepts Langfuse calls."""
|
||||
# Only mock Langfuse API calls
|
||||
if isinstance(url, str) and ("langfuse.com" in url or "langfuse" in url.lower()):
|
||||
print(f"[LANGFUSE MOCK] POST to {url}")
|
||||
return MockLangfuseResponse(status_code=200, json_data={"status": "success"}, url=url)
|
||||
# For non-Langfuse calls, use original method
|
||||
if _original_httpx_post is not None:
|
||||
return _original_httpx_post(self, url, **kwargs)
|
||||
# Fallback: if original not set, create a temporary client for this call
|
||||
import httpx
|
||||
with httpx.Client() as client:
|
||||
return client.post(url, **kwargs)
|
||||
|
||||
|
||||
def create_mock_langfuse_client():
|
||||
"""
|
||||
Monkey-patch httpx.Client.post to intercept Langfuse calls.
|
||||
|
||||
Returns a real httpx.Client instance - the monkey-patch intercepts all calls.
|
||||
"""
|
||||
global _original_httpx_post
|
||||
|
||||
if _original_httpx_post is None:
|
||||
_original_httpx_post = httpx.Client.post
|
||||
httpx.Client.post = _mock_httpx_post # type: ignore
|
||||
print("[LANGFUSE MOCK] Patched httpx.Client.post")
|
||||
|
||||
# Return real client - monkey-patch handles interception
|
||||
return httpx.Client()
|
||||
|
||||
|
||||
def should_use_langfuse_mock() -> bool:
|
||||
"""
|
||||
Determine if Langfuse should run in mock mode.
|
||||
|
||||
Checks the LANGFUSE_MOCK environment variable.
|
||||
|
||||
Returns:
|
||||
bool: True if mock mode should be enabled
|
||||
"""
|
||||
import os
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
mock_mode = os.getenv("LANGFUSE_MOCK", "false")
|
||||
result = str_to_bool(mock_mode)
|
||||
|
||||
# Ensure we return a bool, not None
|
||||
result = bool(result) if result is not None else False
|
||||
|
||||
if result:
|
||||
verbose_logger.info("Langfuse Mock Mode: ENABLED - API calls will be mocked")
|
||||
|
||||
return result
|
||||
Loading…
Add table
Reference in a new issue