mypy

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -37,7 +37,7 @@ def __init__(
         self.model_settings = model_settings
         self.is_streaming = is_streaming
         self._isolation_scope: "Any" = None
-        self._span: "Optional[sentry_sdk.tracing.Span]" = None
+        self._span: "Optional[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]" = None
         self._result: "Any" = None
 
     async def __aenter__(self) -> "Any":

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -314,7 +314,7 @@ def ai_client_span(
 
 
 def update_ai_client_span(
-    span: "sentry_sdk.tracing.Span", model_response: "Any"
+    span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", model_response: "Any"
 ) -> None:
     """Update the AI client span with response data."""
     if not span: