feat: posthog per request api key

This commit is contained in:
Carlos Marchal 2025-10-09 18:32:11 +02:00
parent 7b9897e88f
commit dd560f792e
No known key found for this signature in database
4 changed files with 234 additions and 92 deletions

View File

@ -55,6 +55,26 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
}'
```
### Team-Based Logging
Configure different PostHog credentials per team using the team callback settings:
```bash
curl -X POST 'http://localhost:4000/team/{team_id}/callback' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"callback_name": "posthog",
"callback_type": "success",
"callback_vars": {
"posthog_api_key": "ph_team_specific_key",
"posthog_api_url": "https://custom.posthog.com"
}
}'
```
Now all requests from that team will be logged to their specific PostHog project.
## Usage with LiteLLM Python SDK
### Quick Start
@ -142,6 +162,31 @@ response = client.chat.completions.create(
)
```
#### Per-Request Credentials
You can override PostHog credentials on a per-request basis:
```python
import litellm
litellm.success_callback = ["posthog"]
# Use custom PostHog credentials for this specific request
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "Hello world"}
],
posthog_api_key="ph_custom_project_key",
posthog_api_url="https://custom.posthog.com"
)
```
This is useful when you need to:
- Log different teams/projects to separate PostHog instances
- Use different PostHog projects for staging vs production
- Route logs based on customer or tenant
#### Disable Logging for Specific Calls
Use the `no-log` flag to prevent logging for specific calls:

View File

@ -26,7 +26,7 @@ from litellm.types.integrations.posthog import (
POSTHOG_MAX_BATCH_SIZE,
PostHogEventPayload,
)
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload
class PostHogLogger(CustomBatchLogger):
@ -72,17 +72,19 @@ class PostHogLogger(CustomBatchLogger):
verbose_logger.debug(
"PostHog: Sync logging - Enters logging function for model %s", kwargs
)
api_key, api_url = self._get_credentials_for_request(kwargs)
event_payload = self.create_posthog_event_payload(kwargs)
headers = {
"Content-Type": "application/json",
}
payload = self._create_posthog_payload([event_payload])
payload = self._create_posthog_payload([event_payload], api_key)
capture_url = f"{api_url.rstrip('/')}/batch/"
response = self.sync_client.post(
url=self.capture_url,
url=capture_url,
json=payload,
headers=headers,
)
@ -92,9 +94,9 @@ class PostHogLogger(CustomBatchLogger):
raise Exception(
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
)
verbose_logger.debug("PostHog: Sync event successfully sent")
except Exception as e:
verbose_logger.exception(f"PostHog Sync Layer Error - {str(e)}")
@ -122,9 +124,15 @@ class PostHogLogger(CustomBatchLogger):
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)
self.log_queue.append(event_payload)
# 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..."
)
@ -257,16 +265,42 @@ class PostHogLogger(CustomBatchLogger):
metadata = self._extract_metadata(kwargs)
user_id = self._safe_get(metadata, "user_id")
if user_id:
return str(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 str(trace_id)
return self._safe_uuid()
def _get_credentials_for_request(self, kwargs: Dict[str, Any]) -> tuple[str, 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
@ -282,23 +316,34 @@ class PostHogLogger(CustomBatchLogger):
f"PostHog: Sending batch of {len(self.log_queue)} events"
)
headers = {
"Content-Type": "application/json",
}
# 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"])
payload = self._create_posthog_payload(list(self.log_queue))
# Send each batch to its respective PostHog instance
for (api_key, api_url), events in batches_by_credentials.items():
headers = {
"Content-Type": "application/json",
}
response = await self.async_client.post(
url=self.capture_url,
json=payload,
headers=headers,
)
response.raise_for_status()
payload = self._create_posthog_payload(events, api_key)
capture_url = f"{api_url.rstrip('/')}/batch/"
if response.status_code != 200:
raise Exception(
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
response = await self.async_client.post(
url=capture_url,
json=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}"
)
verbose_logger.debug(
f"PostHog: Batch of {len(self.log_queue)} events successfully sent"
@ -324,8 +369,8 @@ class PostHogLogger(CustomBatchLogger):
def _safe_uuid(self) -> str:
return str(uuid.uuid4())
def _create_posthog_payload(self, events: list) -> Dict[str, Any]:
return {"api_key": self.POSTHOG_API_KEY, "batch": events}
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'):

View File

@ -2222,6 +2222,10 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
arize_space_key: Optional[str]
arize_space_id: Optional[str]
# PostHog dynamic params
posthog_api_key: Optional[str]
posthog_api_url: Optional[str]
# Logging settings
turn_off_message_logging: Optional[bool] # when true will not log messages
litellm_disabled_callbacks: Optional[List[str]]

View File

@ -16,33 +16,36 @@ os.environ["POSTHOG_API_URL"] = "https://app.posthog.com"
def create_standard_logging_payload() -> StandardLoggingPayload:
# Use cast to bypass strict TypedDict requirements for tests
return cast(StandardLoggingPayload, {
"id": "test_id",
"trace_id": "test_trace_id",
"call_type": "completion",
"stream": False,
"response_cost": 0.1,
"status": "success",
"custom_llm_provider": "openai",
"total_tokens": 30,
"prompt_tokens": 20,
"completion_tokens": 10,
"startTime": 1234567890.0,
"endTime": 1234567891.0,
"completionStartTime": 1234567890.5,
"response_time": 1.0,
"model": "gpt-3.5-turbo",
"model_id": "model-123",
"api_base": "https://api.openai.com",
"cache_hit": False,
"saved_cache_cost": 0.0,
"request_tags": [],
"end_user": None,
"messages": [{"role": "user", "content": "Hello, world!"}],
"response": {"choices": [{"message": {"content": "Hi there!"}}]},
"error_str": None,
"model_parameters": {"stream": True},
})
return cast(
StandardLoggingPayload,
{
"id": "test_id",
"trace_id": "test_trace_id",
"call_type": "completion",
"stream": False,
"response_cost": 0.1,
"status": "success",
"custom_llm_provider": "openai",
"total_tokens": 30,
"prompt_tokens": 20,
"completion_tokens": 10,
"startTime": 1234567890.0,
"endTime": 1234567891.0,
"completionStartTime": 1234567890.5,
"response_time": 1.0,
"model": "gpt-3.5-turbo",
"model_id": "model-123",
"api_base": "https://api.openai.com",
"cache_hit": False,
"saved_cache_cost": 0.0,
"request_tags": [],
"end_user": None,
"messages": [{"role": "user", "content": "Hello, world!"}],
"response": {"choices": [{"message": {"content": "Hi there!"}}]},
"error_str": None,
"model_parameters": {"stream": True},
},
)
@pytest.mark.asyncio
@ -94,16 +97,18 @@ async def test_trace_id_fallback_from_standard_logging_object():
posthog_logger = PostHogLogger()
standard_payload = create_standard_logging_payload()
standard_payload["trace_id"] = "test-trace-123"
kwargs = {"standard_logging_object": standard_payload}
event_payload = posthog_logger.create_posthog_event_payload(kwargs)
assert event_payload["properties"]["$ai_trace_id"] == "test-trace-123"
assert event_payload["properties"]["$ai_span_id"] == "test_id" # from standard_payload["id"]
assert (
event_payload["properties"]["$ai_span_id"] == "test_id"
) # from standard_payload["id"]
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_trace_id_uuid_fallback():
"""Test that UUID is generated when no trace_id is available"""
posthog_logger = PostHogLogger()
@ -111,62 +116,62 @@ async def test_trace_id_uuid_fallback():
# Remove trace_id to test fallback
del standard_payload["trace_id"]
del standard_payload["id"]
kwargs = {"standard_logging_object": standard_payload}
event_payload = posthog_logger.create_posthog_event_payload(kwargs)
# Should have generated UUIDs
assert len(event_payload["properties"]["$ai_trace_id"]) == 36 # UUID length
assert len(event_payload["properties"]["$ai_span_id"]) == 36 # UUID length
assert "-" in event_payload["properties"]["$ai_trace_id"] # UUID format
assert len(event_payload["properties"]["$ai_span_id"]) == 36 # UUID length
assert "-" in event_payload["properties"]["$ai_trace_id"] # UUID format
@pytest.mark.asyncio
async def test_distinct_id_fallback_chain():
"""Test the distinct_id fallback priority chain"""
posthog_logger = PostHogLogger()
# Test 1: user_id from metadata (highest priority)
standard_payload = create_standard_logging_payload()
kwargs = {
"standard_logging_object": standard_payload,
"litellm_params": {"metadata": {"user_id": "metadata-user-123"}}
"litellm_params": {"metadata": {"user_id": "metadata-user-123"}},
}
distinct_id = posthog_logger._get_distinct_id(standard_payload, kwargs)
assert distinct_id == "metadata-user-123"
# Test 2: trace_id from standard_logging_object (second priority)
# Test 2: trace_id from standard_logging_object (second priority)
kwargs = {"standard_logging_object": standard_payload} # no metadata
distinct_id = posthog_logger._get_distinct_id(standard_payload, kwargs)
assert distinct_id == "test_trace_id"
# Test 3: end_user from standard_logging_object (third priority)
standard_payload_no_trace = create_standard_logging_payload()
del standard_payload_no_trace["trace_id"]
standard_payload_no_trace["end_user"] = "end-user-456"
distinct_id = posthog_logger._get_distinct_id(standard_payload_no_trace, {})
assert distinct_id == "end-user-456"
# Test 4: UUID fallback (lowest priority)
standard_payload_empty = create_standard_logging_payload()
del standard_payload_empty["trace_id"]
del standard_payload_empty["end_user"]
distinct_id = posthog_logger._get_distinct_id(standard_payload_empty, {})
assert len(distinct_id) == 36 # UUID length
assert "-" in distinct_id # UUID format
assert "-" in distinct_id # UUID format
@pytest.mark.asyncio
async def test_missing_standard_logging_object():
"""Test error handling when standard_logging_object is missing"""
posthog_logger = PostHogLogger()
kwargs = {} # Missing standard_logging_object
with pytest.raises(ValueError, match="standard_logging_object not found in kwargs"):
posthog_logger.create_posthog_event_payload(kwargs)
@ -176,26 +181,26 @@ async def test_custom_metadata_support():
"""Test that custom metadata fields are added directly to properties"""
posthog_logger = PostHogLogger()
standard_payload = create_standard_logging_payload()
kwargs = {
"standard_logging_object": standard_payload,
"litellm_params": {
"metadata": {
"user_id": "user-123", # should be used for distinct_id, not custom property
"project_name": "test_project", # should appear as project_name
"environment": "staging", # should appear as environment
"custom_field": "custom_value" # should appear as custom_field
"environment": "staging", # should appear as environment
"custom_field": "custom_value", # should appear as custom_field
}
}
},
}
event_payload = posthog_logger.create_posthog_event_payload(kwargs)
# Check that custom fields are added directly
assert event_payload["properties"]["project_name"] == "test_project"
assert event_payload["properties"]["environment"] == "staging"
assert event_payload["properties"]["custom_field"] == "custom_value"
# Check that user_id is used for distinct_id, not as custom property
assert event_payload["distinct_id"] == "user-123"
assert "user_id" not in event_payload["properties"]
@ -206,7 +211,7 @@ async def test_custom_metadata_filters_internal_fields():
"""Test that LiteLLM internal fields are filtered out from custom metadata"""
posthog_logger = PostHogLogger()
standard_payload = create_standard_logging_payload()
kwargs = {
"standard_logging_object": standard_payload,
"litellm_params": {
@ -214,17 +219,19 @@ async def test_custom_metadata_filters_internal_fields():
"custom_field": "should_appear",
"endpoint": "/chat/completions", # internal field - should be filtered
"user_api_key_hash": "hash123", # internal field - should be filtered
"headers": {"content-type": "application/json"}, # internal field - should be filtered
"headers": {
"content-type": "application/json"
}, # internal field - should be filtered
"model_info": {"id": "123"}, # internal field - should be filtered
}
}
},
}
event_payload = posthog_logger.create_posthog_event_payload(kwargs)
# Check that custom field appears
assert event_payload["properties"]["custom_field"] == "should_appear"
# Check that internal fields are filtered out
assert "endpoint" not in event_payload["properties"]
assert "user_api_key_hash" not in event_payload["properties"]
@ -237,22 +244,63 @@ async def test_custom_metadata_with_no_metadata():
"""Test that logger handles cases with no metadata gracefully"""
posthog_logger = PostHogLogger()
standard_payload = create_standard_logging_payload()
# Test with no litellm_params
kwargs = {"standard_logging_object": standard_payload}
event_payload = posthog_logger.create_posthog_event_payload(kwargs)
# Should not error and should have standard properties
assert event_payload["event"] == "$ai_generation"
assert event_payload["properties"]["$ai_model"] == "gpt-3.5-turbo"
# Test with empty metadata
kwargs = {
"standard_logging_object": standard_payload,
"litellm_params": {"metadata": {}}
"litellm_params": {"metadata": {}},
}
event_payload = posthog_logger.create_posthog_event_payload(kwargs)
# Should not error and should have standard properties
assert event_payload["event"] == "$ai_generation"
assert event_payload["properties"]["$ai_model"] == "gpt-3.5-turbo"
@pytest.mark.asyncio
async def test_dynamic_credentials():
"""Test that per-request credentials override environment variables"""
from litellm.types.utils import StandardCallbackDynamicParams
posthog_logger = PostHogLogger()
# Test with no dynamic params - should use env vars
kwargs = {}
api_key, api_url = posthog_logger._get_credentials_for_request(kwargs)
assert api_key == "test_key" # from env var
assert api_url == "https://app.posthog.com" # from env var
# Test with dynamic params - should override env vars
standard_callback_dynamic_params = StandardCallbackDynamicParams(
posthog_api_key="dynamic_key", posthog_api_url="https://custom.posthog.com"
)
kwargs = {"standard_callback_dynamic_params": standard_callback_dynamic_params}
api_key, api_url = posthog_logger._get_credentials_for_request(kwargs)
assert api_key == "dynamic_key"
assert api_url == "https://custom.posthog.com"
# Test partial override - only api_key
standard_callback_dynamic_params = StandardCallbackDynamicParams(
posthog_api_key="another_key"
)
kwargs = {"standard_callback_dynamic_params": standard_callback_dynamic_params}
api_key, api_url = posthog_logger._get_credentials_for_request(kwargs)
assert api_key == "another_key"
assert api_url == "https://app.posthog.com" # falls back to env var
# Test partial override - only api_url
standard_callback_dynamic_params = StandardCallbackDynamicParams(
posthog_api_url="https://another.posthog.com"
)
kwargs = {"standard_callback_dynamic_params": standard_callback_dynamic_params}
api_key, api_url = posthog_logger._get_credentials_for_request(kwargs)
assert api_key == "test_key" # falls back to env var
assert api_url == "https://another.posthog.com"