fix(ai): keep the openai.Stream API on sync OpenAI streams

.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.

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,

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("_"):

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(

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.