fix(pydantic-ai): Determine if response is streamed based on the method

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -2,6 +2,7 @@
 from functools import wraps
 
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.consts import SPANDATA
 
 from ..spans import (
     ai_client_span,
@@ -60,6 +61,8 @@ async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any":
         messages, model, model_settings = _extract_span_data(self, ctx)
 
         with ai_client_span(messages, None, model, model_settings) as span:
+            span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+
             result = await original_model_request_run(self, ctx)
 
             # Extract response from result if available
@@ -87,6 +90,8 @@ async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any":
 
             # Create chat span for streaming request
             with ai_client_span(messages, None, model, model_settings) as span:
+                span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+
                 # Call the original stream method
                 async with original_stream_method(self, ctx) as stream:
                     yield stream

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -17,7 +17,6 @@
     _should_send_prompts,
     _get_model_name,
     get_current_agent,
-    get_is_streaming,
 )
 from .utils import (
     _serialize_binary_content_item,
@@ -264,9 +263,6 @@ def ai_client_span(
     _set_agent_data(span, agent)
     _set_model_data(span, model, model_settings)
 
-    # Set streaming flag from contextvar
-    span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming())
-
     # Add available tools if agent is available
     agent_obj = agent or get_current_agent()
     _set_available_tools(span, agent_obj)