import json from typing import TYPE_CHECKING import sentry_sdk from sentry_sdk.ai.utils import ( normalize_message_roles, set_data_normalized, truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import safe_serialize from ..consts import SPAN_ORIGIN from ..utils import ( _get_model_name, _set_agent_data, _set_available_tools, _set_model_data, _should_send_prompts, get_current_agent, get_is_streaming, ) from .utils import ( _serialize_binary_content_item, _serialize_image_url_item, _set_usage_data, ) if TYPE_CHECKING: from typing import Any, Dict, List, Union from pydantic_ai.messages import ModelMessage, SystemPromptPart # type: ignore from sentry_sdk._types import TextPart as SentryTextPart 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 BaseToolReturnPart = None SystemPromptPart = None UserPromptPart = None TextPart = None ThinkingPart = None BinaryContent = None ImageUrl = None def _transform_system_instructions( permanent_instructions: "list[SystemPromptPart]", current_instructions: "list[str]", ) -> "list[SentryTextPart]": text_parts: "list[SentryTextPart]" = [ { "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 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" ) -> None: """Set input messages data on a span.""" if not _should_send_prompts(): return if not messages: return permanent_instructions, current_instructions = _get_system_instructions(messages) if len(permanent_instructions) > 0 or len(current_instructions) > 0: if isinstance(span, StreamedSpan): span.set_attribute( SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, json.dumps( _transform_system_instructions( permanent_instructions, current_instructions ) ), ) else: span.set_data( SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, json.dumps( _transform_system_instructions( permanent_instructions, current_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 and isinstance(part, SystemPromptPart): continue elif ( (TextPart and isinstance(part, TextPart)) or (ThinkingPart and isinstance(part, ThinkingPart)) or (BaseToolCallPart and isinstance(part, BaseToolCallPart)) ): role = "assistant" elif BaseToolReturnPart 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 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 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 and isinstance(item, ImageUrl): content.append(_serialize_image_url_item(item)) elif BinaryContent 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) if formatted_messages: normalized_messages = normalize_message_roles(formatted_messages) client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() messages_data = ( normalized_messages if client.options.get("stream_gen_ai_spans", False) else truncate_and_annotate_messages(normalized_messages, span, scope) ) set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False ) except Exception: # If we fail to format messages, just skip it pass def _set_output_data( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", response: "Any" ) -> None: """Set output data on a span.""" if not _should_send_prompts(): return if not response: return set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) try: # Extract text from ModelResponse if hasattr(response, "parts"): texts = [] tool_calls = [] for part in response.parts: if TextPart and isinstance(part, TextPart) and hasattr(part, "content"): texts.append(part.content) elif BaseToolCallPart and isinstance(part, BaseToolCallPart): tool_call_data = { "type": "function", } if hasattr(part, "tool_name"): tool_call_data["name"] = part.tool_name if hasattr(part, "args"): tool_call_data["arguments"] = safe_serialize(part.args) tool_calls.append(tool_call_data) if texts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) if tool_calls: set_on_span( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) ) except Exception: # If we fail to format output, just skip it pass def ai_client_span( messages: "Any", agent: "Any", model: "Any", model_settings: "Any" ) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": """Create a span for an AI client call (model request). Args: messages: Full conversation history (list of messages) agent: Agent object 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" span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: span = sentry_sdk.traces.start_span( name=f"chat {model_name}", attributes={ "sentry.op": OP.GEN_AI_CHAT, "sentry.origin": SPAN_ORIGIN, SPANDATA.GEN_AI_OPERATION_NAME: "chat", SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(), }, ) else: span = sentry_sdk.start_span( op=OP.GEN_AI_CHAT, name=f"chat {model_name}", origin=SPAN_ORIGIN, ) span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") # Set streaming flag from contextvar span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming()) _set_agent_data(span, agent) _set_model_data(span, model, model_settings) # Add available tools if agent is available agent_obj = agent or get_current_agent() _set_available_tools(span, agent_obj) # Set input messages (full conversation history) if messages: _set_input_messages(span, messages) return span def update_ai_client_span( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", model_response: "Any" ) -> None: """Update the AI client span with response data.""" if not span: return # Set usage data if available if model_response and hasattr(model_response, "usage"): _set_usage_data(span, model_response.usage) # Set output data _set_output_data(span, model_response)