From 3e729e93810bf90c9f7b94849daa2d9c48850b00 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 22 May 2026 12:40:26 +0200 Subject: [PATCH 1/2] fix openai chat completions parse tracking --- .sampo/changesets/wily-shaman-lempo.md | 5 ++ posthog/ai/openai/openai.py | 78 +++++++++++++++++++--- posthog/ai/openai/openai_async.py | 83 ++++++++++++++++++----- posthog/ai/openai/wrapper_utils.py | 19 ++++++ posthog/test/ai/openai/test_openai.py | 92 ++++++++++++++++++++++++++ 5 files changed, 252 insertions(+), 25 deletions(-) create mode 100644 .sampo/changesets/wily-shaman-lempo.md create mode 100644 posthog/ai/openai/wrapper_utils.py diff --git a/.sampo/changesets/wily-shaman-lempo.md b/.sampo/changesets/wily-shaman-lempo.md new file mode 100644 index 00000000..5d5e858d --- /dev/null +++ b/.sampo/changesets/wily-shaman-lempo.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Track OpenAI chat completions parse calls diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 90fd5399..d3fcf39c 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -26,6 +26,7 @@ from posthog.ai.sanitization import sanitize_openai, sanitize_openai_response from posthog.client import Client as PostHogClient from posthog import setup +from posthog.ai.openai.wrapper_utils import warn_on_fallback class OpenAI(openai.OpenAI): @@ -67,6 +68,29 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): self.responses = WrappedResponses(self, self._original_responses) +def _parse_and_track( + wrapper, + posthog_distinct_id: Optional[str], + posthog_trace_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + **kwargs: Any, +): + return call_llm_and_track_usage( + posthog_distinct_id, + wrapper._client._ph_client, + "openai", + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + wrapper._client.base_url, + wrapper._original.parse, + **kwargs, + ) + + class WrappedResponses: """Wrapper for OpenAI responses that tracks usage in PostHog.""" @@ -76,6 +100,7 @@ def __init__(self, client: OpenAI, original_responses): def __getattr__(self, name): """Fallback to original responses object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) def create( @@ -276,16 +301,13 @@ def parse( Returns: The response from OpenAI's responses.parse call. """ - return call_llm_and_track_usage( + return _parse_and_track( + self, posthog_distinct_id, - self._client._ph_client, - "openai", posthog_trace_id, posthog_properties, posthog_privacy_mode, posthog_groups, - self._client.base_url, - self._original.parse, **kwargs, ) @@ -299,6 +321,7 @@ def __init__(self, client: OpenAI, original_chat): def __getattr__(self, name): """Fallback to original chat object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) @property @@ -316,8 +339,42 @@ def __init__(self, client: OpenAI, original_completions): def __getattr__(self, name): """Fallback to original completions object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) + def parse( + self, + posthog_distinct_id: Optional[str] = None, + posthog_trace_id: Optional[str] = None, + posthog_properties: Optional[Dict[str, Any]] = None, + posthog_privacy_mode: bool = False, + posthog_groups: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + """ + Parse an OpenAI chat completion while tracking usage in PostHog. + + Args: + posthog_distinct_id: Optional distinct ID to associate with the usage event. + posthog_trace_id: Optional trace ID. Generated automatically when omitted. + posthog_properties: Additional properties to include with the usage event. + posthog_privacy_mode: Whether to redact captured input and output. + posthog_groups: Optional PostHog groups to associate with the event. + **kwargs: Arguments passed to OpenAI's ``chat.completions.parse`` API. + + Returns: + The parsed response from OpenAI. + """ + return _parse_and_track( + self, + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + **kwargs, + ) + def create( self, posthog_distinct_id: Optional[str] = None, @@ -518,6 +575,7 @@ def __init__(self, client: OpenAI, original_embeddings): def __getattr__(self, name): """Fallback to original embeddings object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) def create( @@ -602,6 +660,7 @@ def __init__(self, client: OpenAI, original_beta): def __getattr__(self, name): """Fallback to original beta object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) @property @@ -619,6 +678,7 @@ def __init__(self, client: OpenAI, original_beta_chat): def __getattr__(self, name): """Fallback to original beta chat object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) @property @@ -636,6 +696,7 @@ def __init__(self, client: OpenAI, original_beta_completions): def __getattr__(self, name): """Fallback to original beta completions object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) def parse( @@ -661,15 +722,12 @@ def parse( Returns: The parsed response from OpenAI. """ - return call_llm_and_track_usage( + return _parse_and_track( + self, posthog_distinct_id, - self._client._ph_client, - "openai", posthog_trace_id, posthog_properties, posthog_privacy_mode, posthog_groups, - self._client.base_url, - self._original.parse, **kwargs, ) diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index cb25e138..dce2a584 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -28,6 +28,7 @@ ) from posthog.ai.sanitization import sanitize_openai, sanitize_openai_response from posthog.client import Client as PostHogClient +from posthog.ai.openai.wrapper_utils import warn_on_fallback class AsyncOpenAI(openai.AsyncOpenAI): @@ -69,6 +70,29 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs): self.responses = WrappedResponses(self, self._original_responses) +async def _parse_and_track( + wrapper, + posthog_distinct_id: Optional[str], + posthog_trace_id: Optional[str], + posthog_properties: Optional[Dict[str, Any]], + posthog_privacy_mode: bool, + posthog_groups: Optional[Dict[str, Any]], + **kwargs: Any, +): + return await call_llm_and_track_usage_async( + posthog_distinct_id, + wrapper._client._ph_client, + "openai", + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + wrapper._client.base_url, + wrapper._original.parse, + **kwargs, + ) + + class WrappedResponses: """Async wrapper for OpenAI responses that tracks usage in PostHog.""" @@ -78,7 +102,7 @@ def __init__(self, client: AsyncOpenAI, original_responses): def __getattr__(self, name): """Fallback to original responses object for any methods we don't explicitly handle.""" - + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) async def create( @@ -305,16 +329,13 @@ async def parse( Returns: The response from OpenAI's responses.parse call. """ - return await call_llm_and_track_usage_async( + return await _parse_and_track( + self, posthog_distinct_id, - self._client._ph_client, - "openai", posthog_trace_id, posthog_properties, posthog_privacy_mode, posthog_groups, - self._client.base_url, - self._original.parse, **kwargs, ) @@ -328,6 +349,7 @@ def __init__(self, client: AsyncOpenAI, original_chat): def __getattr__(self, name): """Fallback to original chat object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) @property @@ -345,8 +367,42 @@ def __init__(self, client: AsyncOpenAI, original_completions): def __getattr__(self, name): """Fallback to original completions object for any methods we don't explicitly handle.""" + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) + async def parse( + self, + posthog_distinct_id: Optional[str] = None, + posthog_trace_id: Optional[str] = None, + posthog_properties: Optional[Dict[str, Any]] = None, + posthog_privacy_mode: bool = False, + posthog_groups: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + """ + Parse an OpenAI chat completion while tracking usage in PostHog. + + Args: + posthog_distinct_id: Optional distinct ID to associate with the usage event. + posthog_trace_id: Optional trace ID. Generated automatically when omitted. + posthog_properties: Additional properties to include with the usage event. + posthog_privacy_mode: Whether to redact captured input and output. + posthog_groups: Optional PostHog groups to associate with the event. + **kwargs: Arguments passed to OpenAI's async ``chat.completions.parse`` API. + + Returns: + The parsed response from OpenAI. + """ + return await _parse_and_track( + self, + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + **kwargs, + ) + async def create( self, posthog_distinct_id: Optional[str] = None, @@ -574,7 +630,7 @@ def __init__(self, client: AsyncOpenAI, original_embeddings): def __getattr__(self, name): """Fallback to original embeddings object for any methods we don't explicitly handle.""" - + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) async def create( @@ -660,7 +716,7 @@ def __init__(self, client: AsyncOpenAI, original_beta): def __getattr__(self, name): """Fallback to original beta object for any methods we don't explicitly handle.""" - + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) @property @@ -678,7 +734,7 @@ def __init__(self, client: AsyncOpenAI, original_beta_chat): def __getattr__(self, name): """Fallback to original beta chat object for any methods we don't explicitly handle.""" - + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) @property @@ -696,7 +752,7 @@ def __init__(self, client: AsyncOpenAI, original_beta_completions): def __getattr__(self, name): """Fallback to original beta completions object for any methods we don't explicitly handle.""" - + warn_on_fallback(self.__class__.__name__, name) return getattr(self._original, name) async def parse( @@ -722,15 +778,12 @@ async def parse( Returns: The parsed response from OpenAI. """ - return await call_llm_and_track_usage_async( + return await _parse_and_track( + self, posthog_distinct_id, - self._client._ph_client, - "openai", posthog_trace_id, posthog_properties, posthog_privacy_mode, posthog_groups, - self._client.base_url, - self._original.parse, **kwargs, ) diff --git a/posthog/ai/openai/wrapper_utils.py b/posthog/ai/openai/wrapper_utils.py new file mode 100644 index 00000000..a62d2abc --- /dev/null +++ b/posthog/ai/openai/wrapper_utils.py @@ -0,0 +1,19 @@ +import logging + + +log = logging.getLogger("posthog") +_fallback_warnings: set[tuple[str, str]] = set() + + +def warn_on_fallback(wrapper_name: str, name: str) -> None: + key = (wrapper_name, name) + if key in _fallback_warnings: + return + + _fallback_warnings.add(key) + log.warning( + "Falling back to unwrapped OpenAI API for %s.%s; PostHog LLM tracking " + "and posthog_* arguments will not be applied.", + wrapper_name, + name, + ) diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 2f35329c..192f4f67 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -1,4 +1,5 @@ import json +import logging import os import time from unittest.mock import AsyncMock, patch @@ -1112,6 +1113,97 @@ def test_responses_parse(mock_client, mock_parsed_response): assert isinstance(props["$ai_latency"], float) +def test_chat_completions_parse(mock_client, mock_openai_response): + with patch( + "openai.resources.chat.completions.Completions.parse", + return_value=mock_openai_response, + ) as mock_parse: + client = OpenAI(api_key="test-key", posthog_client=mock_client) + response = client.chat.completions.parse( + model="gpt-4", + messages=[{"role": "user", "content": "Hello"}], + response_format={"type": "json_object"}, + posthog_distinct_id="test-id", + posthog_properties={"foo": "bar"}, + ) + + assert response == mock_openai_response + assert mock_parse.call_count == 1 + assert "posthog_distinct_id" not in mock_parse.call_args.kwargs + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + + assert call_args["distinct_id"] == "test-id" + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "openai" + assert props["$ai_model"] == "gpt-4" + assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] + assert props["$ai_output_choices"] == [ + { + "role": "assistant", + "content": [{"type": "text", "text": "Test response"}], + } + ] + assert props["$ai_input_tokens"] == 20 + assert props["$ai_output_tokens"] == 10 + assert props["foo"] == "bar" + assert isinstance(props["$ai_latency"], float) + + +@pytest.mark.asyncio +async def test_async_chat_completions_parse(mock_client, mock_openai_response): + mock_parse = AsyncMock(return_value=mock_openai_response) + + with patch( + "openai.resources.chat.completions.AsyncCompletions.parse", new=mock_parse + ): + client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) + response = await client.chat.completions.parse( + model="gpt-4", + messages=[{"role": "user", "content": "Hello"}], + response_format={"type": "json_object"}, + posthog_distinct_id="test-id", + posthog_properties={"foo": "bar"}, + ) + + assert response == mock_openai_response + mock_parse.assert_awaited_once() + assert "posthog_distinct_id" not in mock_parse.call_args.kwargs + assert mock_client.capture.call_count == 1 + + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + + assert call_args["distinct_id"] == "test-id" + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "openai" + assert props["$ai_model"] == "gpt-4" + assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] + assert props["$ai_output_choices"] == [ + { + "role": "assistant", + "content": [{"type": "text", "text": "Test response"}], + } + ] + assert props["$ai_input_tokens"] == 20 + assert props["$ai_output_tokens"] == 10 + assert props["foo"] == "bar" + assert isinstance(props["$ai_latency"], float) + + +def test_fallback_logs_warning(mock_client, caplog): + client = OpenAI(api_key="test-key", posthog_client=mock_client) + + with caplog.at_level(logging.WARNING, logger="posthog"): + with pytest.raises(AttributeError): + client.chat.posthog_unwrapped_test_attribute + + assert "Falling back to unwrapped OpenAI API" in caplog.text + assert "WrappedChat.posthog_unwrapped_test_attribute" in caplog.text + + def test_responses_api_streaming_with_tokens(mock_client): """Test that Responses API streaming properly captures token usage from response.usage.""" from openai.types.responses import ResponseUsage From 6e43d3a61cc29286d156988833567fa9a8c767a1 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Fri, 22 May 2026 12:49:59 +0200 Subject: [PATCH 2/2] address openai wrapper review feedback --- posthog/ai/openai/wrapper_utils.py | 4 +++ posthog/test/ai/openai/test_openai.py | 41 ++++++++++++++++++++++++--- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/posthog/ai/openai/wrapper_utils.py b/posthog/ai/openai/wrapper_utils.py index a62d2abc..886f15ec 100644 --- a/posthog/ai/openai/wrapper_utils.py +++ b/posthog/ai/openai/wrapper_utils.py @@ -5,6 +5,10 @@ _fallback_warnings: set[tuple[str, str]] = set() +def reset_fallback_warnings() -> None: + _fallback_warnings.clear() + + def warn_on_fallback(wrapper_name: str, name: str) -> None: key = (wrapper_name, name) if key in _fallback_warnings: diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 192f4f67..7049fb40 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -38,6 +38,7 @@ from posthog.ai.openai import OpenAI from posthog.ai.openai.openai_async import AsyncOpenAI + from posthog.ai.openai.wrapper_utils import reset_fallback_warnings OPENAI_AVAILABLE = True except ImportError: @@ -1193,15 +1194,47 @@ async def test_async_chat_completions_parse(mock_client, mock_openai_response): assert isinstance(props["$ai_latency"], float) -def test_fallback_logs_warning(mock_client, caplog): - client = OpenAI(api_key="test-key", posthog_client=mock_client) +@pytest.mark.parametrize( + "client_factory, wrapper_accessor, wrapper_name", + [ + (OpenAI, lambda client: client.responses, "WrappedResponses"), + (OpenAI, lambda client: client.chat, "WrappedChat"), + (OpenAI, lambda client: client.chat.completions, "WrappedCompletions"), + (OpenAI, lambda client: client.embeddings, "WrappedEmbeddings"), + (OpenAI, lambda client: client.beta, "WrappedBeta"), + (OpenAI, lambda client: client.beta.chat, "WrappedBetaChat"), + ( + OpenAI, + lambda client: client.beta.chat.completions, + "WrappedBetaCompletions", + ), + (AsyncOpenAI, lambda client: client.responses, "WrappedResponses"), + (AsyncOpenAI, lambda client: client.chat, "WrappedChat"), + (AsyncOpenAI, lambda client: client.chat.completions, "WrappedCompletions"), + (AsyncOpenAI, lambda client: client.embeddings, "WrappedEmbeddings"), + (AsyncOpenAI, lambda client: client.beta, "WrappedBeta"), + (AsyncOpenAI, lambda client: client.beta.chat, "WrappedBetaChat"), + ( + AsyncOpenAI, + lambda client: client.beta.chat.completions, + "WrappedBetaCompletions", + ), + ], +) +def test_fallback_logs_warning( + mock_client, caplog, client_factory, wrapper_accessor, wrapper_name +): + reset_fallback_warnings() + caplog.clear() + client = client_factory(api_key="test-key", posthog_client=mock_client) + wrapper = wrapper_accessor(client) with caplog.at_level(logging.WARNING, logger="posthog"): with pytest.raises(AttributeError): - client.chat.posthog_unwrapped_test_attribute + wrapper.posthog_unwrapped_test_attribute assert "Falling back to unwrapped OpenAI API" in caplog.text - assert "WrappedChat.posthog_unwrapped_test_attribute" in caplog.text + assert f"{wrapper_name}.posthog_unwrapped_test_attribute" in caplog.text def test_responses_api_streaming_with_tokens(mock_client):