From 6f358c0974b067408e5918c8cfb0d80b1b9faf2b Mon Sep 17 00:00:00 2001 From: mo Date: Wed, 14 Oct 2026 20:23:01 +0100 Subject: [PATCH] fix(ai): keep the openai.Stream API on sync OpenAI streams Sync chat completions and Responses streams returned a bare generator, so `with`, `close()` and attributes such as `.response` were lost. Wrap them with _StreamWrapper, as #804 did for Anthropic and #645 did for async. --- .sampo/changesets/openai-sync-stream-api.md | 5 + posthog/ai/openai/openai.py | 5 +- posthog/ai/stream.py | 4 +- posthog/test/ai/openai/test_openai.py | 119 +++++++++++++++++++- posthog/test/ai/utils.py | 23 ++++ 5 files changed, 152 insertions(+), 4 deletions(-) create mode 100644 .sampo/changesets/openai-sync-stream-api.md diff --git a/.sampo/changesets/openai-sync-stream-api.md b/.sampo/changesets/openai-sync-stream-api.md new file mode 100644 index 000000000..8d70d9284 --- /dev/null +++ b/.sampo/changesets/openai-sync-stream-api.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Keep the `openai.Stream` API on sync OpenAI chat completions and Responses streams, so `with`, `close()` and attributes such as `.response` work as they do without PostHog. diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index c1950939c..ea00a197b 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -33,6 +33,7 @@ _ResponsesStreamState, _build_streaming_event_data, ) +from ..stream import _StreamWrapper from ._embeddings import _capture_embedding_event from .wrapper_utils import ( _OpenAIWrapperResource, @@ -185,7 +186,7 @@ def generator(): time.time() - start_time, ) - return generator() + return _StreamWrapper(generator(), stream=response) def _capture_streaming_event( self, @@ -403,7 +404,7 @@ def generator(): time.time() - start_time, ) - return generator() + return _StreamWrapper(generator(), stream=response) def _capture_streaming_event( self, diff --git a/posthog/ai/stream.py b/posthog/ai/stream.py index 2499f6e68..8744c6177 100644 --- a/posthog/ai/stream.py +++ b/posthog/ai/stream.py @@ -32,7 +32,9 @@ def close(self) -> None: try: self._generator.close() finally: - self._stream.close() + close = getattr(self._stream, "close", None) + if close is not None: + close() def __getattr__(self, name: str) -> Any: if name.startswith("_"): diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 0d22f393a..83304a4b1 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -36,7 +36,11 @@ from posthog.ai.openai import OpenAI from posthog.ai.openai.openai_async import AsyncOpenAI from posthog.ai.openai.wrapper_utils import reset_fallback_warnings - from posthog.test.ai.utils import RecordingAsyncStream, make_response_usage + from posthog.test.ai.utils import ( + RecordingAsyncStream, + RecordingStream, + make_response_usage, + ) OPENAI_AVAILABLE = True except ImportError: @@ -2403,6 +2407,119 @@ async def mock_create(self, **kwargs): assert mock_client.capture.call_count == 1 +def test_chat_streaming_supports_with_and_proxies_provider_stream( + mock_client, streaming_tool_call_chunks +): + """Sync chat completions stream=True must keep the `openai.Stream` API: + `with`, and attributes such as `.response`.""" + source = RecordingStream(streaming_tool_call_chunks) + + with patch( + "openai.resources.chat.completions.Completions.create", return_value=source + ): + client = OpenAI(api_key="test-key", posthog_client=mock_client) + + response = client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + posthog_distinct_id="test-id", + ) + + assert response.response == "provider-response" + with response as stream: + chunks = list(stream) + + assert chunks == streaming_tool_call_chunks + assert source.closed is True + assert mock_client.capture.call_count == 1 + call_args = mock_client.capture.call_args[1] + props = call_args["properties"] + assert call_args["event"] == "$ai_generation" + assert props["$ai_provider"] == "openai" + assert props["$ai_model"] == "gpt-4" + + +def test_responses_streaming_supports_with_and_proxies_provider_stream(mock_client): + """Sync responses stream=True must keep the `openai.Stream` API.""" + from unittest.mock import MagicMock + + chunk = MagicMock() + chunk.type = "response.text.delta" + chunk.text = "hello" + source = RecordingStream([chunk]) + + with patch("openai.resources.responses.Responses.create", return_value=source): + client = OpenAI(api_key="test-key", posthog_client=mock_client) + + response = client.responses.create( + model="gpt-4o-mini", + input=[{"role": "user", "content": "Hi"}], + stream=True, + posthog_distinct_id="test-id", + ) + + assert response.response == "provider-response" + with response as stream: + received = list(stream) + + assert received == [chunk] + assert source.closed is True + assert mock_client.capture.call_count == 1 + + +def test_chat_streaming_early_exit_closes_provider_stream( + mock_client, streaming_tool_call_chunks +): + """Breaking out of a sync stream early must close the underlying provider + stream (release the HTTP connection) and still capture the event.""" + source = RecordingStream(streaming_tool_call_chunks) + + with patch( + "openai.resources.chat.completions.Completions.create", return_value=source + ): + client = OpenAI(api_key="test-key", posthog_client=mock_client) + + response = client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + posthog_distinct_id="test-id", + ) + + with response as stream: + for _ in stream: + break + + assert source.closed is True + assert mock_client.capture.call_count == 1 + + +def test_chat_streaming_close_closes_provider_stream( + mock_client, streaming_tool_call_chunks +): + """Calling `close()` mid-stream, outside a `with` block, must close the + provider stream and still capture the event.""" + source = RecordingStream(streaming_tool_call_chunks) + + with patch( + "openai.resources.chat.completions.Completions.create", return_value=source + ): + client = OpenAI(api_key="test-key", posthog_client=mock_client) + + response = client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + posthog_distinct_id="test-id", + ) + next(response) + response.close() + + assert source.closed is True + assert mock_client.capture.call_count == 1 + + def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_openai_response): mock_client.enable_full_ai_capture = True with patch( diff --git a/posthog/test/ai/utils.py b/posthog/test/ai/utils.py index e1ad48f7b..9f27c1f38 100644 --- a/posthog/test/ai/utils.py +++ b/posthog/test/ai/utils.py @@ -41,6 +41,29 @@ def build(model_cls, known): ) +class RecordingStream: + """Mock provider sync stream that is iterable and records when closed. + + Mirrors the real ``openai.Stream``: it supports ``for`` and exposes + ``close()`` plus a ``response`` attribute, so tests can assert both + iteration and that the underlying stream is closed on context exit. + """ + + def __init__(self, items): + self._items = iter(items) + self.closed = False + self.response = "provider-response" + + def __iter__(self): + return self + + def __next__(self): + return next(self._items) + + def close(self): + self.closed = True + + class RecordingAsyncStream: """Mock provider async stream that is iterable and records when closed.