"""Instrumentation capability for OpenTelemetry/Logfire tracing of agent runs.""" from __future__ import annotations import warnings from collections.abc import Awaitable, Callable from dataclasses import KW_ONLY, dataclass, field, replace from typing import TYPE_CHECKING, Any, ClassVar, Literal from opentelemetry.baggage import set_baggage as _otel_set_baggage from opentelemetry.context import attach as _otel_attach, detach as _otel_detach from opentelemetry.trace import StatusCode from pydantic_core import ValidationError, to_json from pydantic_ai import _usage_attribution from pydantic_ai._instrumentation import ( DEFAULT_INSTRUMENTATION_VERSION, InstrumentationNames, MessageJsonCache, get_agent_run_baggage_attributes, get_instructions, has_stale_message_json, open_model_request_span, record_exception as _record_exception, record_uncaught_errors as _record_uncaught_errors, redact_binary_content, safe_to_json, serialize_any, time_to_first_chunk_ctx, ) from pydantic_ai._utils import UNSET, Unset from pydantic_ai.exceptions import ( ApprovalRequired, CallDeferred, MessageHistoryMutatedWarning, ModelRetry, ToolFailedError, ToolRetryError, ) from pydantic_ai.messages import ModelMessage, ModelResponse, RetryPromptPart, ToolCallPart, tool_return_ta from pydantic_ai.tools import ToolDefinition from pydantic_ai.usage import RunUsage from .abstract import ( AbstractCapability, CapabilityOrdering, RawToolArgs, ValidatedToolArgs, WrapModelRequestHandler, WrapOutputProcessHandler, WrapRunHandler, WrapToolExecuteHandler, ) if TYPE_CHECKING: from pydantic_ai._run_context import RunContext from pydantic_ai.models import ModelRequestContext, ModelRequestParameters from pydantic_ai.models.instrumented import InstrumentationSettings from pydantic_ai.output import OutputContext from pydantic_ai.run import AgentRunResult from pydantic_ai.tools import AgentDepsT def _default_settings() -> InstrumentationSettings: """Lazy import to avoid loading the OTel SDK eagerly at module import time.""" from pydantic_ai.models.instrumented import InstrumentationSettings return InstrumentationSettings() @dataclass class Instrumentation(AbstractCapability[Any]): """Capability that instruments agent runs with OpenTelemetry/Logfire tracing. When added to an agent via `capabilities=[Instrumentation(...)]`, this capability creates OpenTelemetry spans for the agent run, model requests, and tool executions. Other capabilities can add attributes to these spans using the OpenTelemetry API (`opentelemetry.trace.get_current_span().set_attribute(key, value)`). """ _safe_at_runtime: ClassVar[bool] = True """Workflow-side only — no toolsets, native tools, or model wrapping introduced — so safe to attach per-run even when a durability capability is bound. Internal flag read by the bundled durable-execution integrations. """ settings: InstrumentationSettings = field(default_factory=lambda: _default_settings()) """OTel/Logfire instrumentation settings. Defaults to `InstrumentationSettings()`, which uses the global `TracerProvider` (typically configured by `logfire.configure()`).""" # Per-run state (set in `for_run`, mutated by `wrap_model_request`). `for_run` # returns a shallow copy via `replace(self)` for per-run isolation. These fields # are updated as the run progresses and assume sequential model requests within # a run — if the agent loop ever issues concurrent model requests, accesses to # these fields would race. _agent_name: str = field(default='agent', repr=False, init=False) _new_message_index: int = field(default=0, repr=False, init=False) _run_usage: RunUsage = field(default_factory=RunUsage, repr=False, init=False) """Usage this run recorded while its span was open, credited by `_usage_attribution`. A nested run's `accumulate` replaces the active accumulator for the length of its own span, so what a delegate records is the delegate's; this holds only what this run recorded itself. """ _last_messages: list[ModelMessage] | None = field(default=None, repr=False, init=False) _last_model_request_parameters: ModelRequestParameters | None = field(default=None, repr=False, init=False) _last_formatted_instructions: str | None | Unset = field(default=UNSET, repr=False, init=False) """Last formatted instructions sent to the model, or `UNSET` before the first request.""" _variable_instructions: bool = field(default=False, repr=False, init=False) """Whether agent-level instructions varied across requests in this run.""" _: KW_ONLY id: str | None = 'instrumentation' """One-off: an agent has a single instrumentation configuration, so the id is fixed by default. Two of them resolve to one via [`combine`][pydantic_ai.capabilities.AbstractCapability.combine], which keeps the last. Pass a distinct `id` to keep both, or `id=None` for derived ids. """ _message_json_cache: MessageJsonCache = field(default_factory=MessageJsonCache, repr=False, init=False) """Per-run cache of input messages' serialized OTel JSON fragments (see `MessageJsonCache`). `for_run`'s `replace(self)` re-runs the factory, so each run starts with an empty cache that's discarded when the run ends.""" # Resolved once from `self.settings.version` in `__post_init__` and preserved across # `dataclasses.replace` calls in `for_run` (which only touches init=True fields). _instrumentation_names: InstrumentationNames = field( default_factory=lambda: InstrumentationNames.for_version(DEFAULT_INSTRUMENTATION_VERSION), repr=False, init=False, ) def __post_init__(self) -> None: self._instrumentation_names = InstrumentationNames.for_version(self.settings.version) def get_ordering(self) -> CapabilityOrdering: return CapabilityOrdering(position='outermost') @classmethod def from_spec( cls, *, include_binary_content: bool = True, include_content: bool = True, include_model_request_parameters: bool = True, version: Literal[2, 3, 4, 5, 6] = DEFAULT_INSTRUMENTATION_VERSION, use_aggregated_usage_attribute_names: bool = True, ) -> Instrumentation: """Build an `Instrumentation` capability from a YAML/JSON spec. Accepts every serializable [`InstrumentationSettings`][pydantic_ai.models.instrumented.InstrumentationSettings] option. The OTel `tracer_provider` and `meter_provider` fields can't be expressed in YAML and default to the global providers (typically configured via `logfire.configure()`). `id` is deliberately not accepted. An agent has one instrumentation configuration -- which is what the class-level default `id` says -- so there is nothing for a spec to name, and two `Instrumentation` capabilities resolve to one rather than colliding. YAML form: capabilities: - Instrumentation: {} # default settings - Instrumentation: version: 2 include_content: false """ from pydantic_ai.models.instrumented import InstrumentationSettings return cls( settings=InstrumentationSettings( include_binary_content=include_binary_content, include_content=include_content, include_model_request_parameters=include_model_request_parameters, version=version, use_aggregated_usage_attribute_names=use_aggregated_usage_attribute_names, ), ) async def for_run(self, ctx: RunContext[Any]) -> Instrumentation: """Return a fresh copy for per-run state isolation.""" inst = replace(self) inst._agent_name = (ctx.agent.name if ctx.agent else None) or 'agent' inst._new_message_index = len(ctx.messages) # Usage this run's span is accountable for, credited by `_usage_attribution` for as long as # the span is open in `wrap_run`; see `_run_span_end_attributes`. inst._run_usage = RunUsage() return inst # ------------------------------------------------------------------ # wrap_run — agent run span # ------------------------------------------------------------------ async def wrap_run( self, ctx: RunContext[AgentDepsT], *, handler: WrapRunHandler, ) -> AgentRunResult[Any]: # `RealtimeSession` owns its session and per-response spans; a second run span here would # duplicate the session's canonical `invoke_agent` span. See the capability-owned span # direction documented in `realtime/_session.py`. if ctx.realtime: return await handler() settings = self.settings names = self._instrumentation_names agent_name = self._agent_name span_attributes: dict[str, Any] = { 'model_name': ctx.model.model_name if ctx.model else 'no-model', 'agent_name': agent_name, 'gen_ai.agent.name': agent_name, 'gen_ai.agent.call.id': ctx.run_id or '', 'gen_ai.conversation.id': ctx.conversation_id or '', 'gen_ai.operation.name': 'invoke_agent', 'logfire.msg': f'{agent_name} run', } if ctx.agent is not None: # pragma: no branch rendered = ctx.agent.render_description(ctx.deps) if rendered is not None: span_attributes['gen_ai.agent.description'] = rendered with ( settings.tracer.start_as_current_span( names.get_agent_run_span_name(agent_name), attributes=span_attributes, record_exception=False, set_status_on_exception=False, ) as span, _record_uncaught_errors(span, include_content=settings.include_content), # Entered with the span and exited with it, so `_run_usage` ends up holding exactly # the usage this run recorded — nested runs report their own on their own spans. _usage_attribution.accumulate(self._run_usage), ): otel_ctx = _otel_set_baggage('gen_ai.agent.name', agent_name) otel_ctx = _otel_set_baggage('gen_ai.agent.call.id', ctx.run_id or '', context=otel_ctx) otel_ctx = _otel_set_baggage('gen_ai.conversation.id', ctx.conversation_id or '', context=otel_ctx) token = _otel_attach(otel_ctx) result: AgentRunResult[Any] | None = None try: result = await handler() if settings.include_content and span.is_recording(): span.set_attribute( 'final_result', ( result.output if isinstance(result.output, str) else safe_to_json(serialize_any(redact_binary_content(result.output, settings))).decode() ), ) return result finally: _otel_detach(token) if span.is_recording(): # Get current messages and metadata from the result (which holds the up-to-date state). # ctx.messages/ctx.metadata may be stale because the run state is mutated during execution. if result is not None: message_history = result.all_messages() metadata = result.metadata else: # On error, use the last messages seen during model requests. message_history = self._last_messages or ctx.messages metadata = ctx.metadata span.set_attributes(self._run_span_end_attributes(ctx, message_history, metadata)) if result is not None: # One O(history) pass per run: turn any silent staleness the per-request # fragment cache may have recorded into a loud signal. Skipped when the run # errored: with warnings configured as errors, warning here in the `finally` # would displace the propagating run exception. if self._message_json_cache and has_stale_message_json( settings, message_history, self._message_json_cache ): warnings.warn( 'In-place mutation of messages already in the history was detected during this run: ' "the `gen_ai.input.messages` attribute recorded on the run's model request spans may " 'not match the messages actually sent to the model. Mutating history messages in ' 'place is not supported; build new message or part objects instead, e.g. via a ' 'history processor.', MessageHistoryMutatedWarning, ) def _run_span_end_attributes( self, ctx: RunContext[Any], message_history: list[ModelMessage], metadata: dict[str, Any] | None, ) -> dict[str, str | int | float | bool]: """Compute the end-of-run span attributes.""" settings = self.settings new_message_index = self._new_message_index last_instructions = get_instructions(message_history, self._last_model_request_parameters) attrs: dict[str, Any] = { 'pydantic_ai.all_messages': safe_to_json( settings.messages_to_otel_messages(list(message_history)) ).decode(), **settings.system_instructions_attributes(last_instructions), } if new_message_index > 0: attrs['pydantic_ai.new_message_index'] = new_message_index if self._variable_instructions: attrs['pydantic_ai.variable_instructions'] = True if metadata is not None: attrs['metadata'] = safe_to_json(serialize_any(redact_binary_content(metadata, settings))).decode() # What this run spent, which is what `gen_ai.aggregated_usage.*` means and what lets the # agent-run spans in a trace be summed without counting a nested run twice. Not `ctx.usage`: # that is the object the caller passed in, accumulated into in place, so it holds the whole # conversation when usage is carried across runs and a delegate's tokens when it is shared. # The per-request `chat` spans are unaffected either way. usage_attrs = settings.aggregated_usage_attributes(self._run_usage) return { **usage_attrs, **attrs, 'logfire.json_schema': to_json( { 'type': 'object', 'properties': { **{k: {'type': 'array'} if isinstance(v, str) else {} for k, v in attrs.items()}, 'final_result': {'type': 'object'}, }, } ).decode(), } # ------------------------------------------------------------------ # wrap_model_request — model request span # ------------------------------------------------------------------ async def wrap_model_request( self, ctx: RunContext[AgentDepsT], *, request_context: ModelRequestContext, handler: WrapModelRequestHandler, ) -> ModelResponse: # Track the latest messages so _run_span_end_attributes has them on error paths # (ctx.messages may be stale because UserPromptNode replaces the list reference). self._last_messages = request_context.messages with open_model_request_span(self.settings, request_context, message_json_cache=self._message_json_cache) as ( finish, prepared_request_context, ): # Stash for `_run_span_end_attributes`: feeding the parameters into # `get_instructions` lets it use the canonical `instruction_parts` source # (which includes prompted-output template instructions and is properly sorted) # instead of falling back to reading `ModelRequest.instructions` from history. self._last_model_request_parameters = prepared_request_context.model_request_parameters # Track whether the fully formatted instructions (including prompted-output schemas) vary across requests. # This does an apples-to-apples comparison of the final payload sent to the model. current_instructions = get_instructions( request_context.messages, prepared_request_context.model_request_parameters ) if not isinstance(self._last_formatted_instructions, Unset): if current_instructions != self._last_formatted_instructions: self._variable_instructions = True self._last_formatted_instructions = current_instructions response = await handler(request_context) # For streaming requests, the agent graph's handler reports TTFT through # `time_to_first_chunk_ctx` (set in the same task, so the value is visible here); # for non-streaming requests this reads the `None` default. finish(response, time_to_first_chunk=time_to_first_chunk_ctx.get()) return response # ------------------------------------------------------------------ # wrap_tool_execute — tool execution span # ------------------------------------------------------------------ async def on_tool_validate_error( self, ctx: RunContext[AgentDepsT], *, call: ToolCallPart, tool_def: ToolDefinition, args: RawToolArgs, error: ValidationError | ModelRetry, ) -> ValidatedToolArgs: """Emit an error span for a tool call whose argument validation failed. Runs only after every other capability has declined to recover the error, so a recovered validation failure produces no span. The span keeps the `execute_tool` operation name so tracing backends group it with other tool spans, and sets `pydantic_ai.tool.failure_stage: 'validation'` to distinguish it from execution failures. With content capture enabled, the span records the retry prompt built from the error as the tool result. That is the exact message the model receives when the agent loop handles the failure; raw-mode callers (e.g. sandboxed dispatch via `handle_call(wrap_validation_errors=False)`) surface the raw exception to the calling code instead, and the recorded prompt is just the rendered description of the failure. """ names = self._instrumentation_names attributes = self._tool_span_attributes(call) # The tool never ran: keep the `execute_tool` operation name so backends find the # span, but say so in the message and mark the failure stage for querying. attributes['logfire.msg'] = f'invalid tool call: {call.tool_name}' attributes[names.tool_failure_stage_attr] = 'validation' with self.settings.tracer.start_as_current_span( names.get_tool_span_name(call.tool_name), attributes=attributes, record_exception=False, set_status_on_exception=False, ) as span: if self.settings.include_content and span.is_recording(): retry = RetryPromptPart.from_error(error, tool_name=call.tool_name, tool_call_id=call.tool_call_id) span.set_attribute(names.tool_result_attr, retry.model_response()) _record_exception(span, error, include_content=self.settings.include_content) span.set_status(StatusCode.ERROR) raise error def _tool_span_attributes(self, call: ToolCallPart) -> dict[str, Any]: """Build the span attributes shared by `wrap_tool_execute` and `wrap_output_process`. Both spans use `gen_ai.operation.name='execute_tool'` and the same `gen_ai.tool.*` attributes — they only differ in how the result is serialized and which exceptions are special-cased, which stays in the call-site `try/except`. """ names = self._instrumentation_names include_content = self.settings.include_content return { 'gen_ai.operation.name': 'execute_tool', 'gen_ai.tool.name': call.tool_name, 'gen_ai.tool.call.id': call.tool_call_id, **({names.tool_arguments_attr: call.args_as_json_str()} if include_content else {}), **get_agent_run_baggage_attributes(), 'logfire.msg': f'running tool: {call.tool_name}', 'logfire.json_schema': to_json( { 'type': 'object', 'properties': { **( { names.tool_arguments_attr: {'type': 'object'}, names.tool_result_attr: {'type': 'object'}, } if include_content else {} ), 'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}, }, } ).decode(), } async def _run_tool_span( self, *, span_name: str, attributes: dict[str, Any], action: Callable[[], Awaitable[Any]], serialize_result: Callable[[Any], str], handle_tool_control_flow: bool = False, ) -> Any: """Open a `gen_ai`-flavoured tool/output span around `action`. Records the serialized result on success (when `include_content` is enabled and the span is recording), records the exception and sets status `ERROR` on failure. When `handle_tool_control_flow` is True, the helper additionally special-cases `CallDeferred`/`ApprovalRequired` (deferrals are control flow, not errors) and records `ToolRetryError`'s retry prompt as the tool result before re-raising. Output-function spans leave that flag off — `ToolRetryError` is treated as a plain error there because the retry prompt is recorded on the surrounding request/agent spans, and `CallDeferred`/`ApprovalRequired` never reach output processing. """ settings = self.settings names = self._instrumentation_names include_content = settings.include_content with settings.tracer.start_as_current_span( span_name, attributes=attributes, record_exception=False, set_status_on_exception=False, ) as span: try: result = await action() except (CallDeferred, ApprovalRequired) as exc: if not handle_tool_control_flow: _record_exception(span, exc, include_content=include_content) span.set_status(StatusCode.ERROR) raise # Deferrals are control flow, not errors: capture the deferral name (and # metadata when available) as span attributes, and only mark the span # ERROR for older instrumentation versions that expected that shape. span.set_attribute(names.tool_deferral_name_attr, type(exc).__name__) if include_content and span.is_recording() and exc.metadata is not None: redacted_metadata = redact_binary_content(exc.metadata, settings) try: metadata_str = to_json(redacted_metadata).decode() except (TypeError, ValueError): metadata_str = repr(redacted_metadata) span.set_attribute(names.tool_deferral_metadata_attr, metadata_str) if settings.version < 5: _record_exception(span, exc, include_content=include_content) span.set_status(StatusCode.ERROR) raise except ToolRetryError as e: if handle_tool_control_flow and include_content and span.is_recording(): # Tool retries are surfaced as model-visible errors; record the prompt # the model will see as the tool result before re-raising. span.set_attribute(names.tool_result_attr, e.tool_retry.model_response()) _record_exception(span, e, include_content=include_content) span.set_status(StatusCode.ERROR) raise except ToolFailedError as e: if handle_tool_control_flow and include_content and span.is_recording(): span.set_attribute(names.tool_result_attr, e.tool_failed.model_response_str(wrap_if_error=False)) _record_exception(span, e, include_content=include_content) span.set_status(StatusCode.ERROR) raise except BaseException as e: _record_exception(span, e, include_content=include_content) span.set_status(StatusCode.ERROR) raise if include_content and span.is_recording(): span.set_attribute( names.tool_result_attr, result if isinstance(result, str) else serialize_result(result), ) return result async def wrap_tool_execute( self, ctx: RunContext[AgentDepsT], *, call: ToolCallPart, tool_def: ToolDefinition, args: ValidatedToolArgs, handler: WrapToolExecuteHandler, ) -> Any: attributes = self._tool_span_attributes(call) if ctx.realtime: # Realtime spans all carry this marker (see `docs/realtime/observability.md`) so # backends can recognize the session tree; the tool span is shared with classic runs, # which stay unmarked. attributes['pydantic_ai.realtime'] = True return await self._run_tool_span( span_name=self._instrumentation_names.get_tool_span_name(call.tool_name), attributes=attributes, action=lambda: handler(args), serialize_result=lambda value: tool_return_ta.dump_json( redact_binary_content(value, self.settings) ).decode(), handle_tool_control_flow=True, ) # ------------------------------------------------------------------ # wrap_output_process — output tool execution span (tool-mode only) # ------------------------------------------------------------------ async def wrap_output_process( self, ctx: RunContext[AgentDepsT], *, output_context: OutputContext, output: Any, handler: WrapOutputProcessHandler, ) -> Any: """Emit a span for output-function execution. Output processing for plain validation (no function) is not span-worthy — the validated value is the model's response itself, no user code ran. We open a span only when an output function will execute, regardless of whether the output arrived via a tool call. The span name reflects the function (or tool name when the function name is unavailable, e.g. union processors). """ if not output_context.has_function: return await handler(output) names = self._instrumentation_names include_content = self.settings.include_content tool_call = output_context.tool_call # Tool-mode output: the registered tool name (e.g. `final_result`) is what the # model called, so use it as the span target. For non-tool output, fall back to # the function name (when known) or a generic placeholder. span_target = tool_call.tool_name if tool_call else (output_context.function_name or 'output_function') attributes: dict[str, Any] = { 'gen_ai.operation.name': 'execute_tool', 'gen_ai.tool.name': span_target, **get_agent_run_baggage_attributes(), 'logfire.msg': f'running output function: {span_target}', } if tool_call is not None and tool_call.tool_call_id: attributes['gen_ai.tool.call.id'] = tool_call.tool_call_id if include_content: attributes[names.tool_arguments_attr] = safe_to_json(redact_binary_content(output, self.settings)).decode() attributes['logfire.json_schema'] = to_json( { 'type': 'object', 'properties': { **( { names.tool_arguments_attr: {'type': 'object'}, names.tool_result_attr: {'type': 'object'}, } if include_content else {} ), 'gen_ai.tool.name': {}, **({'gen_ai.tool.call.id': {}} if tool_call is not None and tool_call.tool_call_id else {}), }, } ).decode() return await self._run_tool_span( span_name=names.get_output_tool_span_name(span_target), attributes=attributes, action=lambda: handler(output), serialize_result=lambda value: safe_to_json( serialize_any(redact_binary_content(value, self.settings)) ).decode(), )