diff --git a/.sampo/changesets/streaming-midstream-error.md b/.sampo/changesets/streaming-midstream-error.md new file mode 100644 index 000000000..b66920190 --- /dev/null +++ b/.sampo/changesets/streaming-midstream-error.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Streamed OpenAI, Anthropic and Gemini generations that raise mid-stream are now captured with `$ai_is_error` and `$ai_error` (and the error's status code as `$ai_http_status`) instead of being recorded as a successful 200. The error is still re-raised to the caller, matching what non-streaming calls already did. diff --git a/posthog/ai/anthropic/anthropic.py b/posthog/ai/anthropic/anthropic.py index a474fefcf..704cacb20 100644 --- a/posthog/ai/anthropic/anthropic.py +++ b/posthog/ai/anthropic/anthropic.py @@ -188,10 +188,14 @@ def _track_streaming_response( accumulator = _AnthropicStreamAccumulator() def generator(): + stream_error: Optional[BaseException] = None try: for event in response: accumulator.consume(event) yield event + except Exception as exc: + stream_error = exc + raise finally: end_time = time.time() latency = end_time - start_time @@ -208,6 +212,7 @@ def generator(): accumulator.content_blocks, accumulator.accumulated_content, stop_reason=accumulator.stop_reason, + error=stream_error, ) return _StreamWrapper(generator(), stream=response) @@ -225,6 +230,7 @@ def _capture_streaming_event( content_blocks: List[StreamingContentBlock], accumulated_content: str, stop_reason: Optional[str] = None, + error: Optional[BaseException] = None, ): from posthog.ai.types import StreamingEventData from posthog.ai.anthropic.anthropic_converter import ( @@ -252,6 +258,7 @@ def _capture_streaming_event( privacy_mode=posthog_privacy_mode, groups=posthog_groups, stop_reason=stop_reason, + error=error, ) # Use the common capture function diff --git a/posthog/ai/anthropic/anthropic_async.py b/posthog/ai/anthropic/anthropic_async.py index 9ac091c64..5207fa6b4 100644 --- a/posthog/ai/anthropic/anthropic_async.py +++ b/posthog/ai/anthropic/anthropic_async.py @@ -188,10 +188,14 @@ def _track_streaming_response( accumulator = _AnthropicStreamAccumulator() async def generator(): + stream_error: Optional[BaseException] = None try: async for event in response: accumulator.consume(event) yield event + except Exception as exc: + stream_error = exc + raise finally: end_time = time.time() latency = end_time - start_time @@ -208,6 +212,7 @@ async def generator(): accumulator.content_blocks, accumulator.accumulated_content, stop_reason=accumulator.stop_reason, + error=stream_error, ) return AsyncStreamWrapper(generator(), stream=response) @@ -225,6 +230,7 @@ async def _capture_streaming_event( content_blocks: List[StreamingContentBlock], accumulated_content: str, stop_reason: Optional[str] = None, + error: Optional[BaseException] = None, ): from posthog.ai.types import StreamingEventData from posthog.ai.anthropic.anthropic_converter import ( @@ -252,6 +258,7 @@ async def _capture_streaming_event( privacy_mode=posthog_privacy_mode, groups=posthog_groups, stop_reason=stop_reason, + error=error, ) # Use the common capture function diff --git a/posthog/ai/gemini/_shared.py b/posthog/ai/gemini/_shared.py index fc3635589..e1036929d 100644 --- a/posthog/ai/gemini/_shared.py +++ b/posthog/ai/gemini/_shared.py @@ -203,6 +203,7 @@ def _capture_streaming_event( latency: float, output: Any, stop_reason: Optional[str] = None, + error: Optional[BaseException] = None, ) -> None: formatted_input = merge_system_prompt( {"contents": contents, **kwargs}, "gemini" @@ -222,6 +223,7 @@ def _capture_streaming_event( privacy_mode=privacy_mode, groups=groups, stop_reason=stop_reason, + error=error, ) capture_streaming_event(self._ph_client, event_data) diff --git a/posthog/ai/gemini/gemini.py b/posthog/ai/gemini/gemini.py index ff74bd270..c5006544a 100644 --- a/posthog/ai/gemini/gemini.py +++ b/posthog/ai/gemini/gemini.py @@ -266,6 +266,7 @@ def generator(): nonlocal usage_stats nonlocal accumulated_content nonlocal stop_reason + stream_error: Optional[BaseException] = None try: for chunk in response: # Extract usage stats from chunk @@ -288,6 +289,9 @@ def generator(): yield chunk + except Exception as exc: + stream_error = exc + raise finally: end_time = time.time() latency = end_time - start_time @@ -305,6 +309,7 @@ def generator(): latency, accumulated_content, stop_reason=stop_reason, + error=stream_error, ) return generator() diff --git a/posthog/ai/gemini/gemini_async.py b/posthog/ai/gemini/gemini_async.py index edede3dca..0580b285b 100644 --- a/posthog/ai/gemini/gemini_async.py +++ b/posthog/ai/gemini/gemini_async.py @@ -253,6 +253,7 @@ async def async_generator(): nonlocal accumulated_content nonlocal stop_reason + stream_error: Optional[BaseException] = None try: async for chunk in response: # Extract usage stats from chunk @@ -275,6 +276,9 @@ async def async_generator(): yield chunk + except Exception as exc: + stream_error = exc + raise finally: end_time = time.time() latency = end_time - start_time @@ -292,6 +296,7 @@ async def async_generator(): latency, accumulated_content, stop_reason=stop_reason, + error=stream_error, ) return AsyncStreamWrapper(async_generator(), stream=response) diff --git a/posthog/ai/openai/_streaming.py b/posthog/ai/openai/_streaming.py index 2ee95ec91..376372687 100644 --- a/posthog/ai/openai/_streaming.py +++ b/posthog/ai/openai/_streaming.py @@ -102,6 +102,7 @@ def _build_streaming_event_data( model_from_response: Optional[str], stop_reason: Optional[str], service_tier: Optional[str] = None, + error: Optional[BaseException] = None, ) -> StreamingEventData: """Build the fields shared by both OpenAI streaming endpoint events.""" @@ -121,4 +122,5 @@ def _build_streaming_event_data( groups=groups, stop_reason=stop_reason, service_tier=service_tier, + error=error, ) diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index ea00a197b..ccdf6e69a 100644 --- a/posthog/ai/openai/openai.py +++ b/posthog/ai/openai/openai.py @@ -170,10 +170,14 @@ def _create_streaming( response = self._original.create(**kwargs) def generator(): + stream_error: Optional[BaseException] = None try: for chunk in response: state.process_chunk(chunk) yield chunk + except Exception as exc: + stream_error = exc + raise finally: self._capture_streaming_event( posthog_distinct_id, @@ -184,6 +188,7 @@ def generator(): kwargs, state, time.time() - start_time, + error=stream_error, ) return _StreamWrapper(generator(), stream=response) @@ -198,6 +203,7 @@ def _capture_streaming_event( kwargs: Dict[str, Any], state: _ResponsesStreamState, latency: float, + error: Optional[BaseException] = None, ): from posthog.ai.utils import capture_streaming_event @@ -216,6 +222,7 @@ def _capture_streaming_event( model_from_response=state.model, stop_reason=state.stop_reason, service_tier=state.service_tier, + error=error, ) capture_streaming_event(self._client._ph_client, event_data) @@ -388,10 +395,14 @@ def _create_streaming( response = self._original.create(**kwargs) def generator(): + stream_error: Optional[BaseException] = None try: for chunk in response: state.process_chunk(chunk) yield chunk + except Exception as exc: + stream_error = exc + raise finally: self._capture_streaming_event( posthog_distinct_id, @@ -402,6 +413,7 @@ def generator(): kwargs, state, time.time() - start_time, + error=stream_error, ) return _StreamWrapper(generator(), stream=response) @@ -416,6 +428,7 @@ def _capture_streaming_event( kwargs: Dict[str, Any], state: _ChatCompletionsStreamState, latency: float, + error: Optional[BaseException] = None, ): from posthog.ai.utils import capture_streaming_event @@ -436,6 +449,7 @@ def _capture_streaming_event( model_from_response=state.model, stop_reason=state.stop_reason, service_tier=state.service_tier, + error=error, ) capture_streaming_event(self._client._ph_client, event_data) diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index dac3eeb4a..07e0de670 100644 --- a/posthog/ai/openai/openai_async.py +++ b/posthog/ai/openai/openai_async.py @@ -171,10 +171,14 @@ async def _create_streaming( response = await self._original.create(**kwargs) async def async_generator(): + stream_error: Optional[BaseException] = None try: async for chunk in response: state.process_chunk(chunk) yield chunk + except Exception as exc: + stream_error = exc + raise finally: await self._capture_streaming_event( posthog_distinct_id, @@ -185,6 +189,7 @@ async def async_generator(): kwargs, state, time.time() - start_time, + error=stream_error, ) return AsyncStreamWrapper(async_generator(), stream=response) @@ -199,6 +204,7 @@ async def _capture_streaming_event( kwargs: Dict[str, Any], state: _ResponsesStreamState, latency: float, + error: Optional[BaseException] = None, ): from posthog.ai.utils import capture_streaming_event @@ -217,6 +223,7 @@ async def _capture_streaming_event( model_from_response=state.model, stop_reason=state.stop_reason, service_tier=state.service_tier, + error=error, ) capture_streaming_event(self._client._ph_client, event_data) @@ -392,10 +399,14 @@ async def _create_streaming( response = await self._original.create(**kwargs) async def async_generator(): + stream_error: Optional[BaseException] = None try: async for chunk in response: state.process_chunk(chunk) yield chunk + except Exception as exc: + stream_error = exc + raise finally: await self._capture_streaming_event( posthog_distinct_id, @@ -406,6 +417,7 @@ async def async_generator(): kwargs, state, time.time() - start_time, + error=stream_error, ) return AsyncStreamWrapper(async_generator(), stream=response) @@ -420,6 +432,7 @@ async def _capture_streaming_event( kwargs: Dict[str, Any], state: _ChatCompletionsStreamState, latency: float, + error: Optional[BaseException] = None, ): from posthog.ai.utils import capture_streaming_event @@ -440,6 +453,7 @@ async def _capture_streaming_event( model_from_response=state.model, stop_reason=state.stop_reason, service_tier=state.service_tier, + error=error, ) capture_streaming_event(self._client._ph_client, event_data) diff --git a/posthog/ai/types.py b/posthog/ai/types.py index 36424b133..1fc5555fa 100644 --- a/posthog/ai/types.py +++ b/posthog/ai/types.py @@ -143,3 +143,4 @@ class StreamingEventData(TypedDict): groups: Optional[Dict[str, Any]] stop_reason: Optional[str] service_tier: NotRequired[Optional[str]] + error: NotRequired[Optional[BaseException]] # raised mid-stream, if any diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index c943bcc28..5f2b956e0 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -837,6 +837,7 @@ def capture_streaming_event( event_data: Standardized streaming event data containing all necessary information """ trace_id = event_data.get("trace_id") or str(uuid.uuid4()) + error = event_data.get("error") # Omitted when the provider never reported a count: absent means unknown, # 0 is a report of nothing. An interrupted stream often reports neither. @@ -865,13 +866,18 @@ def capture_streaming_event( event_data["privacy_mode"], finalize_ai_content(event_data["formatted_output"], ph_client), ), - "$ai_http_status": 200, + "$ai_http_status": ( + getattr(error, "status_code", 0) if error is not None else 200 + ), **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), **({"$ai_output_tokens": output_tokens} if output_tokens is not None else {}), "$ai_latency": event_data["latency"], "$ai_trace_id": trace_id, "$ai_base_url": str(event_data["base_url"]), **(event_data.get("properties") or {}), + **( + {"$ai_is_error": True, "$ai_error": str(error)} if error is not None else {} + ), } warn_if_posthog_ai_gateway(event_data["base_url"]) diff --git a/posthog/test/ai/test_stream_midstream_error.py b/posthog/test/ai/test_stream_midstream_error.py new file mode 100644 index 000000000..1fbc8278f --- /dev/null +++ b/posthog/test/ai/test_stream_midstream_error.py @@ -0,0 +1,269 @@ +"""A stream that raises mid-flight must be captured as an error, not as a 200.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from posthog.ai.types import StreamingEventData +from posthog.ai.utils import capture_streaming_event + + +@pytest.fixture +def mock_client(): + with patch("posthog.client.Client") as mock_client: + mock_client.privacy_mode = False + yield mock_client + + +def _props(mock_client): + assert mock_client.capture.call_count == 1 + return mock_client.capture.call_args[1]["properties"] + + +class _StatusError(RuntimeError): + status_code = 529 + + +def _event_data(**overrides) -> StreamingEventData: + data = StreamingEventData( + provider="openai", + model="gpt-4", + base_url="https://api.openai.com/v1", + kwargs={}, + formatted_input=[], + formatted_output=[], + usage_stats={}, + latency=0.1, + distinct_id="id", + trace_id="trace", + properties=None, + privacy_mode=False, + groups=None, + stop_reason=None, + ) + data.update(overrides) + return data + + +def test_capture_streaming_event_without_error_is_ok(mock_client): + capture_streaming_event(mock_client, _event_data()) + props = _props(mock_client) + assert props["$ai_http_status"] == 200 + assert "$ai_is_error" not in props + + +@pytest.mark.parametrize( + "error, status", [(RuntimeError("boom"), 0), (_StatusError("overloaded"), 529)] +) +def test_capture_streaming_event_with_error(mock_client, error, status): + capture_streaming_event(mock_client, _event_data(error=error)) + props = _props(mock_client) + assert props["$ai_http_status"] == status + assert props["$ai_is_error"] is True + assert props["$ai_error"] == str(error) + + +def test_openai_chat_stream_error(mock_client): + from openai.types.chat import ChatCompletionChunk + from openai.types.chat.chat_completion_chunk import Choice, ChoiceDelta + + from posthog.ai.openai import OpenAI + + chunk = ChatCompletionChunk( + id="c1", + created=1, + model="gpt-4", + object="chat.completion.chunk", + choices=[Choice(index=0, delta=ChoiceDelta(content="partial"))], + ) + + def stream(): + yield chunk + raise RuntimeError("connection reset") + + with patch("openai.resources.chat.completions.Completions.create") as create: + create.return_value = stream() + client = OpenAI(api_key="k", posthog_client=mock_client) + with pytest.raises(RuntimeError, match="connection reset"): + list( + 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_error"] == "connection reset" + assert props["$ai_http_status"] == 0 + + +@pytest.mark.asyncio +async def test_openai_chat_stream_error_async(mock_client): + from openai.types.chat import ChatCompletionChunk + from openai.types.chat.chat_completion_chunk import Choice, ChoiceDelta + + from posthog.ai.openai import AsyncOpenAI + + chunk = ChatCompletionChunk( + id="c1", + created=1, + model="gpt-4", + object="chat.completion.chunk", + choices=[Choice(index=0, delta=ChoiceDelta(content="partial"))], + ) + + async def stream(): + yield chunk + raise RuntimeError("connection reset") + + async def create(self, **kwargs): + return stream() + + with patch("openai.resources.chat.completions.AsyncCompletions.create", new=create): + client = AsyncOpenAI(api_key="k", posthog_client=mock_client) + response = await client.chat.completions.create( + model="gpt-4", messages=[{"role": "user", "content": "hi"}], stream=True + ) + with pytest.raises(RuntimeError, match="connection reset"): + async for _ in response: + pass + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 0 + + +def test_openai_responses_stream_error(mock_client): + from posthog.ai.openai import OpenAI + + def stream(): + yield MagicMock(type="response.created", response=None) + raise RuntimeError("connection reset") + + with patch("openai.resources.responses.Responses.create") as create: + create.return_value = stream() + client = OpenAI(api_key="k", posthog_client=mock_client) + with pytest.raises(RuntimeError, match="connection reset"): + list(client.responses.create(model="gpt-4", input="hi", stream=True)) + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_error"] == "connection reset" + + +def test_anthropic_stream_error(mock_client): + from posthog.ai.anthropic import Anthropic + + def stream(): + yield MagicMock(type="message_start", message=None) + raise _StatusError("overloaded") + + with patch("anthropic.resources.messages.Messages.create") as create: + create.return_value = stream() + client = Anthropic(api_key="k", posthog_client=mock_client) + with pytest.raises(_StatusError): + list( + 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"] == 529 + + +@pytest.mark.asyncio +async def test_anthropic_stream_error_async(mock_client): + from posthog.ai.anthropic import AsyncAnthropic + + async def stream(): + yield MagicMock(type="message_start", message=None) + raise _StatusError("overloaded") + + async def create(**kwargs): + return stream() + + with patch("anthropic.resources.messages.AsyncMessages.create", side_effect=create): + client = AsyncAnthropic(api_key="k", posthog_client=mock_client) + response = await client.messages.create( + model="claude-3-opus-20240229", + messages=[{"role": "user", "content": "hi"}], + max_tokens=5, + stream=True, + ) + with pytest.raises(_StatusError): + async for _ in response: + pass + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 529 + + +def _gemini_chunk(): + chunk = MagicMock() + chunk.text = "partial" + chunk.usage_metadata = None + return chunk + + +def test_gemini_stream_error(mock_client): + pytest.importorskip("google.genai") + from google import genai as google_genai + + from posthog.ai.gemini import Client + + def stream(): + yield _gemini_chunk() + raise RuntimeError("503 overloaded mid-stream") + + with patch.object(google_genai, "Client") as cls: + inst = MagicMock() + cls.return_value = inst + inst.models.generate_content_stream.return_value = stream() + client = Client(api_key="k", posthog_client=mock_client) + with pytest.raises(RuntimeError, match="503"): + list( + 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_error"] == "503 overloaded mid-stream" + assert props["$ai_http_status"] == 0 + + +@pytest.mark.asyncio +async def test_gemini_stream_error_async(mock_client): + pytest.importorskip("google.genai") + from google import genai as google_genai + + from posthog.ai.gemini import AsyncClient + + async def stream(): + yield _gemini_chunk() + raise RuntimeError("503 overloaded mid-stream") + + with patch.object(google_genai, "Client") as cls: + inst = MagicMock() + cls.return_value = inst + inst.aio.models.generate_content_stream = AsyncMock(return_value=stream()) + client = AsyncClient(api_key="k", posthog_client=mock_client) + response = await client.models.generate_content_stream( + model="gemini-2.0-flash", contents=["hi"] + ) + with pytest.raises(RuntimeError, match="503"): + async for _ in response: + pass + + props = _props(mock_client) + assert props["$ai_is_error"] is True + assert props["$ai_http_status"] == 0 diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 5a093b650..693f066e6 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -623,6 +623,7 @@ attribute posthog.ai.types.StreamingContentBlock.thinking: Optional[str] attribute posthog.ai.types.StreamingContentBlock.type: str attribute posthog.ai.types.StreamingEventData.base_url: str attribute posthog.ai.types.StreamingEventData.distinct_id: Optional[str] +attribute posthog.ai.types.StreamingEventData.error: NotRequired[Optional[BaseException]] attribute posthog.ai.types.StreamingEventData.formatted_input: Any attribute posthog.ai.types.StreamingEventData.formatted_output: Any attribute posthog.ai.types.StreamingEventData.groups: Optional[Dict[str, Any]]