mypy

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/spans/agent_workflow.py

@@ -7,10 +7,14 @@
 from ..consts import SPAN_ORIGIN
 
 if TYPE_CHECKING:
+    from typing import Union
+
     import agents
 
 
-def agent_workflow_span(agent: "agents.Agent") -> "sentry_sdk.tracing.Span":
+def agent_workflow_span(
+    agent: "agents.Agent",
+) -> "Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]":
     # Create a transaction or a span if an transaction is already active
     span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
     if span_streaming:

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -21,7 +21,7 @@
 
 def ai_client_span(
     agent: "Agent", get_response_kwargs: "dict[str, Any]"
-) -> "sentry_sdk.tracing.Span":
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
     # TODO-anton: implement other types of operations. Now "chat" is hardcoded.
     # Get model name from agent.model or fall back to request model (for when agent.model is None/default)
     model_name = None

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -17,7 +17,7 @@
 
 def execute_tool_span(
     tool: "agents.Tool", *args: "Any", **kwargs: "Any"
-) -> "sentry_sdk.tracing.Span":
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
     span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
     if span_streaming:
         span = sentry_sdk.traces.start_span(

sentry_sdk/integrations/openai_agents/spans/handoff.py

@@ -26,7 +26,7 @@ def handoff_span(
             # Add conversation ID from agent
             conv_id = getattr(from_agent, "_sentry_conversation_id", None)
             if conv_id:
-                span.set_data(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
+                span.set_attribute(SPANDATA.GEN_AI_CONVERSATION_ID, conv_id)
     else:
         with sentry_sdk.start_span(
             op=OP.GEN_AI_HANDOFF,

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -24,7 +24,7 @@
 
 def invoke_agent_span(
     context: "agents.RunContextWrapper", agent: "agents.Agent", kwargs: "dict[str, Any]"
-) -> "sentry_sdk.tracing.Span":
+) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
     span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
     if span_streaming:
         span = sentry_sdk.traces.start_span(