From 06aaf97e103b5d12e468c1f7118ca4de40756c2b Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Wed, 12 Aug 2026 15:44:44 -0400 Subject: [PATCH] ref(pydantic-ai): Move pydantic-ai object reads behind an extraction layer Introduce _extract.py as the module that concentrates reads of pydantic-ai object internals (private attributes, message part classes, version-dependent shapes) behind typed accessors returning plain data structures. Span modules and patches now consume those accessors, collapsing the duplicated message formatters and blob serializers into one implementation. Review-driven fixes folded in: response access in extract_response_model_name is now exception-safe (AgentRunResult.response can raise), token usage reporting goes through the shared record_token_usage helper, model-name resolution for the chat span name matches gen_ai.request.model resolution, and unknown model settings are skipped instead of raising KeyError. --- .../integrations/pydantic_ai/_extract.py | 488 ++++++++++++++++++ .../pydantic_ai/patches/graph_nodes.py | 30 +- .../integrations/pydantic_ai/patches/tools.py | 11 +- .../pydantic_ai/spans/ai_client.py | 226 +------- .../pydantic_ai/spans/invoke_agent.py | 87 +--- .../integrations/pydantic_ai/spans/utils.py | 79 +-- sentry_sdk/integrations/pydantic_ai/utils.py | 139 ++--- .../pydantic_ai/test_pydantic_ai.py | 26 +- 8 files changed, 585 insertions(+), 501 deletions(-) create mode 100644 sentry_sdk/integrations/pydantic_ai/_extract.py diff --git a/sentry_sdk/integrations/pydantic_ai/_extract.py b/sentry_sdk/integrations/pydantic_ai/_extract.py new file mode 100644 index 0000000000..0ddb7f4104 --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/_extract.py @@ -0,0 +1,488 @@ +"""Typed accessors for reading data off pydantic-ai objects. + +This module concentrates reads of pydantic-ai object internals (including +private attributes and version-dependent shapes) so that upstream library +changes are absorbed here rather than throughout the integration. The one +exception is control-flow state read at the patch points themselves (e.g. +ModelRequestNode._did_stream in patches/graph_nodes.py and Tool.tool_def in +patches/tools.py); everything else consumes the plain data structures +returned here. +""" + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from sentry_sdk._types import BLOB_DATA_SUBSTITUTE +from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX +from sentry_sdk.ai.utils import get_modality_from_mime_type +from sentry_sdk.utils import safe_serialize + +try: + from pydantic_ai.messages import ( + BaseToolCallPart, + BaseToolReturnPart, + BinaryContent, + ImageUrl, + SystemPromptPart, + TextPart, + ThinkingPart, + ) +except ImportError: + # Fallback if these classes are not available + BaseToolCallPart = None # type: ignore[misc,assignment] + BaseToolReturnPart = None # type: ignore[misc,assignment] + BinaryContent = None # type: ignore[misc,assignment] + ImageUrl = None # type: ignore[misc,assignment] + SystemPromptPart = None # type: ignore[misc,assignment] + TextPart = None # type: ignore[misc,assignment] + ThinkingPart = None # type: ignore[misc,assignment] + +if TYPE_CHECKING: + from typing import Any, Dict, List, Optional + + from pydantic_ai.messages import ModelMessage, ModelResponse + from pydantic_ai.messages import SystemPromptPart as SystemPromptPartType + + from sentry_sdk import _types + + +@dataclass +class ModelInfo: + name: "Optional[str]" = None + system: "Optional[str]" = None + settings: "Dict[str, Any]" = field(default_factory=dict) + + +@dataclass +class UsageInfo: + input_tokens: "Optional[int]" = None + cache_read_tokens: "Optional[int]" = None + cache_write_tokens: "Optional[int]" = None + output_tokens: "Optional[int]" = None + total_tokens: "Optional[int]" = None + + +# Model settings that get mirrored onto spans; values are read with dict +# access first because ModelSettings is a TypedDict (dict at runtime). +MODEL_SETTING_NAMES = ( + "max_tokens", + "temperature", + "top_p", + "frequency_penalty", + "presence_penalty", +) + + +def get_model_name(model_obj: "Any") -> "Optional[str]": + """Extract model name from a model object.""" + if not model_obj: + return None + + if hasattr(model_obj, "model_name"): + return model_obj.model_name + elif hasattr(model_obj, "name"): + try: + return model_obj.name() + except Exception: + return str(model_obj) + elif isinstance(model_obj, str): + return model_obj + else: + return str(model_obj) + + +def extract_model_settings(settings: "Any") -> "Dict[str, Any]": + """Extract known model settings as a plain dict of non-None values.""" + extracted: "Dict[str, Any]" = {} + if not settings: + return extracted + + for setting_name in MODEL_SETTING_NAMES: + if isinstance(settings, dict): + value = settings.get(setting_name) + else: + # Fallback for object-style settings + value = getattr(settings, setting_name, None) + if value is not None: + extracted[setting_name] = value + + return extracted + + +def extract_model_info( + model: "Any", model_settings: "Any", agent: "Any" +) -> "ModelInfo": + """Extract model name, provider system, and settings. + + Falls back to the agent's model and model_settings when the explicit + arguments are not provided. + """ + model_obj = model + if not model_obj and agent and hasattr(agent, "model"): + model_obj = agent.model + + info = ModelInfo() + if model_obj: + info.system = getattr(model_obj, "system", None) + info.name = get_model_name(model_obj) + + settings = model_settings + if not settings and agent and hasattr(agent, "model_settings"): + settings = agent.model_settings + info.settings = extract_model_settings(settings) + + return info + + +def extract_agent_name(agent: "Any") -> "Optional[str]": + if agent and hasattr(agent, "name") and agent.name: + return agent.name + return None + + +def extract_available_tools(agent: "Any") -> "Optional[List[Dict[str, Any]]]": + """Extract the agent's available tool definitions from its function toolset.""" + if not agent or not hasattr(agent, "_function_toolset"): + return None + + try: + tools = [] + if hasattr(agent._function_toolset, "tools"): + for tool_name, tool in agent._function_toolset.tools.items(): + tool_info: "Dict[str, Any]" = {"name": tool_name} + + if hasattr(tool, "function_schema"): + schema = tool.function_schema + if getattr(schema, "description", None): + tool_info["description"] = schema.description + + if getattr(schema, "json_schema", None): + tool_info["parameters"] = schema.json_schema + + tools.append(tool_info) + + return tools or None + except Exception: + # If we can't extract tools, just skip it + return None + + +def extract_usage(usage: "Any") -> "Optional[UsageInfo]": + """Extract token usage counts. + + Works with both RequestUsage (single request) and RunUsage (agent run) + objects from pydantic-ai; note the library uses cache_read_tokens / + cache_write_tokens naming. + """ + if usage is None: + return None + + return UsageInfo( + input_tokens=getattr(usage, "input_tokens", None), + cache_read_tokens=getattr(usage, "cache_read_tokens", None), + cache_write_tokens=getattr(usage, "cache_write_tokens", None), + output_tokens=getattr(usage, "output_tokens", None), + total_tokens=getattr(usage, "total_tokens", None), + ) + + +def serialize_image_url_item(item: "Any") -> "Dict[str, Any]": + """Serialize an ImageUrl content item for span data. + + For data URLs containing base64-encoded images, the content is redacted. + For regular HTTP URLs, the URL string is preserved. + """ + url = str(item.url) + data_url_match = DATA_URL_BASE64_REGEX.match(url) + + if data_url_match: + return { + "type": "image", + "content": BLOB_DATA_SUBSTITUTE, + } + + return { + "type": "image", + "content": url, + } + + +def serialize_binary_content_item(item: "Any") -> "Dict[str, Any]": + """Serialize a BinaryContent item for span data, redacting the blob data.""" + return { + "type": "blob", + "modality": get_modality_from_mime_type(item.media_type), + "mime_type": item.media_type, + "content": BLOB_DATA_SUBSTITUTE, + } + + +def _collect_system_instructions( + messages: "List[ModelMessage]", +) -> "tuple[List[SystemPromptPartType], List[str]]": + permanent_instructions = [] + current_instructions = [] + + for msg in messages: + if hasattr(msg, "parts"): + for part in msg.parts: + if SystemPromptPart is not None and isinstance(part, SystemPromptPart): + permanent_instructions.append(part) + + if hasattr(msg, "instructions") and msg.instructions is not None: + current_instructions.append(msg.instructions) + + return permanent_instructions, current_instructions + + +def extract_system_instructions( + messages: "List[ModelMessage]", +) -> "List[_types.TextPart]": + """Extract permanent and per-request system instructions as text parts.""" + permanent_instructions, current_instructions = _collect_system_instructions( + messages + ) + + text_parts: "List[_types.TextPart]" = [ + { + "type": "text", + "content": instruction.content, + } + for instruction in permanent_instructions + ] + + text_parts.extend( + { + "type": "text", + "content": instruction, + } + for instruction in current_instructions + ) + + return text_parts + + +def extract_request_messages(messages: "Any") -> "List[Dict[str, Any]]": + """Extract a conversation history as gen_ai-format message dicts. + + System prompt parts are skipped; they are reported separately via + extract_system_instructions(). + """ + formatted_messages = [] + + for msg in messages: + if not hasattr(msg, "parts"): + continue + + for part in msg.parts: + role = "user" + # Use isinstance checks with proper base classes + if SystemPromptPart is not None and isinstance(part, SystemPromptPart): + continue + elif ( + (TextPart is not None and isinstance(part, TextPart)) + or (ThinkingPart is not None and isinstance(part, ThinkingPart)) + or (BaseToolCallPart is not None and isinstance(part, BaseToolCallPart)) + ): + role = "assistant" + elif BaseToolReturnPart is not None and isinstance( + part, BaseToolReturnPart + ): + role = "tool" + + content: "List[Dict[str, Any] | str]" = [] + tool_calls = None + tool_call_id = None + + # Handle ToolCallPart (assistant requesting tool use) + if BaseToolCallPart is not None and isinstance(part, BaseToolCallPart): + tool_call_data = {} + if hasattr(part, "tool_name"): + tool_call_data["name"] = part.tool_name + if hasattr(part, "args"): + tool_call_data["arguments"] = safe_serialize(part.args) + if tool_call_data: + tool_calls = [tool_call_data] + # Handle ToolReturnPart (tool result) + elif BaseToolReturnPart is not None and isinstance( + part, BaseToolReturnPart + ): + if hasattr(part, "tool_name"): + tool_call_id = part.tool_name + if hasattr(part, "content"): + content.append({"type": "text", "text": str(part.content)}) + # Handle regular content + elif hasattr(part, "content"): + if isinstance(part.content, str): + content.append({"type": "text", "text": part.content}) + elif isinstance(part.content, list): + for item in part.content: + if isinstance(item, str): + content.append({"type": "text", "text": item}) + elif ImageUrl is not None and isinstance(item, ImageUrl): + content.append(serialize_image_url_item(item)) + elif BinaryContent is not None and isinstance( + item, BinaryContent + ): + content.append(serialize_binary_content_item(item)) + else: + content.append(safe_serialize(item)) + else: + content.append({"type": "text", "text": str(part.content)}) + # Add message if we have content or tool calls + if content or tool_calls: + message: "Dict[str, Any]" = {"role": role} + if content: + message["content"] = content + if tool_calls: + message["tool_calls"] = tool_calls + if tool_call_id: + message["tool_call_id"] = tool_call_id + formatted_messages.append(message) + + return formatted_messages + + +def extract_response_parts( + response: "ModelResponse", +) -> "List[_types.TextPart | _types.ReasoningPart | _types.ToolCallPart]": + """Extract a model response's parts as gen_ai-format output parts.""" + parts: "List[_types.TextPart | _types.ReasoningPart | _types.ToolCallPart]" = [] + + if not hasattr(response, "parts"): + return parts + + for part in response.parts: + if ( + TextPart is not None + and isinstance(part, TextPart) + and hasattr(part, "content") + ): + parts.append({"type": "text", "content": part.content}) + + elif ThinkingPart is not None and isinstance(part, ThinkingPart): + parts.append( + { + "type": "reasoning", + "content": part.content, + } + ) + + elif BaseToolCallPart is not None and isinstance(part, BaseToolCallPart): + tool_part: "_types.ToolCallPart" = {"type": "tool_call"} + if hasattr(part, "tool_name"): + tool_part["name"] = part.tool_name + if hasattr(part, "args"): + tool_part["arguments"] = safe_serialize(part.args) + parts.append(tool_part) + + return parts + + +def extract_agent_prompt_messages( + agent: "Any", user_prompt: "Any" +) -> "List[Dict[str, Any]]": + """Extract an agent's static system prompts plus the user prompt as + gen_ai-format message dicts.""" + messages: "List[Dict[str, Any]]" = [] + + # Add system prompts (both system_prompt and instructions) + system_texts = [] + + if agent: + system_prompts = getattr(agent, "_system_prompts", None) or [] + for prompt in system_prompts: + if isinstance(prompt, str): + system_texts.append(prompt) + + instructions = getattr(agent, "_instructions", None) + if instructions: + if isinstance(instructions, str): + system_texts.append(instructions) + elif isinstance(instructions, (list, tuple)): + for instr in instructions: + if isinstance(instr, str): + system_texts.append(instr) + elif callable(instr): + # Skip dynamic/callable instructions + pass + + for system_text in system_texts: + messages.append( + { + "content": [{"text": system_text, "type": "text"}], + "role": "system", + } + ) + + if user_prompt: + if isinstance(user_prompt, str): + messages.append( + { + "content": [{"text": user_prompt, "type": "text"}], + "role": "user", + } + ) + elif isinstance(user_prompt, list): + # Handle list of user content + content = [] + for item in user_prompt: + if isinstance(item, str): + content.append({"text": item, "type": "text"}) + elif ImageUrl is not None and isinstance(item, ImageUrl): + content.append(serialize_image_url_item(item)) + elif BinaryContent is not None and isinstance(item, BinaryContent): + content.append(serialize_binary_content_item(item)) + if content: + messages.append( + { + "content": content, + "role": "user", + } + ) + + return messages + + +def extract_response_model_name(result: "Any") -> "Optional[str]": + """Extract the responding model's name from an agent run result.""" + try: + # Accessing .response can itself raise (e.g. AgentRunResult raises + # ValueError when the run produced no ModelResponse), so the access + # must live inside the try block. + response = result.response + if hasattr(response, "model_name") and response.model_name: + return response.model_name + except Exception: + # If response access fails, continue without the model name + pass + return None + + +def extract_graph_request_data(node: "Any", ctx: "Any") -> "tuple[List[Any], Any, Any]": + """Extract (messages, model, model_settings) from a ModelRequestNode and + its graph context, for the legacy (pre-hooks) instrumentation path.""" + model = None + model_settings = None + if hasattr(ctx, "deps"): + model = getattr(ctx.deps, "model", None) + model_settings = getattr(ctx.deps, "model_settings", None) + + # Build full message list: history + current request + messages = [] + if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"): + messages.extend(ctx.state.message_history) + + current_request = getattr(node, "request", None) + if current_request: + messages.append(current_request) + + return messages, model, model_settings + + +def extract_tool_call_args(call: "Any") -> "Dict[str, Any]": + """Extract a tool call's arguments as a dict.""" + try: + return call.args_as_dict() + except Exception: + return call.args if isinstance(call.args, dict) else {} diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py index a177628773..bb4ace697d 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py @@ -3,6 +3,7 @@ from sentry_sdk.integrations import DidNotEnable +from .._extract import extract_graph_request_data from ..spans import ( ai_client_span, update_ai_client_span, @@ -21,31 +22,6 @@ from pydantic_ai.messages import ModelResponse -def _extract_span_data(node: "Any", ctx: "Any") -> "tuple[list[Any], Any, Any]": - """Extract common data needed for creating chat spans. - - Returns: - Tuple of (messages, model, model_settings) - """ - # Extract model and settings from context - model = None - model_settings = None - if hasattr(ctx, "deps"): - model = getattr(ctx.deps, "model", None) - model_settings = getattr(ctx.deps, "model_settings", None) - - # Build full message list: history + current request - messages = [] - if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"): - messages.extend(ctx.state.message_history) - - current_request = getattr(node, "request", None) - if current_request: - messages.append(current_request) - - return messages, model, model_settings - - def _patch_graph_nodes() -> None: """ Patches the graph node execution to create appropriate spans. @@ -67,7 +43,7 @@ async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": if did_stream or cached_result is not None: return await original_model_request_run(self, ctx) - messages, model, model_settings = _extract_span_data(self, ctx) + messages, model, model_settings = extract_graph_request_data(self, ctx) with ai_client_span(messages, None, model, model_settings) as span: result = await original_model_request_run(self, ctx) @@ -101,7 +77,7 @@ async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any": yield stream return - messages, model, model_settings = _extract_span_data(self, ctx) + messages, model, model_settings = extract_graph_request_data(self, ctx) # Create chat span for streaming request with ai_client_span(messages, None, model, model_settings) as span: diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index 958d729fbb..3cffcf3826 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -6,6 +6,7 @@ from sentry_sdk.integrations import DidNotEnable from sentry_sdk.utils import capture_internal_exceptions, reraise +from .._extract import extract_tool_call_args from ..spans import execute_tool_span, update_execute_tool_span from ..utils import _capture_exception, get_current_agent @@ -52,10 +53,7 @@ async def wrapped_execute_tool_call( agent = get_current_agent() if agent and tool: - try: - args_dict = call.args_as_dict() - except Exception: - args_dict = call.args if isinstance(call.args, dict) else {} + args_dict = extract_tool_call_args(call) # Create execute_tool span # Nesting is handled by isolation_scope() to ensure proper parent-child relationships @@ -125,10 +123,7 @@ async def wrapped_call_tool( agent = get_current_agent() if agent and tool: - try: - args_dict = call.args_as_dict() - except Exception: - args_dict = call.args if isinstance(call.args, dict) else {} + args_dict = extract_tool_call_args(call) # Create execute_tool span # Nesting is handled by isolation_scope() to ensure proper parent-child relationships diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 2991889059..6b795cfa50 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -13,11 +13,15 @@ has_span_streaming_enabled, should_truncate_gen_ai_input, ) -from sentry_sdk.utils import safe_serialize +from .._extract import ( + extract_model_info, + extract_request_messages, + extract_response_parts, + extract_system_instructions, +) from ..consts import SPAN_ORIGIN from ..utils import ( - _get_model_name, _set_agent_data, _set_available_tools, _set_model_data, @@ -25,84 +29,15 @@ get_current_agent, get_is_streaming, ) -from .utils import ( - _serialize_binary_content_item, - _serialize_image_url_item, - _set_usage_data, -) +from .utils import _set_usage_data if TYPE_CHECKING: - from typing import Any, Dict, List, Optional, Union + from typing import Any, Optional, Union - from pydantic_ai.messages import ModelMessage, ModelResponse, SystemPromptPart + from pydantic_ai.messages import ModelResponse - from sentry_sdk import _types from sentry_sdk.traces import StreamedSpan -try: - from pydantic_ai.messages import ( - BaseToolCallPart, - BaseToolReturnPart, - BinaryContent, - ImageUrl, - SystemPromptPart, - TextPart, - ThinkingPart, - UserPromptPart, - ) -except ImportError: - # Fallback if these classes are not available - BaseToolCallPart = None # type: ignore[misc,assignment] - BaseToolReturnPart = None # type: ignore[misc,assignment] - SystemPromptPart = None # type: ignore[misc,assignment] - UserPromptPart = None # type: ignore[misc,assignment] - TextPart = None # type: ignore[misc,assignment] - ThinkingPart = None # type: ignore[misc,assignment] - BinaryContent = None # type: ignore[misc,assignment] - ImageUrl = None # type: ignore[misc,assignment] - ThinkingPart = None # type: ignore[misc,assignment] - - -def _transform_system_instructions( - permanent_instructions: "list[SystemPromptPart]", - current_instructions: "list[str]", -) -> "list[_types.TextPart]": - text_parts: "list[_types.TextPart]" = [ - { - "type": "text", - "content": instruction.content, - } - for instruction in permanent_instructions - ] - - text_parts.extend( - { - "type": "text", - "content": instruction, - } - for instruction in current_instructions - ) - - return text_parts - - -def _get_system_instructions( - messages: "list[ModelMessage]", -) -> "tuple[list[SystemPromptPart], list[str]]": - permanent_instructions = [] - current_instructions = [] - - for msg in messages: - if hasattr(msg, "parts"): - for part in msg.parts: - if SystemPromptPart is not None and isinstance(part, SystemPromptPart): - permanent_instructions.append(part) - - if hasattr(msg, "instructions") and msg.instructions is not None: - current_instructions.append(msg.instructions) - - return permanent_instructions, current_instructions - def _set_input_messages( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any" @@ -114,97 +49,16 @@ def _set_input_messages( if not messages: return - permanent_instructions, current_instructions = _get_system_instructions(messages) - if len(permanent_instructions) > 0 or len(current_instructions) > 0: + system_instructions = extract_system_instructions(messages) + if system_instructions: _set_span_data_attribute( span, SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps( - _transform_system_instructions( - permanent_instructions, current_instructions - ) - ), + json.dumps(system_instructions), ) try: - formatted_messages = [] - - for msg in messages: - if hasattr(msg, "parts"): - for part in msg.parts: - role = "user" - # Use isinstance checks with proper base classes - if SystemPromptPart is not None and isinstance( - part, SystemPromptPart - ): - continue - elif ( - (TextPart is not None and isinstance(part, TextPart)) - or (ThinkingPart is not None and isinstance(part, ThinkingPart)) - or ( - BaseToolCallPart is not None - and isinstance(part, BaseToolCallPart) - ) - ): - role = "assistant" - elif BaseToolReturnPart is not None and isinstance( - part, BaseToolReturnPart - ): - role = "tool" - - content: "List[Dict[str, Any] | str]" = [] - tool_calls = None - tool_call_id = None - - # Handle ToolCallPart (assistant requesting tool use) - if BaseToolCallPart is not None and isinstance( - part, BaseToolCallPart - ): - tool_call_data = {} - if hasattr(part, "tool_name"): - tool_call_data["name"] = part.tool_name - if hasattr(part, "args"): - tool_call_data["arguments"] = safe_serialize(part.args) - if tool_call_data: - tool_calls = [tool_call_data] - # Handle ToolReturnPart (tool result) - elif BaseToolReturnPart is not None and isinstance( - part, BaseToolReturnPart - ): - if hasattr(part, "tool_name"): - tool_call_id = part.tool_name - if hasattr(part, "content"): - content.append({"type": "text", "text": str(part.content)}) - # Handle regular content - elif hasattr(part, "content"): - if isinstance(part.content, str): - content.append({"type": "text", "text": part.content}) - elif isinstance(part.content, list): - for item in part.content: - if isinstance(item, str): - content.append({"type": "text", "text": item}) - elif ImageUrl is not None and isinstance( - item, ImageUrl - ): - content.append(_serialize_image_url_item(item)) - elif BinaryContent is not None and isinstance( - item, BinaryContent - ): - content.append(_serialize_binary_content_item(item)) - else: - content.append(safe_serialize(item)) - else: - content.append({"type": "text", "text": str(part.content)}) - # Add message if we have content or tool calls - if content or tool_calls: - message: "Dict[str, Any]" = {"role": role} - if content: - message["content"] = content - if tool_calls: - message["tool_calls"] = tool_calls - if tool_call_id: - message["tool_call_id"] = tool_call_id - formatted_messages.append(message) + formatted_messages = extract_request_messages(messages) if formatted_messages: normalized_messages = normalize_message_roles(formatted_messages) @@ -240,42 +94,13 @@ def _set_output_data( ) try: - if hasattr(response, "parts"): - parts: "list[Union[_types.TextPart, _types.ReasoningPart, _types.ToolCallPart]]" = [] - - for part in response.parts: - if ( - TextPart is not None - and isinstance(part, TextPart) - and hasattr(part, "content") - ): - parts.append({"type": "text", "content": part.content}) - - elif ThinkingPart is not None and isinstance(part, ThinkingPart): - parts.append( - { - "type": "reasoning", - "content": part.content, - } - ) - - elif BaseToolCallPart is not None and isinstance( - part, BaseToolCallPart - ): - tool_part: "_types.ToolCallPart" = {"type": "tool_call"} - if hasattr(part, "tool_name"): - tool_part["name"] = part.tool_name - if hasattr(part, "args"): - tool_part["arguments"] = safe_serialize(part.args) - parts.append(tool_part) - - if parts: - _set_span_data_attribute( - span, - SPANDATA.GEN_AI_OUTPUT_MESSAGES, - json.dumps([{"role": "assistant", "parts": parts}]), - ) - + parts = extract_response_parts(response) + if parts: + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps([{"role": "assistant", "parts": parts}]), + ) except Exception: # If we fail to format output, just skip it pass @@ -292,12 +117,11 @@ def ai_client_span( model: Model object model_settings: Model settings """ - # Determine model name for span name - model_obj = model - if agent and hasattr(agent, "model"): - model_obj = agent.model - - model_name = _get_model_name(model_obj) or "unknown" + # Determine model name for span name, resolving the same way as + # _set_model_data so the span name and gen_ai.request.model agree + model_name = ( + extract_model_info(model, None, agent or get_current_agent()).name or "unknown" + ) span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index f5cc5c4f97..bf896b3a6d 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -14,6 +14,7 @@ should_truncate_gen_ai_input, ) +from .._extract import extract_agent_prompt_messages, extract_response_model_name from ..consts import SPAN_ORIGIN from ..utils import ( _set_agent_data, @@ -21,22 +22,12 @@ _set_model_data, _should_send_prompts, ) -from .utils import ( - _serialize_binary_content_item, - _serialize_image_url_item, -) if TYPE_CHECKING: from typing import Any, Union from sentry_sdk.traces import StreamedSpan -try: - from pydantic_ai.messages import BinaryContent, ImageUrl -except ImportError: - BinaryContent = None # type: ignore[misc,assignment] - ImageUrl = None # type: ignore[misc,assignment] - def invoke_agent_span( user_prompt: "Any", @@ -76,66 +67,7 @@ def invoke_agent_span( # Add user prompt and system prompts if available and prompts are enabled if _should_send_prompts(): - messages = [] - - # Add system prompts (both instructions and system_prompt) - system_texts = [] - - if agent: - # Check for system_prompt - system_prompts = getattr(agent, "_system_prompts", None) or [] - for prompt in system_prompts: - if isinstance(prompt, str): - system_texts.append(prompt) - - # Check for instructions (stored in _instructions) - instructions = getattr(agent, "_instructions", None) - if instructions: - if isinstance(instructions, str): - system_texts.append(instructions) - elif isinstance(instructions, (list, tuple)): - for instr in instructions: - if isinstance(instr, str): - system_texts.append(instr) - elif callable(instr): - # Skip dynamic/callable instructions - pass - - # Add all system texts as system messages - for system_text in system_texts: - messages.append( - { - "content": [{"text": system_text, "type": "text"}], - "role": "system", - } - ) - - # Add user prompt - if user_prompt: - if isinstance(user_prompt, str): - messages.append( - { - "content": [{"text": user_prompt, "type": "text"}], - "role": "user", - } - ) - elif isinstance(user_prompt, list): - # Handle list of user content - content = [] - for item in user_prompt: - if isinstance(item, str): - content.append({"text": item, "type": "text"}) - elif ImageUrl is not None and isinstance(item, ImageUrl): - content.append(_serialize_image_url_item(item)) - elif BinaryContent is not None and isinstance(item, BinaryContent): - content.append(_serialize_binary_content_item(item)) - if content: - messages.append( - { - "content": content, - "role": "user", - } - ) + messages = extract_agent_prompt_messages(agent, user_prompt) if messages: normalized_messages = normalize_message_roles(messages) @@ -171,13 +103,8 @@ def update_invoke_agent_span( ) # Set model name from response if available - if hasattr(result, "response"): - try: - response = result.response - if hasattr(response, "model_name") and response.model_name: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name - ) - except Exception: - # If response access fails, continue without setting model name - pass + response_model_name = extract_response_model_name(result) + if response_model_name: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model_name + ) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py index ae4a899f4b..98947fdd5f 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/utils.py @@ -3,50 +3,18 @@ from typing import TYPE_CHECKING import sentry_sdk -from sentry_sdk._types import BLOB_DATA_SUBSTITUTE -from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX -from sentry_sdk.ai.utils import _set_span_data_attribute, get_modality_from_mime_type -from sentry_sdk.consts import SPANDATA +from sentry_sdk.ai.monitoring import record_token_usage + +from .._extract import extract_usage if TYPE_CHECKING: - from typing import Any, Dict, Union + from typing import Union from pydantic_ai.usage import RequestUsage, RunUsage from sentry_sdk.traces import StreamedSpan -def _serialize_image_url_item(item: "Any") -> "Dict[str, Any]": - """Serialize an ImageUrl content item for span data. - - For data URLs containing base64-encoded images, the content is redacted. - For regular HTTP URLs, the URL string is preserved. - """ - url = str(item.url) - data_url_match = DATA_URL_BASE64_REGEX.match(url) - - if data_url_match: - return { - "type": "image", - "content": BLOB_DATA_SUBSTITUTE, - } - - return { - "type": "image", - "content": url, - } - - -def _serialize_binary_content_item(item: "Any") -> "Dict[str, Any]": - """Serialize a BinaryContent item for span data, redacting the blob data.""" - return { - "type": "blob", - "modality": get_modality_from_mime_type(item.media_type), - "mime_type": item.media_type, - "content": BLOB_DATA_SUBSTITUTE, - } - - def _set_usage_data( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", usage: "Union[RequestUsage, RunUsage]", @@ -60,34 +28,15 @@ def _set_usage_data( span: The Sentry span to set data on. usage: RequestUsage or RunUsage object containing token usage information. """ - if usage is None: + usage_info = extract_usage(usage) + if usage_info is None: return - if hasattr(usage, "input_tokens") and usage.input_tokens is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens - ) - - # Pydantic AI uses cache_read_tokens (not input_tokens_cached) - if hasattr(usage, "cache_read_tokens") and usage.cache_read_tokens is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens - ) - - # Pydantic AI uses cache_write_tokens (not input_tokens_cache_write) - if hasattr(usage, "cache_write_tokens") and usage.cache_write_tokens is not None: - _set_span_data_attribute( - span, - SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, - usage.cache_write_tokens, - ) - - if hasattr(usage, "output_tokens") and usage.output_tokens is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens - ) - - if hasattr(usage, "total_tokens") and usage.total_tokens is not None: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens - ) + record_token_usage( + span, + input_tokens=usage_info.input_tokens, + input_tokens_cached=usage_info.cache_read_tokens, + input_tokens_cache_write=usage_info.cache_write_tokens, + output_tokens=usage_info.output_tokens, + total_tokens=usage_info.total_tokens, + ) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 560e4715fb..cf24caa6da 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -7,12 +7,23 @@ from sentry_sdk.scope import should_send_default_pii from sentry_sdk.utils import event_from_exception, safe_serialize +from ._extract import extract_agent_name, extract_available_tools, extract_model_info + if TYPE_CHECKING: - from typing import Any, Optional, Union + from typing import Any, Union from sentry_sdk.traces import StreamedSpan +_MODEL_SETTINGS_SPANDATA = { + "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, + "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, + "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, + "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, +} + + # Store the current agent context in a contextvar for re-entrant safety # Using a list as a stack to support nested agent calls _agent_context_stack: "ContextVar[list[dict[str, Any]]]" = ContextVar( @@ -81,38 +92,9 @@ def _set_agent_data( agent: Agent object (can be None, will try to get from contextvar if not provided) """ # Extract agent name from agent object or contextvar - agent_obj = agent - if not agent_obj: - # Try to get from contextvar - agent_obj = get_current_agent() - - if agent_obj and hasattr(agent_obj, "name") and agent_obj.name: - _set_span_data_attribute(span, SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) - - -def _get_model_name(model_obj: "Any") -> "Optional[str]": - """Extract model name from a model object. - - Args: - model_obj: Model object to extract name from - - Returns: - Model name string or None if not found - """ - if not model_obj: - return None - - if hasattr(model_obj, "model_name"): - return model_obj.model_name - elif hasattr(model_obj, "name"): - try: - return model_obj.name() - except Exception: - return str(model_obj) - elif isinstance(model_obj, str): - return model_obj - else: - return str(model_obj) + agent_name = extract_agent_name(agent or get_current_agent()) + if agent_name: + _set_span_data_attribute(span, SPANDATA.GEN_AI_AGENT_NAME, agent_name) def _set_model_data( @@ -127,51 +109,18 @@ def _set_model_data( model: Model object (can be None, will try to get from agent if not provided) model_settings: Model settings (can be None, will try to get from agent if not provided) """ - # Try to get agent from contextvar if we need it - agent_obj = get_current_agent() - - # Extract model information - model_obj = model - if not model_obj and agent_obj and hasattr(agent_obj, "model"): - model_obj = agent_obj.model - - if model_obj: - # Set system from model - if hasattr(model_obj, "system"): - _set_span_data_attribute(span, SPANDATA.GEN_AI_SYSTEM, model_obj.system) - - # Set model name - model_name = _get_model_name(model_obj) - if model_name: - _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model_name) - - # Extract model settings - settings = model_settings - if not settings and agent_obj and hasattr(agent_obj, "model_settings"): - settings = agent_obj.model_settings - - if settings: - settings_map = { - "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, - "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, - "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, - "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, - "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, - } - - # ModelSettings is a TypedDict (dict at runtime), so use dict access - if isinstance(settings, dict): - for setting_name, spandata_key in settings_map.items(): - value = settings.get(setting_name) - if value is not None: - _set_span_data_attribute(span, spandata_key, value) - else: - # Fallback for object-style settings - for setting_name, spandata_key in settings_map.items(): - if hasattr(settings, setting_name): - value = getattr(settings, setting_name) - if value is not None: - _set_span_data_attribute(span, spandata_key, value) + model_info = extract_model_info(model, model_settings, get_current_agent()) + + if model_info.system is not None: + _set_span_data_attribute(span, SPANDATA.GEN_AI_SYSTEM, model_info.system) + + if model_info.name: + _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model_info.name) + + for setting_name, value in model_info.settings.items(): + spandata_key = _MODEL_SETTINGS_SPANDATA.get(setting_name) + if spandata_key is not None: + _set_span_data_attribute(span, spandata_key, value) def _set_available_tools( @@ -183,35 +132,11 @@ def _set_available_tools( span: The span to set data on agent: Agent object with _function_toolset attribute """ - if not agent or not hasattr(agent, "_function_toolset"): - return - - try: - tools = [] - # Get tools from the function toolset - if hasattr(agent._function_toolset, "tools"): - for tool_name, tool in agent._function_toolset.tools.items(): - tool_info = {"name": tool_name} - - # Add description from function_schema if available - if hasattr(tool, "function_schema"): - schema = tool.function_schema - if getattr(schema, "description", None): - tool_info["description"] = schema.description - - # Add parameters from json_schema - if getattr(schema, "json_schema", None): - tool_info["parameters"] = schema.json_schema - - tools.append(tool_info) - - if tools: - _set_span_data_attribute( - span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) - ) - except Exception: - # If we can't extract tools, just skip it - pass + tools = extract_available_tools(agent) + if tools: + _set_span_data_attribute( + span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) def _capture_exception(exc: "Any", handled: bool = False) -> None: diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index fc6ccc5088..bcb0d36ab2 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -2304,7 +2304,7 @@ async def test_model_name_extraction_with_callable(sentry_init, capture_items): """ from unittest.mock import MagicMock - from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name + from sentry_sdk.integrations.pydantic_ai._extract import get_model_name sentry_init( integrations=[PydanticAIIntegration()], @@ -2318,7 +2318,7 @@ async def test_model_name_extraction_with_callable(sentry_init, capture_items): mock_model.name = lambda: "custom-model-name" # Get model name - should call the callable name() - result = _get_model_name(mock_model) + result = get_model_name(mock_model) # Should return the result from callable assert result == "custom-model-name" @@ -2331,7 +2331,7 @@ async def test_model_name_extraction_fallback_to_str(sentry_init, capture_items) """ from unittest.mock import MagicMock - from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name + from sentry_sdk.integrations.pydantic_ai._extract import get_model_name sentry_init( integrations=[PydanticAIIntegration()], @@ -2345,7 +2345,7 @@ async def test_model_name_extraction_fallback_to_str(sentry_init, capture_items) del mock_model.model_name # Get model name - should fall back to str() - result = _get_model_name(mock_model) + result = get_model_name(mock_model) # Should return string representation assert result is not None @@ -3331,11 +3331,11 @@ async def test_set_input_messages_without_prompts(sentry_init, capture_items): @pytest.mark.asyncio async def test_get_model_name_with_exception_in_callable(sentry_init, capture_items): """ - Test that _get_model_name handles exceptions in name() callable. + Test that get_model_name handles exceptions in name() callable. """ from unittest.mock import MagicMock - from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name + from sentry_sdk.integrations.pydantic_ai._extract import get_model_name sentry_init( integrations=[PydanticAIIntegration()], @@ -3347,7 +3347,7 @@ async def test_get_model_name_with_exception_in_callable(sentry_init, capture_it mock_model.name = MagicMock(side_effect=Exception("Error")) # Should fall back to str() - result = _get_model_name(mock_model) + result = get_model_name(mock_model) # Should return something (str fallback) assert result is not None @@ -3356,9 +3356,9 @@ async def test_get_model_name_with_exception_in_callable(sentry_init, capture_it @pytest.mark.asyncio async def test_get_model_name_with_string_model(sentry_init, capture_items): """ - Test that _get_model_name handles string models. + Test that get_model_name handles string models. """ - from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name + from sentry_sdk.integrations.pydantic_ai._extract import get_model_name sentry_init( integrations=[PydanticAIIntegration()], @@ -3366,7 +3366,7 @@ async def test_get_model_name_with_string_model(sentry_init, capture_items): ) # Pass a string as model - result = _get_model_name("gpt-4") + result = get_model_name("gpt-4") # Should return the string assert result == "gpt-4" @@ -3375,9 +3375,9 @@ async def test_get_model_name_with_string_model(sentry_init, capture_items): @pytest.mark.asyncio async def test_get_model_name_with_none(sentry_init, capture_items): """ - Test that _get_model_name handles None model. + Test that get_model_name handles None model. """ - from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name + from sentry_sdk.integrations.pydantic_ai._extract import get_model_name sentry_init( integrations=[PydanticAIIntegration()], @@ -3385,7 +3385,7 @@ async def test_get_model_name_with_none(sentry_init, capture_items): ) # Pass None - result = _get_model_name(None) + result = get_model_name(None) # Should return None assert result is None