diff --git a/.sampo/changesets/somber-duchess-joukahainen.md b/.sampo/changesets/somber-duchess-joukahainen.md new file mode 100644 index 00000000..d84256a9 --- /dev/null +++ b/.sampo/changesets/somber-duchess-joukahainen.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Fix async OpenAI streaming captures to include token usage and other generation properties emitted by synchronous streams. diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 57764415..5abc23e9 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -16,9 +16,9 @@ from posthog.ai.utils import ( call_llm_and_track_usage_async, _capture_ai_event, - extract_available_tool_calls, + extract_available_tool_calls as extract_available_tool_calls, finalize_ai_content, - get_model_params, + get_model_params as get_model_params, merge_usage_stats, with_privacy_mode, ) @@ -232,7 +232,6 @@ async def async_generator(): usage_stats, latency, output, - extract_available_tool_calls("openai", kwargs), model_from_response, stop_reason=stop_reason, ) @@ -250,75 +249,36 @@ async def _capture_streaming_event( usage_stats: TokenUsage, latency: float, output: Any, - available_tool_calls: Optional[List[Dict[str, Any]]] = None, model_from_response: Optional[str] = None, stop_reason: Optional[str] = None, ): - if posthog_trace_id is None: - posthog_trace_id = str(uuid.uuid4()) + from posthog.ai.types import StreamingEventData + from posthog.ai.utils import capture_streaming_event + + formatted_input = format_openai_streaming_input(kwargs, "responses") # Use model from kwargs, fallback to model from response model = kwargs.get("model") or model_from_response or "unknown" - event_properties = { - "$ai_provider": "openai", - "$ai_model": model, - "$ai_model_parameters": get_model_params(kwargs), - "$ai_input": with_privacy_mode( - self._client._ph_client, - posthog_privacy_mode, - finalize_ai_content( - format_openai_streaming_input(kwargs, "responses"), - self._client._ph_client, - ), - ), - "$ai_output_choices": with_privacy_mode( - self._client._ph_client, - posthog_privacy_mode, - finalize_ai_content( - format_openai_streaming_output(output, "responses"), - self._client._ph_client, - ), - ), - "$ai_http_status": 200, - "$ai_input_tokens": usage_stats.get("input_tokens", 0), - "$ai_output_tokens": usage_stats.get("output_tokens", 0), - "$ai_cache_read_input_tokens": usage_stats.get( - "cache_read_input_tokens", 0 - ), - "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0), - "$ai_latency": latency, - "$ai_trace_id": posthog_trace_id, - "$ai_base_url": str(self._client.base_url), - **(posthog_properties or {}), - } - - # Add web search count if present - web_search_count = usage_stats.get("web_search_count") - if ( - web_search_count is not None - and isinstance(web_search_count, int) - and web_search_count > 0 - ): - event_properties["$ai_web_search_count"] = web_search_count - - if stop_reason is not None: - event_properties["$ai_stop_reason"] = stop_reason - - if available_tool_calls: - event_properties["$ai_tools"] = available_tool_calls + event_data = StreamingEventData( + provider="openai", + model=model, + base_url=str(self._client.base_url), + kwargs=kwargs, + formatted_input=formatted_input, + formatted_output=format_openai_streaming_output(output, "responses"), + usage_stats=usage_stats, + latency=latency, + distinct_id=posthog_distinct_id, + trace_id=posthog_trace_id, + properties=posthog_properties, + privacy_mode=posthog_privacy_mode, + groups=posthog_groups, + stop_reason=stop_reason, + ) - if posthog_distinct_id is None: - event_properties["$process_person_profile"] = False - - if hasattr(self._client._ph_client, "capture"): - _capture_ai_event( - self._client._ph_client, - "$ai_generation", - distinct_id=posthog_distinct_id or posthog_trace_id, - properties=event_properties, - groups=posthog_groups, - ) + # Use the common capture function + capture_streaming_event(self._client._ph_client, event_data) async def parse( self, @@ -557,7 +517,6 @@ async def async_generator(): latency, accumulated_content, tool_calls_list, - extract_available_tool_calls("openai", kwargs), model_from_response, stop_reason=stop_reason, ) @@ -576,76 +535,36 @@ async def _capture_streaming_event( latency: float, output: Any, tool_calls: Optional[List[Dict[str, Any]]] = None, - available_tool_calls: Optional[List[Dict[str, Any]]] = None, model_from_response: Optional[str] = None, stop_reason: Optional[str] = None, ): - if posthog_trace_id is None: - posthog_trace_id = str(uuid.uuid4()) + from posthog.ai.types import StreamingEventData + from posthog.ai.utils import capture_streaming_event + + formatted_input = format_openai_streaming_input(kwargs, "chat") # Use model from kwargs, fallback to model from response model = kwargs.get("model") or model_from_response or "unknown" - event_properties = { - "$ai_provider": "openai", - "$ai_model": model, - "$ai_model_parameters": get_model_params(kwargs), - "$ai_input": with_privacy_mode( - self._client._ph_client, - posthog_privacy_mode, - finalize_ai_content( - format_openai_streaming_input(kwargs, "chat"), - self._client._ph_client, - ), - ), - "$ai_output_choices": with_privacy_mode( - self._client._ph_client, - posthog_privacy_mode, - finalize_ai_content( - format_openai_streaming_output(output, "chat", tool_calls), - self._client._ph_client, - ), - ), - "$ai_http_status": 200, - "$ai_input_tokens": usage_stats.get("input_tokens", 0), - "$ai_output_tokens": usage_stats.get("output_tokens", 0), - "$ai_cache_read_input_tokens": usage_stats.get( - "cache_read_input_tokens", 0 - ), - "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0), - "$ai_latency": latency, - "$ai_trace_id": posthog_trace_id, - "$ai_base_url": str(self._client.base_url), - **(posthog_properties or {}), - } - - # Add web search count if present - web_search_count = usage_stats.get("web_search_count") - - if ( - web_search_count is not None - and isinstance(web_search_count, int) - and web_search_count > 0 - ): - event_properties["$ai_web_search_count"] = web_search_count - - if stop_reason is not None: - event_properties["$ai_stop_reason"] = stop_reason - - if available_tool_calls: - event_properties["$ai_tools"] = available_tool_calls - - if posthog_distinct_id is None: - event_properties["$process_person_profile"] = False + event_data = StreamingEventData( + provider="openai", + model=model, + base_url=str(self._client.base_url), + kwargs=kwargs, + formatted_input=formatted_input, + formatted_output=format_openai_streaming_output(output, "chat", tool_calls), + usage_stats=usage_stats, + latency=latency, + distinct_id=posthog_distinct_id, + trace_id=posthog_trace_id, + properties=posthog_properties, + privacy_mode=posthog_privacy_mode, + groups=posthog_groups, + stop_reason=stop_reason, + ) - if hasattr(self._client._ph_client, "capture"): - _capture_ai_event( - self._client._ph_client, - "$ai_generation", - distinct_id=posthog_distinct_id or posthog_trace_id, - properties=event_properties, - groups=posthog_groups, - ) + # Use the common capture function + capture_streaming_event(self._client._ph_client, event_data) class WrappedEmbeddings(_OpenAIWrapperResource): diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 057fec63..3910b879 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -790,7 +790,6 @@ def capture_streaming_event( # Add optional token fields # For Anthropic, always include cache fields even if 0 (backward compatibility) - # For others, only include if present and non-zero if event_data["provider"] == "anthropic": # Anthropic always includes cache fields cache_read = event_data["usage_stats"].get("cache_read_input_tokens", 0) @@ -798,7 +797,6 @@ def capture_streaming_event( event_properties["$ai_cache_read_input_tokens"] = cache_read event_properties["$ai_cache_creation_input_tokens"] = cache_creation else: - # Other providers only include if non-zero optional_token_fields = [ "cache_read_input_tokens", "cache_creation_input_tokens", @@ -807,8 +805,20 @@ def capture_streaming_event( for field in optional_token_fields: value = event_data["usage_stats"].get(field) - if value is not None and isinstance(value, int) and value > 0: - event_properties[f"$ai_{field}"] = value + property_name = f"$ai_{field}" + + # OpenAI async streams historically included these fields even when 0. + # Keep those defaults in the shared path so they are not mistaken for + # caller-supplied token passthrough properties. + if event_data["provider"] == "openai" and field in { + "cache_read_input_tokens", + "reasoning_tokens", + }: + event_properties.setdefault( + property_name, event_data["usage_stats"].get(field, 0) + ) + elif value is not None and isinstance(value, int) and value > 0: + event_properties[property_name] = value cache_reporting_exclusive = event_data["usage_stats"].get( "cache_reporting_exclusive" diff --git a/posthog/test/ai/openai/conftest.py b/posthog/test/ai/openai/conftest.py new file mode 100644 index 00000000..8afcc467 --- /dev/null +++ b/posthog/test/ai/openai/conftest.py @@ -0,0 +1,126 @@ +"""Fixtures shared by the OpenAI test modules. + +Defined here rather than in a test module so pytest supplies them by discovery. Importing them +between test files bound the names in the importing module and tripped Ruff F811. +""" + +from unittest.mock import patch + +import pytest +from openai.types.chat.chat_completion_chunk import ( + ChatCompletionChunk, + ChoiceDelta, + ChoiceDeltaToolCall, + ChoiceDeltaToolCallFunction, +) +from openai.types.chat.chat_completion_chunk import ( + Choice as ChoiceChunk, +) +from openai.types.completion_usage import CompletionUsage + + +@pytest.fixture +def mock_client(): + with patch("posthog.client.Client") as mock_client: + mock_client.privacy_mode = False + yield mock_client + + +@pytest.fixture +def streaming_tool_call_chunks(): + return [ + ChatCompletionChunk( + id="chunk1", + model="gpt-4", + object="chat.completion.chunk", + created=1234567890, + choices=[ + ChoiceChunk( + index=0, + delta=ChoiceDelta( + role="assistant", + tool_calls=[ + ChoiceDeltaToolCall( + index=0, + id="call_abc123", + type="function", + function=ChoiceDeltaToolCallFunction( + name="get_weather", + arguments='{"location": "', + ), + ) + ], + ), + finish_reason=None, + ) + ], + ), + ChatCompletionChunk( + id="chunk2", + model="gpt-4", + object="chat.completion.chunk", + created=1234567891, + choices=[ + ChoiceChunk( + index=0, + delta=ChoiceDelta( + tool_calls=[ + ChoiceDeltaToolCall( + index=0, + id="call_abc123", + type="function", + function=ChoiceDeltaToolCallFunction( + arguments='San Francisco"', + ), + ) + ], + ), + finish_reason=None, + ) + ], + ), + ChatCompletionChunk( + id="chunk3", + model="gpt-4", + object="chat.completion.chunk", + created=1234567892, + choices=[ + ChoiceChunk( + index=0, + delta=ChoiceDelta( + tool_calls=[ + ChoiceDeltaToolCall( + index=0, + id="call_abc123", + type="function", + function=ChoiceDeltaToolCallFunction( + arguments=', "unit": "celsius"}', + ), + ) + ], + ), + finish_reason=None, + ) + ], + ), + ChatCompletionChunk( + id="chunk4", + model="gpt-4", + object="chat.completion.chunk", + created=1234567893, + choices=[ + ChoiceChunk( + index=0, + delta=ChoiceDelta( + content="The weather in San Francisco is 15°C.", + ), + finish_reason=None, + ) + ], + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=15, + total_tokens=35, + ), + ), + ] diff --git a/posthog/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py new file mode 100644 index 00000000..f2066b93 --- /dev/null +++ b/posthog/test/ai/openai/test_async_parity.py @@ -0,0 +1,90 @@ +""" +End-to-end regression test for the openai async streaming property gap. + +Drives the real wrapped clients through a mocked OpenAI stream, sync and async, and asserts +the async twin emits the same $ai_generation properties the sync twin does. Reuses posthog's +own fixtures so the stream shape is theirs, not mine. + +FAILS on main. Would pass if the async twin called posthog.ai.utils.capture_streaming_event +the way the sync twin, anthropic and gemini all do. + + cp test_async_parity.py repo-posthog/posthog/test/ai/openai/test_async_parity.py + .venv-posthog/bin/python -m pytest repo-posthog/posthog/test/ai/openai/test_async_parity.py -v +""" + +from unittest.mock import patch + +import pytest + +from posthog.ai.openai import AsyncOpenAI, OpenAI + +TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": {}, + }, + } +] +MESSAGES = [{"role": "user", "content": "What's the weather in San Francisco?"}] + + +def _sync_props(mock_client, chunks): + with patch("openai.resources.chat.completions.Completions.create") as create: + create.return_value = chunks + client = OpenAI(api_key="test-key", posthog_client=mock_client) + list( + client.chat.completions.create( + model="gpt-4", + messages=MESSAGES, + tools=TOOLS, + stream=True, + posthog_distinct_id="test-id", + ) + ) + return mock_client.capture.call_args[1]["properties"] + + +async def _async_props(mock_client, chunks): + async def create(self, **kwargs): + async def it(): + for chunk in chunks: + yield chunk + + return it() + + with patch("openai.resources.chat.completions.AsyncCompletions.create", new=create): + client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) + stream = await client.chat.completions.create( + model="gpt-4", + messages=MESSAGES, + tools=TOOLS, + stream=True, + posthog_distinct_id="test-id", + ) + async for _ in stream: + pass + return mock_client.capture.call_args[1]["properties"] + + +@pytest.mark.asyncio +async def test_async_streaming_emits_the_same_properties_as_sync( + mock_client, streaming_tool_call_chunks +): + sync_props = _sync_props(mock_client, streaming_tool_call_chunks) + mock_client.capture.reset_mock() + async_props = await _async_props(mock_client, streaming_tool_call_chunks) + + # Guard: if the sync side stopped emitting these, the comparison below is vacuous. + assert "$ai_usage" in sync_props + assert sync_props["$ai_tokens_source"] == "sdk" + assert async_props["$ai_tokens_source"] == sync_props["$ai_tokens_source"] + assert async_props["$ai_cache_read_input_tokens"] == 0 + assert async_props["$ai_reasoning_tokens"] == 0 + + missing = sorted(set(sync_props) - set(async_props)) + assert missing == [], ( + f"the async openai streaming path drops {missing} that the sync path sends" + ) diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 83d4e58f..4b0bec5f 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -9,13 +9,11 @@ try: from openai.types.chat import ChatCompletion, ChatCompletionMessage from openai.types.chat.chat_completion import Choice - from openai.types.chat.chat_completion_chunk import ChatCompletionChunk - from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk from openai.types.chat.chat_completion_chunk import ( + ChatCompletionChunk, ChoiceDelta, - ChoiceDeltaToolCall, - ChoiceDeltaToolCallFunction, ) + from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk from openai.types.chat.chat_completion_message_tool_call import ( ChatCompletionMessageToolCall, Function, @@ -24,11 +22,11 @@ from openai.types.create_embedding_response import CreateEmbeddingResponse, Usage from openai.types.embedding import Embedding from openai.types.responses import ( + ParsedResponse, Response, + ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText, - ResponseFunctionToolCall, - ParsedResponse, ) from openai.types.responses.parsed_response import ( ParsedResponseOutputMessage, @@ -50,13 +48,6 @@ ) -@pytest.fixture -def mock_client(): - with patch("posthog.client.Client") as mock_client: - mock_client.privacy_mode = False - yield mock_client - - @pytest.fixture def mock_openai_response(): return ChatCompletion( @@ -248,106 +239,6 @@ def mock_openai_response_with_null_token_details(): ) -@pytest.fixture -def streaming_tool_call_chunks(): - return [ - ChatCompletionChunk( - id="chunk1", - model="gpt-4", - object="chat.completion.chunk", - created=1234567890, - choices=[ - ChoiceChunk( - index=0, - delta=ChoiceDelta( - role="assistant", - tool_calls=[ - ChoiceDeltaToolCall( - index=0, - id="call_abc123", - type="function", - function=ChoiceDeltaToolCallFunction( - name="get_weather", - arguments='{"location": "', - ), - ) - ], - ), - finish_reason=None, - ) - ], - ), - ChatCompletionChunk( - id="chunk2", - model="gpt-4", - object="chat.completion.chunk", - created=1234567891, - choices=[ - ChoiceChunk( - index=0, - delta=ChoiceDelta( - tool_calls=[ - ChoiceDeltaToolCall( - index=0, - id="call_abc123", - type="function", - function=ChoiceDeltaToolCallFunction( - arguments='San Francisco"', - ), - ) - ], - ), - finish_reason=None, - ) - ], - ), - ChatCompletionChunk( - id="chunk3", - model="gpt-4", - object="chat.completion.chunk", - created=1234567892, - choices=[ - ChoiceChunk( - index=0, - delta=ChoiceDelta( - tool_calls=[ - ChoiceDeltaToolCall( - index=0, - id="call_abc123", - type="function", - function=ChoiceDeltaToolCallFunction( - arguments=', "unit": "celsius"}', - ), - ) - ], - ), - finish_reason=None, - ) - ], - ), - ChatCompletionChunk( - id="chunk4", - model="gpt-4", - object="chat.completion.chunk", - created=1234567893, - choices=[ - ChoiceChunk( - index=0, - delta=ChoiceDelta( - content="The weather in San Francisco is 15°C.", - ), - finish_reason=None, - ) - ], - usage=CompletionUsage( - prompt_tokens=20, - completion_tokens=15, - total_tokens=35, - ), - ), - ] - - @pytest.fixture def mock_openai_response_with_tool_calls(): return ChatCompletion(