diff --git a/.sampo/changesets/stream-connect-error.md b/.sampo/changesets/stream-connect-error.md new file mode 100644 index 000000000..9f684067b --- /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 704cacb20..0267eb84f 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 5207fa6b4..ce08ad0f2 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 dad87da9a..4d06c554c 100644 --- a/posthog/ai/gemini/gemini.py +++ b/posthog/ai/gemini/gemini.py @@ -256,11 +256,31 @@ 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} - 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 6c70efd9b..bef6e3fe3 100644 --- a/posthog/ai/gemini/gemini_async.py +++ b/posthog/ai/gemini/gemini_async.py @@ -240,13 +240,31 @@ 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} - 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 1bdcdeb6f..39b86f58c 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 @@ -395,7 +409,21 @@ def _create_streaming( **(stream_options if isinstance(stream_options, dict) else {}), "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 1282e1ddf..ac34f9b88 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 @@ -399,7 +413,21 @@ async def _create_streaming( **(stream_options if isinstance(stream_options, dict) else {}), "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 1fbc8278f..7958b307d 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