From bcd579706bdbe1c6c5ce21c542dc04b2dcebf2c9 Mon Sep 17 00:00:00 2001 From: Cohen Karnell <7363269+ckarnell@users.noreply.github.com> Date: Sun, 9 Aug 2026 16:07:22 -0400 Subject: [PATCH 1/6] fix(ai): route async OpenAI streaming through capture_streaming_event Both async OpenAI streaming paths built the $ai_generation property dict inline, while sync OpenAI and both Anthropic and Gemini twins delegate to posthog.ai.utils.capture_streaming_event. Async streaming therefore dropped $ai_usage, $ai_tokens_source, $ai_instructions on the Responses path, and the AI Gateway double-billing warning. Routing through the shared helper rather than adding the keys inline, since the duplication is what allowed the drift. Adds a parity test that compares sync and async captured properties and fails on main. --- posthog/ai/openai/openai_async.py | 166 ++++++-------------- posthog/test/ai/openai/test_async_parity.py | 93 +++++++++++ 2 files changed, 137 insertions(+), 122 deletions(-) create mode 100644 posthog/test/ai/openai/test_async_parity.py diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 57764415..e059e1d7 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -18,7 +18,6 @@ _capture_ai_event, extract_available_tool_calls, finalize_ai_content, - get_model_params, merge_usage_stats, with_privacy_mode, ) @@ -254,71 +253,33 @@ async def _capture_streaming_event( 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, @@ -580,72 +541,33 @@ async def _capture_streaming_event( 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/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py new file mode 100644 index 00000000..f6771aaf --- /dev/null +++ b/posthog/test/ai/openai/test_async_parity.py @@ -0,0 +1,93 @@ +""" +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 OpenAI, AsyncOpenAI + +# fixtures mock_client and streaming_tool_call_chunks come from test_openai.py in this dir +from posthog.test.ai.openai.test_openai import ( # noqa: F401 + mock_client, + streaming_tool_call_chunks, +) + +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 "$ai_tokens_source" in sync_props + + missing = sorted(set(sync_props) - set(async_props)) + assert missing == [], ( + f"the async openai streaming path drops {missing} that the sync path sends" + ) From 6e9e243542e18d48370f573cdbfa996d1a611968 Mon Sep 17 00:00:00 2001 From: Cohen Karnell <7363269+ckarnell@users.noreply.github.com> Date: Mon, 10 Aug 2026 04:08:09 -0400 Subject: [PATCH 2/6] Fix CI: silence F811 on imported fixtures, refresh the API snapshot The test module imports mock_client and streaming_tool_call_chunks from test_openai.py and then takes them as parameters, which ruff reads as redefinition. No other test file in the repo imports fixtures across modules, so rather than move them into a conftest and restructure the test layout, the three call sites carry a targeted noqa. The public API snapshot dropped the get_model_params alias from openai_async, which this PR removed by routing streaming through capture_streaming_event. Regenerated with the repo's own script. --- posthog/test/ai/openai/test_async_parity.py | 9 +++++---- references/public_api_snapshot.txt | 1 - 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/posthog/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py index f6771aaf..1d9d956e 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -16,7 +16,7 @@ import pytest -from posthog.ai.openai import OpenAI, AsyncOpenAI +from posthog.ai.openai import AsyncOpenAI, OpenAI # fixtures mock_client and streaming_tool_call_chunks come from test_openai.py in this dir from posthog.test.ai.openai.test_openai import ( # noqa: F401 @@ -37,7 +37,7 @@ MESSAGES = [{"role": "user", "content": "What's the weather in San Francisco?"}] -def _sync_props(mock_client, chunks): +def _sync_props(mock_client, chunks): # noqa: F811 with patch("openai.resources.chat.completions.Completions.create") as create: create.return_value = chunks client = OpenAI(api_key="test-key", posthog_client=mock_client) @@ -53,7 +53,7 @@ def _sync_props(mock_client, chunks): return mock_client.capture.call_args[1]["properties"] -async def _async_props(mock_client, chunks): +async def _async_props(mock_client, chunks): # noqa: F811 async def create(self, **kwargs): async def it(): for chunk in chunks: @@ -77,7 +77,8 @@ async def it(): @pytest.mark.asyncio async def test_async_streaming_emits_the_same_properties_as_sync( - mock_client, streaming_tool_call_chunks + mock_client, # noqa: F811 + streaming_tool_call_chunks, # noqa: F811 ): sync_props = _sync_props(mock_client, streaming_tool_call_chunks) mock_client.capture.reset_mock() diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 34307b03..8f9b00a4 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -174,7 +174,6 @@ alias posthog.ai.openai.openai_async.extract_openai_usage_from_chunk -> posthog. alias posthog.ai.openai.openai_async.finalize_ai_content -> posthog.ai.utils.finalize_ai_content alias posthog.ai.openai.openai_async.format_openai_streaming_input -> posthog.ai.openai.openai_converter.format_openai_streaming_input alias posthog.ai.openai.openai_async.format_openai_streaming_output -> posthog.ai.openai.openai_converter.format_openai_streaming_output -alias posthog.ai.openai.openai_async.get_model_params -> posthog.ai.utils.get_model_params alias posthog.ai.openai.openai_async.merge_provider_override -> posthog.ai.openai.wrapper_utils.merge_provider_override alias posthog.ai.openai.openai_async.merge_usage_stats -> posthog.ai.utils.merge_usage_stats alias posthog.ai.openai.openai_async.setup -> posthog.setup From ade11ac613dc219216b1c85410ee72505fba021a Mon Sep 17 00:00:00 2001 From: Cohen Karnell <7363269+ckarnell@users.noreply.github.com> Date: Mon, 10 Aug 2026 04:15:34 -0400 Subject: [PATCH 3/6] Move the shared OpenAI fixtures into conftest.py Replaces the noqa suppressions from the previous commit, per review. The two fixtures now live in conftest.py so pytest supplies them by discovery, which means no test module imports or rebinds the names and Ruff F811 has nothing to fire on. mock_client and streaming_tool_call_chunks moved out of test_openai.py unchanged, with the imports they need. Two chunk-type imports that only the moved fixture used are dropped from test_openai.py. posthog/test/ai is 420 passed, 88 skipped, the same as before the move. --- posthog/test/ai/openai/conftest.py | 126 ++++++++++++++++++++ posthog/test/ai/openai/test_async_parity.py | 13 +- posthog/test/ai/openai/test_openai.py | 117 +----------------- 3 files changed, 133 insertions(+), 123 deletions(-) create mode 100644 posthog/test/ai/openai/conftest.py 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 index 1d9d956e..01e05cb1 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -18,12 +18,6 @@ from posthog.ai.openai import AsyncOpenAI, OpenAI -# fixtures mock_client and streaming_tool_call_chunks come from test_openai.py in this dir -from posthog.test.ai.openai.test_openai import ( # noqa: F401 - mock_client, - streaming_tool_call_chunks, -) - TOOLS = [ { "type": "function", @@ -37,7 +31,7 @@ MESSAGES = [{"role": "user", "content": "What's the weather in San Francisco?"}] -def _sync_props(mock_client, chunks): # noqa: F811 +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) @@ -53,7 +47,7 @@ def _sync_props(mock_client, chunks): # noqa: F811 return mock_client.capture.call_args[1]["properties"] -async def _async_props(mock_client, chunks): # noqa: F811 +async def _async_props(mock_client, chunks): async def create(self, **kwargs): async def it(): for chunk in chunks: @@ -77,8 +71,7 @@ async def it(): @pytest.mark.asyncio async def test_async_streaming_emits_the_same_properties_as_sync( - mock_client, # noqa: F811 - streaming_tool_call_chunks, # noqa: F811 + mock_client, streaming_tool_call_chunks ): sync_props = _sync_props(mock_client, streaming_tool_call_chunks) mock_client.capture.reset_mock() 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( From ea6ebc0980f7f20b087d6620f1b1f76992c38d2f Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Mon, 10 Aug 2026 10:30:05 +0200 Subject: [PATCH 4/6] chore: add OpenAI async streaming changeset --- .sampo/changesets/somber-duchess-joukahainen.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .sampo/changesets/somber-duchess-joukahainen.md 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. From b35d5b9816ab7504670a46ac222cb0883e25793f Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Tue, 11 Aug 2026 13:14:11 +0200 Subject: [PATCH 5/6] fix(ai): address async streaming review feedback --- posthog/ai/openai/openai_async.py | 23 ++++++++++++++------- posthog/test/ai/openai/test_async_parity.py | 2 ++ references/public_api_snapshot.txt | 1 + 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index e059e1d7..4031fd37 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -16,8 +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 as get_model_params, merge_usage_stats, with_privacy_mode, ) @@ -231,7 +232,6 @@ async def async_generator(): usage_stats, latency, output, - extract_available_tool_calls("openai", kwargs), model_from_response, stop_reason=stop_reason, ) @@ -249,7 +249,6 @@ 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, ): @@ -272,7 +271,13 @@ async def _capture_streaming_event( latency=latency, distinct_id=posthog_distinct_id, trace_id=posthog_trace_id, - properties=posthog_properties, + properties={ + "$ai_cache_read_input_tokens": usage_stats.get( + "cache_read_input_tokens", 0 + ), + "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0), + **(posthog_properties or {}), + }, privacy_mode=posthog_privacy_mode, groups=posthog_groups, stop_reason=stop_reason, @@ -518,7 +523,6 @@ async def async_generator(): latency, accumulated_content, tool_calls_list, - extract_available_tool_calls("openai", kwargs), model_from_response, stop_reason=stop_reason, ) @@ -537,7 +541,6 @@ 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, ): @@ -560,7 +563,13 @@ async def _capture_streaming_event( latency=latency, distinct_id=posthog_distinct_id, trace_id=posthog_trace_id, - properties=posthog_properties, + properties={ + "$ai_cache_read_input_tokens": usage_stats.get( + "cache_read_input_tokens", 0 + ), + "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0), + **(posthog_properties or {}), + }, privacy_mode=posthog_privacy_mode, groups=posthog_groups, stop_reason=stop_reason, diff --git a/posthog/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py index 01e05cb1..97f2dc5d 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -80,6 +80,8 @@ async def test_async_streaming_emits_the_same_properties_as_sync( # Guard: if the sync side stopped emitting these, the comparison below is vacuous. assert "$ai_usage" in sync_props assert "$ai_tokens_source" in sync_props + 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 == [], ( diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 8f9b00a4..34307b03 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -174,6 +174,7 @@ alias posthog.ai.openai.openai_async.extract_openai_usage_from_chunk -> posthog. alias posthog.ai.openai.openai_async.finalize_ai_content -> posthog.ai.utils.finalize_ai_content alias posthog.ai.openai.openai_async.format_openai_streaming_input -> posthog.ai.openai.openai_converter.format_openai_streaming_input alias posthog.ai.openai.openai_async.format_openai_streaming_output -> posthog.ai.openai.openai_converter.format_openai_streaming_output +alias posthog.ai.openai.openai_async.get_model_params -> posthog.ai.utils.get_model_params alias posthog.ai.openai.openai_async.merge_provider_override -> posthog.ai.openai.wrapper_utils.merge_provider_override alias posthog.ai.openai.openai_async.merge_usage_stats -> posthog.ai.utils.merge_usage_stats alias posthog.ai.openai.openai_async.setup -> posthog.setup From 049e213da8d94e929668fe5efdeb8cf7940b153b Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Tue, 11 Aug 2026 13:31:55 +0200 Subject: [PATCH 6/6] fix(ai): preserve async streaming token source --- posthog/ai/openai/openai_async.py | 16 ++-------------- posthog/ai/utils.py | 18 ++++++++++++++---- posthog/test/ai/openai/test_async_parity.py | 3 ++- 3 files changed, 18 insertions(+), 19 deletions(-) diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 4031fd37..5abc23e9 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -271,13 +271,7 @@ async def _capture_streaming_event( latency=latency, distinct_id=posthog_distinct_id, trace_id=posthog_trace_id, - properties={ - "$ai_cache_read_input_tokens": usage_stats.get( - "cache_read_input_tokens", 0 - ), - "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0), - **(posthog_properties or {}), - }, + properties=posthog_properties, privacy_mode=posthog_privacy_mode, groups=posthog_groups, stop_reason=stop_reason, @@ -563,13 +557,7 @@ async def _capture_streaming_event( latency=latency, distinct_id=posthog_distinct_id, trace_id=posthog_trace_id, - properties={ - "$ai_cache_read_input_tokens": usage_stats.get( - "cache_read_input_tokens", 0 - ), - "$ai_reasoning_tokens": usage_stats.get("reasoning_tokens", 0), - **(posthog_properties or {}), - }, + properties=posthog_properties, privacy_mode=posthog_privacy_mode, groups=posthog_groups, stop_reason=stop_reason, 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/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py index 97f2dc5d..f2066b93 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -79,7 +79,8 @@ async def test_async_streaming_emits_the_same_properties_as_sync( # Guard: if the sync side stopped emitting these, the comparison below is vacuous. assert "$ai_usage" in sync_props - assert "$ai_tokens_source" 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