Skip to content

fix(ai): route async OpenAI streaming through capture_streaming_event #861

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marandaneto merged 7 commits into from
Aug 12, 2026
Merged

fix(ai): route async OpenAI streaming through capture_streaming_event #861

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
bcd5797
fix(ai): route async OpenAI streaming through capture_streaming_event
ckarnell Aug 9, 2026
6e9e243
Fix CI: silence F811 on imported fixtures, refresh the API snapshot
ckarnell Aug 10, 2026
ade11ac
Move the shared OpenAI fixtures into conftest.py
ckarnell Aug 10, 2026
ea6ebc0
chore: add OpenAI async streaming changeset
marandaneto Aug 10, 2026
b35d5b9
fix(ai): address async streaming review feedback
marandaneto Aug 11, 2026
049e213
fix(ai): preserve async streaming token source
marandaneto Aug 11, 2026
ab2a639
Merge branch 'main' into fix/openai-async-streaming-parity
marandaneto Aug 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .sampo/changesets/somber-duchess-joukahainen.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
pypi/posthog: patch
---

Fix async OpenAI streaming captures to include token usage and other generation properties emitted by synchronous streams.
173 changes: 46 additions & 127 deletions posthog/ai/openai/openai_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Comment thread
marandaneto marked this conversation as resolved.
get_model_params as get_model_params,
merge_usage_stats,
with_privacy_mode,
)
Expand Down Expand Up @@ -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,
)
Expand All @@ -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(
Comment thread
marandaneto marked this conversation as resolved.
"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)
Comment thread
marandaneto marked this conversation as resolved.

async def parse(
self,
Expand Down Expand Up @@ -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,
)
Expand All @@ -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):
Expand Down
18 changes: 14 additions & 4 deletions posthog/ai/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -790,15 +790,13 @@ 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)
cache_creation = event_data["usage_stats"].get("cache_creation_input_tokens", 0)
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",
Expand All @@ -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"
Expand Down
126 changes: 126 additions & 0 deletions posthog/test/ai/openai/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
),
),
]
Loading