From 434a8d3ac0c707e05d2926a8c92a6ab272ad1728 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 21 Aug 2026 16:06:36 +0200 Subject: [PATCH 1/2] refactor(openai): extract streaming state processors --- posthog/ai/openai/openai.py | 202 ++++---------------- posthog/ai/openai/openai_async.py | 185 ++++-------------- posthog/ai/openai/streaming.py | 114 +++++++++++ posthog/test/ai/openai/test_async_parity.py | 75 ++++++++ 4 files changed, 255 insertions(+), 321 deletions(-) create mode 100644 posthog/ai/openai/streaming.py diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 35bb0c2df..24cbc6f33 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -1,8 +1,6 @@ import time import uuid -from typing import Any, Dict, List, Optional - -from posthog.ai.types import TokenUsage +from typing import Any, Dict, Optional try: import openai @@ -14,19 +12,20 @@ from posthog.ai.utils import ( call_llm_and_track_usage, _capture_ai_event, - extract_available_tool_calls, finalize_ai_content, - merge_usage_stats, with_privacy_mode, ) from posthog.ai.openai.openai_converter import ( - extract_openai_usage_from_chunk, - extract_openai_content_from_chunk, - extract_openai_tool_calls_from_chunk, - accumulate_openai_tool_calls, + format_openai_streaming_input, + format_openai_streaming_output, ) from posthog.client import Client as PostHogClient from posthog import setup +from posthog.ai.openai.streaming import ( + _ChatCompletionsStreamState, + _ResponsesStreamState, + _build_streaming_event_data, +) from posthog.ai.openai.wrapper_utils import ( _OpenAIWrapperResource, merge_provider_override, @@ -168,55 +167,15 @@ def _create_streaming( **kwargs: Any, ): start_time = time.time() - usage_stats: TokenUsage = TokenUsage() - final_content: List[Any] = [] - model_from_response: Optional[str] = None - stop_reason: Optional[str] = None + state = _ResponsesStreamState() response = self._original.create(**kwargs) def generator(): - nonlocal usage_stats - nonlocal final_content - nonlocal model_from_response - nonlocal stop_reason - try: for chunk in response: - # Extract model from response object in chunk (for stored prompts) - if hasattr(chunk, "response") and chunk.response: - if model_from_response is None and hasattr( - chunk.response, "model" - ): - model_from_response = chunk.response.model - - # Extract usage stats from chunk - chunk_usage = extract_openai_usage_from_chunk(chunk, "responses") - - if chunk_usage: - merge_usage_stats(usage_stats, chunk_usage) - - content = extract_openai_content_from_chunk(chunk, "responses") - - if content is not None: - final_content.extend(content) - - # Capture stop reason from response.completed event - if ( - hasattr(chunk, "type") - and chunk.type == "response.completed" - and hasattr(chunk, "response") - and chunk.response - ): - chunk_status = getattr(chunk.response, "status", None) - if chunk_status is not None: - stop_reason = chunk_status - + state.process_chunk(chunk) yield chunk - finally: - end_time = time.time() - latency = end_time - start_time - output = final_content self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, @@ -224,12 +183,8 @@ def generator(): posthog_privacy_mode, posthog_groups, kwargs, - usage_stats, - latency, - output, - None, # Responses API doesn't have tools - model_from_response, - stop_reason=stop_reason, + state, + time.time() - start_time, ) return generator() @@ -242,43 +197,26 @@ def _capture_streaming_event( posthog_privacy_mode: bool, posthog_groups: Optional[Dict[str, Any]], kwargs: Dict[str, Any], - usage_stats: TokenUsage, + state: _ResponsesStreamState, latency: float, - output: Any, - available_tool_calls: Optional[List[Dict[str, Any]]] = None, - model_from_response: Optional[str] = None, - stop_reason: Optional[str] = None, ): - from posthog.ai.types import StreamingEventData - from posthog.ai.openai.openai_converter import ( - format_openai_streaming_input, - format_openai_streaming_output, - ) 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_data = StreamingEventData( - provider="openai", - model=model, - base_url=str(self._client.base_url), + event_data = _build_streaming_event_data( + base_url=self._client.base_url, kwargs=kwargs, - formatted_input=formatted_input, - formatted_output=format_openai_streaming_output(output, "responses"), - usage_stats=usage_stats, + formatted_input=format_openai_streaming_input(kwargs, "responses"), + formatted_output=format_openai_streaming_output(state.output, "responses"), + usage_stats=state.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, + model_from_response=state.model, + stop_reason=state.stop_reason, ) - - # Use the common capture function capture_streaming_event(self._client._ph_client, event_data) def parse( @@ -443,69 +381,18 @@ def _create_streaming( **kwargs: Any, ): start_time = time.time() - usage_stats: TokenUsage = TokenUsage() - accumulated_content: List[Any] = [] - accumulated_tool_calls: Dict[int, Dict[str, Any]] = {} - model_from_response: Optional[str] = None - stop_reason: Optional[str] = None + state = _ChatCompletionsStreamState() if "stream_options" not in kwargs: kwargs["stream_options"] = {} kwargs["stream_options"]["include_usage"] = True response = self._original.create(**kwargs) def generator(): - nonlocal usage_stats - nonlocal accumulated_content - nonlocal accumulated_tool_calls - nonlocal model_from_response - nonlocal stop_reason - try: for chunk in response: - # Extract model from chunk (Chat Completions chunks have model field) - if model_from_response is None and hasattr(chunk, "model"): - model_from_response = chunk.model - - # Extract usage stats from chunk - chunk_usage = extract_openai_usage_from_chunk(chunk, "chat") - - if chunk_usage: - merge_usage_stats(usage_stats, chunk_usage) - - # Extract content from chunk - content = extract_openai_content_from_chunk(chunk, "chat") - - if content is not None: - accumulated_content.append(content) - - # Extract and accumulate tool calls from chunk - chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk) - if chunk_tool_calls: - accumulate_openai_tool_calls( - accumulated_tool_calls, chunk_tool_calls - ) - - # Capture stop reason from chunk - if ( - hasattr(chunk, "choices") - and chunk.choices - and getattr(chunk.choices[0], "finish_reason", None) is not None - ): - stop_reason = chunk.choices[0].finish_reason - + state.process_chunk(chunk) yield chunk - finally: - end_time = time.time() - latency = end_time - start_time - - # Convert accumulated tool calls dict to list - tool_calls_list = ( - list(accumulated_tool_calls.values()) - if accumulated_tool_calls - else None - ) - self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, @@ -513,13 +400,8 @@ def generator(): posthog_privacy_mode, posthog_groups, kwargs, - usage_stats, - latency, - accumulated_content, - tool_calls_list, - extract_available_tool_calls("openai", kwargs), - model_from_response, - stop_reason=stop_reason, + state, + time.time() - start_time, ) return generator() @@ -532,44 +414,28 @@ def _capture_streaming_event( posthog_privacy_mode: bool, posthog_groups: Optional[Dict[str, Any]], kwargs: Dict[str, Any], - usage_stats: TokenUsage, + state: _ChatCompletionsStreamState, 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, ): - from posthog.ai.types import StreamingEventData - from posthog.ai.openai.openai_converter import ( - format_openai_streaming_input, - format_openai_streaming_output, - ) 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_data = StreamingEventData( - provider="openai", - model=model, - base_url=str(self._client.base_url), + event_data = _build_streaming_event_data( + base_url=self._client.base_url, kwargs=kwargs, - formatted_input=formatted_input, - formatted_output=format_openai_streaming_output(output, "chat", tool_calls), - usage_stats=usage_stats, + formatted_input=format_openai_streaming_input(kwargs, "chat"), + formatted_output=format_openai_streaming_output( + state.output, "chat", state.tool_calls + ), + usage_stats=state.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, + model_from_response=state.model, + stop_reason=state.stop_reason, ) - - # Use the common capture function capture_streaming_event(self._client._ph_client, event_data) diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 5abc23e9b..723ea9b1c 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -1,6 +1,6 @@ import time import uuid -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional from posthog.ai.stream import AsyncStreamWrapper from posthog.ai.types import TokenUsage @@ -16,21 +16,20 @@ from posthog.ai.utils import ( call_llm_and_track_usage_async, _capture_ai_event, - 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, ) from posthog.ai.openai.openai_converter import ( - extract_openai_usage_from_chunk, - extract_openai_content_from_chunk, - extract_openai_tool_calls_from_chunk, - accumulate_openai_tool_calls, format_openai_streaming_input, format_openai_streaming_output, ) from posthog.client import Client as PostHogClient +from posthog.ai.openai.streaming import ( + _ChatCompletionsStreamState, + _ResponsesStreamState, + _build_streaming_event_data, +) from posthog.ai.openai.wrapper_utils import ( _OpenAIWrapperResource, merge_provider_override, @@ -172,56 +171,15 @@ async def _create_streaming( **kwargs: Any, ): start_time = time.time() - usage_stats: TokenUsage = TokenUsage() - final_content: List[Any] = [] - model_from_response: Optional[str] = None - stop_reason: Optional[str] = None + state = _ResponsesStreamState() response = await self._original.create(**kwargs) async def async_generator(): - nonlocal usage_stats - nonlocal final_content - nonlocal model_from_response - nonlocal stop_reason - try: async for chunk in response: - # Extract model from response object in chunk (for stored prompts) - if hasattr(chunk, "response") and chunk.response: - if model_from_response is None and hasattr( - chunk.response, "model" - ): - model_from_response = chunk.response.model - - # Extract usage stats from chunk - chunk_usage = extract_openai_usage_from_chunk(chunk, "responses") - - if chunk_usage: - merge_usage_stats(usage_stats, chunk_usage) - - content = extract_openai_content_from_chunk(chunk, "responses") - - if content is not None: - final_content.extend(content) - - # Capture stop reason from response.completed event - if ( - hasattr(chunk, "type") - and chunk.type == "response.completed" - and hasattr(chunk, "response") - and chunk.response - ): - chunk_status = getattr(chunk.response, "status", None) - if chunk_status is not None: - stop_reason = chunk_status - + state.process_chunk(chunk) yield chunk - finally: - end_time = time.time() - latency = end_time - start_time - output = final_content - await self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, @@ -229,11 +187,8 @@ async def async_generator(): posthog_privacy_mode, posthog_groups, kwargs, - usage_stats, - latency, - output, - model_from_response, - stop_reason=stop_reason, + state, + time.time() - start_time, ) return AsyncStreamWrapper(async_generator(), stream=response) @@ -246,38 +201,26 @@ async def _capture_streaming_event( posthog_privacy_mode: bool, posthog_groups: Optional[Dict[str, Any]], kwargs: Dict[str, Any], - usage_stats: TokenUsage, + state: _ResponsesStreamState, latency: float, - output: Any, - model_from_response: Optional[str] = None, - stop_reason: Optional[str] = None, ): - 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_data = StreamingEventData( - provider="openai", - model=model, - base_url=str(self._client.base_url), + event_data = _build_streaming_event_data( + base_url=self._client.base_url, kwargs=kwargs, - formatted_input=formatted_input, - formatted_output=format_openai_streaming_output(output, "responses"), - usage_stats=usage_stats, + formatted_input=format_openai_streaming_input(kwargs, "responses"), + formatted_output=format_openai_streaming_output(state.output, "responses"), + usage_stats=state.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, + model_from_response=state.model, + stop_reason=state.stop_reason, ) - - # Use the common capture function capture_streaming_event(self._client._ph_client, event_data) async def parse( @@ -444,11 +387,7 @@ async def _create_streaming( **kwargs: Any, ): start_time = time.time() - usage_stats: TokenUsage = TokenUsage() - accumulated_content: List[Any] = [] - accumulated_tool_calls: Dict[int, Dict[str, Any]] = {} - model_from_response: Optional[str] = None - stop_reason: Optional[str] = None + state = _ChatCompletionsStreamState() if "stream_options" not in kwargs: kwargs["stream_options"] = {} @@ -456,56 +395,11 @@ async def _create_streaming( response = await self._original.create(**kwargs) async def async_generator(): - nonlocal usage_stats - nonlocal accumulated_content - nonlocal accumulated_tool_calls - nonlocal model_from_response - nonlocal stop_reason - try: async for chunk in response: - # Extract model from chunk (Chat Completions chunks have model field) - if model_from_response is None and hasattr(chunk, "model"): - model_from_response = chunk.model - - # Extract usage stats from chunk - chunk_usage = extract_openai_usage_from_chunk(chunk, "chat") - if chunk_usage: - merge_usage_stats(usage_stats, chunk_usage) - - # Extract content from chunk - content = extract_openai_content_from_chunk(chunk, "chat") - if content is not None: - accumulated_content.append(content) - - # Extract and accumulate tool calls from chunk - chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk) - if chunk_tool_calls: - accumulate_openai_tool_calls( - accumulated_tool_calls, chunk_tool_calls - ) - - # Capture stop reason from chunk - if ( - hasattr(chunk, "choices") - and chunk.choices - and getattr(chunk.choices[0], "finish_reason", None) is not None - ): - stop_reason = chunk.choices[0].finish_reason - + state.process_chunk(chunk) yield chunk - finally: - end_time = time.time() - latency = end_time - start_time - - # Convert accumulated tool calls dict to list - tool_calls_list = ( - list(accumulated_tool_calls.values()) - if accumulated_tool_calls - else None - ) - await self._capture_streaming_event( posthog_distinct_id, posthog_trace_id, @@ -513,12 +407,8 @@ async def async_generator(): posthog_privacy_mode, posthog_groups, kwargs, - usage_stats, - latency, - accumulated_content, - tool_calls_list, - model_from_response, - stop_reason=stop_reason, + state, + time.time() - start_time, ) return AsyncStreamWrapper(async_generator(), stream=response) @@ -531,39 +421,28 @@ async def _capture_streaming_event( posthog_privacy_mode: bool, posthog_groups: Optional[Dict[str, Any]], kwargs: Dict[str, Any], - usage_stats: TokenUsage, + state: _ChatCompletionsStreamState, latency: float, - output: Any, - tool_calls: Optional[List[Dict[str, Any]]] = None, - model_from_response: Optional[str] = None, - stop_reason: Optional[str] = None, ): - 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_data = StreamingEventData( - provider="openai", - model=model, - base_url=str(self._client.base_url), + event_data = _build_streaming_event_data( + base_url=self._client.base_url, kwargs=kwargs, - formatted_input=formatted_input, - formatted_output=format_openai_streaming_output(output, "chat", tool_calls), - usage_stats=usage_stats, + formatted_input=format_openai_streaming_input(kwargs, "chat"), + formatted_output=format_openai_streaming_output( + state.output, "chat", state.tool_calls + ), + usage_stats=state.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, + model_from_response=state.model, + stop_reason=state.stop_reason, ) - - # Use the common capture function capture_streaming_event(self._client._ph_client, event_data) diff --git a/posthog/ai/openai/streaming.py b/posthog/ai/openai/streaming.py new file mode 100644 index 000000000..3f2f38e74 --- /dev/null +++ b/posthog/ai/openai/streaming.py @@ -0,0 +1,114 @@ +"""Sync-neutral state accumulation for OpenAI streaming endpoints.""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from ..types import StreamingEventData, TokenUsage +from ..utils import merge_usage_stats +from .openai_converter import ( + accumulate_openai_tool_calls, + extract_openai_content_from_chunk, + extract_openai_tool_calls_from_chunk, + extract_openai_usage_from_chunk, +) + + +@dataclass +class _ResponsesStreamState: + """Accumulates state specific to a Responses API stream.""" + + usage_stats: TokenUsage = field(default_factory=lambda: TokenUsage()) + output: List[Any] = field(default_factory=list) + model: Optional[str] = None + stop_reason: Optional[str] = None + + def process_chunk(self, chunk: Any) -> None: + response = getattr(chunk, "response", None) + if response and self.model is None and hasattr(response, "model"): + self.model = response.model + + chunk_usage = extract_openai_usage_from_chunk(chunk, "responses") + if chunk_usage: + merge_usage_stats(self.usage_stats, chunk_usage) + + content = extract_openai_content_from_chunk(chunk, "responses") + if content is not None: + self.output.extend(content) + + if getattr(chunk, "type", None) == "response.completed" and response: + status = getattr(response, "status", None) + if status is not None: + self.stop_reason = status + + +@dataclass +class _ChatCompletionsStreamState: + """Accumulates state specific to a Chat Completions stream.""" + + usage_stats: TokenUsage = field(default_factory=lambda: TokenUsage()) + output: List[Any] = field(default_factory=list) + _tool_calls: Dict[int, Dict[str, Any]] = field(default_factory=dict) + model: Optional[str] = None + stop_reason: Optional[str] = None + + def process_chunk(self, chunk: Any) -> None: + if self.model is None and hasattr(chunk, "model"): + self.model = chunk.model + + chunk_usage = extract_openai_usage_from_chunk(chunk, "chat") + if chunk_usage: + merge_usage_stats(self.usage_stats, chunk_usage) + + content = extract_openai_content_from_chunk(chunk, "chat") + if content is not None: + self.output.append(content) + + chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk) + if chunk_tool_calls: + accumulate_openai_tool_calls(self._tool_calls, chunk_tool_calls) + + choices = getattr(chunk, "choices", None) + if choices: + finish_reason = getattr(choices[0], "finish_reason", None) + if finish_reason is not None: + self.stop_reason = finish_reason + + @property + def tool_calls(self) -> Optional[List[Dict[str, Any]]]: + return list(self._tool_calls.values()) if self._tool_calls else None + + +def _build_streaming_event_data( + *, + base_url: Any, + kwargs: Dict[str, Any], + formatted_input: Any, + formatted_output: Any, + usage_stats: TokenUsage, + latency: float, + distinct_id: Optional[str], + trace_id: Optional[str], + properties: Optional[Dict[str, Any]], + privacy_mode: bool, + groups: Optional[Dict[str, Any]], + model_from_response: Optional[str], + stop_reason: Optional[str], +) -> StreamingEventData: + """Build the fields shared by both OpenAI streaming endpoint events.""" + + return StreamingEventData( + provider="openai", + model=kwargs.get("model") or model_from_response or "unknown", + base_url=str(base_url), + kwargs=kwargs, + formatted_input=formatted_input, + formatted_output=formatted_output, + usage_stats=usage_stats, + latency=latency, + distinct_id=distinct_id, + trace_id=trace_id, + properties=properties, + privacy_mode=privacy_mode, + groups=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 f2066b93c..6601197ee 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -12,11 +12,13 @@ .venv-posthog/bin/python -m pytest repo-posthog/posthog/test/ai/openai/test_async_parity.py -v """ +from types import SimpleNamespace from unittest.mock import patch import pytest from posthog.ai.openai import AsyncOpenAI, OpenAI +from posthog.test.ai.utils import make_response_usage TOOLS = [ { @@ -88,3 +90,76 @@ async def test_async_streaming_emits_the_same_properties_as_sync( assert missing == [], ( f"the async openai streaming path drops {missing} that the sync path sends" ) + + +@pytest.mark.asyncio +async def test_responses_streaming_properties_have_sync_async_parity(mock_client): + response = SimpleNamespace( + model="gpt-4o-response", + status="completed", + usage=make_response_usage(11, 7, 18, cached_tokens=3), + output=[ + SimpleNamespace( + type="message", + role="assistant", + content=[SimpleNamespace(type="output_text", text="hello")], + ) + ], + ) + chunk = SimpleNamespace(type="response.completed", response=response) + request = { + "input": [{"role": "user", "content": "Hi"}], + "stream": True, + "posthog_distinct_id": "test-id", + "posthog_trace_id": "shared-trace", + "posthog_provider_override": "groq", + } + + with patch( + "openai.resources.responses.Responses.create", return_value=iter([chunk]) + ): + client = OpenAI(api_key="test-key", posthog_client=mock_client) + list(client.responses.create(**request)) + sync_props = mock_client.capture.call_args.kwargs["properties"] + + async def create(self, **kwargs): + async def chunks(): + yield chunk + + return chunks() + + mock_client.capture.reset_mock() + with patch("openai.resources.responses.AsyncResponses.create", new=create): + client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) + stream = await client.responses.create(**request) + async for _ in stream: + pass + async_props = mock_client.capture.call_args.kwargs["properties"] + + sync_without_latency = {k: v for k, v in sync_props.items() if k != "$ai_latency"} + async_without_latency = {k: v for k, v in async_props.items() if k != "$ai_latency"} + assert async_without_latency == sync_without_latency + assert async_props["$ai_model"] == "gpt-4o-response" + assert async_props["$ai_stop_reason"] == "completed" + assert async_props["$ai_provider"] == "groq" + + +def test_sync_stream_close_after_early_exit_captures_partial_state( + mock_client, streaming_tool_call_chunks +): + with patch( + "openai.resources.chat.completions.Completions.create", + return_value=iter(streaming_tool_call_chunks), + ): + client = OpenAI(api_key="test-key", posthog_client=mock_client) + stream = client.chat.completions.create( + model="gpt-4", + messages=MESSAGES, + stream=True, + posthog_distinct_id="test-id", + ) + assert next(stream) == streaming_tool_call_chunks[0] + stream.close() + + assert mock_client.capture.call_count == 1 + assert mock_client.capture.call_args.kwargs["properties"]["$ai_model"] == "gpt-4" From 36fbc4d432ba82afb5e2c9a13379a355ef73c09e Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Sat, 22 Aug 2026 08:53:35 +0200 Subject: [PATCH 2/2] fix(openai): preserve public API aliases --- .../ai/openai/{streaming.py => _streaming.py} | 0 posthog/ai/openai/openai.py | 22 +++++++++++++------ posthog/ai/openai/openai_async.py | 8 ++++++- 3 files changed, 22 insertions(+), 8 deletions(-) rename posthog/ai/openai/{streaming.py => _streaming.py} (100%) diff --git a/posthog/ai/openai/streaming.py b/posthog/ai/openai/_streaming.py similarity index 100% rename from posthog/ai/openai/streaming.py rename to posthog/ai/openai/_streaming.py diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 24cbc6f33..3e8234b6a 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -2,6 +2,8 @@ import uuid from typing import Any, Dict, Optional +from posthog.ai.types import TokenUsage as TokenUsage + try: import openai except ImportError: @@ -12,16 +14,22 @@ from posthog.ai.utils import ( call_llm_and_track_usage, _capture_ai_event, + extract_available_tool_calls as extract_available_tool_calls, finalize_ai_content, + merge_usage_stats as merge_usage_stats, with_privacy_mode, ) from posthog.ai.openai.openai_converter import ( - format_openai_streaming_input, - format_openai_streaming_output, + accumulate_openai_tool_calls as accumulate_openai_tool_calls, + extract_openai_content_from_chunk as extract_openai_content_from_chunk, + extract_openai_tool_calls_from_chunk as extract_openai_tool_calls_from_chunk, + extract_openai_usage_from_chunk as extract_openai_usage_from_chunk, + format_openai_streaming_input as _format_openai_streaming_input, + format_openai_streaming_output as _format_openai_streaming_output, ) from posthog.client import Client as PostHogClient from posthog import setup -from posthog.ai.openai.streaming import ( +from posthog.ai.openai._streaming import ( _ChatCompletionsStreamState, _ResponsesStreamState, _build_streaming_event_data, @@ -205,8 +213,8 @@ def _capture_streaming_event( event_data = _build_streaming_event_data( base_url=self._client.base_url, kwargs=kwargs, - formatted_input=format_openai_streaming_input(kwargs, "responses"), - formatted_output=format_openai_streaming_output(state.output, "responses"), + formatted_input=_format_openai_streaming_input(kwargs, "responses"), + formatted_output=_format_openai_streaming_output(state.output, "responses"), usage_stats=state.usage_stats, latency=latency, distinct_id=posthog_distinct_id, @@ -422,8 +430,8 @@ def _capture_streaming_event( event_data = _build_streaming_event_data( base_url=self._client.base_url, kwargs=kwargs, - formatted_input=format_openai_streaming_input(kwargs, "chat"), - formatted_output=format_openai_streaming_output( + formatted_input=_format_openai_streaming_input(kwargs, "chat"), + formatted_output=_format_openai_streaming_output( state.output, "chat", state.tool_calls ), usage_stats=state.usage_stats, diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 723ea9b1c..7d61a9215 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -16,16 +16,22 @@ from posthog.ai.utils import ( call_llm_and_track_usage_async, _capture_ai_event, + extract_available_tool_calls as extract_available_tool_calls, finalize_ai_content, get_model_params as get_model_params, + merge_usage_stats as merge_usage_stats, with_privacy_mode, ) from posthog.ai.openai.openai_converter import ( + accumulate_openai_tool_calls as accumulate_openai_tool_calls, + extract_openai_content_from_chunk as extract_openai_content_from_chunk, + extract_openai_tool_calls_from_chunk as extract_openai_tool_calls_from_chunk, + extract_openai_usage_from_chunk as extract_openai_usage_from_chunk, format_openai_streaming_input, format_openai_streaming_output, ) from posthog.client import Client as PostHogClient -from posthog.ai.openai.streaming import ( +from posthog.ai.openai._streaming import ( _ChatCompletionsStreamState, _ResponsesStreamState, _build_streaming_event_data,