diff --git a/sentry_sdk/integrations/pydantic_ai/__init__.py b/sentry_sdk/integrations/pydantic_ai/__init__.py index db21861e71..58a2a7bfbc 100644 --- a/sentry_sdk/integrations/pydantic_ai/__init__.py +++ b/sentry_sdk/integrations/pydantic_ai/__init__.py @@ -31,6 +31,13 @@ def register_hooks(hooks: "Hooks") -> None: """ Creates hooks for chat model calls and register the hooks by adding the hooks to the `capabilities` argument passed to `Agent.__init__()`. + + The chat span opened in on_request is stored in the run's `RunContext.metadata` + dict, which pydantic-ai shares by reference between the hooks of one run. This + keeps span pairing correct per run (even for overlapping runs in one task) and + covers every entry point that fires request hooks (including `Agent.iter()`, + which the Agent.run/run_stream wrappers never see). It requires seeding a + metadata dict in `patched_init` below when the user did not provide one. """ @hooks.on.before_model_request @@ -41,12 +48,17 @@ async def on_request( if not isinstance(run_context_metadata, dict): return request_context - span = ai_client_span( - messages=request_context.messages, - agent=None, - model=request_context.model, - model_settings=request_context.model_settings, - ) + span = None + with capture_internal_exceptions(): + span = ai_client_span( + messages=request_context.messages, + agent=None, + model=request_context.model, + model_settings=request_context.model_settings, + ) + + if span is None: + return request_context run_context_metadata["_sentry_span"] = span span.__enter__() @@ -68,7 +80,8 @@ async def on_response( if span is None: return response - update_ai_client_span(span, response) + with capture_internal_exceptions(): + update_ai_client_span(span, response) span.__exit__(None, None, None) return response @@ -116,15 +129,15 @@ class PydanticAIIntegration(Integration): Typical interaction with the library: 1. The user creates an Agent instance with configuration, including system instructions sent to every model call. 2. The user calls `Agent.run()` or `Agent.run_stream()` to start an agent run. The latter can be used to incrementally receive progress. - - Each run invocation has `RunContext` objects that are passed to the library hooks. 3. In a loop, the agent repeatedly calls the model, maintaining a conversation history that includes previous messages and tool results, which is passed to each call. Internally, Pydantic AI maintains an execution graph in which ModelRequestNode are responsible for model calls, including retries. - Hooks using the decorators provided by `pydantic_ai.capabilities` create and manage spans for model calls when these hooks are available (newer library versions). - The span is created in `on_request` and stored in the metadata of the `RunContext` object shared with `on_response` and `on_error`. + Hooks using the decorators provided by `pydantic_ai.capabilities` create and manage spans for model calls when these hooks are available (newer library versions); + older versions are instrumented by patching the graph nodes directly (see patches/graph_nodes.py). - The metadata dictionary on the RunContext instance is initialized with `{"_sentry_span": None}` in the `_create_run_wrapper()` and `_create_streaming_wrapper()` wrappers that - instrument `Agent.run()` and `Agent.run_stream()`, respectively. A non-empty dictionary is required for the metadata object to be a shared reference between hooks. + The wrappers around `Agent.run()` and `Agent.run_stream()` track each in-flight run on a contextvar stack (see _run_context.py); the tool patches and span + helpers read the current agent from there. The request hooks pair each chat span with its model request through the run's `RunContext.metadata` dict + (see register_hooks), which stays correct per run and also covers entry points the wrappers don't instrument, such as `Agent.iter()`. """ identifier = "pydantic_ai" diff --git a/sentry_sdk/integrations/pydantic_ai/_extract.py b/sentry_sdk/integrations/pydantic_ai/_extract.py index 0ddb7f4104..ff0a85ae65 100644 --- a/sentry_sdk/integrations/pydantic_ai/_extract.py +++ b/sentry_sdk/integrations/pydantic_ai/_extract.py @@ -15,6 +15,7 @@ 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.consts import SPANDATA from sentry_sdk.utils import safe_serialize try: @@ -53,24 +54,15 @@ class ModelInfo: 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", -) +# Single source of truth for which model settings get mirrored onto spans +# and which span attribute each one maps to. +MODEL_SETTINGS_TO_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, +} def get_model_name(model_obj: "Any") -> "Optional[str]": @@ -97,7 +89,7 @@ def extract_model_settings(settings: "Any") -> "Dict[str, Any]": if not settings: return extracted - for setting_name in MODEL_SETTING_NAMES: + for setting_name in MODEL_SETTINGS_TO_SPANDATA: if isinstance(settings, dict): value = settings.get(setting_name) else: @@ -167,8 +159,8 @@ def extract_available_tools(agent: "Any") -> "Optional[List[Dict[str, Any]]]": return None -def extract_usage(usage: "Any") -> "Optional[UsageInfo]": - """Extract token usage counts. +def extract_usage_kwargs(usage: "Any") -> "Optional[Dict[str, Optional[int]]]": + """Extract token usage counts as record_token_usage keyword arguments. Works with both RequestUsage (single request) and RunUsage (agent run) objects from pydantic-ai; note the library uses cache_read_tokens / @@ -177,13 +169,13 @@ def extract_usage(usage: "Any") -> "Optional[UsageInfo]": 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), - ) + return { + "input_tokens": getattr(usage, "input_tokens", None), + "input_tokens_cached": getattr(usage, "cache_read_tokens", None), + "input_tokens_cache_write": 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]": diff --git a/sentry_sdk/integrations/pydantic_ai/_run_context.py b/sentry_sdk/integrations/pydantic_ai/_run_context.py new file mode 100644 index 0000000000..4253ac50da --- /dev/null +++ b/sentry_sdk/integrations/pydantic_ai/_run_context.py @@ -0,0 +1,63 @@ +"""Run-scoped state shared between the agent wrappers and the model/tool +instrumentation. + +One agent run corresponds to one AgentRun on the contextvar stack. The stack +makes nested agent calls re-entrant safe, and the context manager guarantees +push/pop pairing. +""" + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Iterator, Optional + + +@dataclass +class AgentRun: + """State for one in-flight agent run.""" + + agent: "Any" + is_streaming: bool = False + + +_agent_run_stack: "ContextVar[tuple[AgentRun, ...]]" = ContextVar( + "pydantic_ai_agent_run_stack", default=() +) + + +def current_agent_run() -> "Optional[AgentRun]": + stack = _agent_run_stack.get() + return stack[-1] if stack else None + + +def get_current_agent() -> "Any": + run = current_agent_run() + return run.agent if run is not None else None + + +def get_is_streaming() -> bool: + run = current_agent_run() + return run.is_streaming if run is not None else False + + +@contextmanager +def agent_run_scope(agent: "Any", is_streaming: bool = False) -> "Iterator[AgentRun]": + """Track an agent run on the contextvar stack for the duration of the + with block.""" + run = AgentRun(agent=agent, is_streaming=is_streaming) + token = _agent_run_stack.set(_agent_run_stack.get() + (run,)) + try: + yield run + finally: + try: + _agent_run_stack.reset(token) + except (LookupError, ValueError): + # A streaming run's context manager can be exited in a different + # asyncio task (and therefore a different Context) than it was + # entered in, in which case the token cannot be reset. The stack + # entry only lives in the entering task's context copy, so there + # is nothing to clean up. + pass diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 864c83a506..cde48326a0 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -1,4 +1,5 @@ import sys +from contextlib import ExitStack from functools import wraps from typing import TYPE_CHECKING @@ -6,8 +7,9 @@ from sentry_sdk.integrations import DidNotEnable from sentry_sdk.utils import capture_internal_exceptions, reraise +from .._run_context import agent_run_scope from ..spans import invoke_agent_span, update_invoke_agent_span -from ..utils import _capture_exception, pop_agent, push_agent +from ..utils import _capture_exception try: from pydantic_ai.agent import Agent @@ -18,6 +20,14 @@ from typing import Any, Callable, Optional, Union +def _extract_run_params( + args: "tuple[Any, ...]", kwargs: "dict[str, Any]" +) -> "tuple[Any, Any, Any]": + """Extract (user_prompt, model, model_settings) from a run call.""" + user_prompt = kwargs.get("user_prompt") or (args[0] if args else None) + return user_prompt, kwargs.get("model"), kwargs.get("model_settings") + + class _StreamingContextManagerWrapper: """Wrapper for streaming methods that return async context managers.""" @@ -28,39 +38,37 @@ def __init__( user_prompt: "Any", model: "Any", model_settings: "Any", - is_streaming: bool = True, ) -> None: self.agent = agent self.original_ctx_manager = original_ctx_manager self.user_prompt = user_prompt self.model = model self.model_settings = model_settings - self.is_streaming = is_streaming - self._isolation_scope: "Any" = None + self._contexts: "Optional[ExitStack]" = None self._span: "Optional[Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]]" = None self._result: "Any" = None async def __aenter__(self) -> "Any": - # Set up isolation scope and invoke_agent span - self._isolation_scope = sentry_sdk.isolation_scope() - self._isolation_scope.__enter__() - - # Create invoke_agent span (will be closed in __aexit__) - self._span = invoke_agent_span( - self.user_prompt, - self.agent, - self.model, - self.model_settings, - self.is_streaming, - ) - self._span.__enter__() - - # Push agent to contextvar stack after span is successfully created and entered - # This ensures proper pairing with pop_agent() in __aexit__ even if exceptions occur - push_agent(self.agent, self.is_streaming) + # Isolation scope, invoke_agent span, and run-context tracking are all + # owned by one ExitStack so they unwind together in __aexit__ (or + # right here if entering the original context manager fails). + with ExitStack() as contexts: + contexts.enter_context(sentry_sdk.isolation_scope()) + span = invoke_agent_span( + self.user_prompt, + self.agent, + self.model, + self.model_settings, + is_streaming=True, + ) + contexts.enter_context(span) + self._span = span + contexts.enter_context(agent_run_scope(self.agent, is_streaming=True)) + + result = await self.original_ctx_manager.__aenter__() + + self._contexts = contexts.pop_all() - # Enter the original context manager - result = await self.original_ctx_manager.__aenter__() self._result = result return result @@ -73,27 +81,13 @@ async def __aexit__(self, exc_type: "Any", exc_val: "Any", exc_tb: "Any") -> Non if exc_type is None and self._result and self._span is not None: update_invoke_agent_span(self._span, self._result) finally: - # Pop agent from contextvar stack - pop_agent() + if self._contexts is not None: + self._contexts.__exit__(exc_type, exc_val, exc_tb) - # Clean up invoke span - if self._span: - self._span.__exit__(exc_type, exc_val, exc_tb) - # Clean up isolation scope - if self._isolation_scope: - self._isolation_scope.__exit__(exc_type, exc_val, exc_tb) - - -def _create_run_wrapper( - original_func: "Callable[..., Any]", is_streaming: bool = False -) -> "Callable[..., Any]": +def _create_run_wrapper(original_func: "Callable[..., Any]") -> "Callable[..., Any]": """ Wraps the Agent.run method to create an invoke_agent span. - - Args: - original_func: The original run method - is_streaming: Whether this is a streaming method (for future use) """ from sentry_sdk.integrations.pydantic_ai import ( PydanticAIIntegration, @@ -101,42 +95,30 @@ def _create_run_wrapper( @wraps(original_func) async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - # Isolate each workflow so that when agents are run in asyncio tasks they - # don't touch each other's scopes - with sentry_sdk.isolation_scope(): - # Extract parameters for the span - user_prompt = kwargs.get("user_prompt") or (args[0] if args else None) - model = kwargs.get("model") - model_settings = kwargs.get("model_settings") + user_prompt, model, model_settings = _extract_run_params(args, kwargs) - if PydanticAIIntegration.using_request_hooks: - metadata = kwargs.get("metadata") - if metadata is None: - kwargs["metadata"] = {"_sentry_span": None} + if PydanticAIIntegration.using_request_hooks: + if kwargs.get("metadata") is None: + kwargs["metadata"] = {"_sentry_span": None} - # Create invoke_agent span + # Isolate each workflow so that when agents are run in asyncio tasks + # they don't touch each other's scopes + with sentry_sdk.isolation_scope(): with invoke_agent_span( - user_prompt, self, model, model_settings, is_streaming + user_prompt, self, model, model_settings, is_streaming=False ) as span: - # Push agent to contextvar stack after span is successfully created and entered - # This ensures proper pairing with pop_agent() in finally even if exceptions occur - push_agent(self, is_streaming) - - try: - result = await original_func(self, *args, **kwargs) - - # Update span with result - update_invoke_agent_span(span, result) - - return result - except Exception as exc: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - _capture_exception(exc) - reraise(*exc_info) - finally: - # Pop agent from contextvar stack - pop_agent() + with agent_run_scope(self, is_streaming=False): + try: + result = await original_func(self, *args, **kwargs) + + update_invoke_agent_span(span, result) + + return result + except Exception as exc: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + _capture_exception(exc) + reraise(*exc_info) return wrapper @@ -153,14 +135,10 @@ def _create_streaming_wrapper( @wraps(original_func) def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - # Extract parameters for the span - user_prompt = kwargs.get("user_prompt") or (args[0] if args else None) - model = kwargs.get("model") - model_settings = kwargs.get("model_settings") + user_prompt, model, model_settings = _extract_run_params(args, kwargs) if PydanticAIIntegration.using_request_hooks: - metadata = kwargs.get("metadata") - if metadata is None: + if kwargs.get("metadata") is None: kwargs["metadata"] = {"_sentry_span": None} # Call original function to get the context manager @@ -173,7 +151,6 @@ def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": user_prompt=user_prompt, model=model, model_settings=model_settings, - is_streaming=True, ) return wrapper @@ -183,8 +160,8 @@ def _patch_agent_run() -> None: """ Patches the Agent run methods to create spans for agent execution. - This patches both non-streaming (run, run_sync) and streaming - (run_stream, run_stream_events) methods. + This patches both the non-streaming (run) and streaming (run_stream) + entry points; run_sync delegates to run. """ # Store original methods @@ -192,7 +169,7 @@ def _patch_agent_run() -> None: original_run_stream = Agent.run_stream # Wrap and apply patches for non-streaming methods - Agent.run = _create_run_wrapper(original_run, is_streaming=False) # type: ignore[method-assign] + Agent.run = _create_run_wrapper(original_run) # type: ignore[method-assign] # Wrap and apply patches for streaming methods Agent.run_stream = _create_streaming_wrapper(original_run_stream) # type: ignore[method-assign] diff --git a/sentry_sdk/integrations/pydantic_ai/patches/tools.py b/sentry_sdk/integrations/pydantic_ai/patches/tools.py index 3cffcf3826..81c7cf80e7 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/tools.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/tools.py @@ -7,8 +7,9 @@ from sentry_sdk.utils import capture_internal_exceptions, reraise from .._extract import extract_tool_call_args +from .._run_context import get_current_agent from ..spans import execute_tool_span, update_execute_tool_span -from ..utils import _capture_exception, get_current_agent +from ..utils import _capture_exception if TYPE_CHECKING: from typing import Any diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 6b795cfa50..4720a33726 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -20,14 +20,13 @@ extract_response_parts, extract_system_instructions, ) +from .._run_context import get_current_agent, get_is_streaming from ..consts import SPAN_ORIGIN from ..utils import ( _set_agent_data, _set_available_tools, _set_model_data, _should_send_prompts, - get_current_agent, - get_is_streaming, ) from .utils import _set_usage_data @@ -49,15 +48,15 @@ def _set_input_messages( if not messages: return - system_instructions = extract_system_instructions(messages) - if system_instructions: - _set_span_data_attribute( - span, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps(system_instructions), - ) - try: + system_instructions = extract_system_instructions(messages) + if system_instructions: + _set_span_data_attribute( + span, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps(system_instructions), + ) + formatted_messages = extract_request_messages(messages) if formatted_messages: @@ -117,11 +116,10 @@ def ai_client_span( model: Model object model_settings: Model settings """ - # 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" - ) + # Resolve the agent once so the span name and every attribute derived + # below (gen_ai.request.model, agent data, available tools) agree + agent_obj = agent or get_current_agent() + model_name = extract_model_info(model, model_settings, agent_obj).name or "unknown" span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) if span_streaming: @@ -145,11 +143,10 @@ def ai_client_span( # 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) + _set_agent_data(span, agent_obj) + _set_model_data(span, model, model_settings, agent=agent_obj) # 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) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index bf896b3a6d..fb795339b0 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -62,7 +62,7 @@ def invoke_agent_span( span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") _set_agent_data(span, agent) - _set_model_data(span, model, model_settings) + _set_model_data(span, model, model_settings, agent=agent) _set_available_tools(span, agent) # Add user prompt and system prompts if available and prompts are enabled diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py index 98947fdd5f..2b23a459c6 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/utils.py @@ -5,7 +5,7 @@ import sentry_sdk from sentry_sdk.ai.monitoring import record_token_usage -from .._extract import extract_usage +from .._extract import extract_usage_kwargs if TYPE_CHECKING: from typing import Union @@ -28,15 +28,8 @@ def _set_usage_data( span: The Sentry span to set data on. usage: RequestUsage or RunUsage object containing token usage information. """ - usage_info = extract_usage(usage) - if usage_info is None: + usage_kwargs = extract_usage_kwargs(usage) + if usage_kwargs is None: return - 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, - ) + record_token_usage(span, **usage_kwargs) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index cf24caa6da..b4840603b9 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -1,4 +1,3 @@ -from contextvars import ContextVar from typing import TYPE_CHECKING import sentry_sdk @@ -7,7 +6,13 @@ 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 +from ._extract import ( + MODEL_SETTINGS_TO_SPANDATA, + extract_agent_name, + extract_available_tools, + extract_model_info, +) +from ._run_context import get_current_agent if TYPE_CHECKING: from typing import Any, Union @@ -15,53 +20,6 @@ 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( - "pydantic_ai_agent_context_stack", default=[] -) - - -def push_agent(agent: "Any", is_streaming: bool = False) -> 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() -> 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() -> "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() -> 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() -> bool: """ Check if prompts should be sent to Sentry. @@ -101,6 +59,7 @@ def _set_model_data( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", model: "Any", model_settings: "Any", + agent: "Any" = None, ) -> None: """Set model-related data on a span. @@ -108,8 +67,10 @@ def _set_model_data( 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) + agent: Agent to fall back to for model and settings (defaults to the + agent of the current run) """ - model_info = extract_model_info(model, model_settings, get_current_agent()) + model_info = extract_model_info(model, model_settings, agent or get_current_agent()) if model_info.system is not None: _set_span_data_attribute(span, SPANDATA.GEN_AI_SYSTEM, model_info.system) @@ -118,7 +79,7 @@ def _set_model_data( _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) + spandata_key = MODEL_SETTINGS_TO_SPANDATA.get(setting_name) if spandata_key is not None: _set_span_data_attribute(span, spandata_key, value)