diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 653a9b392f..3d0f9c0dae 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -91,6 +91,7 @@ from ..capabilities.abstract import ( from ..capabilities.combined import bind_capabilities_tier from ..capabilities.hooks import EventT, Hooks, OnEventHookFunc from ..capabilities.instrumentation import Instrumentation as InstrumentationCap +from ..capabilities.wrapper import WrapperCapability from ..models.instrumented import InstrumentationSettings, InstrumentedModel from ..native_tools import AbstractNativeTool from ..native_tools._tool_search import ToolSearchTool @@ -1666,12 +1667,9 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): else: instrumentation_settings = self._resolve_instrumentation_settings() - if instrumentation_settings is not None: - tracer = instrumentation_settings.tracer - instrumentation_cap: InstrumentationCap | None = InstrumentationCap(settings=instrumentation_settings) - else: - tracer = NoOpTracer() - instrumentation_cap = None + instrumentation_cap = ( + InstrumentationCap(settings=instrumentation_settings) if instrumentation_settings is not None else None + ) # Allocated here rather than with the graph deps below, so the context `for_run` receives # shares the very mappings the run fills at setup. A capability that holds on to that @@ -1694,11 +1692,6 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): usage_limits=usage_limits, prompt=user_prompt, messages=state.message_history, - tracer=tracer, - trace_include_content=instrumentation_settings is not None and instrumentation_settings.include_content, - instrumentation_version=instrumentation_settings.version - if instrumentation_settings - else DEFAULT_INSTRUMENTATION_VERSION, run_step=0, pending_messages=state.pending_messages, run_id=state.run_id, @@ -1706,6 +1699,14 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): _cancellation=cancellation, ) + # An explicit `Instrumentation` capability (agent- or call-level) replaces the one injected from + # `instrumentation_settings` (see `_resolve_run_capabilities`), so `for_run` hooks and metadata + # factories are shown the settings of the one that will actually instrument the run. + _set_run_context_instrumentation( + initial_ctx, + _run_instrumentation_settings([base_capability, *extra_capabilities], default=instrumentation_settings), + ) + # Resolve run metadata up front so capability and toolset `for_run` hooks # can see it on `RunContext.metadata`. Metadata factories receive the # `initial_ctx` above (no `tool_manager` / `validation_context` yet); they @@ -1731,6 +1732,10 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): base_is_override=base_is_override, ) run_capability = resolved_caps.run_capability + # Read back off the resolved tree, which also sees an `Instrumentation` that only a `for_run` + # contributed, so every later step's `RunContext` describes what the run's spans use. + instrumentation_settings = _run_instrumentation_settings([run_capability]) + _set_run_context_instrumentation(initial_ctx, instrumentation_settings) capabilities_dict = resolved_caps.capabilities cap_instructions = resolved_caps.instructions cap_native_tools = resolved_caps.native_tools @@ -1891,7 +1896,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): native_tools=cap_native_tools, tool_manager=tool_manager, display_banner=display_banner, - tracer=tracer, + tracer=initial_ctx.tracer, get_instructions=get_instructions, instrumentation_settings=instrumentation_settings, cancellation=cancellation, @@ -3473,40 +3478,16 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): # The session-level `realtime` span and per-response `chat` spans are hand-managed by # `RealtimeSession`; run-lifecycle hooks have no per-exchange boundary on which to build them. - # Drive them from the settings that - # will actually win: an explicit `Instrumentation` capability's (agent- or call-level) over the + # Drive them, and the `RunContext` (as `iter` does), from the settings that will actually win: + # an explicit `Instrumentation` capability's (agent- or call-level) over the # `instrument=`-derived ones, matching the precedence `_resolve_run_capabilities` applies to the - # tool spans. - # The winner is *selected* here, not combined: `combine` settles duplicates *within* one - # layer and never runs across the agent-to-run boundary, so the session reads the - # configuration off the instance `_resolve_run_capabilities` would keep -- the last - # explicit `Instrumentation` in application order (a call-level one supersedes the - # agent-level one, and within a layer the guarded combine resolves to the last instance's - # values). Taking the first match would drive the session spans from settings the - # effective configuration had already turned off. Id guards apply once, where the layers - # actually combine. - instrumentation_layers = [self._effective_root_capability(), *extra_capabilities] - explicit_instrumentations = [ - leaf - for layer in instrumentation_layers - for leaf in leaf_capabilities(layer) - if isinstance(leaf, InstrumentationCap) - ] - explicit_instrumentation = explicit_instrumentations[-1] if explicit_instrumentations else None - session_instrumentation_settings = ( - explicit_instrumentation.settings if explicit_instrumentation is not None else instrumentation_settings + # tool spans. Resolve metadata after, and before `for_run`, so capability/toolset hooks see it + # (same ordering as the graph run). + base_capability, base_is_override = self._base_run_capability() + _set_run_context_instrumentation( + run_context, + _run_instrumentation_settings([base_capability, *extra_capabilities], default=instrumentation_settings), ) - # Mirror `iter`'s `RunContext`: expose the resolved tracer (a `NoOpTracer` when uninstrumented) - # and content-tracing flag, and resolve metadata before `for_run` so capability/toolset hooks see - # it (same ordering as the graph run). - run_context.tracer = ( - session_instrumentation_settings.tracer if session_instrumentation_settings is not None else NoOpTracer() - ) - run_context.trace_include_content = ( - session_instrumentation_settings is not None and session_instrumentation_settings.include_content - ) - if session_instrumentation_settings is not None: - run_context.instrumentation_version = session_instrumentation_settings.version run_context.metadata = self._get_metadata(run_context, metadata) # Resolve the capability layers and extract their contributions, exactly as `run`/`iter` do via @@ -3514,7 +3495,6 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): # Realtime keeps its own surroundings: no `InstrumentedModel` unwrap, once-only model settings # (below), and the `_keep_native` drop plus the shared native ↔ local-tool swap (below). Keep # this in sync with the `iter` call site. - base_capability, base_is_override = self._base_run_capability() resolved_caps = await self._resolve_run_capabilities( run_context, base_capability=base_capability, @@ -3524,6 +3504,10 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]): base_is_override=base_is_override, ) run_capability = resolved_caps.run_capability + # Read back off the resolved tree, as `iter` does, so an `Instrumentation` only a `for_run` + # contributed drives the session's spans and context too. + session_instrumentation_settings = _run_instrumentation_settings([run_capability]) + _set_run_context_instrumentation(run_context, session_instrumentation_settings) # Deferred capabilities load in a session the same way they do in a graph run: the catalog # renders into the connect-time instructions and the loaded instructions come back as the # `load_capability` tool's own result, which every provider supports. What no provider @@ -4381,6 +4365,39 @@ _AUTO_INJECT_CAPABILITY_TYPES: tuple[type[AbstractCapability[Any]], ...] = ( """Infrastructure capabilities auto-injected when not already present.""" +def _run_instrumentation_settings( + capabilities: Sequence[AbstractCapability[Any]], *, default: InstrumentationSettings | None = None +) -> InstrumentationSettings | None: + """The settings of the `Instrumentation` capability that instruments a run with these capabilities. + + That's the last one in application order, or `default` when there is none. `_resolve_run_capabilities` + selects rather than combines across the agent-to-run boundary (`combine` only settles duplicates + *within* a layer), so a call-level one supersedes the agent-level one, and within a layer the guarded + combine keeps the last instance's values. Taking the first match would describe settings the effective + configuration had already turned off. + + Before resolution, `default` is the `instrument=`-derived settings the run injects only when no explicit + capability is present. After resolution, the injected capability is part of the tree, so no default applies. + A wrapper around a single capability is itself a leaf, so it's unwrapped: that's the shape a capability + function's `Instrumentation` takes once resolved, and the wrapper delegates its hooks to it. + """ + instrumentations: list[InstrumentationCap] = [] + for capability in capabilities: + for leaf in leaf_capabilities(capability): + while isinstance(leaf, WrapperCapability): + leaf = leaf.wrapped + if isinstance(leaf, InstrumentationCap): + instrumentations.append(leaf) + return instrumentations[-1].settings if instrumentations else default + + +def _set_run_context_instrumentation(ctx: RunContext[Any], settings: InstrumentationSettings | None) -> None: + """Point `ctx`'s tracer, content flag and instrumentation version at `settings`; uninstrumented if `None`.""" + ctx.tracer = settings.tracer if settings is not None else NoOpTracer() + ctx.trace_include_content = settings is not None and settings.include_content + ctx.instrumentation_version = settings.version if settings is not None else DEFAULT_INSTRUMENTATION_VERSION + + def _registered_capability_count(capability: AbstractCapability[Any]) -> int: """Count the capabilities the user registered, for a `pydantic_ai._display` banner. diff --git a/tests/test_display.py b/tests/test_display.py index 5ba9cbf54d..cea3971ebf 100644 --- a/tests/test_display.py +++ b/tests/test_display.py @@ -18,7 +18,7 @@ import pydantic_ai import pydantic_ai._display as _display from pydantic_ai import Agent, ModelMessage, ModelRequest, UserPromptPart, __version__ from pydantic_ai.agent import _registered_capability_count # pyright: ignore[reportPrivateUsage] -from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.capabilities import AbstractCapability, Instrumentation from pydantic_ai.models.test import TestModel from pydantic_ai.toolsets import FunctionToolset @@ -916,6 +916,15 @@ def test_instrumented_agent_run_is_silent(monkeypatch: pytest.MonkeyPatch, stder assert stderr.getvalue() == '' +def test_run_instrumented_by_an_explicit_capability_is_silent(monkeypatch: pytest.MonkeyPatch, stderr: TTYStream): + """The banner defers to the `Instrumentation` capability that instruments the run, not just `instrument`.""" + monkeypatch.setattr(sys, 'stderr', stderr) + + Agent(TestModel()).run_sync('hello', capabilities=[Instrumentation()]) + + assert stderr.getvalue() == '' + + def test_an_instrumented_run_leaves_the_banner_for_another_agent(monkeypatch: pytest.MonkeyPatch, stderr: TTYStream): """Instrumentation is the agent's, not the process's, so it doesn't speak for the ones after it.""" monkeypatch.setattr(sys, 'stderr', stderr) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 4f29a80be7..5b83885dff 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -55,6 +55,8 @@ with try_import(): # `opentelemetry-sdk` arrives with the `logfire` extra, so it is not importable in the # `pydantic-ai-slim` / `pydantic-evals` install groups either. from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.sdk.trace.sampling import ALWAYS_OFF try: @@ -3916,6 +3918,105 @@ def test_instrument_all_skipped_when_capability_already_present( Agent.instrument_all(False) +@dataclass +class _SeenInstrumentation: + tracer: Any + include_content: bool + version: int + + +def _seen(ctx: RunContext[Any]) -> _SeenInstrumentation: + return _SeenInstrumentation(ctx.tracer, ctx.trace_include_content, ctx.instrumentation_version) + + +@dataclass +class _ObserveForRun(AbstractCapability[Any]): + """Records what a `for_run` hook sees of the run's instrumentation.""" + + seen: list[_SeenInstrumentation] + + async def for_run(self, ctx: RunContext[Any]) -> _ObserveForRun: + self.seen.append(_seen(ctx)) + return self + + +@pytest.mark.skipif(not logfire_installed, reason='logfire not installed') +@pytest.mark.parametrize('explicit_on', [None, 'agent', 'run']) +def test_run_context_instrumentation_follows_the_capability_that_instruments_the_run( + explicit_on: Literal['agent', 'run'] | None, +) -> None: + """An explicit `Instrumentation` capability replaces the `instrument_all()` one, so `RunContext` describes it. + + Its `tracer`, `trace_include_content` and `instrumentation_version` are what the run's spans use, both + in `for_run` hooks and in the context a tool gets later in the run. Without one, `instrument_all()`'s + settings still apply. + """ + implicit = InstrumentationSettings(include_content=True, tracer_provider=TracerProvider()) + explicit = InstrumentationSettings(include_content=False, version=6, tracer_provider=TracerProvider()) + expected = implicit if explicit_on is None else explicit + + for_run_seen: list[_SeenInstrumentation] = [] + agent_capabilities: list[AbstractCapability[Any]] = [_ObserveForRun(for_run_seen)] + run_capabilities: list[AbstractCapability[Any]] = [] + if explicit_on == 'agent': + agent_capabilities.append(Instrumentation(settings=explicit)) + elif explicit_on == 'run': + run_capabilities.append(Instrumentation(settings=explicit)) + + Agent.instrument_all(implicit) + try: + agent = Agent(TestModel(), capabilities=agent_capabilities) + tool_seen: list[_SeenInstrumentation] = [] + + @agent.tool + def peek(ctx: RunContext[Any]) -> str: + tool_seen.append(_seen(ctx)) + return 'ok' + + agent.run_sync('hi', capabilities=run_capabilities) + finally: + Agent.instrument_all(False) + + assert implicit.tracer is not explicit.tracer + expected_seen = _SeenInstrumentation(expected.tracer, expected.include_content, expected.version) + assert for_run_seen == [expected_seen] + assert tool_seen == [expected_seen] + + +@pytest.mark.skipif(not logfire_installed, reason='logfire not installed') +def test_run_context_instrumentation_follows_a_capability_function_instrumentation() -> None: + """An `Instrumentation` only a capability function's `for_run` returns still instruments the run, so the + context the run's tools get describes it too, even though it couldn't be known before `for_run`.""" + implicit_provider = TracerProvider() + implicit_exporter = InMemorySpanExporter() + implicit_provider.add_span_processor(SimpleSpanProcessor(implicit_exporter)) + dynamic_provider = TracerProvider() + dynamic_exporter = InMemorySpanExporter() + dynamic_provider.add_span_processor(SimpleSpanProcessor(dynamic_exporter)) + dynamic = InstrumentationSettings(include_content=False, version=6, tracer_provider=dynamic_provider) + + Agent.instrument_all(InstrumentationSettings(include_content=True, tracer_provider=implicit_provider)) + try: + agent = Agent(TestModel()) + tool_seen: list[_SeenInstrumentation] = [] + + @agent.tool + def peek(ctx: RunContext[Any]) -> str: + tool_seen.append(_seen(ctx)) + return 'ok' + + agent.run_sync('hi', capabilities=[lambda ctx: Instrumentation(settings=dynamic)]) + finally: + Agent.instrument_all(False) + + assert tool_seen == [_SeenInstrumentation(dynamic.tracer, False, 6)] + # The capability function's `Instrumentation` is the one that opened the run's spans. + assert implicit_exporter.get_finished_spans() == () + assert [span.name for span in dynamic_exporter.get_finished_spans()] == snapshot( + ['chat test', 'execute_tool peek', 'chat test', 'invoke_agent agent'] + ) + + @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') def test_agent_with_user_provided_instrumented_model( get_logfire_summary: Callable[[], LogfireSummary],