518 lines
19 KiB
Python
518 lines
19 KiB
Python
"""
|
|
PostHog Integration - sends LLM analytics events to PostHog
|
|
|
|
Follows PostHog's LLM Analytics format: https://posthog.com/docs/llm-analytics/manual-capture
|
|
|
|
async_log_success_event: stores batch of events in memory and flushes to PostHog
|
|
async_log_failure_event: logs failed LLM calls with error information
|
|
|
|
For batching specific details see CustomBatchLogger class
|
|
"""
|
|
|
|
import asyncio
|
|
import atexit
|
|
import os
|
|
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.litellm_core_utils.safe_json_dumps import safe_dumps
|
|
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,
|
|
httpxSpecialProvider,
|
|
)
|
|
from litellm.types.integrations.posthog import (
|
|
POSTHOG_MAX_BATCH_SIZE,
|
|
PostHogEventPayload,
|
|
)
|
|
from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload
|
|
|
|
|
|
class PostHogLogger(CustomBatchLogger):
|
|
def __init__(self, **kwargs):
|
|
"""
|
|
Initializes the PostHog logger, checks if the correct env variables are set
|
|
|
|
Required environment variables:
|
|
`POSTHOG_API_KEY` - your PostHog API key
|
|
`POSTHOG_API_URL` - your PostHog API URL (defaults to https://app.posthog.com)
|
|
"""
|
|
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=<>'")
|
|
|
|
self.async_client = get_async_httpx_client(
|
|
llm_provider=httpxSpecialProvider.LoggingCallback
|
|
)
|
|
self.sync_client = _get_httpx_client()
|
|
|
|
self.POSTHOG_API_KEY = os.getenv("POSTHOG_API_KEY")
|
|
posthog_api_url = os.getenv("POSTHOG_API_URL", "https://us.i.posthog.com")
|
|
self.posthog_host = posthog_api_url.rstrip("/")
|
|
self.capture_url = f"{self.posthog_host}/batch/"
|
|
|
|
self._async_initialized = False
|
|
self.flush_lock = None
|
|
self.log_queue = []
|
|
|
|
# Register cleanup handler to flush internal queue on exit
|
|
atexit.register(self._flush_on_exit)
|
|
|
|
super().__init__(
|
|
**kwargs, flush_lock=None, batch_size=POSTHOG_MAX_BATCH_SIZE
|
|
)
|
|
|
|
except Exception as e:
|
|
verbose_logger.exception(
|
|
f"PostHog: Got exception on init PostHog client {str(e)}"
|
|
)
|
|
raise e
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
verbose_logger.debug(
|
|
"PostHog: Sync logging - Enters logging function for model %s", kwargs
|
|
)
|
|
|
|
api_key, api_url = self._get_credentials_for_request(kwargs)
|
|
if api_key is None or api_url is None:
|
|
raise Exception("PostHog credentials not found in kwargs")
|
|
event_payload = self.create_posthog_event_payload(kwargs)
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload = self._create_posthog_payload([event_payload], api_key)
|
|
capture_url = f"{api_url.rstrip('/')}/batch/"
|
|
|
|
response = self.sync_client.post(
|
|
url=capture_url,
|
|
content=safe_dumps(payload),
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
|
|
if response.status_code != 200:
|
|
raise Exception(
|
|
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
|
|
)
|
|
|
|
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)}")
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
verbose_logger.debug(
|
|
"PostHog: Async logging - Enters logging function for model %s", kwargs
|
|
)
|
|
self._ensure_async_setup() # Lazy initialization
|
|
await self._log_async_event(kwargs, response_obj, start_time, end_time)
|
|
except Exception as e:
|
|
verbose_logger.exception(f"PostHog Layer Error - {str(e)}")
|
|
pass
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
verbose_logger.debug(
|
|
"PostHog: Async logging - Enters logging function for model %s", kwargs
|
|
)
|
|
self._ensure_async_setup() # Lazy initialization
|
|
await self._log_async_event(kwargs, response_obj, start_time, end_time)
|
|
except Exception as e:
|
|
verbose_logger.exception(f"PostHog Layer Error - {str(e)}")
|
|
pass
|
|
|
|
async def _log_async_event(
|
|
self, kwargs, response_obj=None, start_time=0.0, end_time=0.0
|
|
):
|
|
# Note: response_obj, start_time, end_time not used - all data comes from kwargs
|
|
api_key, api_url = self._get_credentials_for_request(kwargs)
|
|
event_payload = self.create_posthog_event_payload(kwargs)
|
|
|
|
# Store event with its credentials for batch sending
|
|
self.log_queue.append(
|
|
{"event": event_payload, "api_key": api_key, "api_url": api_url}
|
|
)
|
|
verbose_logger.debug(
|
|
f"PostHog, event added to queue. Will flush in {self.flush_interval} seconds..."
|
|
)
|
|
|
|
if len(self.log_queue) >= self.batch_size:
|
|
await self.flush_queue()
|
|
|
|
def create_posthog_event_payload(
|
|
self, kwargs: Dict[str, Any]
|
|
) -> PostHogEventPayload:
|
|
"""
|
|
Helper function to create a PostHog event payload for logging
|
|
|
|
Args:
|
|
kwargs (Dict[str, Any]): request kwargs containing standard_logging_object
|
|
|
|
Returns:
|
|
PostHogEventPayload: defined in types.py
|
|
"""
|
|
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
|
|
"standard_logging_object", None
|
|
)
|
|
if standard_logging_object is None:
|
|
raise ValueError("standard_logging_object not found in kwargs")
|
|
|
|
call_type = standard_logging_object.get("call_type", "")
|
|
event_name = "$ai_embedding" if call_type == "embedding" else "$ai_generation"
|
|
|
|
properties = self._create_posthog_properties(
|
|
standard_logging_object=standard_logging_object,
|
|
kwargs=kwargs,
|
|
event_name=event_name,
|
|
)
|
|
|
|
distinct_id = self._get_distinct_id(standard_logging_object, kwargs)
|
|
|
|
return PostHogEventPayload(
|
|
event=event_name,
|
|
properties=properties,
|
|
distinct_id=distinct_id,
|
|
)
|
|
|
|
def _create_posthog_properties(
|
|
self,
|
|
standard_logging_object: StandardLoggingPayload,
|
|
kwargs: Dict[str, Any],
|
|
event_name: str,
|
|
) -> Dict[str, Any]:
|
|
"""Create PostHog properties following LLM Analytics spec"""
|
|
properties = {}
|
|
|
|
# Core model information
|
|
properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "")
|
|
properties["$ai_provider"] = self._safe_get(
|
|
standard_logging_object, "custom_llm_provider", ""
|
|
)
|
|
|
|
# Input/Output data
|
|
messages = self._safe_get(standard_logging_object, "messages")
|
|
if messages is not None:
|
|
properties["$ai_input"] = messages
|
|
|
|
if event_name == "$ai_generation":
|
|
response = self._safe_get(standard_logging_object, "response")
|
|
if response is not None:
|
|
properties["$ai_output_choices"] = response
|
|
|
|
# Token information
|
|
properties["$ai_input_tokens"] = self._safe_get(
|
|
standard_logging_object, "prompt_tokens", 0
|
|
)
|
|
if event_name == "$ai_generation":
|
|
properties["$ai_output_tokens"] = self._safe_get(
|
|
standard_logging_object, "completion_tokens", 0
|
|
)
|
|
|
|
# Cost and performance
|
|
response_cost = self._safe_get(standard_logging_object, "response_cost")
|
|
if response_cost is not None:
|
|
properties["$ai_total_cost_usd"] = response_cost
|
|
|
|
properties["$ai_latency"] = self._safe_get(
|
|
standard_logging_object, "response_time", 0.0
|
|
)
|
|
|
|
# Error handling
|
|
if self._safe_get(standard_logging_object, "status") == "failure":
|
|
properties["$ai_is_error"] = True
|
|
error_str = self._safe_get(standard_logging_object, "error_str")
|
|
if error_str is not None:
|
|
properties["$ai_error"] = error_str
|
|
|
|
# Add trace properties
|
|
self._add_trace_properties(properties, kwargs)
|
|
|
|
# Add custom metadata fields
|
|
self._add_custom_metadata_properties(properties, kwargs)
|
|
|
|
return properties
|
|
|
|
def _add_trace_properties(self, properties: Dict[str, Any], kwargs: Dict[str, Any]):
|
|
standard_logging_object = self._safe_get(kwargs, "standard_logging_object", {})
|
|
|
|
trace_id = self._safe_get(
|
|
standard_logging_object, "trace_id", self._safe_uuid()
|
|
)
|
|
properties["$ai_trace_id"] = trace_id
|
|
|
|
span_id = self._safe_get(standard_logging_object, "id", self._safe_uuid())
|
|
properties["$ai_span_id"] = span_id
|
|
|
|
metadata = self._extract_metadata(kwargs)
|
|
parent_id = metadata.get("parent_run_id") or metadata.get("parent_id")
|
|
if parent_id:
|
|
properties["$ai_parent_id"] = parent_id
|
|
|
|
def _add_custom_metadata_properties(
|
|
self, properties: Dict[str, Any], kwargs: Dict[str, Any]
|
|
):
|
|
"""Add custom metadata fields to PostHog properties"""
|
|
metadata = self._extract_metadata(kwargs)
|
|
if not isinstance(metadata, dict):
|
|
return
|
|
|
|
litellm_internal_fields = {
|
|
"endpoint",
|
|
"caching_groups",
|
|
"user_api_key_hash",
|
|
"user_api_key_alias",
|
|
"user_api_key_team_id",
|
|
"user_api_key_user_id",
|
|
"user_api_key_org_id",
|
|
"user_api_key_team_alias",
|
|
"user_api_key_end_user_id",
|
|
"user_api_key_user_email",
|
|
"user_api_key",
|
|
"user_api_end_user_max_budget",
|
|
"litellm_api_version",
|
|
"global_max_parallel_requests",
|
|
"user_api_key_team_max_budget",
|
|
"user_api_key_team_spend",
|
|
"user_api_key_spend",
|
|
"user_api_key_max_budget",
|
|
"user_api_key_model_max_budget",
|
|
"user_api_key_metadata",
|
|
"headers",
|
|
"litellm_parent_otel_span",
|
|
"requester_ip_address",
|
|
"model_group",
|
|
"model_group_size",
|
|
"deployment",
|
|
"model_info",
|
|
"api_base",
|
|
"caching_groups",
|
|
"hidden_params",
|
|
"parent_run_id",
|
|
"parent_id",
|
|
"user_id",
|
|
}
|
|
|
|
for key, value in metadata.items():
|
|
if key not in litellm_internal_fields:
|
|
properties[key] = value
|
|
|
|
def _get_distinct_id(
|
|
self, standard_logging_object: StandardLoggingPayload, kwargs: Dict[str, Any]
|
|
) -> str:
|
|
metadata = self._extract_metadata(kwargs)
|
|
user_id = self._safe_get(metadata, "user_id")
|
|
if user_id:
|
|
return str(user_id)
|
|
end_user = self._safe_get(standard_logging_object, "end_user")
|
|
if end_user:
|
|
return str(end_user)
|
|
trace_id = self._safe_get(standard_logging_object, "trace_id")
|
|
if trace_id:
|
|
return str(trace_id)
|
|
|
|
return self._safe_uuid()
|
|
|
|
def _get_credentials_for_request(
|
|
self, kwargs: Dict[str, Any]
|
|
) -> Tuple[Optional[str], Optional[str]]:
|
|
"""
|
|
Get PostHog credentials for this request.
|
|
|
|
Checks for per-request credentials in standard_callback_dynamic_params,
|
|
falls back to instance defaults from environment variables.
|
|
|
|
Args:
|
|
kwargs: Request kwargs containing standard_callback_dynamic_params
|
|
|
|
Returns:
|
|
tuple[str, str]: (api_key, api_url)
|
|
"""
|
|
standard_callback_dynamic_params: Optional[
|
|
StandardCallbackDynamicParams
|
|
] = kwargs.get("standard_callback_dynamic_params", None)
|
|
|
|
if standard_callback_dynamic_params is not None:
|
|
api_key = (
|
|
standard_callback_dynamic_params.get("posthog_api_key")
|
|
or self.POSTHOG_API_KEY
|
|
)
|
|
api_url = (
|
|
standard_callback_dynamic_params.get("posthog_api_url")
|
|
or self.posthog_host
|
|
)
|
|
else:
|
|
api_key = self.POSTHOG_API_KEY
|
|
api_url = self.posthog_host
|
|
|
|
return api_key, api_url
|
|
|
|
async def async_send_batch(self):
|
|
"""
|
|
Sends the in memory logs queue to PostHog API
|
|
|
|
Raises:
|
|
Raises a NON Blocking verbose_logger.exception if an error occurs
|
|
"""
|
|
try:
|
|
if not self.log_queue:
|
|
return
|
|
|
|
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] = {}
|
|
for item in self.log_queue:
|
|
key = (item["api_key"], item["api_url"])
|
|
if key not in batches_by_credentials:
|
|
batches_by_credentials[key] = []
|
|
batches_by_credentials[key].append(item["event"])
|
|
|
|
# Send each batch to its respective PostHog instance
|
|
for (api_key, api_url), events in batches_by_credentials.items():
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload = self._create_posthog_payload(events, api_key)
|
|
capture_url = f"{api_url.rstrip('/')}/batch/"
|
|
|
|
response = await self.async_client.post(
|
|
url=capture_url,
|
|
content=safe_dumps(payload),
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
|
|
if response.status_code != 200:
|
|
raise Exception(
|
|
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
|
|
)
|
|
|
|
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)}")
|
|
|
|
def _ensure_async_setup(self):
|
|
if not self._async_initialized:
|
|
try:
|
|
self.flush_lock = asyncio.Lock()
|
|
asyncio.create_task(self.periodic_flush())
|
|
self._async_initialized = True
|
|
verbose_logger.debug("PostHog: Async components initialized")
|
|
except Exception as e:
|
|
verbose_logger.error(
|
|
f"PostHog: Failed to initialize async components: {str(e)}"
|
|
)
|
|
raise
|
|
|
|
def _extract_metadata(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
|
litellm_params = kwargs.get("litellm_params", {}) or {}
|
|
return litellm_params.get("metadata", {}) or {}
|
|
|
|
def _safe_uuid(self) -> str:
|
|
return str(uuid.uuid4())
|
|
|
|
def _create_posthog_payload(self, events: list, api_key: str) -> Dict[str, Any]:
|
|
return {"api_key": api_key, "batch": events}
|
|
|
|
def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any:
|
|
if obj is None or not hasattr(obj, "get"):
|
|
return default
|
|
return obj.get(key, default)
|
|
|
|
def _flush_on_exit(self):
|
|
"""
|
|
Flush remaining events from internal log_queue before process exit.
|
|
Called automatically via atexit handler.
|
|
|
|
This works in conjunction with GLOBAL_LOGGING_WORKER's atexit handler:
|
|
1. GLOBAL_LOGGING_WORKER atexit invokes pending callbacks
|
|
2. Callbacks add events to this logger's internal log_queue
|
|
3. This atexit handler flushes the internal queue to PostHog
|
|
"""
|
|
if not self.log_queue:
|
|
return
|
|
|
|
verbose_logger.debug(
|
|
f"PostHog: Flushing {len(self.log_queue)} remaining events on exit"
|
|
)
|
|
|
|
try:
|
|
# Group events by credentials (same logic as async_send_batch)
|
|
batches_by_credentials: Dict[Tuple[str, str], list] = {}
|
|
for item in self.log_queue:
|
|
key = (item["api_key"], item["api_url"])
|
|
if key not in batches_by_credentials:
|
|
batches_by_credentials[key] = []
|
|
batches_by_credentials[key].append(item["event"])
|
|
|
|
# Send each batch synchronously using sync_client
|
|
for (api_key, api_url), events in batches_by_credentials.items():
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload = self._create_posthog_payload(events, api_key)
|
|
capture_url = f"{api_url.rstrip('/')}/batch/"
|
|
|
|
response = self.sync_client.post(
|
|
url=capture_url,
|
|
content=safe_dumps(payload),
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
|
|
if response.status_code != 200:
|
|
verbose_logger.error(
|
|
f"PostHog: Failed to flush on exit - status {response.status_code}"
|
|
)
|
|
|
|
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:
|
|
verbose_logger.error(f"PostHog: Error flushing events on exit: {str(e)}")
|