feat: posthog per request api key
This commit is contained in:
parent
7b9897e88f
commit
dd560f792e
@ -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:
|
||||
|
||||
@ -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'):
|
||||
|
||||
@ -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]]
|
||||
|
||||
@ -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"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user