Skip to content
189 changes: 182 additions & 7 deletions sentry_sdk/integrations/openai_agents/spans/ai_client.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,47 @@
import json
from typing import TYPE_CHECKING

import sentry_sdk
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing_utils import has_span_streaming_enabled
from sentry_sdk.utils import safe_serialize

try:
from agents import (
CodeInterpreterTool,
ComputerTool,
FileSearchTool,
FunctionTool,
HostedMCPTool,
ImageGenerationTool,
LocalShellTool,
WebSearchTool,
)
except ImportError:
raise DidNotEnable("OpenAI Agents not installed")

try:
from agents import ApplyPatchTool, ShellTool
except ImportError:
ShellTool = None
ApplyPatchTool = None

try:
from agents import ToolSearchTool
except ImportError:
ToolSearchTool = None

try:
from agents import CustomTool
except ImportError:
CustomTool = None

try:
from agents import ProgrammaticToolCallingTool
except ImportError:
ProgrammaticToolCallingTool = None


from ..consts import SPAN_ORIGIN
from ..utils import (
Expand All @@ -17,7 +54,146 @@
if TYPE_CHECKING:
from typing import Any, Optional, Union

from agents import Agent
from agents import Agent, Tool

from sentry_sdk._types import ToolDefinition


def _transform_tool_definitions(tools: "list[Tool]") -> "list[ToolDefinition]":
"""
Transform tool definitions to the schema used by the "gen_ai.tool.definitions" attribute.
Includes special handling for tools where the type includes a name, description or parameters.

Note: These are generated from the same OpenAPI spec as Responses API tools.
"""
if not isinstance(tools, list):
return []

tool_definitions: "list[ToolDefinition]" = []
for tool in tools:
if isinstance(tool, FunctionTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
"description": tool.description,
"parameters": tool.params_json_schema,
}
)
continue

if isinstance(tool, WebSearchTool):
tool_definitions.append(
{
"type": "web_search",
"name": tool.name,
}
)
continue

if isinstance(tool, FileSearchTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if isinstance(tool, ComputerTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if CustomTool is not None and isinstance(tool, CustomTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
"description": tool.description,
}
)
continue

if isinstance(tool, HostedMCPTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if ApplyPatchTool is not None and isinstance(tool, ApplyPatchTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if ShellTool is not None and isinstance(tool, ShellTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if isinstance(tool, ImageGenerationTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if isinstance(tool, CodeInterpreterTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if isinstance(tool, LocalShellTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if ToolSearchTool is not None and isinstance(tool, ToolSearchTool):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

if ProgrammaticToolCallingTool is not None and isinstance(
tool, ProgrammaticToolCallingTool
):
tool_definitions.append(
{
"type": "function",
"name": tool.name,
}
)
continue

return tool_definitions


def ai_client_span(
Expand Down Expand Up @@ -56,11 +232,10 @@ def ai_client_span(

_set_agent_data(span, agent)

if len(agent.tools) > 0:
set_on_span(
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
safe_serialize([vars(tool) for tool in agent.tools]),
)
set_on_span(
SPANDATA.GEN_AI_TOOL_DEFINITIONS,
json.dumps(_transform_tool_definitions(agent.tools)),
Comment thread
sentry-warden[bot] marked this conversation as resolved.
)

_set_input_data(span, get_response_kwargs)

Expand Down
10 changes: 0 additions & 10 deletions sentry_sdk/integrations/openai_agents/spans/invoke_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,6 @@ def invoke_agent_span(
SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent",
},
)

set_on_span = span.set_attribute
else:
start_span_function = get_start_span_function()
span = start_span_function(
Expand All @@ -51,8 +49,6 @@ def invoke_agent_span(

span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")

set_on_span = span.set_data

if should_send_default_pii():
messages = []
if agent.instructions:
Expand Down Expand Up @@ -101,12 +97,6 @@ def invoke_agent_span(

_set_agent_data(span, agent)

if len(agent.tools) > 0:
set_on_span(
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
safe_serialize([vars(tool) for tool in agent.tools]),
)

return span


Expand Down
Loading
Loading