Ultra File Manager

Edit: /opt/hc_python/lib/python3.12/site-packages/sentry_sdk/integrations/pydantic_ai/utils.py

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)

Ukuran: 7.0 KB