fix(ai): capture an error event when a streaming call fails to connect

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

posthog/ai/anthropic/anthropic.py

@@ -6,6 +6,7 @@
         "Please install the Anthropic SDK to use this feature: 'pip install anthropic'"
     )
 
+import logging
 import time
 import uuid
 from typing import Any, Dict, List, Optional
@@ -31,6 +32,8 @@
 from posthog.client import Client as PostHogClient
 from posthog import setup
 
+_log = logging.getLogger("posthog")
+
 
 class Anthropic(anthropic.Anthropic):
     """
@@ -162,7 +165,28 @@ def _create_streaming(
         **kwargs: Any,
     ):
         start_time = time.time()
-        response = super().create(**kwargs)
+        try:
+            response = super().create(**kwargs)
+        except Exception as exc:
+            accumulator = _AnthropicStreamAccumulator()
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
         return self._track_streaming_response(
             response,
             posthog_distinct_id,

posthog/ai/anthropic/anthropic_async.py

@@ -6,11 +6,13 @@
         "Please install the Anthropic SDK to use this feature: 'pip install anthropic'"
     )
 
+import logging
 import time
 import uuid
 from typing import Any, Dict, List, Optional
 
 from posthog import setup
+
 from ..stream import AsyncStreamWrapper as AsyncStreamWrapper
 from ..types import (
     StreamingContentBlock as StreamingContentBlock,
@@ -31,6 +33,8 @@
 )
 from posthog.client import Client as PostHogClient
 
+_log = logging.getLogger("posthog")
+
 
 class AsyncAnthropic(anthropic.AsyncAnthropic):
     """
@@ -162,7 +166,28 @@ 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()
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
         return self._track_streaming_response(
             response,
             posthog_distinct_id,

posthog/ai/gemini/gemini.py

@@ -1,3 +1,4 @@
+import logging
 import time
 from typing import Any, Dict, Optional
 
@@ -29,6 +30,8 @@
     format_gemini_streaming_output as format_gemini_streaming_output,
 )
 
+_log = logging.getLogger("posthog")
+
 
 class Client:
     """
@@ -257,11 +260,36 @@ 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:
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error, and the
+                # fixed text avoids echoing prompt or error content into logs.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
 
         def generator():
             nonlocal usage_stats

posthog/ai/gemini/gemini_async.py

@@ -1,3 +1,4 @@
+import logging
 import time
 from typing import Any, Dict, Optional
 
@@ -29,6 +30,8 @@
     format_gemini_streaming_output as format_gemini_streaming_output,
 )
 
+_log = logging.getLogger("posthog")
+
 
 class AsyncClient:
     """
@@ -241,13 +244,36 @@ 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:
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error, and the
+                # fixed text avoids echoing prompt or error content into logs.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
 
         async def async_generator():
             nonlocal usage_stats

posthog/ai/openai/openai.py

@@ -1,3 +1,4 @@
+import logging
 import time
 import uuid
 from typing import TYPE_CHECKING as _TYPE_CHECKING, Any, Dict, Optional
@@ -41,6 +42,8 @@
     merge_provider_override,
 )
 
+_log = logging.getLogger("posthog")
+
 
 class OpenAI(openai.OpenAI):
     """
@@ -167,7 +170,26 @@ def _create_streaming(
     ):
         start_time = time.time()
         state = _ResponsesStreamState()
-        response = self._original.create(**kwargs)
+        try:
+            response = self._original.create(**kwargs)
+        except Exception as exc:
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error, and the
+                # fixed text avoids echoing prompt or error content into logs.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
 
         def generator():
             stream_error: Optional[BaseException] = None
@@ -395,7 +417,26 @@ 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:
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error, and the
+                # fixed text avoids echoing prompt or error content into logs.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
 
         def generator():
             stream_error: Optional[BaseException] = None

posthog/ai/openai/openai_async.py

@@ -1,3 +1,4 @@
+import logging
 import time
 import uuid
 from typing import TYPE_CHECKING as _TYPE_CHECKING, Any, Dict, Optional
@@ -42,6 +43,8 @@
     merge_provider_override,
 )
 
+_log = logging.getLogger("posthog")
+
 
 class AsyncOpenAI(openai.AsyncOpenAI):
     """
@@ -168,7 +171,26 @@ 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:
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error, and the
+                # fixed text avoids echoing prompt or error content into logs.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
 
         async def async_generator():
             stream_error: Optional[BaseException] = None
@@ -399,7 +421,26 @@ 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:
+            try:
+                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,
+                )
+            except Exception:
+                # Telemetry must never replace the provider's own error, and the
+                # fixed text avoids echoing prompt or error content into logs.
+                _log.warning("[PostHog] Failed to capture streaming connect error")
+            raise
 
         async def async_generator():
             stream_error: Optional[BaseException] = None