sentry-sdk 0.18.0__py2.py3-none-any.whl → 2.46.0__py2.py3-none-any.whl

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -0,0 +1,246 @@
+import sentry_sdk
+from sentry_sdk.ai.utils import set_data_normalized
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.utils import safe_serialize
+
+from ..consts import SPAN_ORIGIN
+from ..utils import (
+    _set_agent_data,
+    _set_available_tools,
+    _set_model_data,
+    _should_send_prompts,
+    _get_model_name,
+    get_current_agent,
+    get_is_streaming,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, List, Dict
+    from pydantic_ai.usage import RequestUsage  # type: ignore
+
+try:
+    from pydantic_ai.messages import (  # type: ignore
+        BaseToolCallPart,
+        BaseToolReturnPart,
+        SystemPromptPart,
+        UserPromptPart,
+        TextPart,
+        ThinkingPart,
+    )
+except ImportError:
+    # Fallback if these classes are not available
+    BaseToolCallPart = None
+    BaseToolReturnPart = None
+    SystemPromptPart = None
+    UserPromptPart = None
+    TextPart = None
+    ThinkingPart = None
+
+
+def _set_usage_data(span, usage):
+    # type: (sentry_sdk.tracing.Span, RequestUsage) -> None
+    """Set token usage data on a span."""
+    if usage is None:
+        return
+
+    if hasattr(usage, "input_tokens") and usage.input_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens)
+
+    if hasattr(usage, "output_tokens") and usage.output_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens)
+
+    if hasattr(usage, "total_tokens") and usage.total_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)
+
+
+def _set_input_messages(span, messages):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set input messages data on a span."""
+    if not _should_send_prompts():
+        return
+
+    if not messages:
+        return
+
+    try:
+        formatted_messages = []
+        system_prompt = None
+
+        # Extract system prompt from any ModelRequest with instructions
+        for msg in messages:
+            if hasattr(msg, "instructions") and msg.instructions:
+                system_prompt = msg.instructions
+                break
+
+        # Add system prompt as first message if present
+        if system_prompt:
+            formatted_messages.append(
+                {"role": "system", "content": [{"type": "text", "text": system_prompt}]}
+            )
+
+        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):
+                        role = "system"
+                    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 = []  # type: 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})
+                                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 = {"role": role}  # type: Dict[str, Any]
+                        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:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, formatted_messages, unpack=False
+            )
+    except Exception:
+        # If we fail to format messages, just skip it
+        pass
+
+
+def _set_output_data(span, response):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set output data on a span."""
+    if not _should_send_prompts():
+        return
+
+    if not response:
+        return
+
+    set_data_normalized(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:
+                span.set_data(
+                    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, agent, model, model_settings):
+    # type: (Any, Any, Any, Any) -> sentry_sdk.tracing.Span
+    """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 = 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_agent_data(span, agent)
+    _set_model_data(span, model, model_settings)
+
+    # Set streaming flag from contextvar
+    span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming())
+
+    # 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, model_response):
+    # type: (sentry_sdk.tracing.Span, 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)

sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py

@@ -0,0 +1,49 @@
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.utils import safe_serialize
+
+from ..consts import SPAN_ORIGIN
+from ..utils import _set_agent_data, _should_send_prompts
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+
+def execute_tool_span(tool_name, tool_args, agent, tool_type="function"):
+    # type: (str, Any, Any, str) -> sentry_sdk.tracing.Span
+    """Create a span for tool execution.
+
+    Args:
+        tool_name: The name of the tool being executed
+        tool_args: The arguments passed to the tool
+        agent: The agent executing the tool
+        tool_type: The type of tool ("function" for regular tools, "mcp" for MCP services)
+    """
+    span = sentry_sdk.start_span(
+        op=OP.GEN_AI_EXECUTE_TOOL,
+        name=f"execute_tool {tool_name}",
+        origin=SPAN_ORIGIN,
+    )
+
+    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
+    span.set_data(SPANDATA.GEN_AI_TOOL_TYPE, tool_type)
+    span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name)
+
+    _set_agent_data(span, agent)
+
+    if _should_send_prompts() and tool_args is not None:
+        span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args))
+
+    return span
+
+
+def update_execute_tool_span(span, result):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Update the execute tool span with the result."""
+    if not span:
+        return
+
+    if _should_send_prompts() and result is not None:
+        span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result))

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

@@ -0,0 +1,112 @@
+import sentry_sdk
+from sentry_sdk.ai.utils import get_start_span_function, set_data_normalized
+from sentry_sdk.consts import OP, SPANDATA
+
+from ..consts import SPAN_ORIGIN
+from ..utils import (
+    _set_agent_data,
+    _set_available_tools,
+    _set_model_data,
+    _should_send_prompts,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+def invoke_agent_span(user_prompt, agent, model, model_settings, is_streaming=False):
+    # type: (Any, Any, Any, Any, bool) -> sentry_sdk.tracing.Span
+    """Create a span for invoking the agent."""
+    # Determine agent name for span
+    name = "agent"
+    if agent and getattr(agent, "name", None):
+        name = agent.name
+
+    span = get_start_span_function()(
+        op=OP.GEN_AI_INVOKE_AGENT,
+        name=f"invoke_agent {name}",
+        origin=SPAN_ORIGIN,
+    )
+
+    span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+
+    _set_agent_data(span, agent)
+    _set_model_data(span, model, model_settings)
+    _set_available_tools(span, agent)
+
+    # 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"})
+                if content:
+                    messages.append(
+                        {
+                            "content": content,
+                            "role": "user",
+                        }
+                    )
+
+        if messages:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False
+            )
+
+    return span
+
+
+def update_invoke_agent_span(span, output):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Update and close the invoke agent span."""
+    if span and _should_send_prompts() and output:
+        set_data_normalized(
+            span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False
+        )

sentry_sdk/integrations/pydantic_ai/utils.py

@@ -0,0 +1,223 @@
+import sentry_sdk
+from contextvars import ContextVar
+from sentry_sdk.consts import SPANDATA
+from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
+from sentry_sdk.utils import event_from_exception, safe_serialize
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+
+# 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("pydantic_ai_agent_context_stack", default=[])  # type: ContextVar[list[dict[str, Any]]]
+
+
+def push_agent(agent, is_streaming=False):
+    # type: (Any, bool) -> None
+    """Push an agent context onto the stack along with its streaming flag."""
+    stack = _agent_context_stack.get().copy()
+    stack.append({"agent": agent, "is_streaming": is_streaming})
+    _agent_context_stack.set(stack)
+
+
+def pop_agent():
+    # type: () -> None
+    """Pop an agent context from the stack."""
+    stack = _agent_context_stack.get().copy()
+    if stack:
+        stack.pop()
+    _agent_context_stack.set(stack)
+
+
+def get_current_agent():
+    # type: () -> Any
+    """Get the current agent from the contextvar stack."""
+    stack = _agent_context_stack.get()
+    if stack:
+        return stack[-1]["agent"]
+    return None
+
+
+def get_is_streaming():
+    # type: () -> bool
+    """Get the streaming flag from the contextvar stack."""
+    stack = _agent_context_stack.get()
+    if stack:
+        return stack[-1].get("is_streaming", False)
+    return False
+
+
+def _should_send_prompts():
+    # type: () -> bool
+    """
+    Check if prompts should be sent to Sentry.
+
+    This checks both send_default_pii and the include_prompts integration setting.
+    """
+    if not should_send_default_pii():
+        return False
+
+    from . import PydanticAIIntegration
+
+    # Get the integration instance from the client
+    integration = sentry_sdk.get_client().get_integration(PydanticAIIntegration)
+
+    if integration is None:
+        return False
+
+    return getattr(integration, "include_prompts", False)
+
+
+def _set_agent_data(span, agent):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set agent-related data on a span.
+
+    Args:
+        span: The span to set data on
+        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:
+        span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name)
+
+
+def _get_model_name(model_obj):
+    # type: (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)
+
+
+def _set_model_data(span, model, model_settings):
+    # type: (sentry_sdk.tracing.Span, Any, Any) -> None
+    """Set model-related data on a span.
+
+    Args:
+        span: The span to set data on
+        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"):
+            span.set_data(SPANDATA.GEN_AI_SYSTEM, model_obj.system)
+
+        # Set model name
+        model_name = _get_model_name(model_obj)
+        if model_name:
+            span.set_data(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:
+                    span.set_data(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:
+                        span.set_data(spandata_key, value)
+
+
+def _set_available_tools(span, agent):
+    # type: (sentry_sdk.tracing.Span, Any) -> None
+    """Set available tools data on a span from an agent's function toolset.
+
+    Args:
+        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:
+            span.set_data(
+                SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools)
+            )
+    except Exception:
+        # If we can't extract tools, just skip it
+        pass
+
+
+def _capture_exception(exc):
+    # type: (Any) -> None
+    set_span_errored()
+
+    event, hint = event_from_exception(
+        exc,
+        client_options=sentry_sdk.get_client().options,
+        mechanism={"type": "pydantic_ai", "handled": False},
+    )
+    sentry_sdk.capture_event(event, hint=hint)