pass response model directly to update_ai_client_span

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -56,8 +56,6 @@ async def wrapped_get_response(*args, **kwargs):
 
                 response_model = getattr(agent, "_sentry_raw_response_model", None)
                 if response_model:
-                    result._sentry_response_model = response_model
-
                     agent_span = getattr(agent, "_sentry_agent_span", None)
                     if agent_span:
                         agent_span.set_data(
@@ -66,7 +64,7 @@ async def wrapped_get_response(*args, **kwargs):
 
                     delattr(agent, "_sentry_raw_response_model")
 
-                update_ai_client_span(span, agent, kwargs, result)
+                update_ai_client_span(span, agent, kwargs, result, response_model)
 
             return result
 

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -14,7 +14,7 @@
 
 if TYPE_CHECKING:
     from agents import Agent
-    from typing import Any
+    from typing import Any, Optional
 
 
 def ai_client_span(agent, get_response_kwargs):
@@ -35,12 +35,14 @@ def ai_client_span(agent, get_response_kwargs):
     return span
 
 
-def update_ai_client_span(span, agent, get_response_kwargs, result):
-    # type: (sentry_sdk.tracing.Span, Agent, dict[str, Any], Any) -> None
+def update_ai_client_span(
+    span, agent, get_response_kwargs, result, response_model=None
+):
+    # type: (sentry_sdk.tracing.Span, Agent, dict[str, Any], Any, Optional[str]) -> None
     _set_usage_data(span, result.usage)
     _set_output_data(span, result)
     _create_mcp_execute_tool_spans(span, result)
 
     # Set response model if captured from raw response
-    if hasattr(result, "_sentry_response_model") and result._sentry_response_model:
-        span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, result._sentry_response_model)
+    if response_model is not None:
+        span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)