diff --git a/sentry_sdk/integrations/openai_agents/__init__.py b/sentry_sdk/integrations/openai_agents/__init__.py index 94c5ecc7a5..806a3bc1ac 100644 --- a/sentry_sdk/integrations/openai_agents/__init__.py +++ b/sentry_sdk/integrations/openai_agents/__init__.py @@ -112,7 +112,7 @@ async def new_wrapped_run_single_turn( *args: "Any", **kwargs: "Any" ) -> "SingleStepResult": return await _run_single_turn( - run_loop.run_single_turn, *args, **kwargs + run_loop.run_single_turn, use_run_hooks, *args, **kwargs ) agents.run.run_single_turn = new_wrapped_run_single_turn @@ -125,6 +125,7 @@ async def new_wrapped_run_single_turn_streamed( ) -> "SingleStepResult": return await _run_single_turn_streamed( original_run_single_turn_streamed, + use_run_hooks, *args, **kwargs, ) @@ -207,7 +208,9 @@ def old_wrapped_get_model( async def old_wrapped_run_single_turn( cls: "agents.Runner", *args: "Any", **kwargs: "Any" ) -> "SingleStepResult": - return await _run_single_turn(original_run_single_turn, *args, **kwargs) + return await _run_single_turn( + original_run_single_turn, use_run_hooks, *args, **kwargs + ) agents.run.AgentRunner._run_single_turn = classmethod( old_wrapped_run_single_turn @@ -220,7 +223,7 @@ async def old_wrapped_run_single_turn_streamed( cls: "agents.Runner", *args: "Any", **kwargs: "Any" ) -> "SingleStepResult": return await _run_single_turn_streamed( - original_run_single_turn_streamed, *args, **kwargs + original_run_single_turn_streamed, use_run_hooks, *args, **kwargs ) agents.run.AgentRunner._run_single_turn_streamed = classmethod( diff --git a/sentry_sdk/integrations/openai_agents/patches/agent_run.py b/sentry_sdk/integrations/openai_agents/patches/agent_run.py index 71883b2eef..ed2579903e 100644 --- a/sentry_sdk/integrations/openai_agents/patches/agent_run.py +++ b/sentry_sdk/integrations/openai_agents/patches/agent_run.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from typing import Any, Awaitable, Callable, Optional, Union + from agents import TResponseInputItem from agents.run_internal.run_steps import SingleStepResult from sentry_sdk.tracing import Span @@ -49,7 +50,7 @@ def _maybe_start_agent_span( context_wrapper: "agents.RunContextWrapper", agent: "agents.Agent", should_run_agent_start_hooks: bool, - span_kwargs: "dict[str, Any]", + turn_input: "Optional[list[TResponseInputItem]]", is_streaming: bool = False, ) -> "Optional[Union[Span, StreamedSpan]]": """ @@ -68,14 +69,14 @@ def _maybe_start_agent_span( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent + span=span, usage=context_wrapper.usage, agent=agent ) span.__exit__(None, None, None) delattr(context_wrapper, "_sentry_agent_span") # Store the agent on the context wrapper so we can access it later context_wrapper._sentry_current_agent = agent - span = invoke_agent_span(context_wrapper, agent, span_kwargs) + span = invoke_agent_span(agent, turn_input) context_wrapper._sentry_agent_span = span agent._sentry_agent_span = span @@ -92,6 +93,7 @@ def _maybe_start_agent_span( async def _run_single_turn( original_run_single_turn: "Callable[..., Awaitable[SingleStepResult]]", + use_run_hooks: "bool", *args: "Any", **kwargs: "Any", ) -> "SingleStepResult": @@ -107,32 +109,48 @@ async def _run_single_turn( if bindings is not None else kwargs.get("agent") ) - context_wrapper = kwargs.get("context_wrapper") - should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks", False) - span = _maybe_start_agent_span( - context_wrapper, agent, should_run_agent_start_hooks, kwargs - ) + context_wrapper: "agents.RunContextWrapper[Any]" = kwargs.get("context_wrapper") + if not use_run_hooks: + should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks", False) - if ( - span is None - or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) - or (not isinstance(span, StreamedSpan) and span.timestamp is not None) - ): - return await original_run_single_turn(*args, **kwargs) + span = _maybe_start_agent_span( + context_wrapper, + agent, + should_run_agent_start_hooks, + kwargs.get("input"), + ) + + if ( + span is None + or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) + or (not isinstance(span, StreamedSpan) and span.timestamp is not None) + ): + return await original_run_single_turn(*args, **kwargs) try: result = await original_run_single_turn(*args, **kwargs) except Exception: exc_info = sys.exc_info() with capture_internal_exceptions(): - span = getattr(context_wrapper, "_sentry_agent_span", None) - if span: - update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent - ) - span.__exit__(*exc_info) - delattr(context_wrapper, "_sentry_agent_span") + if use_run_hooks: + run_hooks = kwargs.get("hooks") + if run_hooks is not None: + span = getattr(run_hooks, "_sentry_invoke_agent_span", None) + if span is not None: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + del run_hooks._sentry_invoke_agent_span + span.__exit__(*exc_info) + else: + span = getattr(context_wrapper, "_sentry_agent_span", None) + if span: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + span.__exit__(*exc_info) + delattr(context_wrapper, "_sentry_agent_span") reraise(*exc_info) return result @@ -140,6 +158,7 @@ async def _run_single_turn( async def _run_single_turn_streamed( original_run_single_turn_streamed: "Callable[..., Awaitable[SingleStepResult]]", + use_run_hooks: "bool", *args: "Any", **kwargs: "Any", ) -> "SingleStepResult": @@ -181,42 +200,59 @@ async def _run_single_turn_streamed( args[1] if len(args) > 1 else kwargs.get("bindings", kwargs.get("agent")) ) agent = getattr(agent_or_bindings, "public_agent", agent_or_bindings) - context_wrapper = args[3] if len(args) > 3 else kwargs.get("context_wrapper") - should_run_agent_start_hooks = bool( - args[5] if len(args) > 5 else kwargs.get("should_run_agent_start_hooks", False) + + context_wrapper: "agents.RunContextWrapper[Any]" = ( + args[3] if len(args) > 3 else kwargs.get("context_wrapper") ) + if not use_run_hooks: + should_run_agent_start_hooks = bool( + args[5] + if len(args) > 5 + else kwargs.get("should_run_agent_start_hooks", False) + ) - span_kwargs: "dict[str, Any]" = {} - if streamed_result and hasattr(streamed_result, "input"): - span_kwargs["original_input"] = streamed_result.input + span_kwargs: "dict[str, Any]" = {} + if streamed_result and hasattr(streamed_result, "input"): + span_kwargs["original_input"] = streamed_result.input - span = _maybe_start_agent_span( - context_wrapper, - agent, - should_run_agent_start_hooks, - span_kwargs, - is_streaming=True, - ) + span = _maybe_start_agent_span( + context_wrapper, + agent, + should_run_agent_start_hooks, + getattr(streamed_result, "input", None), + is_streaming=True, + ) - if ( - span is None - or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) - or (not isinstance(span, StreamedSpan) and span.timestamp is not None) - ): - return await original_run_single_turn_streamed(*args, **kwargs) + if ( + span is None + or (isinstance(span, StreamedSpan) and span.end_timestamp is not None) + or (not isinstance(span, StreamedSpan) and span.timestamp is not None) + ): + return await original_run_single_turn_streamed(*args, **kwargs) try: result = await original_run_single_turn_streamed(*args, **kwargs) except Exception: exc_info = sys.exc_info() with capture_internal_exceptions(): - span = getattr(context_wrapper, "_sentry_agent_span", None) - if span: - update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent - ) - span.__exit__(*exc_info) - delattr(context_wrapper, "_sentry_agent_span") + if use_run_hooks: + run_hooks = args[2] if len(args) > 2 else kwargs.get("hooks") + if run_hooks is not None: + span = getattr(run_hooks, "_sentry_invoke_agent_span", None) + if span is not None: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + del run_hooks._sentry_invoke_agent_span + span.__exit__(*exc_info) + else: + span = getattr(context_wrapper, "_sentry_agent_span", None) + if span: + update_invoke_agent_span( + span=span, usage=context_wrapper.usage, agent=agent + ) + span.__exit__(*exc_info) + delattr(context_wrapper, "_sentry_agent_span") _close_streaming_workflow_span(agent) reraise(*exc_info) @@ -266,7 +302,7 @@ async def _execute_handoffs( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent + span=span, usage=context_wrapper.usage, agent=agent ) span.__exit__(*exc_info) delattr(context_wrapper, "_sentry_agent_span") @@ -274,7 +310,7 @@ async def _execute_handoffs( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: - update_invoke_agent_span(span=span, context=context_wrapper, agent=agent) + update_invoke_agent_span(span=span, usage=context_wrapper.usage, agent=agent) span.__exit__(None, None, None) delattr(context_wrapper, "_sentry_agent_span") @@ -315,7 +351,10 @@ async def _execute_final_output( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent, output=final_output + span=span, + usage=context_wrapper.usage, + agent=agent, + output=final_output, ) span.__exit__(*exc_info) delattr(context_wrapper, "_sentry_agent_span") @@ -324,7 +363,7 @@ async def _execute_final_output( span = getattr(context_wrapper, "_sentry_agent_span", None) if span: update_invoke_agent_span( - span=span, context=context_wrapper, agent=agent, output=final_output + span=span, usage=context_wrapper.usage, agent=agent, output=final_output ) span.__exit__(None, None, None) delattr(context_wrapper, "_sentry_agent_span") diff --git a/sentry_sdk/integrations/openai_agents/patches/runner.py b/sentry_sdk/integrations/openai_agents/patches/runner.py index f091c5dfb9..e3ac7301c1 100644 --- a/sentry_sdk/integrations/openai_agents/patches/runner.py +++ b/sentry_sdk/integrations/openai_agents/patches/runner.py @@ -11,6 +11,7 @@ from ..spans import ( agent_workflow_span, execute_tool_span, + invoke_agent_span, update_execute_tool_span, update_invoke_agent_span, ) @@ -27,7 +28,7 @@ if TYPE_CHECKING: from typing import Any, AsyncIterator, Callable - from agents import Agent, Tool, ToolContext + from agents import Agent, AgentHookContext, Tool, ToolContext TContext = TypeVar("TContext") @@ -35,11 +36,42 @@ class _SentryRunHooks(RunHooks[TContext]): # type: ignore[misc] """ - Responsible for creating and managing Execute Tool spans. These spans are - stored on the ToolContext reference that is shared between `on_tool_start()` - and `on_tool_end()` + Responsible for creating Execute Tool and Invoke Agent spans. Both + spans can also be closed by hooks, and Execute Tool spans are exclusively + closed in the hook. + + Execute Tool spans are stored on the ToolContext reference that is shared + between `on_tool_start()` and `on_tool_end()`. + + Invoke Agent spans are stored as an instance variable on the hooks + instance. This works because agent invocations in a run are always + sequential. """ + async def on_agent_start( + self, + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + ) -> "None": + self._sentry_invoke_agent_span = invoke_agent_span(agent, context.turn_input) + + async def on_agent_end( + self, + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + output: "Any", + ) -> "None": + span = getattr(self, "_sentry_invoke_agent_span", None) + if span is not None: + update_invoke_agent_span( + span=span, + usage=context.usage, + agent=agent, + output=output, + ) + del self._sentry_invoke_agent_span + span.__exit__(None, None, None) + async def on_tool_start( self, context: "ToolContext[TContext]", @@ -88,11 +120,33 @@ def _patch_run_hooks(hooks: "RunHooks[TContext]") -> None: if is_already_patched: return + original_on_agent_start = hooks.on_agent_start + original_on_agent_end = hooks.on_agent_end + original_on_tool_start = hooks.on_tool_start original_on_tool_end = hooks.on_tool_end sentry_hooks = _SentryRunHooks() # type: ignore[var-annotated] + @wraps(original_on_agent_start) + async def on_agent_start( + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + ) -> "None": + with capture_internal_exceptions(): + await sentry_hooks.on_agent_start(context, agent) + await original_on_agent_start(context, agent) + + @wraps(original_on_agent_end) + async def on_agent_end( + context: "AgentHookContext[TContext]", + agent: "Agent[TContext]", + output: "Any", + ) -> "None": + with capture_internal_exceptions(): + await sentry_hooks.on_agent_end(context, agent, output) + await original_on_agent_end(context, agent, output) + @wraps(original_on_tool_start) async def on_tool_start( context: "ToolContext[TContext]", agent: "Agent[TContext]", tool: "Tool" @@ -113,6 +167,10 @@ async def on_tool_end( await original_on_tool_end(context, agent, tool, result) hooks._sentry_is_patched = True + + hooks.on_agent_start = on_agent_start + hooks.on_agent_end = on_agent_end + hooks.on_tool_start = on_tool_start hooks.on_tool_end = on_tool_end @@ -136,7 +194,8 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": if hooks is not None: _patch_run_hooks(hooks=hooks) else: - kwargs["hooks"] = _SentryRunHooks() + hooks = _SentryRunHooks() + kwargs["hooks"] = hooks # Isolate each workflow so that when agents are run in asyncio tasks they # don't touch each other's scopes @@ -175,7 +234,20 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": _capture_exception(exc) context_wrapper = getattr(exc.run_data, "context_wrapper", None) - if context_wrapper is not None: + if context_wrapper is not None and use_run_hooks: + invoke_agent_span = getattr( + hooks, "_sentry_invoke_agent_span", None + ) + + if invoke_agent_span is not None: + update_invoke_agent_span( + span=invoke_agent_span, + usage=context_wrapper.usage, + agent=agent, + ) + del hooks._sentry_invoke_agent_span + invoke_agent_span.__exit__(*exc_info) + elif context_wrapper is not None: invoke_agent_span = getattr( context_wrapper, "_sentry_agent_span", None ) @@ -192,7 +264,7 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": ): update_invoke_agent_span( span=invoke_agent_span, - context=context_wrapper, + usage=context_wrapper.usage, agent=agent, ) @@ -208,20 +280,21 @@ async def wrapper(*args: "Any", **kwargs: "Any") -> "Any": _capture_exception(exc) reraise(*exc_info) - invoke_agent_span = getattr( - run_result.context_wrapper, "_sentry_agent_span", None - ) - if not invoke_agent_span: - return run_result - - update_invoke_agent_span( - span=invoke_agent_span, - context=run_result.context_wrapper, - agent=agent, - ) - - invoke_agent_span.__exit__(None, None, None) - delattr(run_result.context_wrapper, "_sentry_agent_span") + if not use_run_hooks: + invoke_agent_span = getattr( + run_result.context_wrapper, "_sentry_agent_span", None + ) + if not invoke_agent_span: + return run_result + + update_invoke_agent_span( + span=invoke_agent_span, + usage=run_result.context_wrapper.usage, + agent=agent, + ) + + invoke_agent_span.__exit__(None, None, None) + delattr(run_result.context_wrapper, "_sentry_agent_span") return run_result return wrapper diff --git a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py index c21145ac4a..06a8fe0c88 100644 --- a/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py +++ b/sentry_sdk/integrations/openai_agents/spans/invoke_agent.py @@ -20,13 +20,15 @@ from ..utils import _set_agent_data, _set_usage_data if TYPE_CHECKING: - from typing import Any, Union + from typing import Any, Optional, Union import agents + from agents import TResponseInputItem + from agents.usage import Usage def invoke_agent_span( - context: "agents.RunContextWrapper", agent: "agents.Agent", kwargs: "dict[str, Any]" + agent: "agents.Agent", turn_input: "Optional[list[TResponseInputItem]]" ) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: @@ -64,12 +66,11 @@ def invoke_agent_span( } ) - original_input = kwargs.get("original_input") - if original_input is not None: + if turn_input is not None: message = ( - original_input - if isinstance(original_input, str) - else safe_serialize(original_input) + turn_input + if isinstance(turn_input, str) + else safe_serialize(turn_input) ) messages.append( { @@ -102,13 +103,12 @@ def invoke_agent_span( def update_invoke_agent_span( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", - context: "agents.RunContextWrapper", + usage: "Optional[Usage]", agent: "agents.Agent", output: "Any" = None, ) -> None: - # Add aggregated usage data from context_wrapper - if hasattr(context, "usage"): - _set_usage_data(span, context.usage) + if usage is not None: + _set_usage_data(span, usage) if should_send_default_pii(): set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False)