From 7aaa93c9203f85e1426704bc62037395cab957f9 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Thu, 29 Oct 2026 06:32:57 -0400 Subject: [PATCH 1/2] fix(ai): capture an error event when a streaming call fails to connect Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Alex Castillo --- .sampo/changesets/stream-connect-error.md | 5 + posthog/ai/anthropic/anthropic.py | 19 +- posthog/ai/anthropic/anthropic_async.py | 19 +- posthog/ai/gemini/gemini.py | 22 ++- posthog/ai/gemini/gemini_async.py | 24 ++- posthog/ai/openai/openai.py | 32 +++- posthog/ai/openai/openai_async.py | 32 +++- .../test/ai/test_stream_midstream_error.py | 162 ++++++++++++++++++ 8 files changed, 305 insertions(+), 10 deletions(-) create mode 100644 .sampo/changesets/stream-connect-error.md diff --git a/.sampo/changesets/stream-connect-error.md b/.sampo/changesets/stream-connect-error.md new file mode 100644 index 00000000..9f684067 --- /dev/null +++ b/.sampo/changesets/stream-connect-error.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Capture an `$ai_generation` error event when `create(stream=True)` fails before the stream starts (bad key, rate limit, invalid model) in the OpenAI chat, OpenAI Responses, Anthropic and Gemini wrappers, sync and async. Previously only errors raised mid-stream were captured. diff --git a/posthog/ai/anthropic/anthropic.py b/posthog/ai/anthropic/anthropic.py index 704cacb2..0267eb84 100644 --- a/posthog/ai/anthropic/anthropic.py +++ b/posthog/ai/anthropic/anthropic.py @@ -162,7 +162,24 @@ def _create_streaming( **kwargs: Any, ): start_time = time.time() - response = super().create(**kwargs) + try: + response = super().create(**kwargs) + except Exception as exc: + accumulator = _AnthropicStreamAccumulator() + self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + accumulator.usage_stats, + time.time() - start_time, + accumulator.content_blocks, + accumulator.accumulated_content, + error=exc, + ) + raise return self._track_streaming_response( response, posthog_distinct_id, diff --git a/posthog/ai/anthropic/anthropic_async.py b/posthog/ai/anthropic/anthropic_async.py index 5207fa6b..ce08ad0f 100644 --- a/posthog/ai/anthropic/anthropic_async.py +++ b/posthog/ai/anthropic/anthropic_async.py @@ -162,7 +162,24 @@ async def _create_streaming( **kwargs: Any, ): start_time = time.time() - response = await super().create(**kwargs) + try: + response = await super().create(**kwargs) + except Exception as exc: + accumulator = _AnthropicStreamAccumulator() + await self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + accumulator.usage_stats, + time.time() - start_time, + accumulator.content_blocks, + accumulator.accumulated_content, + error=exc, + ) + raise return self._track_streaming_response( response, posthog_distinct_id, diff --git a/posthog/ai/gemini/gemini.py b/posthog/ai/gemini/gemini.py index dad87da9..cacf8e4f 100644 --- a/posthog/ai/gemini/gemini.py +++ b/posthog/ai/gemini/gemini.py @@ -260,7 +260,27 @@ def _generate_content_streaming( stop_reason: Optional[str] = None kwargs_without_stream = {"model": model, "contents": contents, **kwargs} - response = self._client.models.generate_content_stream(**kwargs_without_stream) + try: + response = self._client.models.generate_content_stream( + **kwargs_without_stream + ) + except Exception as exc: + self._capture_streaming_event( + model, + contents, + distinct_id, + trace_id, + properties, + privacy_mode, + groups, + kwargs, + usage_stats, + time.time() - start_time, + accumulated_content, + stop_reason=stop_reason, + error=exc, + ) + raise def generator(): nonlocal usage_stats diff --git a/posthog/ai/gemini/gemini_async.py b/posthog/ai/gemini/gemini_async.py index 6c70efd9..44e971e2 100644 --- a/posthog/ai/gemini/gemini_async.py +++ b/posthog/ai/gemini/gemini_async.py @@ -244,9 +244,27 @@ async def _generate_content_streaming( stop_reason: Optional[str] = None kwargs_without_stream = {"model": model, "contents": contents, **kwargs} - response = await self._client.aio.models.generate_content_stream( - **kwargs_without_stream - ) + try: + response = await self._client.aio.models.generate_content_stream( + **kwargs_without_stream + ) + except Exception as exc: + self._capture_streaming_event( + model, + contents, + distinct_id, + trace_id, + properties, + privacy_mode, + groups, + kwargs, + usage_stats, + time.time() - start_time, + accumulated_content, + stop_reason=stop_reason, + error=exc, + ) + raise async def async_generator(): nonlocal usage_stats diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index ccdf6e69..07358c0b 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -167,7 +167,21 @@ def _create_streaming( ): start_time = time.time() state = _ResponsesStreamState() - response = self._original.create(**kwargs) + try: + response = self._original.create(**kwargs) + except Exception as exc: + self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + state, + time.time() - start_time, + error=exc, + ) + raise def generator(): stream_error: Optional[BaseException] = None @@ -392,7 +406,21 @@ def _create_streaming( if "stream_options" not in kwargs: kwargs["stream_options"] = {} kwargs["stream_options"]["include_usage"] = True - response = self._original.create(**kwargs) + try: + response = self._original.create(**kwargs) + except Exception as exc: + self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + state, + time.time() - start_time, + error=exc, + ) + raise def generator(): stream_error: Optional[BaseException] = None diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index 07e0de67..e3ff433c 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -168,7 +168,21 @@ async def _create_streaming( ): start_time = time.time() state = _ResponsesStreamState() - response = await self._original.create(**kwargs) + try: + response = await self._original.create(**kwargs) + except Exception as exc: + await self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + state, + time.time() - start_time, + error=exc, + ) + raise async def async_generator(): stream_error: Optional[BaseException] = None @@ -396,7 +410,21 @@ async def _create_streaming( if "stream_options" not in kwargs: kwargs["stream_options"] = {} kwargs["stream_options"]["include_usage"] = True - response = await self._original.create(**kwargs) + try: + response = await self._original.create(**kwargs) + except Exception as exc: + await self._capture_streaming_event( + posthog_distinct_id, + posthog_trace_id, + posthog_properties, + posthog_privacy_mode, + posthog_groups, + kwargs, + state, + time.time() - start_time, + error=exc, + ) + raise async def async_generator(): stream_error: Optional[BaseException] = None diff --git a/posthog/test/ai/test_stream_midstream_error.py b/posthog/test/ai/test_stream_midstream_error.py index 1fbc8278..7958b307 100644 --- a/posthog/test/ai/test_stream_midstream_error.py +++ b/posthog/test/ai/test_stream_midstream_error.py @@ -267,3 +267,165 @@ async def stream(): props = _props(mock_client) assert props["$ai_is_error"] is True assert props["$ai_http_status"] == 0 + + +class _ConnectError(RuntimeError): + status_code = 401 + + +def test_openai_chat_stream_connect_error(mock_client): + from posthog.ai.openai import OpenAI + + with patch("openai.resources.chat.completions.Completions.create") as create: + create.side_effect = _ConnectError("bad key") + client = OpenAI(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + stream=True, + ) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 + + +@pytest.mark.asyncio +async def test_openai_chat_stream_connect_error_async(mock_client): + from posthog.ai.openai import AsyncOpenAI + + async def create(**kwargs): + raise _ConnectError("bad key") + + with patch( + "openai.resources.chat.completions.AsyncCompletions.create", side_effect=create + ): + client = AsyncOpenAI(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + await client.chat.completions.create( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + stream=True, + ) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 + + +def test_openai_responses_stream_connect_error(mock_client): + from posthog.ai.openai import OpenAI + + with patch("openai.resources.responses.Responses.create") as create: + create.side_effect = _ConnectError("bad key") + client = OpenAI(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + client.responses.create(model="gpt-4", input="hi", stream=True) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 + + +@pytest.mark.asyncio +async def test_openai_responses_stream_connect_error_async(mock_client): + from posthog.ai.openai import AsyncOpenAI + + async def create(**kwargs): + raise _ConnectError("bad key") + + with patch("openai.resources.responses.AsyncResponses.create", side_effect=create): + client = AsyncOpenAI(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + await client.responses.create(model="gpt-4", input="hi", stream=True) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 + + +def test_anthropic_stream_connect_error(mock_client): + from posthog.ai.anthropic import Anthropic + + with patch("anthropic.resources.messages.Messages.create") as create: + create.side_effect = _ConnectError("bad key") + client = Anthropic(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "hi"}], + max_tokens=5, + stream=True, + ) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 + + +@pytest.mark.asyncio +async def test_anthropic_stream_connect_error_async(mock_client): + from posthog.ai.anthropic import AsyncAnthropic + + async def create(**kwargs): + raise _ConnectError("bad key") + + with patch("anthropic.resources.messages.AsyncMessages.create", side_effect=create): + client = AsyncAnthropic(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + await client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "hi"}], + max_tokens=5, + stream=True, + ) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 + + +def test_gemini_stream_connect_error(mock_client): + pytest.importorskip("google.genai") + from google import genai as google_genai + + from posthog.ai.gemini import Client + + with patch.object(google_genai, "Client") as cls: + inst = MagicMock() + cls.return_value = inst + inst.models.generate_content_stream.side_effect = _ConnectError("bad key") + client = Client(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + client.models.generate_content_stream( + model="gemini-2.0-flash", contents=["hi"] + ) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 + + +@pytest.mark.asyncio +async def test_gemini_stream_connect_error_async(mock_client): + pytest.importorskip("google.genai") + from google import genai as google_genai + + from posthog.ai.gemini import AsyncClient + + with patch.object(google_genai, "Client") as cls: + inst = MagicMock() + cls.return_value = inst + inst.aio.models.generate_content_stream = AsyncMock( + side_effect=_ConnectError("bad key") + ) + client = AsyncClient(api_key="k", posthog_client=mock_client) + with pytest.raises(_ConnectError): + await client.models.generate_content_stream( + model="gemini-2.0-flash", contents=["hi"] + ) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 401 From 29325a412843f29c24a975d9ea265c25ed8268cd Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Thu, 29 Oct 2026 21:13:23 -0400 Subject: [PATCH 2/2] fix(ai): annotate the Gemini streaming accumulator for mypy --- posthog/ai/gemini/gemini.py | 2 +- posthog/ai/gemini/gemini_async.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/posthog/ai/gemini/gemini.py b/posthog/ai/gemini/gemini.py index cacf8e4f..4d06c554 100644 --- a/posthog/ai/gemini/gemini.py +++ b/posthog/ai/gemini/gemini.py @@ -256,7 +256,7 @@ def _generate_content_streaming( ): start_time = time.time() usage_stats: TokenUsage = TokenUsage() - accumulated_content = [] + accumulated_content: list[Any] = [] stop_reason: Optional[str] = None kwargs_without_stream = {"model": model, "contents": contents, **kwargs} diff --git a/posthog/ai/gemini/gemini_async.py b/posthog/ai/gemini/gemini_async.py index 44e971e2..bef6e3fe 100644 --- a/posthog/ai/gemini/gemini_async.py +++ b/posthog/ai/gemini/gemini_async.py @@ -240,7 +240,7 @@ async def _generate_content_streaming( ): start_time = time.time() usage_stats: TokenUsage = TokenUsage() - accumulated_content = [] + accumulated_content: list[Any] = [] stop_reason: Optional[str] = None kwargs_without_stream = {"model": model, "contents": contents, **kwargs}