mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
feat: add PostHog mock client support
- Create posthog_mock_client.py using factory pattern - Integrate mock client into PostHogLogger with mock mode detection - Add verbose logging for mock mode initialization and batch operations - Enable mock mode via POSTHOG_MOCK environment variable
This commit is contained in:
parent
e5bea6dd28
commit
05b39a02fb
2 changed files with 61 additions and 7 deletions
|
|
@ -17,6 +17,10 @@ from typing import Any, Dict, Optional, Tuple
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.integrations.posthog_mock_client import (
|
||||
should_use_posthog_mock,
|
||||
create_mock_posthog_client,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
|
|
@ -40,6 +44,12 @@ class PostHogLogger(CustomBatchLogger):
|
|||
"""
|
||||
try:
|
||||
verbose_logger.debug("PostHog: in init posthog logger")
|
||||
|
||||
self.is_mock_mode = should_use_posthog_mock()
|
||||
if self.is_mock_mode:
|
||||
create_mock_posthog_client()
|
||||
verbose_logger.debug("[POSTHOG MOCK] PostHog logger initialized in mock mode")
|
||||
|
||||
if os.getenv("POSTHOG_API_KEY", None) is None:
|
||||
raise Exception("POSTHOG_API_KEY is not set, set 'POSTHOG_API_KEY=<>'")
|
||||
|
||||
|
|
@ -100,7 +110,10 @@ class PostHogLogger(CustomBatchLogger):
|
|||
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
|
||||
)
|
||||
|
||||
verbose_logger.debug("PostHog: Sync event successfully sent")
|
||||
if self.is_mock_mode:
|
||||
verbose_logger.debug("[POSTHOG MOCK] Sync event successfully mocked")
|
||||
else:
|
||||
verbose_logger.debug("PostHog: Sync event successfully sent")
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"PostHog Sync Layer Error - {str(e)}")
|
||||
|
|
@ -320,6 +333,9 @@ class PostHogLogger(CustomBatchLogger):
|
|||
verbose_logger.debug(
|
||||
f"PostHog: Sending batch of {len(self.log_queue)} events"
|
||||
)
|
||||
|
||||
if self.is_mock_mode:
|
||||
verbose_logger.debug("[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted")
|
||||
|
||||
# Group events by credentials for batch sending
|
||||
batches_by_credentials: Dict[tuple[str, str], list] = {}
|
||||
|
|
@ -350,9 +366,12 @@ class PostHogLogger(CustomBatchLogger):
|
|||
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"PostHog: Batch of {len(self.log_queue)} events successfully sent"
|
||||
)
|
||||
if self.is_mock_mode:
|
||||
verbose_logger.debug(f"[POSTHOG MOCK] Batch of {len(self.log_queue)} events successfully mocked")
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
f"PostHog: Batch of {len(self.log_queue)} events successfully sent"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"PostHog Error sending batch API - {str(e)}")
|
||||
|
||||
|
|
@ -429,9 +448,14 @@ class PostHogLogger(CustomBatchLogger):
|
|||
f"PostHog: Failed to flush on exit - status {response.status_code}"
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"PostHog: Successfully flushed {len(self.log_queue)} events on exit"
|
||||
)
|
||||
if self.is_mock_mode:
|
||||
verbose_logger.debug(
|
||||
f"[POSTHOG MOCK] Successfully flushed {len(self.log_queue)} events on exit"
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
f"PostHog: Successfully flushed {len(self.log_queue)} events on exit"
|
||||
)
|
||||
self.log_queue.clear()
|
||||
|
||||
except Exception as e:
|
||||
|
|
|
|||
30
litellm/integrations/posthog_mock_client.py
Normal file
30
litellm/integrations/posthog_mock_client.py
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
"""
|
||||
Mock httpx client for PostHog integration testing.
|
||||
|
||||
This module intercepts PostHog API calls and returns successful mock responses,
|
||||
allowing full code execution without making actual network calls.
|
||||
|
||||
Usage:
|
||||
Set POSTHOG_MOCK=true in environment variables or config to enable mock mode.
|
||||
"""
|
||||
|
||||
from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory
|
||||
|
||||
# Create mock client using factory
|
||||
_config = MockClientConfig(
|
||||
name="POSTHOG",
|
||||
env_var="POSTHOG_MOCK",
|
||||
default_latency_ms=100,
|
||||
default_status_code=200,
|
||||
default_json_data={"status": "success"},
|
||||
url_matchers=[
|
||||
".posthog.com",
|
||||
"posthog.com",
|
||||
"us.i.posthog.com",
|
||||
"app.posthog.com",
|
||||
],
|
||||
patch_async_handler=True,
|
||||
patch_sync_client=True,
|
||||
)
|
||||
|
||||
create_mock_posthog_client, should_use_posthog_mock = create_mock_client_factory(_config)
|
||||
Loading…
Add table
Reference in a new issue