Skip to content

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

Open
mosafariuk wants to merge 1 commit into
PostHog:mainfrom
mosafariuk:fix/openai-sync-stream-api
Open

mosafariuk wants to merge 1 commit into
PostHog:mainfrom
mosafariuk:fix/openai-sync-stream-api

Conversation

@mosafariuk

Copy link
Copy Markdown

💡 Motivation and Context

With stream=True, the sync posthog.ai.openai.OpenAI wrapper returns a bare generator from chat.completions.create and responses.create, not something that behaves like openai.Stream. So the wrapper is no longer a drop-in replacement for streaming:

  • with client.chat.completions.create(..., stream=True) as stream: raises AttributeError: __enter__ (LangChain's ChatOpenAI streams this way)
  • stream.response (HTTP status and headers) is gone
  • stream.close() only closes the generator. The HTTP connection stays open until it's garbage-collected

This was already fixed for async streams in #645 (AsyncStreamWrapper) and for sync Anthropic in #804 (_StreamWrapper). Sync OpenAI chat completions and Responses were the remaining paths still returning generator(). This PR wraps them with the same _StreamWrapper, so with, close() and attributes such as .response reach the provider stream, and the PostHog event is still captured once.

It also makes _StreamWrapper.close() skip a provider stream that has no close(), the same check AsyncStreamWrapper already does. openai.Stream always has one, but without the check the existing test_sync_stream_close_after_early_exit_captures_partial_state test, which mocks the stream as iter([...]), would fail. With the check it passes unchanged.

Sync Gemini is unchanged on purpose: google-genai's own generate_content_stream() returns a plain generator, so returning a generator there already matches the provider API.

💚 How did you test it?

  • New regression tests in posthog/test/ai/openai/test_openai.py, modelled on the fix(ai): support async with on async streaming responses (Fixes #393) #645 async tests, using a new sync RecordingStream helper next to RecordingAsyncStream:
    • chat and Responses streams support with, expose .response, close the provider stream on exit, and capture exactly one $ai_generation event
    • breaking out early inside with closes the provider stream and still captures the event
    • calling close() mid-stream outside a with block closes the provider stream and still captures the event
    • all four fail on main and pass with this change
  • Full suite: 4611 passed, 16 skipped. ruff format --check, ruff check, mypy | mypy-baseline filter and make public_api_check are clean (no public API change, since _StreamWrapper is internal).
  • End to end with a real openai.Stream: a script points the real OpenAI client at an httpx.MockTransport that serves SSE chunks. I compared plain openai.OpenAI with the PostHog wrapper and checked that with and .response work, the HTTP response is_closed after a full with iteration and after an early close(), and one event is captured in each case. I ran it on openai 2.54 (the repo's lockfile) and reproduced the original failure on 3.31.

📝 Checklist

  • I reviewed the submitted code.
  • I added tests to verify the changes.
  • I updated the docs if needed.
  • No breaking change or entry added to the changelog.

If releasing new changes

  • Ran sampo add to generate a changeset file

🤖 Agent context

Autonomy: Human-driven (agent-assisted)

  • An AI coding agent helped with the investigation and the first pass. It diffed the sync and async wrappers to find the paths fix(ai): preserve Anthropic messages stream API #804 missed, then wrote the patch, tests and repro script. I reviewed and ran all of it.
  • I left sync Gemini alone because its native stream is a plain generator. I added the close() check instead of changing the existing parity test, so current behaviour stays as it is.

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 PostHog#804 did for Anthropic and PostHog#645 did for async.
@mosafariuk
mosafariuk requested a review from a team as a code owner October 14, 2026 19:44

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant