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/6] 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/6] 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, From deeffe48615e28796f6506c79dfda73bc2ae11ca Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 21 Aug 2026 16:16:29 +0200 Subject: [PATCH 3/6] refactor(openai): centralize resource and embedding policy # Conflicts: # posthog/ai/openai/openai.py # posthog/ai/openai/openai_async.py --- posthog/ai/openai/embeddings.py | 49 +++++++++++ posthog/ai/openai/openai.py | 83 ++++++------------ posthog/ai/openai/openai_async.py | 85 ++++++------------- posthog/ai/openai/openai_providers.py | 52 ++---------- posthog/ai/openai/wrapper_utils.py | 11 ++- posthog/test/ai/openai/test_async_parity.py | 49 ++++++++++- .../test/ai/openai/test_resource_wrapping.py | 81 ++++++++++++++++++ 7 files changed, 241 insertions(+), 169 deletions(-) create mode 100644 posthog/ai/openai/embeddings.py create mode 100644 posthog/test/ai/openai/test_resource_wrapping.py diff --git a/posthog/ai/openai/embeddings.py b/posthog/ai/openai/embeddings.py new file mode 100644 index 000000000..834e8bc57 --- /dev/null +++ b/posthog/ai/openai/embeddings.py @@ -0,0 +1,49 @@ +from typing import Any, Dict, Optional + +from ..utils import _capture_ai_event, finalize_ai_content, with_privacy_mode + + +def _capture_embedding_event( + *, + posthog_client: Any, + base_url: Any, + response: Any, + request_kwargs: Dict[str, Any], + latency: float, + distinct_id: Optional[str], + trace_id: str, + properties: Optional[Dict[str, Any]], + privacy_mode: bool, + groups: Optional[Dict[str, Any]], +) -> None: + """Build and capture telemetry shared by sync and async embedding wrappers.""" + usage = getattr(response, "usage", None) + input_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0 + + event_properties = { + "$ai_provider": "openai", + "$ai_model": request_kwargs.get("model"), + "$ai_input": with_privacy_mode( + posthog_client, + privacy_mode, + finalize_ai_content(request_kwargs.get("input"), posthog_client), + ), + "$ai_http_status": 200, + "$ai_input_tokens": input_tokens, + "$ai_latency": latency, + "$ai_trace_id": trace_id, + "$ai_base_url": str(base_url), + **(properties or {}), + } + + if distinct_id is None: + event_properties["$process_person_profile"] = False + + if hasattr(posthog_client, "capture"): + _capture_ai_event( + posthog_client, + "$ai_embedding", + distinct_id=distinct_id or trace_id, + properties=event_properties, + groups=groups, + ) diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 3e8234b6a..365fee0a4 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -34,8 +34,10 @@ _ResponsesStreamState, _build_streaming_event_data, ) -from posthog.ai.openai.wrapper_utils import ( +from .embeddings import _capture_embedding_event +from .wrapper_utils import ( _OpenAIWrapperResource, + _wrap_openai_resources, merge_provider_override, ) @@ -59,24 +61,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): super().__init__(**kwargs) self._ph_client = posthog_client or setup() - # Store original objects after parent initialization (only if they exist) - self._original_chat = getattr(self, "chat", None) - self._original_embeddings = getattr(self, "embeddings", None) - self._original_beta = getattr(self, "beta", None) - self._original_responses = getattr(self, "responses", None) - - # Replace with wrapped versions (only if originals exist) - if self._original_chat is not None: - self.chat = WrappedChat(self, self._original_chat) - - if self._original_embeddings is not None: - self.embeddings = WrappedEmbeddings(self, self._original_embeddings) - - if self._original_beta is not None: - self.beta = WrappedBeta(self, self._original_beta) - - if self._original_responses is not None: - self.responses = WrappedResponses(self, self._original_responses) + _wrap_openai_resources(self, _SYNC_RESOURCE_WRAPPERS) def _parse_and_track( @@ -491,46 +476,18 @@ def create( response = self._original.create(**kwargs) end_time = time.time() - # Extract usage statistics if available - usage_stats = {} - if hasattr(response, "usage") and response.usage: - usage_stats = { - "prompt_tokens": getattr(response.usage, "prompt_tokens", 0), - "total_tokens": getattr(response.usage, "total_tokens", 0), - } - - latency = end_time - start_time - - # Build the event properties - event_properties = { - "$ai_provider": "openai", - "$ai_model": kwargs.get("model"), - "$ai_input": with_privacy_mode( - self._client._ph_client, - posthog_privacy_mode, - finalize_ai_content(kwargs.get("input"), self._client._ph_client), - ), - "$ai_http_status": 200, - "$ai_input_tokens": usage_stats.get("prompt_tokens", 0), - "$ai_latency": latency, - "$ai_trace_id": posthog_trace_id, - "$ai_base_url": str(self._client.base_url), - **(posthog_properties or {}), - } - - if posthog_distinct_id is None: - event_properties["$process_person_profile"] = False - - # Send capture event for embeddings - if hasattr(self._client._ph_client, "capture"): - _capture_ai_event( - self._client._ph_client, - "$ai_embedding", - distinct_id=posthog_distinct_id or posthog_trace_id, - properties=event_properties, - groups=posthog_groups, - ) - + _capture_embedding_event( + posthog_client=self._client._ph_client, + base_url=self._client.base_url, + response=response, + request_kwargs=kwargs, + latency=end_time - start_time, + distinct_id=posthog_distinct_id, + trace_id=posthog_trace_id, + properties=posthog_properties, + privacy_mode=posthog_privacy_mode, + groups=posthog_groups, + ) return response @@ -594,3 +551,11 @@ def parse( posthog_provider_override, **kwargs, ) + + +_SYNC_RESOURCE_WRAPPERS = { + "chat": WrappedChat, + "embeddings": WrappedEmbeddings, + "beta": WrappedBeta, + "responses": WrappedResponses, +} diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 7d61a9215..dc4c039e3 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -3,7 +3,6 @@ from typing import Any, Dict, Optional from posthog.ai.stream import AsyncStreamWrapper -from posthog.ai.types import TokenUsage try: import openai @@ -36,8 +35,10 @@ _ResponsesStreamState, _build_streaming_event_data, ) -from posthog.ai.openai.wrapper_utils import ( +from .embeddings import _capture_embedding_event +from .wrapper_utils import ( _OpenAIWrapperResource, + _wrap_openai_resources, merge_provider_override, ) @@ -61,24 +62,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): super().__init__(**kwargs) self._ph_client = posthog_client or setup() - # Store original objects after parent initialization (only if they exist) - self._original_chat = getattr(self, "chat", None) - self._original_embeddings = getattr(self, "embeddings", None) - self._original_beta = getattr(self, "beta", None) - self._original_responses = getattr(self, "responses", None) - - # Replace with wrapped versions (only if originals exist) - if self._original_chat is not None: - self.chat = WrappedChat(self, self._original_chat) - - if self._original_embeddings is not None: - self.embeddings = WrappedEmbeddings(self, self._original_embeddings) - - if self._original_beta is not None: - self.beta = WrappedBeta(self, self._original_beta) - - if self._original_responses is not None: - self.responses = WrappedResponses(self, self._original_responses) + _wrap_openai_resources(self, _ASYNC_RESOURCE_WRAPPERS) async def _parse_and_track( @@ -496,47 +480,18 @@ async def create( response = await self._original.create(**kwargs) end_time = time.time() - # Extract usage statistics if available - usage_stats: TokenUsage = TokenUsage() - - if hasattr(response, "usage") and response.usage: - usage_stats = TokenUsage( - input_tokens=getattr(response.usage, "prompt_tokens", 0), - output_tokens=getattr(response.usage, "completion_tokens", 0), - ) - - latency = end_time - start_time - - # Build the event properties - event_properties = { - "$ai_provider": "openai", - "$ai_model": kwargs.get("model"), - "$ai_input": with_privacy_mode( - self._client._ph_client, - posthog_privacy_mode, - finalize_ai_content(kwargs.get("input"), self._client._ph_client), - ), - "$ai_http_status": 200, - "$ai_input_tokens": usage_stats.get("input_tokens", 0), - "$ai_latency": latency, - "$ai_trace_id": posthog_trace_id, - "$ai_base_url": str(self._client.base_url), - **(posthog_properties or {}), - } - - if posthog_distinct_id is None: - event_properties["$process_person_profile"] = False - - # Send capture event for embeddings - if hasattr(self._client._ph_client, "capture"): - _capture_ai_event( - self._client._ph_client, - "$ai_embedding", - distinct_id=posthog_distinct_id or posthog_trace_id, - properties=event_properties, - groups=posthog_groups, - ) - + _capture_embedding_event( + posthog_client=self._client._ph_client, + base_url=self._client.base_url, + response=response, + request_kwargs=kwargs, + latency=end_time - start_time, + distinct_id=posthog_distinct_id, + trace_id=posthog_trace_id, + properties=posthog_properties, + privacy_mode=posthog_privacy_mode, + groups=posthog_groups, + ) return response @@ -600,3 +555,11 @@ async def parse( posthog_provider_override, **kwargs, ) + + +_ASYNC_RESOURCE_WRAPPERS = { + "chat": WrappedChat, + "embeddings": WrappedEmbeddings, + "beta": WrappedBeta, + "responses": WrappedResponses, +} diff --git a/posthog/ai/openai/openai_providers.py b/posthog/ai/openai/openai_providers.py index d86d49988..46ce5e058 100644 --- a/posthog/ai/openai/openai_providers.py +++ b/posthog/ai/openai/openai_providers.py @@ -5,16 +5,9 @@ "Please install the Open AI SDK to use this feature: 'pip install openai'" ) -from posthog.ai.openai.openai import ( - WrappedBeta, - WrappedChat, - WrappedEmbeddings, - WrappedResponses, -) -from posthog.ai.openai.openai_async import WrappedBeta as AsyncWrappedBeta -from posthog.ai.openai.openai_async import WrappedChat as AsyncWrappedChat -from posthog.ai.openai.openai_async import WrappedEmbeddings as AsyncWrappedEmbeddings -from posthog.ai.openai.openai_async import WrappedResponses as AsyncWrappedResponses +from .openai import _SYNC_RESOURCE_WRAPPERS +from .openai_async import _ASYNC_RESOURCE_WRAPPERS +from .wrapper_utils import _wrap_openai_resources from typing import Optional from posthog.client import Client as PostHogClient @@ -39,24 +32,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): super().__init__(**kwargs) self._ph_client = posthog_client or setup() - # Store original objects after parent initialization (only if they exist) - self._original_chat = getattr(self, "chat", None) - self._original_embeddings = getattr(self, "embeddings", None) - self._original_beta = getattr(self, "beta", None) - self._original_responses = getattr(self, "responses", None) - - # Replace with wrapped versions (only if originals exist) - if self._original_chat is not None: - self.chat = WrappedChat(self, self._original_chat) - - if self._original_embeddings is not None: - self.embeddings = WrappedEmbeddings(self, self._original_embeddings) - - if self._original_beta is not None: - self.beta = WrappedBeta(self, self._original_beta) - - if self._original_responses is not None: - self.responses = WrappedResponses(self, self._original_responses) + _wrap_openai_resources(self, _SYNC_RESOURCE_WRAPPERS) class AsyncAzureOpenAI(openai.AsyncAzureOpenAI): @@ -77,22 +53,4 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): super().__init__(**kwargs) self._ph_client = posthog_client or setup() - # Store original objects after parent initialization (only if they exist) - self._original_chat = getattr(self, "chat", None) - self._original_embeddings = getattr(self, "embeddings", None) - self._original_beta = getattr(self, "beta", None) - self._original_responses = getattr(self, "responses", None) - - # Replace with wrapped versions (only if originals exist) - if self._original_chat is not None: - self.chat = AsyncWrappedChat(self, self._original_chat) - - if self._original_embeddings is not None: - self.embeddings = AsyncWrappedEmbeddings(self, self._original_embeddings) - - if self._original_beta is not None: - self.beta = AsyncWrappedBeta(self, self._original_beta) - - # Only add responses if available (newer OpenAI versions) - if self._original_responses is not None: - self.responses = AsyncWrappedResponses(self, self._original_responses) + _wrap_openai_resources(self, _ASYNC_RESOURCE_WRAPPERS) diff --git a/posthog/ai/openai/wrapper_utils.py b/posthog/ai/openai/wrapper_utils.py index 5cefab814..509f2d94d 100644 --- a/posthog/ai/openai/wrapper_utils.py +++ b/posthog/ai/openai/wrapper_utils.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Dict, Optional +from typing import Any, Dict, Mapping, Optional log = logging.getLogger("posthog") @@ -33,6 +33,15 @@ def merge_provider_override( return {**(posthog_properties or {}), "$ai_provider": posthog_provider_override} +def _wrap_openai_resources(client: Any, wrappers: Mapping[str, type]) -> None: + """Save and replace available SDK resources using an explicit wrapper mapping.""" + for resource_name, wrapper_type in wrappers.items(): + original = getattr(client, resource_name, None) + setattr(client, f"_original_{resource_name}", original) + if original is not None: + setattr(client, resource_name, wrapper_type(client, original)) + + def reset_fallback_warnings() -> None: _fallback_warnings.clear() diff --git a/posthog/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py index 6601197ee..850b2a885 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -13,7 +13,7 @@ """ from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest @@ -144,6 +144,53 @@ async def chunks(): assert async_props["$ai_provider"] == "groq" +@pytest.mark.asyncio +async def test_embedding_telemetry_has_sync_async_parity(mock_client): + response = SimpleNamespace(usage=SimpleNamespace(prompt_tokens=12, total_tokens=12)) + request = { + "model": "text-embedding-3-small", + "input": "private input", + "posthog_trace_id": "shared-trace", + "posthog_properties": {"custom": "value"}, + "posthog_privacy_mode": True, + "posthog_groups": {"company": "test-company"}, + "posthog_provider_override": "azure", + } + provider_request = { + "model": "text-embedding-3-small", + "input": "private input", + } + + with patch( + "openai.resources.embeddings.Embeddings.create", return_value=response + ) as sync_create: + client = OpenAI(api_key="test-key", posthog_client=mock_client) + assert client.embeddings.create(**request) is response + sync_create.assert_called_once_with(**provider_request) + sync_capture = mock_client.capture.call_args + + async_create = AsyncMock(return_value=response) + mock_client.capture.reset_mock() + with patch("openai.resources.embeddings.AsyncEmbeddings.create", new=async_create): + client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) + assert await client.embeddings.create(**request) is response + async_create.assert_awaited_once_with(**provider_request) + async_capture = mock_client.capture.call_args + + sync_props = sync_capture.kwargs["properties"] + async_props = async_capture.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_input"] is None + assert async_props["$ai_input_tokens"] == 12 + assert async_props["$ai_provider"] == "azure" + assert async_props["$process_person_profile"] is False + assert async_capture.kwargs["distinct_id"] == "shared-trace" + assert async_capture.kwargs["groups"] == {"company": "test-company"} + + def test_sync_stream_close_after_early_exit_captures_partial_state( mock_client, streaming_tool_call_chunks ): diff --git a/posthog/test/ai/openai/test_resource_wrapping.py b/posthog/test/ai/openai/test_resource_wrapping.py new file mode 100644 index 000000000..3e1117c46 --- /dev/null +++ b/posthog/test/ai/openai/test_resource_wrapping.py @@ -0,0 +1,81 @@ +from unittest.mock import MagicMock + +import pytest +from openai.resources.beta import AsyncBeta, Beta +from openai.resources.chat import AsyncChat, Chat +from openai.resources.embeddings import AsyncEmbeddings, Embeddings +from openai.resources.responses import AsyncResponses, Responses + +from posthog.ai.openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI +from posthog.ai.openai.openai import ( + WrappedBeta, + WrappedChat, + WrappedEmbeddings, + WrappedResponses, +) +from posthog.ai.openai.openai_async import ( + WrappedBeta as AsyncWrappedBeta, +) +from posthog.ai.openai.openai_async import ( + WrappedChat as AsyncWrappedChat, +) +from posthog.ai.openai.openai_async import ( + WrappedEmbeddings as AsyncWrappedEmbeddings, +) +from posthog.ai.openai.openai_async import ( + WrappedResponses as AsyncWrappedResponses, +) + +_SYNC_WRAPPERS = { + "chat": WrappedChat, + "embeddings": WrappedEmbeddings, + "beta": WrappedBeta, + "responses": WrappedResponses, +} +_ASYNC_WRAPPERS = { + "chat": AsyncWrappedChat, + "embeddings": AsyncWrappedEmbeddings, + "beta": AsyncWrappedBeta, + "responses": AsyncWrappedResponses, +} +_SYNC_RESOURCES = { + "chat": Chat, + "embeddings": Embeddings, + "beta": Beta, + "responses": Responses, +} +_ASYNC_RESOURCES = { + "chat": AsyncChat, + "embeddings": AsyncEmbeddings, + "beta": AsyncBeta, + "responses": AsyncResponses, +} +_AZURE_KWARGS = { + "api_key": "test-key", + "azure_endpoint": "https://example.openai.azure.com", + "api_version": "2024-02-01", +} + + +@pytest.mark.parametrize( + "client_type, client_kwargs, wrappers, resource_types", + [ + (OpenAI, {"api_key": "test-key"}, _SYNC_WRAPPERS, _SYNC_RESOURCES), + (AsyncOpenAI, {"api_key": "test-key"}, _ASYNC_WRAPPERS, _ASYNC_RESOURCES), + (AzureOpenAI, _AZURE_KWARGS, _SYNC_WRAPPERS, _SYNC_RESOURCES), + (AsyncAzureOpenAI, _AZURE_KWARGS, _ASYNC_WRAPPERS, _ASYNC_RESOURCES), + ], +) +def test_client_resources_are_discovered_and_wrapped( + client_type, client_kwargs, wrappers, resource_types +): + client = client_type(posthog_client=MagicMock(), **client_kwargs) + + for resource_name, wrapper_type in wrappers.items(): + wrapped = getattr(client, resource_name) + original = getattr(client, f"_original_{resource_name}") + + assert type(wrapped) is wrapper_type + assert type(original) is resource_types[resource_name] + assert wrapped._client is client + assert wrapped._original is original From 7ba992516d8d240ea70912a9836b6bf02fe914d4 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Sat, 22 Aug 2026 09:10:03 +0200 Subject: [PATCH 4/6] fix(openai): preserve wrapper public API --- .../openai/{embeddings.py => _embeddings.py} | 0 posthog/ai/openai/openai.py | 16 ++++++--- posthog/ai/openai/openai_async.py | 17 +++++++--- posthog/ai/openai/openai_providers.py | 33 +++++++++++++++++-- 4 files changed, 53 insertions(+), 13 deletions(-) rename posthog/ai/openai/{embeddings.py => _embeddings.py} (100%) diff --git a/posthog/ai/openai/embeddings.py b/posthog/ai/openai/_embeddings.py similarity index 100% rename from posthog/ai/openai/embeddings.py rename to posthog/ai/openai/_embeddings.py diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 365fee0a4..a9bbb7b0e 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -1,6 +1,6 @@ import time import uuid -from typing import Any, Dict, Optional +from typing import TYPE_CHECKING as _TYPE_CHECKING, Any, Dict, Optional from posthog.ai.types import TokenUsage as TokenUsage @@ -13,11 +13,10 @@ 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, + finalize_ai_content as finalize_ai_content, merge_usage_stats as merge_usage_stats, - with_privacy_mode, + with_privacy_mode as with_privacy_mode, ) from posthog.ai.openai.openai_converter import ( accumulate_openai_tool_calls as accumulate_openai_tool_calls, @@ -34,7 +33,7 @@ _ResponsesStreamState, _build_streaming_event_data, ) -from .embeddings import _capture_embedding_event +from ._embeddings import _capture_embedding_event from .wrapper_utils import ( _OpenAIWrapperResource, _wrap_openai_resources, @@ -63,6 +62,13 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): _wrap_openai_resources(self, _SYNC_RESOURCE_WRAPPERS) + # Keep dynamically installed resources visible to API and type inspection. + if _TYPE_CHECKING: + self.chat = WrappedChat(self, self._original_chat) + self.embeddings = WrappedEmbeddings(self, self._original_embeddings) + self.beta = WrappedBeta(self, self._original_beta) + self.responses = WrappedResponses(self, self._original_responses) + def _parse_and_track( wrapper, diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index dc4c039e3..0cb41a039 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -1,8 +1,9 @@ import time import uuid -from typing import Any, Dict, Optional +from typing import TYPE_CHECKING as _TYPE_CHECKING, Any, Dict, Optional from posthog.ai.stream import AsyncStreamWrapper +from posthog.ai.types import TokenUsage as TokenUsage try: import openai @@ -14,12 +15,11 @@ from posthog import setup 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, + finalize_ai_content as finalize_ai_content, get_model_params as get_model_params, merge_usage_stats as merge_usage_stats, - with_privacy_mode, + with_privacy_mode as with_privacy_mode, ) from posthog.ai.openai.openai_converter import ( accumulate_openai_tool_calls as accumulate_openai_tool_calls, @@ -35,7 +35,7 @@ _ResponsesStreamState, _build_streaming_event_data, ) -from .embeddings import _capture_embedding_event +from ._embeddings import _capture_embedding_event from .wrapper_utils import ( _OpenAIWrapperResource, _wrap_openai_resources, @@ -64,6 +64,13 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): _wrap_openai_resources(self, _ASYNC_RESOURCE_WRAPPERS) + # Keep dynamically installed resources visible to API and type inspection. + if _TYPE_CHECKING: + self.chat = WrappedChat(self, self._original_chat) + self.embeddings = WrappedEmbeddings(self, self._original_embeddings) + self.beta = WrappedBeta(self, self._original_beta) + self.responses = WrappedResponses(self, self._original_responses) + async def _parse_and_track( wrapper, diff --git a/posthog/ai/openai/openai_providers.py b/posthog/ai/openai/openai_providers.py index 46ce5e058..08d8c8b35 100644 --- a/posthog/ai/openai/openai_providers.py +++ b/posthog/ai/openai/openai_providers.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING as _TYPE_CHECKING, Optional + try: import openai except ImportError: @@ -5,10 +7,21 @@ "Please install the Open AI SDK to use this feature: 'pip install openai'" ) -from .openai import _SYNC_RESOURCE_WRAPPERS -from .openai_async import _ASYNC_RESOURCE_WRAPPERS +from .openai import ( + WrappedBeta as WrappedBeta, + WrappedChat as WrappedChat, + WrappedEmbeddings as WrappedEmbeddings, + WrappedResponses as WrappedResponses, + _SYNC_RESOURCE_WRAPPERS, +) +from .openai_async import ( + WrappedBeta as AsyncWrappedBeta, + WrappedChat as AsyncWrappedChat, + WrappedEmbeddings as AsyncWrappedEmbeddings, + WrappedResponses as AsyncWrappedResponses, + _ASYNC_RESOURCE_WRAPPERS, +) from .wrapper_utils import _wrap_openai_resources -from typing import Optional from posthog.client import Client as PostHogClient from posthog import setup @@ -34,6 +47,13 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): _wrap_openai_resources(self, _SYNC_RESOURCE_WRAPPERS) + # Keep dynamically installed resources visible to API and type inspection. + if _TYPE_CHECKING: + self.chat = WrappedChat(self, self._original_chat) + self.embeddings = WrappedEmbeddings(self, self._original_embeddings) + self.beta = WrappedBeta(self, self._original_beta) + self.responses = WrappedResponses(self, self._original_responses) + class AsyncAzureOpenAI(openai.AsyncAzureOpenAI): """ @@ -54,3 +74,10 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): self._ph_client = posthog_client or setup() _wrap_openai_resources(self, _ASYNC_RESOURCE_WRAPPERS) + + # Keep dynamically installed resources visible to API and type inspection. + if _TYPE_CHECKING: + self.chat = AsyncWrappedChat(self, self._original_chat) + self.embeddings = AsyncWrappedEmbeddings(self, self._original_embeddings) + self.beta = AsyncWrappedBeta(self, self._original_beta) + self.responses = AsyncWrappedResponses(self, self._original_responses) From b910864905e691515a65e9974cd1fc99a78b2b55 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Tue, 25 Aug 2026 09:26:41 +0200 Subject: [PATCH 5/6] test(openai): use resource wrapper mappings --- .../test/ai/openai/test_resource_wrapping.py | 55 +++++++------------ 1 file changed, 21 insertions(+), 34 deletions(-) diff --git a/posthog/test/ai/openai/test_resource_wrapping.py b/posthog/test/ai/openai/test_resource_wrapping.py index 3e1117c46..02bf7b094 100644 --- a/posthog/test/ai/openai/test_resource_wrapping.py +++ b/posthog/test/ai/openai/test_resource_wrapping.py @@ -7,37 +7,9 @@ from openai.resources.responses import AsyncResponses, Responses from posthog.ai.openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI -from posthog.ai.openai.openai import ( - WrappedBeta, - WrappedChat, - WrappedEmbeddings, - WrappedResponses, -) -from posthog.ai.openai.openai_async import ( - WrappedBeta as AsyncWrappedBeta, -) -from posthog.ai.openai.openai_async import ( - WrappedChat as AsyncWrappedChat, -) -from posthog.ai.openai.openai_async import ( - WrappedEmbeddings as AsyncWrappedEmbeddings, -) -from posthog.ai.openai.openai_async import ( - WrappedResponses as AsyncWrappedResponses, -) +from posthog.ai.openai.openai import _SYNC_RESOURCE_WRAPPERS +from posthog.ai.openai.openai_async import _ASYNC_RESOURCE_WRAPPERS -_SYNC_WRAPPERS = { - "chat": WrappedChat, - "embeddings": WrappedEmbeddings, - "beta": WrappedBeta, - "responses": WrappedResponses, -} -_ASYNC_WRAPPERS = { - "chat": AsyncWrappedChat, - "embeddings": AsyncWrappedEmbeddings, - "beta": AsyncWrappedBeta, - "responses": AsyncWrappedResponses, -} _SYNC_RESOURCES = { "chat": Chat, "embeddings": Embeddings, @@ -60,10 +32,25 @@ @pytest.mark.parametrize( "client_type, client_kwargs, wrappers, resource_types", [ - (OpenAI, {"api_key": "test-key"}, _SYNC_WRAPPERS, _SYNC_RESOURCES), - (AsyncOpenAI, {"api_key": "test-key"}, _ASYNC_WRAPPERS, _ASYNC_RESOURCES), - (AzureOpenAI, _AZURE_KWARGS, _SYNC_WRAPPERS, _SYNC_RESOURCES), - (AsyncAzureOpenAI, _AZURE_KWARGS, _ASYNC_WRAPPERS, _ASYNC_RESOURCES), + ( + OpenAI, + {"api_key": "test-key"}, + _SYNC_RESOURCE_WRAPPERS, + _SYNC_RESOURCES, + ), + ( + AsyncOpenAI, + {"api_key": "test-key"}, + _ASYNC_RESOURCE_WRAPPERS, + _ASYNC_RESOURCES, + ), + (AzureOpenAI, _AZURE_KWARGS, _SYNC_RESOURCE_WRAPPERS, _SYNC_RESOURCES), + ( + AsyncAzureOpenAI, + _AZURE_KWARGS, + _ASYNC_RESOURCE_WRAPPERS, + _ASYNC_RESOURCES, + ), ], ) def test_client_resources_are_discovered_and_wrapped( From 425fec91b088cdb7daa285bda1bf61db9d114d8a Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Tue, 25 Aug 2026 09:47:36 +0200 Subject: [PATCH 6/6] refactor(openai): declare wrapped resource types --- posthog/ai/openai/openai.py | 13 +++++------ posthog/ai/openai/openai_async.py | 13 +++++------ posthog/ai/openai/openai_providers.py | 26 ++++++++++------------ references/public_api_snapshot.txt | 32 +++++++++++++-------------- 4 files changed, 40 insertions(+), 44 deletions(-) diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index a9bbb7b0e..9a046831c 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -48,6 +48,12 @@ class OpenAI(openai.OpenAI): _ph_client: PostHogClient + if _TYPE_CHECKING: + chat: "WrappedChat" + embeddings: "WrappedEmbeddings" + beta: "WrappedBeta" + responses: "WrappedResponses" + def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): """ Args: @@ -62,13 +68,6 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): _wrap_openai_resources(self, _SYNC_RESOURCE_WRAPPERS) - # Keep dynamically installed resources visible to API and type inspection. - if _TYPE_CHECKING: - self.chat = WrappedChat(self, self._original_chat) - self.embeddings = WrappedEmbeddings(self, self._original_embeddings) - self.beta = WrappedBeta(self, self._original_beta) - self.responses = WrappedResponses(self, self._original_responses) - def _parse_and_track( wrapper, diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 0cb41a039..e7a955cc1 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -50,6 +50,12 @@ class AsyncOpenAI(openai.AsyncOpenAI): _ph_client: PostHogClient + if _TYPE_CHECKING: + chat: "WrappedChat" + embeddings: "WrappedEmbeddings" + beta: "WrappedBeta" + responses: "WrappedResponses" + def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): """ Args: @@ -64,13 +70,6 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): _wrap_openai_resources(self, _ASYNC_RESOURCE_WRAPPERS) - # Keep dynamically installed resources visible to API and type inspection. - if _TYPE_CHECKING: - self.chat = WrappedChat(self, self._original_chat) - self.embeddings = WrappedEmbeddings(self, self._original_embeddings) - self.beta = WrappedBeta(self, self._original_beta) - self.responses = WrappedResponses(self, self._original_responses) - async def _parse_and_track( wrapper, diff --git a/posthog/ai/openai/openai_providers.py b/posthog/ai/openai/openai_providers.py index 08d8c8b35..7c0ce04b6 100644 --- a/posthog/ai/openai/openai_providers.py +++ b/posthog/ai/openai/openai_providers.py @@ -34,6 +34,12 @@ class AzureOpenAI(openai.AzureOpenAI): _ph_client: PostHogClient + if _TYPE_CHECKING: + chat: "WrappedChat" + embeddings: "WrappedEmbeddings" + beta: "WrappedBeta" + responses: "WrappedResponses" + def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): """ Args: @@ -47,13 +53,6 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): _wrap_openai_resources(self, _SYNC_RESOURCE_WRAPPERS) - # Keep dynamically installed resources visible to API and type inspection. - if _TYPE_CHECKING: - self.chat = WrappedChat(self, self._original_chat) - self.embeddings = WrappedEmbeddings(self, self._original_embeddings) - self.beta = WrappedBeta(self, self._original_beta) - self.responses = WrappedResponses(self, self._original_responses) - class AsyncAzureOpenAI(openai.AsyncAzureOpenAI): """ @@ -62,6 +61,12 @@ class AsyncAzureOpenAI(openai.AsyncAzureOpenAI): _ph_client: PostHogClient + if _TYPE_CHECKING: + chat: "AsyncWrappedChat" + embeddings: "AsyncWrappedEmbeddings" + beta: "AsyncWrappedBeta" + responses: "AsyncWrappedResponses" + def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): """ Args: @@ -74,10 +79,3 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): self._ph_client = posthog_client or setup() _wrap_openai_resources(self, _ASYNC_RESOURCE_WRAPPERS) - - # Keep dynamically installed resources visible to API and type inspection. - if _TYPE_CHECKING: - self.chat = AsyncWrappedChat(self, self._original_chat) - self.embeddings = AsyncWrappedEmbeddings(self, self._original_embeddings) - self.beta = AsyncWrappedBeta(self, self._original_beta) - self.responses = AsyncWrappedResponses(self, self._original_responses) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 54ca98d0d..3a414ae9f 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -403,28 +403,28 @@ attribute posthog.ai.langchain.callbacks.SpanMetadata.latency: float attribute posthog.ai.langchain.callbacks.SpanMetadata.name: str attribute posthog.ai.langchain.callbacks.SpanMetadata.start_time: float attribute posthog.ai.langchain.callbacks.log = logging.getLogger('posthog') -attribute posthog.ai.openai.openai.OpenAI.beta = WrappedBeta(self, self._original_beta) -attribute posthog.ai.openai.openai.OpenAI.chat = WrappedChat(self, self._original_chat) -attribute posthog.ai.openai.openai.OpenAI.embeddings = WrappedEmbeddings(self, self._original_embeddings) -attribute posthog.ai.openai.openai.OpenAI.responses = WrappedResponses(self, self._original_responses) +attribute posthog.ai.openai.openai.OpenAI.beta: WrappedBeta +attribute posthog.ai.openai.openai.OpenAI.chat: WrappedChat +attribute posthog.ai.openai.openai.OpenAI.embeddings: WrappedEmbeddings +attribute posthog.ai.openai.openai.OpenAI.responses: WrappedResponses attribute posthog.ai.openai.openai.WrappedBeta.chat attribute posthog.ai.openai.openai.WrappedBetaChat.completions attribute posthog.ai.openai.openai.WrappedChat.completions -attribute posthog.ai.openai.openai_async.AsyncOpenAI.beta = WrappedBeta(self, self._original_beta) -attribute posthog.ai.openai.openai_async.AsyncOpenAI.chat = WrappedChat(self, self._original_chat) -attribute posthog.ai.openai.openai_async.AsyncOpenAI.embeddings = WrappedEmbeddings(self, self._original_embeddings) -attribute posthog.ai.openai.openai_async.AsyncOpenAI.responses = WrappedResponses(self, self._original_responses) +attribute posthog.ai.openai.openai_async.AsyncOpenAI.beta: WrappedBeta +attribute posthog.ai.openai.openai_async.AsyncOpenAI.chat: WrappedChat +attribute posthog.ai.openai.openai_async.AsyncOpenAI.embeddings: WrappedEmbeddings +attribute posthog.ai.openai.openai_async.AsyncOpenAI.responses: WrappedResponses attribute posthog.ai.openai.openai_async.WrappedBeta.chat attribute posthog.ai.openai.openai_async.WrappedBetaChat.completions attribute posthog.ai.openai.openai_async.WrappedChat.completions -attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.beta = AsyncWrappedBeta(self, self._original_beta) -attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.chat = AsyncWrappedChat(self, self._original_chat) -attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.embeddings = AsyncWrappedEmbeddings(self, self._original_embeddings) -attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.responses = AsyncWrappedResponses(self, self._original_responses) -attribute posthog.ai.openai.openai_providers.AzureOpenAI.beta = WrappedBeta(self, self._original_beta) -attribute posthog.ai.openai.openai_providers.AzureOpenAI.chat = WrappedChat(self, self._original_chat) -attribute posthog.ai.openai.openai_providers.AzureOpenAI.embeddings = WrappedEmbeddings(self, self._original_embeddings) -attribute posthog.ai.openai.openai_providers.AzureOpenAI.responses = WrappedResponses(self, self._original_responses) +attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.beta: AsyncWrappedBeta +attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.chat: AsyncWrappedChat +attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.embeddings: AsyncWrappedEmbeddings +attribute posthog.ai.openai.openai_providers.AsyncAzureOpenAI.responses: AsyncWrappedResponses +attribute posthog.ai.openai.openai_providers.AzureOpenAI.beta: WrappedBeta +attribute posthog.ai.openai.openai_providers.AzureOpenAI.chat: WrappedChat +attribute posthog.ai.openai.openai_providers.AzureOpenAI.embeddings: WrappedEmbeddings +attribute posthog.ai.openai.openai_providers.AzureOpenAI.responses: WrappedResponses attribute posthog.ai.openai.wrapper_utils.log = logging.getLogger('posthog') attribute posthog.ai.openai_agents.processor.log = logging.getLogger('posthog') attribute posthog.ai.otel.spans.AI_SPAN_PREFIXES = ('gen_ai.', 'llm.', 'ai.', 'traceloop.')