From af73a61bd8387f9bf47576e7583aa85834864e50 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 18:41:56 +1000 Subject: [PATCH 01/25] feat: add gen_ai.agent.name attribute to agent run logging --- pydantic_ai_slim/pydantic_ai/agent/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index b70f541262..95282d7b43 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -649,6 +649,7 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None: attributes={ 'model_name': model_used.model_name if model_used else 'no-model', 'agent_name': agent_name, + 'gen_ai.agent.name': agent_name, 'logfire.msg': f'{agent_name} run', }, ) From 6454dd48e5371e1c4099c3d6f2de6f1cf228ed44 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:02:46 +1000 Subject: [PATCH 02/25] feat: update span naming convention for agent invocations in v3 instrumentation --- pydantic_ai_slim/pydantic_ai/agent/__init__.py | 4 +++- pydantic_ai_slim/pydantic_ai/models/instrumented.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 95282d7b43..536a0b514c 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -645,7 +645,9 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None: agent_name = self.name or 'agent' run_span = tracer.start_span( - 'agent run', + f'invoke_agent {agent_name}' + if instrumentation_settings and instrumentation_settings.version > 2 + else 'agent run', attributes={ 'model_name': model_used.model_name if model_used else 'no-model', 'agent_name': agent_name, diff --git a/pydantic_ai_slim/pydantic_ai/models/instrumented.py b/pydantic_ai_slim/pydantic_ai/models/instrumented.py index 57dc235da6..0016c2cb65 100644 --- a/pydantic_ai_slim/pydantic_ai/models/instrumented.py +++ b/pydantic_ai_slim/pydantic_ai/models/instrumented.py @@ -90,7 +90,7 @@ class InstrumentationSettings: event_mode: Literal['attributes', 'logs'] = 'attributes' include_binary_content: bool = True include_content: bool = True - version: Literal[1, 2] = 1 + version: Literal[1, 2, 3] = 1 def __init__( self, @@ -99,7 +99,7 @@ def __init__( meter_provider: MeterProvider | None = None, include_binary_content: bool = True, include_content: bool = True, - version: Literal[1, 2] = 2, + version: Literal[1, 2, 3] = 2, event_mode: Literal['attributes', 'logs'] = 'attributes', event_logger_provider: EventLoggerProvider | None = None, ): From 785a928a77c5d6094359eb74c4cbb4b37f78a534 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:17:24 +1000 Subject: [PATCH 03/25] refactor: extract default instrumentation version into a constant and use it consistently --- pydantic_ai_slim/pydantic_ai/models/instrumented.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/instrumented.py b/pydantic_ai_slim/pydantic_ai/models/instrumented.py index 0016c2cb65..052bd4fd58 100644 --- a/pydantic_ai_slim/pydantic_ai/models/instrumented.py +++ b/pydantic_ai_slim/pydantic_ai/models/instrumented.py @@ -72,6 +72,9 @@ def instrument_model(model: Model, instrument: InstrumentationSettings | bool) - return model +DEFAULT_INSTRUMENTATION_VERSION = 2 + + @dataclass(init=False) class InstrumentationSettings: """Options for instrumenting models and agents with OpenTelemetry. @@ -90,7 +93,7 @@ class InstrumentationSettings: event_mode: Literal['attributes', 'logs'] = 'attributes' include_binary_content: bool = True include_content: bool = True - version: Literal[1, 2, 3] = 1 + version: Literal[1, 2, 3] = DEFAULT_INSTRUMENTATION_VERSION def __init__( self, @@ -99,7 +102,7 @@ def __init__( meter_provider: MeterProvider | None = None, include_binary_content: bool = True, include_content: bool = True, - version: Literal[1, 2, 3] = 2, + version: Literal[1, 2, 3] = DEFAULT_INSTRUMENTATION_VERSION, event_mode: Literal['attributes', 'logs'] = 'attributes', event_logger_provider: EventLoggerProvider | None = None, ): From 723809f27002d1f8701445b8dd8be8130fe43493 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:24:31 +1000 Subject: [PATCH 04/25] feat: add instrumentation version field to RunContext for tracking settings --- pydantic_ai_slim/pydantic_ai/_run_context.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pydantic_ai_slim/pydantic_ai/_run_context.py b/pydantic_ai_slim/pydantic_ai/_run_context.py index 428697b59a..7383c9fbb5 100644 --- a/pydantic_ai_slim/pydantic_ai/_run_context.py +++ b/pydantic_ai_slim/pydantic_ai/_run_context.py @@ -36,6 +36,8 @@ class RunContext(Generic[AgentDepsT]): """The tracer to use for tracing the run.""" trace_include_content: bool = False """Whether to include the content of the messages in the trace.""" + instrumentation_version: int | None = None + """Instrumentation settings version, if instrumentation is enabled.""" retries: dict[str, int] = field(default_factory=dict) """Number of retries for each tool so far.""" tool_call_id: str | None = None From 42268f08f5070a8c1f3fefa4b352e47ac0e157db Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:32:28 +1000 Subject: [PATCH 05/25] feat: add default instrumentation version constant and update RunContext defaults --- pydantic_ai_slim/pydantic_ai/_instrumentation.py | 2 ++ pydantic_ai_slim/pydantic_ai/_run_context.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) create mode 100644 pydantic_ai_slim/pydantic_ai/_instrumentation.py diff --git a/pydantic_ai_slim/pydantic_ai/_instrumentation.py b/pydantic_ai_slim/pydantic_ai/_instrumentation.py new file mode 100644 index 0000000000..342352ca44 --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/_instrumentation.py @@ -0,0 +1,2 @@ +DEFAULT_INSTRUMENTATION_VERSION = 2 +"""Default instrumentation version for `InstrumentationSettings`.""" diff --git a/pydantic_ai_slim/pydantic_ai/_run_context.py b/pydantic_ai_slim/pydantic_ai/_run_context.py index 7383c9fbb5..df2a4c1b5a 100644 --- a/pydantic_ai_slim/pydantic_ai/_run_context.py +++ b/pydantic_ai_slim/pydantic_ai/_run_context.py @@ -8,6 +8,8 @@ from opentelemetry.trace import NoOpTracer, Tracer from typing_extensions import TypeVar +from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION + from . import _utils, messages as _messages if TYPE_CHECKING: @@ -36,7 +38,7 @@ class RunContext(Generic[AgentDepsT]): """The tracer to use for tracing the run.""" trace_include_content: bool = False """Whether to include the content of the messages in the trace.""" - instrumentation_version: int | None = None + instrumentation_version: int = DEFAULT_INSTRUMENTATION_VERSION """Instrumentation settings version, if instrumentation is enabled.""" retries: dict[str, int] = field(default_factory=dict) """Number of retries for each tool so far.""" From 14eedd46db317229a852217415eba48ef181e432 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:35:13 +1000 Subject: [PATCH 06/25] feat: add instrumentation version to run context using settings or default --- pydantic_ai_slim/pydantic_ai/_agent_graph.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pydantic_ai_slim/pydantic_ai/_agent_graph.py b/pydantic_ai_slim/pydantic_ai/_agent_graph.py index 5f89155092..b8ab970696 100644 --- a/pydantic_ai_slim/pydantic_ai/_agent_graph.py +++ b/pydantic_ai_slim/pydantic_ai/_agent_graph.py @@ -16,6 +16,7 @@ from typing_extensions import TypeVar, assert_never from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx # type: ignore +from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION from pydantic_ai._tool_manager import ToolManager from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, is_async_callable, run_in_executor from pydantic_ai.builtin_tools import AbstractBuiltinTool @@ -704,6 +705,9 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT tracer=ctx.deps.tracer, trace_include_content=ctx.deps.instrumentation_settings is not None and ctx.deps.instrumentation_settings.include_content, + instrumentation_version=ctx.deps.instrumentation_settings.version + if ctx.deps.instrumentation_settings + else DEFAULT_INSTRUMENTATION_VERSION, run_step=ctx.state.run_step, tool_call_approved=ctx.state.run_step == 0, ) From 16e764ec4a805c4f7120defe1d667afe2d0f604d Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:40:28 +1000 Subject: [PATCH 07/25] refactor: move DEFAULT_INSTRUMENTATION_VERSION constant to _instrumentation module --- pydantic_ai_slim/pydantic_ai/models/instrumented.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/instrumented.py b/pydantic_ai_slim/pydantic_ai/models/instrumented.py index 052bd4fd58..17064a8e1e 100644 --- a/pydantic_ai_slim/pydantic_ai/models/instrumented.py +++ b/pydantic_ai_slim/pydantic_ai/models/instrumented.py @@ -21,6 +21,8 @@ from opentelemetry.util.types import AttributeValue from pydantic import TypeAdapter +from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION + from .. import _otel_messages from .._run_context import RunContext from ..messages import ( @@ -72,9 +74,6 @@ def instrument_model(model: Model, instrument: InstrumentationSettings | bool) - return model -DEFAULT_INSTRUMENTATION_VERSION = 2 - - @dataclass(init=False) class InstrumentationSettings: """Options for instrumenting models and agents with OpenTelemetry. From 59d04cc13328cae6bd85e26a6467e021d9722934 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:40:47 +1000 Subject: [PATCH 08/25] feat: add instrumentation version to tool execution tracing --- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index 5e27c2584e..eba6fb4ae7 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -115,6 +115,7 @@ async def handle_call( wrap_validation_errors, self.ctx.tracer, self.ctx.trace_include_content, + self.ctx.instrumentation_version, usage_limits, ) @@ -203,7 +204,8 @@ async def _call_tool_traced( allow_partial: bool, wrap_validation_errors: bool, tracer: Tracer, - include_content: bool = False, + include_content: bool, + instrumentation_version: int, usage_limits: UsageLimits | None = None, ) -> Any: """See .""" @@ -220,8 +222,12 @@ async def _call_tool_traced( 'properties': { **( { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, + 'gen_ai.tool.call.arguments' if instrumentation_version > 2 else 'tool_arguments': { + 'type': 'object' + }, + 'gen_ai.tool.call.result' if instrumentation_version > 2 else 'tool_response': { + 'type': 'object' + }, } if include_content else {} @@ -232,7 +238,10 @@ async def _call_tool_traced( } ), } - with tracer.start_as_current_span('running tool', attributes=span_attributes) as span: + with tracer.start_as_current_span( + f'execute_tool {call.tool_name}' if instrumentation_version > 2 else 'running tool', + attributes=span_attributes, + ) as span: try: tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits) except ToolRetryError as e: From 233f496b4a5b128efdcafe6346f6e810512156e9 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:55:22 +1000 Subject: [PATCH 09/25] feat: add InstrumentationConfig to manage versioned span names and attributes --- .../pydantic_ai/_instrumentation.py | 75 +++++++++++++++++++ pydantic_ai_slim/pydantic_ai/_tool_manager.py | 19 +++-- .../pydantic_ai/agent/__init__.py | 9 ++- 3 files changed, 90 insertions(+), 13 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_instrumentation.py b/pydantic_ai_slim/pydantic_ai/_instrumentation.py index 342352ca44..2e432f7246 100644 --- a/pydantic_ai_slim/pydantic_ai/_instrumentation.py +++ b/pydantic_ai_slim/pydantic_ai/_instrumentation.py @@ -1,2 +1,77 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing_extensions import Self + DEFAULT_INSTRUMENTATION_VERSION = 2 """Default instrumentation version for `InstrumentationSettings`.""" + + +@dataclass(frozen=True) +class InstrumentationConfig: + """Configuration for instrumentation span names and attributes based on version.""" + + # Agent run span configuration + agent_run_span_name: str + agent_name_attr: str + + # Tool execution span configuration + tool_span_name: str + tool_arguments_attr: str + tool_result_attr: str + + @classmethod + def for_version(cls, version: int) -> Self: + """Create instrumentation configuration for a specific version. + + Args: + version: The instrumentation version (1, 2, or 3+) + + Returns: + InstrumentationConfig instance with version-appropriate settings + """ + if version <= 2: + return cls( + agent_run_span_name='agent run', + agent_name_attr='agent_name', + tool_span_name='running tool', + tool_arguments_attr='tool_arguments', + tool_result_attr='tool_response', + ) + else: + return cls( + agent_run_span_name='invoke_agent', + agent_name_attr='gen_ai.agent.name', + tool_span_name='execute_tool', # Will be formatted with tool name + tool_arguments_attr='gen_ai.tool.call.arguments', + tool_result_attr='gen_ai.tool.call.result', + ) + + def get_tool_span_name(self, tool_name: str) -> str: + """Get the formatted tool span name. + + Args: + tool_name: Name of the tool being executed + + Returns: + Formatted span name + """ + if self.tool_span_name == 'execute_tool': + return f'execute_tool {tool_name}' + return self.tool_span_name + + def get_agent_run_span_name(self, agent_name: str) -> str: + """Get the formatted agent span name. + + Args: + agent_name: Name of the agent being executed + + Returns: + Formatted span name + """ + if self.agent_run_span_name == 'invoke_agent': + return f'invoke_agent {agent_name}' + return self.agent_run_span_name diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index eba6fb4ae7..efbfca9337 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -12,6 +12,7 @@ from typing_extensions import assert_never from . import messages as _messages +from ._instrumentation import InstrumentationConfig from ._run_context import AgentDepsT, RunContext from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior from .messages import ToolCallPart @@ -209,11 +210,13 @@ async def _call_tool_traced( usage_limits: UsageLimits | None = None, ) -> Any: """See .""" + instrumentation_config = InstrumentationConfig.for_version(instrumentation_version) + span_attributes = { 'gen_ai.tool.name': call.tool_name, # NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai 'gen_ai.tool.call.id': call.tool_call_id, - **({'tool_arguments': call.args_as_json_str()} if include_content else {}), + **({instrumentation_config.tool_arguments_attr: call.args_as_json_str()} if include_content else {}), 'logfire.msg': f'running tool: {call.tool_name}', # add the JSON schema so these attributes are formatted nicely in Logfire 'logfire.json_schema': json.dumps( @@ -222,12 +225,8 @@ async def _call_tool_traced( 'properties': { **( { - 'gen_ai.tool.call.arguments' if instrumentation_version > 2 else 'tool_arguments': { - 'type': 'object' - }, - 'gen_ai.tool.call.result' if instrumentation_version > 2 else 'tool_response': { - 'type': 'object' - }, + instrumentation_config.tool_arguments_attr: {'type': 'object'}, + instrumentation_config.tool_result_attr: {'type': 'object'}, } if include_content else {} @@ -239,7 +238,7 @@ async def _call_tool_traced( ), } with tracer.start_as_current_span( - f'execute_tool {call.tool_name}' if instrumentation_version > 2 else 'running tool', + instrumentation_config.get_tool_span_name(call.tool_name), attributes=span_attributes, ) as span: try: @@ -247,12 +246,12 @@ async def _call_tool_traced( except ToolRetryError as e: part = e.tool_retry if include_content and span.is_recording(): - span.set_attribute('tool_response', part.model_response()) + span.set_attribute(instrumentation_config.tool_result_attr, part.model_response()) raise e if include_content and span.is_recording(): span.set_attribute( - 'tool_response', + instrumentation_config.tool_result_attr, tool_result if isinstance(tool_result, str) else _messages.tool_return_ta.dump_json(tool_result).decode(), diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 536a0b514c..83fe87631a 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -14,6 +14,7 @@ from pydantic.json_schema import GenerateJsonSchema from typing_extensions import Self, TypeVar, deprecated +from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION, InstrumentationConfig from pydantic_graph import Graph from .. import ( @@ -644,10 +645,12 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None: ) agent_name = self.name or 'agent' + instrumentation_config = InstrumentationConfig.for_version( + instrumentation_settings.version if instrumentation_settings else DEFAULT_INSTRUMENTATION_VERSION + ) + run_span = tracer.start_span( - f'invoke_agent {agent_name}' - if instrumentation_settings and instrumentation_settings.version > 2 - else 'agent run', + instrumentation_config.get_agent_run_span_name(agent_name), attributes={ 'model_name': model_used.model_name if model_used else 'no-model', 'agent_name': agent_name, From f375a3c19f85e405c37f85f7568a1d3c52ab9d92 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 19:59:15 +1000 Subject: [PATCH 10/25] fix: add gen_ai.agent.name attribute to fallback model test snapshots --- tests/models/test_fallback.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/models/test_fallback.py b/tests/models/test_fallback.py index 484a73ac37..9e0172b438 100644 --- a/tests/models/test_fallback.py +++ b/tests/models/test_fallback.py @@ -170,6 +170,7 @@ def test_first_failed_instrumented(capfire: CaptureLogfire) -> None: 'attributes': { 'model_name': 'fallback:function:failure_response:,function:success_response:', 'agent_name': 'agent', + 'gen_ai.agent.name': 'agent', 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -269,6 +270,7 @@ async def test_first_failed_instrumented_stream(capfire: CaptureLogfire) -> None 'attributes': { 'model_name': 'fallback:function::failure_response_stream,function::success_response_stream', 'agent_name': 'agent', + 'gen_ai.agent.name': 'agent', 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 50, @@ -376,6 +378,7 @@ def test_all_failed_instrumented(capfire: CaptureLogfire) -> None: 'attributes': { 'model_name': 'fallback:function:failure_response:,function:failure_response:', 'agent_name': 'agent', + 'gen_ai.agent.name': 'agent', 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'pydantic_ai.all_messages': [{'role': 'user', 'parts': [{'type': 'text', 'content': 'hello'}]}], From 7fe2c180348054771bf9d5253eb418893069eec1 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 20:00:54 +1000 Subject: [PATCH 11/25] fix: add gen_ai.agent.name field to logfire test assertions --- tests/test_logfire.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 583537f3c5..b9412bd762 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -123,6 +123,7 @@ async def my_ret(x: int) -> str: { 'model_name': 'test', 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'final_result': '{"my_ret":"1"}', @@ -176,6 +177,7 @@ async def my_ret(x: int) -> str: { 'model_name': 'test', 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 103, @@ -435,6 +437,7 @@ class MyOutput: { 'model_name': 'test', 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -530,6 +533,7 @@ class MyOutput: { 'model_name': 'test', 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -620,6 +624,7 @@ class MyOutput: { 'model_name': 'test', 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -710,6 +715,7 @@ class MyOutput: { 'model_name': 'test', 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', 'logfire.msg': 'my_agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, @@ -931,6 +937,7 @@ async def test_feedback(capfire: CaptureLogfire) -> None: 'attributes': { 'model_name': 'test', 'agent_name': 'agent', + 'gen_ai.agent.name': 'agent', 'logfire.msg': 'agent run', 'logfire.span_type': 'span', 'gen_ai.usage.input_tokens': 51, From 09db4de30415d805d6d29cafe55fb009839149b8 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Sun, 28 Sep 2025 20:32:51 +1000 Subject: [PATCH 12/25] feat: add support for version 3 instrumentation settings in logfire tests --- tests/test_logfire.py | 796 ++++++++++++++++++++++++++++++------------ 1 file changed, 573 insertions(+), 223 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index b9412bd762..519b386ff1 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -2,7 +2,7 @@ from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import Any, Literal import pytest from dirty_equals import IsInt, IsJson, IsList @@ -77,6 +77,7 @@ def get_summary() -> LogfireSummary: InstrumentationSettings(version=1, event_mode='attributes'), InstrumentationSettings(version=1, event_mode='logs'), InstrumentationSettings(version=2), + InstrumentationSettings(version=3), ], ) def test_logfire( @@ -98,27 +99,102 @@ async def my_ret(x: int) -> str: assert summary.traces == [] return - assert summary.traces == snapshot( - [ + if isinstance(instrument, InstrumentationSettings) and instrument.version == 3: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'message': 'chat test'}, + { + 'id': 2, + 'message': 'running 1 tool', + 'children': [ + {'id': 3, 'message': 'running tool: my_ret'}, + ], + }, + {'id': 4, 'message': 'chat test'}, + ], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'message': 'chat test'}, + { + 'id': 2, + 'message': 'running 1 tool', + 'children': [ + {'id': 3, 'message': 'running tool: my_ret'}, + ], + }, + {'id': 4, 'message': 'chat test'}, + ], + } + ] + ) + + if isinstance(instrument, InstrumentationSettings) and instrument.version == 3: + assert summary.attributes[0] == snapshot( { - 'id': 0, - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'message': 'chat test'}, - { - 'id': 2, - 'message': 'running 1 tool', - 'children': [ - {'id': 3, 'message': 'running tool: my_ret'}, - ], - }, - {'id': 4, 'message': 'chat test'}, - ], + 'model_name': 'test', + 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', + 'logfire.msg': 'my_agent run', + 'logfire.span_type': 'span', + 'final_result': '{"my_ret":"1"}', + 'gen_ai.usage.input_tokens': 103, + 'gen_ai.usage.output_tokens': 12, + 'pydantic_ai.all_messages': IsJson( + snapshot( + [ + {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}, + { + 'role': 'assistant', + 'parts': [ + { + 'type': 'tool_call', + 'id': IsStr(), + 'name': 'my_ret', + 'arguments': {'x': 0}, + } + ], + }, + { + 'role': 'user', + 'parts': [ + { + 'type': 'tool_call_response', + 'id': IsStr(), + 'name': 'my_ret', + 'result': '1', + } + ], + }, + {'role': 'assistant', 'parts': [{'type': 'text', 'content': '{"my_ret":"1"}'}]}, + ] + ) + ), + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'pydantic_ai.all_messages': {'type': 'array'}, + 'final_result': {'type': 'object'}, + }, + } + ) + ), } - ] - ) - - if instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 2: + ) + elif instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 2: assert summary.attributes[0] == snapshot( { 'model_name': 'test', @@ -413,10 +489,7 @@ async def my_ret(x: int) -> str: @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') @pytest.mark.parametrize( 'instrument', - [ - InstrumentationSettings(version=1), - InstrumentationSettings(version=2), - ], + [InstrumentationSettings(version=1), InstrumentationSettings(version=2), InstrumentationSettings(version=3)], ) def test_instructions_with_structured_output( get_logfire_summary: Callable[[], LogfireSummary], instrument: InstrumentationSettings @@ -528,7 +601,62 @@ class MyOutput: ] ) ) - else: + elif instrument.version == 3: + assert summary.attributes[0] == snapshot( + { + 'model_name': 'test', + 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', + 'logfire.msg': 'my_agent run', + 'logfire.span_type': 'span', + 'final_result': '{"content": "a"}', + 'gen_ai.usage.input_tokens': 51, + 'gen_ai.usage.output_tokens': 5, + 'pydantic_ai.all_messages': IsJson( + snapshot( + [ + {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}, + { + 'role': 'assistant', + 'parts': [ + { + 'type': 'tool_call', + 'id': IsStr(), + 'name': 'final_result', + 'arguments': {'content': 'a'}, + } + ], + }, + { + 'role': 'user', + 'parts': [ + { + 'type': 'tool_call_response', + 'id': IsStr(), + 'name': 'final_result', + 'result': 'Final result processed.', + } + ], + }, + ] + ) + ), + 'gen_ai.system_instructions': '[{"type": "text", "content": "Here are some instructions"}]', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'pydantic_ai.all_messages': {'type': 'array'}, + 'gen_ai.system_instructions': {'type': 'array'}, + 'final_result': {'type': 'object'}, + }, + } + ) + ), + } + ) + elif instrument.version == 2: assert summary.attributes[0] == snapshot( { 'model_name': 'test', @@ -696,140 +824,263 @@ class MyOutput: @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') -def test_instructions_with_structured_output_exclude_content_v2( +@pytest.mark.parametrize('version', [2, 3]) +def test_instructions_with_structured_output_exclude_content_v2_v3( get_logfire_summary: Callable[[], LogfireSummary], + version: Literal[2, 3], ) -> None: @dataclass class MyOutput: content: str - settings: InstrumentationSettings = InstrumentationSettings(include_content=False, version=2) + settings: InstrumentationSettings = InstrumentationSettings(include_content=False, version=version) my_agent = Agent(model=TestModel(), instructions='Here are some instructions', instrument=settings) result = my_agent.run_sync('Hello', output_type=MyOutput) - assert result.output == snapshot(MyOutput(content='a')) + assert result.output == MyOutput(content='a') summary = get_logfire_summary() - assert summary.attributes[0] == snapshot( - { - 'model_name': 'test', - 'agent_name': 'my_agent', - 'gen_ai.agent.name': 'my_agent', - 'logfire.msg': 'my_agent run', - 'logfire.span_type': 'span', - 'gen_ai.usage.input_tokens': 51, - 'gen_ai.usage.output_tokens': 5, - 'pydantic_ai.all_messages': IsJson( - snapshot( - [ - {'role': 'user', 'parts': [{'type': 'text'}]}, + if version == 2: + assert summary.attributes[0] == snapshot( + { + 'model_name': 'test', + 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', + 'logfire.msg': 'my_agent run', + 'logfire.span_type': 'span', + 'gen_ai.usage.input_tokens': 51, + 'gen_ai.usage.output_tokens': 5, + 'pydantic_ai.all_messages': IsJson( + snapshot( + [ + {'role': 'user', 'parts': [{'type': 'text'}]}, + { + 'role': 'assistant', + 'parts': [ + { + 'type': 'tool_call', + 'id': IsStr(), + 'name': 'final_result', + } + ], + }, + { + 'role': 'user', + 'parts': [ + { + 'type': 'tool_call_response', + 'id': IsStr(), + 'name': 'final_result', + } + ], + }, + ] + ) + ), + 'logfire.json_schema': IsJson( + snapshot( { - 'role': 'assistant', - 'parts': [ - { - 'type': 'tool_call', - 'id': IsStr(), - 'name': 'final_result', - } - ], - }, + 'type': 'object', + 'properties': { + 'pydantic_ai.all_messages': {'type': 'array'}, + 'final_result': {'type': 'object'}, + }, + } + ) + ), + } + ) + chat_span_attributes = summary.attributes[1] + assert chat_span_attributes == snapshot( + { + 'gen_ai.operation.name': 'chat', + 'gen_ai.system': 'test', + 'gen_ai.request.model': 'test', + 'model_request_parameters': IsJson( + snapshot( { - 'role': 'user', - 'parts': [ + 'function_tools': [], + 'builtin_tools': [], + 'output_mode': 'tool', + 'output_object': None, + 'output_tools': [ { - 'type': 'tool_call_response', - 'id': IsStr(), 'name': 'final_result', + 'parameters_json_schema': { + 'properties': {'content': {'type': 'string'}}, + 'required': ['content'], + 'title': 'MyOutput', + 'type': 'object', + }, + 'description': 'The final response which ends this conversation', + 'outer_typed_dict_key': None, + 'strict': None, + 'sequential': False, + 'kind': 'output', + 'metadata': None, } ], - }, - ] - ) - ), - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'pydantic_ai.all_messages': {'type': 'array'}, - 'final_result': {'type': 'object'}, - }, - } - ) - ), - } - ) - chat_span_attributes = summary.attributes[1] - assert chat_span_attributes == snapshot( - { - 'gen_ai.operation.name': 'chat', - 'gen_ai.system': 'test', - 'gen_ai.request.model': 'test', - 'model_request_parameters': IsJson( - snapshot( - { - 'function_tools': [], - 'builtin_tools': [], - 'output_mode': 'tool', - 'output_object': None, - 'output_tools': [ + 'allow_text_output': False, + } + ) + ), + 'logfire.span_type': 'span', + 'logfire.msg': 'chat test', + 'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])), + 'gen_ai.output.messages': IsJson( + snapshot( + [ { - 'name': 'final_result', - 'parameters_json_schema': { - 'properties': {'content': {'type': 'string'}}, - 'required': ['content'], - 'title': 'MyOutput', - 'type': 'object', - }, - 'description': 'The final response which ends this conversation', - 'outer_typed_dict_key': None, - 'strict': None, - 'sequential': False, - 'kind': 'output', - 'metadata': None, + 'role': 'assistant', + 'parts': [ + { + 'type': 'tool_call', + 'id': IsStr(), + 'name': 'final_result', + } + ], } - ], - 'allow_text_output': False, - } - ) - ), - 'logfire.span_type': 'span', - 'logfire.msg': 'chat test', - 'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])), - 'gen_ai.output.messages': IsJson( - snapshot( - [ + ] + ) + ), + 'logfire.json_schema': IsJson( + snapshot( { - 'role': 'assistant', - 'parts': [ + 'type': 'object', + 'properties': { + 'gen_ai.input.messages': {'type': 'array'}, + 'gen_ai.output.messages': {'type': 'array'}, + 'model_request_parameters': {'type': 'object'}, + }, + } + ) + ), + 'gen_ai.usage.input_tokens': 51, + 'gen_ai.usage.output_tokens': 5, + 'gen_ai.response.model': 'test', + } + ) + else: + assert summary.attributes[0] == snapshot( + { + 'model_name': 'test', + 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', + 'logfire.msg': 'my_agent run', + 'logfire.span_type': 'span', + 'gen_ai.usage.input_tokens': 51, + 'gen_ai.usage.output_tokens': 5, + 'pydantic_ai.all_messages': IsJson( + snapshot( + [ + {'role': 'user', 'parts': [{'type': 'text'}]}, + { + 'role': 'assistant', + 'parts': [ + { + 'type': 'tool_call', + 'id': IsStr(), + 'name': 'final_result', + } + ], + }, + { + 'role': 'user', + 'parts': [ + { + 'type': 'tool_call_response', + 'id': IsStr(), + 'name': 'final_result', + } + ], + }, + ] + ) + ), + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'pydantic_ai.all_messages': {'type': 'array'}, + 'final_result': {'type': 'object'}, + }, + } + ) + ), + } + ) + chat_span_attributes = summary.attributes[1] + assert chat_span_attributes == snapshot( + { + 'gen_ai.operation.name': 'chat', + 'gen_ai.system': 'test', + 'gen_ai.request.model': 'test', + 'model_request_parameters': IsJson( + snapshot( + { + 'function_tools': [], + 'builtin_tools': [], + 'output_mode': 'tool', + 'output_object': None, + 'output_tools': [ { - 'type': 'tool_call', - 'id': IsStr(), 'name': 'final_result', + 'parameters_json_schema': { + 'properties': {'content': {'type': 'string'}}, + 'required': ['content'], + 'title': 'MyOutput', + 'type': 'object', + }, + 'description': 'The final response which ends this conversation', + 'outer_typed_dict_key': None, + 'strict': None, + 'sequential': False, + 'kind': 'output', + 'metadata': None, } ], + 'allow_text_output': False, } - ] - ) - ), - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.input.messages': {'type': 'array'}, - 'gen_ai.output.messages': {'type': 'array'}, - 'model_request_parameters': {'type': 'object'}, - }, - } - ) - ), - 'gen_ai.usage.input_tokens': 51, - 'gen_ai.usage.output_tokens': 5, - 'gen_ai.response.model': 'test', - } - ) + ) + ), + 'logfire.span_type': 'span', + 'logfire.msg': 'chat test', + 'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])), + 'gen_ai.output.messages': IsJson( + snapshot( + [ + { + 'role': 'assistant', + 'parts': [ + { + 'type': 'tool_call', + 'id': IsStr(), + 'name': 'final_result', + } + ], + } + ] + ) + ), + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.input.messages': {'type': 'array'}, + 'gen_ai.output.messages': {'type': 'array'}, + 'model_request_parameters': {'type': 'object'}, + }, + } + ) + ), + 'gen_ai.usage.input_tokens': 51, + 'gen_ai.usage.output_tokens': 5, + 'gen_ai.response.model': 'test', + } + ) def test_instrument_all(): @@ -991,15 +1242,28 @@ async def test_feedback(capfire: CaptureLogfire) -> None: @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') -@pytest.mark.parametrize('include_content,tool_error', [(True, False), (True, True), (False, False), (False, True)]) +@pytest.mark.parametrize( + 'include_content,tool_error,version', + [ + (True, False, 2), + (True, True, 2), + (False, False, 2), + (False, True, 2), + (True, False, 3), + (True, True, 3), + (False, False, 3), + (False, True, 3), + ], +) def test_include_tool_args_span_attributes( get_logfire_summary: Callable[[], LogfireSummary], include_content: bool, tool_error: bool, + version: Literal[2, 3], ) -> None: """Test that tool arguments are included/excluded in span attributes based on instrumentation settings.""" - instrumentation_settings = InstrumentationSettings(include_content=include_content) + instrumentation_settings = InstrumentationSettings(include_content=include_content, version=version) test_model = TestModel(seed=42) my_agent = Agent(model=test_model, instrument=instrumentation_settings) @@ -1024,101 +1288,187 @@ async def add_numbers(x: int, y: int) -> int: ) if include_content: - if tool_error: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"x":42,"y":42}', - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'tool_response': """\ + if version == 3: + if tool_error: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'gen_ai.tool.call.arguments': '{"x":42,"y":42}', + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.call.arguments': {'type': 'object'}, + 'gen_ai.tool.call.result': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'gen_ai.tool.call.result': """\ Tool error Fix the errors and try again.\ """, - 'logfire.level_num': 17, - } - ) + 'logfire.level_num': 17, + } + ) + else: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'gen_ai.tool.call.arguments': '{"x":42,"y":42}', + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.call.arguments': {'type': 'object'}, + 'gen_ai.tool.call.result': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'gen_ai.tool.call.result': '84', + } + ) else: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"x":42,"y":42}', - 'tool_response': '84', - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - } - ) + if tool_error: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"x":42,"y":42}', + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'tool_response': """\ +Tool error + +Fix the errors and try again.\ +""", + 'logfire.level_num': 17, + } + ) + else: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"x":42,"y":42}', + 'tool_response': '84', + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + } + ) else: - if tool_error: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'logfire.level_num': 17, - } - ) + if version == 3: + if tool_error: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + {'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}} + ) + ), + 'logfire.span_type': 'span', + 'logfire.level_num': 17, + } + ) + else: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + {'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}} + ) + ), + 'logfire.span_type': 'span', + } + ) else: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - } - ) + if tool_error: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'logfire.level_num': 17, + } + ) + else: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + } + ) class WeatherInfo(BaseModel): From eb1229ff5c5cc85640bf9507b52eb710eb061030 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 17:01:08 +1000 Subject: [PATCH 13/25] fix: add name field to SpanSummary and update test fixtures with span names --- tests/test_logfire.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 519b386ff1..307009c5fd 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -32,6 +32,7 @@ class SpanSummary(TypedDict): id: int + name: str message: str children: NotRequired[list[SpanSummary]] @@ -50,7 +51,9 @@ def __init__(self, capfire: CaptureLogfire): id_counter = 0 for span in spans: tid = span['context']['trace_id'], span['context']['span_id'] - span_lookup[tid] = span_summary = SpanSummary(id=id_counter, message=span['attributes']['logfire.msg']) + span_lookup[tid] = span_summary = SpanSummary( + id=id_counter, name=span['name'], message=span['attributes']['logfire.msg'] + ) self.attributes[id_counter] = span['attributes'] id_counter += 1 if parent := span['parent']: @@ -104,17 +107,19 @@ async def my_ret(x: int) -> str: [ { 'id': 0, + 'name': 'invoke_agent my_agent', 'message': 'my_agent run', 'children': [ - {'id': 1, 'message': 'chat test'}, + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, { 'id': 2, + 'name': 'running tools', 'message': 'running 1 tool', 'children': [ - {'id': 3, 'message': 'running tool: my_ret'}, + {'id': 3, 'name': 'execute_tool my_ret', 'message': 'running tool: my_ret'}, ], }, - {'id': 4, 'message': 'chat test'}, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, ], } ] @@ -124,17 +129,19 @@ async def my_ret(x: int) -> str: [ { 'id': 0, + 'name': 'agent run', 'message': 'my_agent run', 'children': [ - {'id': 1, 'message': 'chat test'}, + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, { 'id': 2, + 'name': 'running tools', 'message': 'running 1 tool', 'children': [ - {'id': 3, 'message': 'running tool: my_ret'}, + {'id': 3, 'name': 'running tool', 'message': 'running tool: my_ret'}, ], }, - {'id': 4, 'message': 'chat test'}, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, ], } ] From 872fe1aadb6109cc0fa303284ba51493247b6b0b Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 17:21:13 +1000 Subject: [PATCH 14/25] feat: add output tool span configuration and improve logfire schema attributes --- .../pydantic_ai/_instrumentation.py | 30 +++++++++--- pydantic_ai_slim/pydantic_ai/_output.py | 38 ++++++++++----- tests/test_logfire.py | 48 +++++++++++++++++++ 3 files changed, 98 insertions(+), 18 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_instrumentation.py b/pydantic_ai_slim/pydantic_ai/_instrumentation.py index 2e432f7246..4e56f8c91a 100644 --- a/pydantic_ai_slim/pydantic_ai/_instrumentation.py +++ b/pydantic_ai_slim/pydantic_ai/_instrumentation.py @@ -23,6 +23,9 @@ class InstrumentationConfig: tool_arguments_attr: str tool_result_attr: str + # Output Tool execution span configuration + output_tool_span_name: str + @classmethod def for_version(cls, version: int) -> Self: """Create instrumentation configuration for a specific version. @@ -40,6 +43,7 @@ def for_version(cls, version: int) -> Self: tool_span_name='running tool', tool_arguments_attr='tool_arguments', tool_result_attr='tool_response', + output_tool_span_name='running output function', ) else: return cls( @@ -48,8 +52,22 @@ def for_version(cls, version: int) -> Self: tool_span_name='execute_tool', # Will be formatted with tool name tool_arguments_attr='gen_ai.tool.call.arguments', tool_result_attr='gen_ai.tool.call.result', + output_tool_span_name='execute_tool', ) + def get_agent_run_span_name(self, agent_name: str) -> str: + """Get the formatted agent span name. + + Args: + agent_name: Name of the agent being executed + + Returns: + Formatted span name + """ + if self.agent_run_span_name == 'invoke_agent': + return f'invoke_agent {agent_name}' + return self.agent_run_span_name + def get_tool_span_name(self, tool_name: str) -> str: """Get the formatted tool span name. @@ -63,15 +81,15 @@ def get_tool_span_name(self, tool_name: str) -> str: return f'execute_tool {tool_name}' return self.tool_span_name - def get_agent_run_span_name(self, agent_name: str) -> str: - """Get the formatted agent span name. + def get_output_tool_span_name(self, tool_name: str) -> str: + """Get the formatted output tool span name. Args: - agent_name: Name of the agent being executed + tool_name: Name of the tool being executed Returns: Formatted span name """ - if self.agent_run_span_name == 'invoke_agent': - return f'invoke_agent {agent_name}' - return self.agent_run_span_name + if self.output_tool_span_name == 'execute_tool': + return f'execute_tool {tool_name}' + return self.output_tool_span_name diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index bd0245ae83..22318b5758 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -11,6 +11,8 @@ from pydantic_core import SchemaValidator, to_json from typing_extensions import Self, TypedDict, TypeVar, assert_never +from pydantic_ai._instrumentation import InstrumentationConfig + from . import _function_schema, _utils, messages as _messages from ._run_context import AgentDepsT, RunContext from .exceptions import ModelRetry, ToolRetryError, UserError @@ -94,6 +96,7 @@ async def execute_traced_output_function( ToolRetryError: When wrap_validation_errors is True and a ModelRetry is caught ModelRetry: When wrap_validation_errors is False and a ModelRetry occurs """ + instrumentation_config = InstrumentationConfig.for_version(run_context.instrumentation_version) # Set up span attributes tool_name = run_context.tool_name or getattr(function_schema.function, '__name__', 'output_function') attributes = { @@ -103,18 +106,29 @@ async def execute_traced_output_function( if run_context.tool_call_id: attributes['gen_ai.tool.call.id'] = run_context.tool_call_id if run_context.trace_include_content: - attributes['tool_arguments'] = to_json(args).decode() - attributes['logfire.json_schema'] = json.dumps( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - }, - } - ) + attributes[instrumentation_config.tool_arguments_attr] = to_json(args).decode() + + attributes['logfire.json_schema'] = json.dumps( + { + 'type': 'object', + 'properties': { + **( + { + instrumentation_config.tool_arguments_attr: {'type': 'object'}, + instrumentation_config.tool_result_attr: {'type': 'object'}, + } + if run_context.trace_include_content + else {} + ), + 'gen_ai.tool.name': {}, + **({'gen_ai.tool.call.id': {}} if run_context.tool_call_id else {}), + }, + } + ) - with run_context.tracer.start_as_current_span('running output function', attributes=attributes) as span: + with run_context.tracer.start_as_current_span( + instrumentation_config.get_output_tool_span_name(tool_name), attributes=attributes + ) as span: try: output = await function_schema.call(args, run_context) except ModelRetry as r: @@ -134,7 +148,7 @@ async def execute_traced_output_function( from .models.instrumented import InstrumentedModel span.set_attribute( - 'tool_response', + instrumentation_config.tool_result_attr, output if isinstance(output, str) else json.dumps(InstrumentedModel.serialize_any(output)), ) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 307009c5fd..09c56acb2d 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1525,6 +1525,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1539,6 +1541,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', } ) @@ -1586,6 +1591,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1600,6 +1607,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', } ) @@ -1655,6 +1665,8 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1674,6 +1686,8 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1690,6 +1704,9 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'logfire.msg': 'running output function: final_result', 'gen_ai.tool.call.id': IsStr(), + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', 'logfire.level_num': 17, }, @@ -1697,6 +1714,9 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'logfire.msg': 'running output function: final_result', 'gen_ai.tool.call.id': IsStr(), + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', }, ] @@ -1743,6 +1763,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1757,6 +1779,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'get_weather', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: get_weather', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', } ) @@ -1809,6 +1834,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1823,6 +1850,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', } ) @@ -1876,6 +1906,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1890,6 +1922,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', } ) @@ -1938,6 +1973,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1952,6 +1989,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), 'logfire.span_type': 'span', } ) @@ -2004,6 +2044,7 @@ def call_text_response(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, }, } ) @@ -2017,6 +2058,7 @@ def call_text_response(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: { 'gen_ai.tool.name': 'upcase_text', 'logfire.msg': 'running output function: upcase_text', + 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', } ) @@ -2074,6 +2116,7 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, }, } ) @@ -2087,6 +2130,7 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: { 'gen_ai.tool.name': 'upcase_text', 'logfire.msg': 'running output function: upcase_text', + 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', } ) @@ -2144,6 +2188,7 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, }, } ) @@ -2162,6 +2207,7 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, }, } ) @@ -2177,12 +2223,14 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: { 'gen_ai.tool.name': 'get_weather_with_retry', 'logfire.msg': 'running output function: get_weather_with_retry', + 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', 'logfire.level_num': 17, }, { 'gen_ai.tool.name': 'get_weather_with_retry', 'logfire.msg': 'running output function: get_weather_with_retry', + 'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})), 'logfire.span_type': 'span', }, ] From 85d4a5d7f7d3e0c011c54aca6b74561949862f83 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 17:35:23 +1000 Subject: [PATCH 15/25] refactor: consolidate version 2 and 3 test assertions into shared logic --- tests/test_logfire.py | 541 +++++++++++------------------------------- 1 file changed, 133 insertions(+), 408 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 09c56acb2d..c37ecec277 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -147,61 +147,7 @@ async def my_ret(x: int) -> str: ] ) - if isinstance(instrument, InstrumentationSettings) and instrument.version == 3: - assert summary.attributes[0] == snapshot( - { - 'model_name': 'test', - 'agent_name': 'my_agent', - 'gen_ai.agent.name': 'my_agent', - 'logfire.msg': 'my_agent run', - 'logfire.span_type': 'span', - 'final_result': '{"my_ret":"1"}', - 'gen_ai.usage.input_tokens': 103, - 'gen_ai.usage.output_tokens': 12, - 'pydantic_ai.all_messages': IsJson( - snapshot( - [ - {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}, - { - 'role': 'assistant', - 'parts': [ - { - 'type': 'tool_call', - 'id': IsStr(), - 'name': 'my_ret', - 'arguments': {'x': 0}, - } - ], - }, - { - 'role': 'user', - 'parts': [ - { - 'type': 'tool_call_response', - 'id': IsStr(), - 'name': 'my_ret', - 'result': '1', - } - ], - }, - {'role': 'assistant', 'parts': [{'type': 'text', 'content': '{"my_ret":"1"}'}]}, - ] - ) - ), - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'pydantic_ai.all_messages': {'type': 'array'}, - 'final_result': {'type': 'object'}, - }, - } - ) - ), - } - ) - elif instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 2: + if instrument is True or (isinstance(instrument, InstrumentationSettings) and instrument.version in (2, 3)): assert summary.attributes[0] == snapshot( { 'model_name': 'test', @@ -608,7 +554,7 @@ class MyOutput: ] ) ) - elif instrument.version == 3: + elif instrument.version in (2, 3): assert summary.attributes[0] == snapshot( { 'model_name': 'test', @@ -663,61 +609,6 @@ class MyOutput: ), } ) - elif instrument.version == 2: - assert summary.attributes[0] == snapshot( - { - 'model_name': 'test', - 'agent_name': 'my_agent', - 'gen_ai.agent.name': 'my_agent', - 'logfire.msg': 'my_agent run', - 'logfire.span_type': 'span', - 'gen_ai.usage.input_tokens': 51, - 'gen_ai.usage.output_tokens': 5, - 'pydantic_ai.all_messages': IsJson( - snapshot( - [ - {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}, - { - 'role': 'assistant', - 'parts': [ - { - 'type': 'tool_call', - 'id': IsStr(), - 'name': 'final_result', - 'arguments': {'content': 'a'}, - } - ], - }, - { - 'role': 'user', - 'parts': [ - { - 'type': 'tool_call_response', - 'id': IsStr(), - 'name': 'final_result', - 'result': 'Final result processed.', - } - ], - }, - ] - ) - ), - 'final_result': '{"content": "a"}', - 'gen_ai.system_instructions': '[{"type": "text", "content": "Here are some instructions"}]', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'pydantic_ai.all_messages': {'type': 'array'}, - 'gen_ai.system_instructions': {'type': 'array'}, - 'final_result': {'type': 'object'}, - }, - } - ) - ), - } - ) assert chat_span_attributes['gen_ai.input.messages'] == IsJson( snapshot([{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}]) @@ -848,246 +739,126 @@ class MyOutput: assert result.output == MyOutput(content='a') summary = get_logfire_summary() - if version == 2: - assert summary.attributes[0] == snapshot( - { - 'model_name': 'test', - 'agent_name': 'my_agent', - 'gen_ai.agent.name': 'my_agent', - 'logfire.msg': 'my_agent run', - 'logfire.span_type': 'span', - 'gen_ai.usage.input_tokens': 51, - 'gen_ai.usage.output_tokens': 5, - 'pydantic_ai.all_messages': IsJson( - snapshot( - [ - {'role': 'user', 'parts': [{'type': 'text'}]}, - { - 'role': 'assistant', - 'parts': [ - { - 'type': 'tool_call', - 'id': IsStr(), - 'name': 'final_result', - } - ], - }, - { - 'role': 'user', - 'parts': [ - { - 'type': 'tool_call_response', - 'id': IsStr(), - 'name': 'final_result', - } - ], - }, - ] - ) - ), - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'pydantic_ai.all_messages': {'type': 'array'}, - 'final_result': {'type': 'object'}, - }, - } - ) - ), - } - ) - chat_span_attributes = summary.attributes[1] - assert chat_span_attributes == snapshot( - { - 'gen_ai.operation.name': 'chat', - 'gen_ai.system': 'test', - 'gen_ai.request.model': 'test', - 'model_request_parameters': IsJson( - snapshot( + # Version 2 and 3 have identical snapshots for this test case + assert summary.attributes[0] == snapshot( + { + 'model_name': 'test', + 'agent_name': 'my_agent', + 'gen_ai.agent.name': 'my_agent', + 'logfire.msg': 'my_agent run', + 'logfire.span_type': 'span', + 'gen_ai.usage.input_tokens': 51, + 'gen_ai.usage.output_tokens': 5, + 'pydantic_ai.all_messages': IsJson( + snapshot( + [ + {'role': 'user', 'parts': [{'type': 'text'}]}, { - 'function_tools': [], - 'builtin_tools': [], - 'output_mode': 'tool', - 'output_object': None, - 'output_tools': [ + 'role': 'assistant', + 'parts': [ { + 'type': 'tool_call', + 'id': IsStr(), 'name': 'final_result', - 'parameters_json_schema': { - 'properties': {'content': {'type': 'string'}}, - 'required': ['content'], - 'title': 'MyOutput', - 'type': 'object', - }, - 'description': 'The final response which ends this conversation', - 'outer_typed_dict_key': None, - 'strict': None, - 'sequential': False, - 'kind': 'output', - 'metadata': None, } ], - 'allow_text_output': False, - } - ) - ), - 'logfire.span_type': 'span', - 'logfire.msg': 'chat test', - 'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])), - 'gen_ai.output.messages': IsJson( - snapshot( - [ - { - 'role': 'assistant', - 'parts': [ - { - 'type': 'tool_call', - 'id': IsStr(), - 'name': 'final_result', - } - ], - } - ] - ) - ), - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.input.messages': {'type': 'array'}, - 'gen_ai.output.messages': {'type': 'array'}, - 'model_request_parameters': {'type': 'object'}, - }, - } - ) - ), - 'gen_ai.usage.input_tokens': 51, - 'gen_ai.usage.output_tokens': 5, - 'gen_ai.response.model': 'test', - } - ) - else: - assert summary.attributes[0] == snapshot( - { - 'model_name': 'test', - 'agent_name': 'my_agent', - 'gen_ai.agent.name': 'my_agent', - 'logfire.msg': 'my_agent run', - 'logfire.span_type': 'span', - 'gen_ai.usage.input_tokens': 51, - 'gen_ai.usage.output_tokens': 5, - 'pydantic_ai.all_messages': IsJson( - snapshot( - [ - {'role': 'user', 'parts': [{'type': 'text'}]}, - { - 'role': 'assistant', - 'parts': [ - { - 'type': 'tool_call', - 'id': IsStr(), - 'name': 'final_result', - } - ], - }, - { - 'role': 'user', - 'parts': [ - { - 'type': 'tool_call_response', - 'id': IsStr(), - 'name': 'final_result', - } - ], - }, - ] - ) - ), - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'pydantic_ai.all_messages': {'type': 'array'}, - 'final_result': {'type': 'object'}, - }, - } - ) - ), - } - ) - chat_span_attributes = summary.attributes[1] - assert chat_span_attributes == snapshot( - { - 'gen_ai.operation.name': 'chat', - 'gen_ai.system': 'test', - 'gen_ai.request.model': 'test', - 'model_request_parameters': IsJson( - snapshot( + }, { - 'function_tools': [], - 'builtin_tools': [], - 'output_mode': 'tool', - 'output_object': None, - 'output_tools': [ + 'role': 'user', + 'parts': [ { + 'type': 'tool_call_response', + 'id': IsStr(), 'name': 'final_result', - 'parameters_json_schema': { - 'properties': {'content': {'type': 'string'}}, - 'required': ['content'], - 'title': 'MyOutput', - 'type': 'object', - }, - 'description': 'The final response which ends this conversation', - 'outer_typed_dict_key': None, - 'strict': None, - 'sequential': False, - 'kind': 'output', - 'metadata': None, } ], - 'allow_text_output': False, - } - ) - ), - 'logfire.span_type': 'span', - 'logfire.msg': 'chat test', - 'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])), - 'gen_ai.output.messages': IsJson( - snapshot( - [ + }, + ] + ) + ), + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'pydantic_ai.all_messages': {'type': 'array'}, + 'final_result': {'type': 'object'}, + }, + } + ) + ), + } + ) + chat_span_attributes = summary.attributes[1] + assert chat_span_attributes == snapshot( + { + 'gen_ai.operation.name': 'chat', + 'gen_ai.system': 'test', + 'gen_ai.request.model': 'test', + 'model_request_parameters': IsJson( + snapshot( + { + 'function_tools': [], + 'builtin_tools': [], + 'output_mode': 'tool', + 'output_object': None, + 'output_tools': [ { - 'role': 'assistant', - 'parts': [ - { - 'type': 'tool_call', - 'id': IsStr(), - 'name': 'final_result', - } - ], + 'name': 'final_result', + 'parameters_json_schema': { + 'properties': {'content': {'type': 'string'}}, + 'required': ['content'], + 'title': 'MyOutput', + 'type': 'object', + }, + 'description': 'The final response which ends this conversation', + 'outer_typed_dict_key': None, + 'strict': None, + 'sequential': False, + 'kind': 'output', + 'metadata': None, } - ] - ) - ), - 'logfire.json_schema': IsJson( - snapshot( + ], + 'allow_text_output': False, + } + ) + ), + 'logfire.span_type': 'span', + 'logfire.msg': 'chat test', + 'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])), + 'gen_ai.output.messages': IsJson( + snapshot( + [ { - 'type': 'object', - 'properties': { - 'gen_ai.input.messages': {'type': 'array'}, - 'gen_ai.output.messages': {'type': 'array'}, - 'model_request_parameters': {'type': 'object'}, - }, + 'role': 'assistant', + 'parts': [ + { + 'type': 'tool_call', + 'id': IsStr(), + 'name': 'final_result', + } + ], } - ) - ), - 'gen_ai.usage.input_tokens': 51, - 'gen_ai.usage.output_tokens': 5, - 'gen_ai.response.model': 'test', - } - ) + ] + ) + ), + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.input.messages': {'type': 'array'}, + 'gen_ai.output.messages': {'type': 'array'}, + 'model_request_parameters': {'type': 'object'}, + }, + } + ) + ), + 'gen_ai.usage.input_tokens': 51, + 'gen_ai.usage.output_tokens': 5, + 'gen_ai.response.model': 'test', + } + ) def test_instrument_all(): @@ -1404,78 +1175,32 @@ async def add_numbers(x: int, y: int) -> int: } ) else: - if version == 3: - if tool_error: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - {'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}} - ) - ), - 'logfire.span_type': 'span', - 'logfire.level_num': 17, - } - ) - else: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - {'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}} - ) - ), - 'logfire.span_type': 'span', - } - ) + # Version 2 and 3 have identical snapshots when include_content=False + if tool_error: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), + 'logfire.span_type': 'span', + 'logfire.level_num': 17, + } + ) else: - if tool_error: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'logfire.level_num': 17, - } - ) - else: - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - } - ) + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + ), + 'logfire.span_type': 'span', + } + ) class WeatherInfo(BaseModel): From 926258d0828e000a2007fdd17ce7f497a181050f Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 18:56:30 +1000 Subject: [PATCH 16/25] test: add trace assertions for different instrumentation versions in logfire tests --- tests/test_logfire.py | 554 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 500 insertions(+), 54 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index c37ecec277..5da4c1b8bc 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -148,6 +148,50 @@ async def my_ret(x: int) -> str: ) if instrument is True or (isinstance(instrument, InstrumentationSettings) and instrument.version in (2, 3)): + if instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 2: + # default instrumentation settings + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [{'id': 3, 'name': 'running tool', 'message': 'running tool: my_ret'}], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + ], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + {'id': 3, 'name': 'execute_tool my_ret', 'message': 'running tool: my_ret'} + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + ], + } + ] + ) + assert summary.attributes[0] == snapshot( { 'model_name': 'test', @@ -555,6 +599,29 @@ class MyOutput: ) ) elif instrument.version in (2, 3): + if instrument.version == 2: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + } + ] + ) + assert summary.attributes[0] == snapshot( { 'model_name': 'test', @@ -739,6 +806,30 @@ class MyOutput: assert result.output == MyOutput(content='a') summary = get_logfire_summary() + + if version == 2: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + } + ] + ) + # Version 2 and 3 have identical snapshots for this test case assert summary.attributes[0] == snapshot( { @@ -1068,6 +1159,43 @@ async def add_numbers(x: int, y: int) -> int: if include_content: if version == 3: if tool_error: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + { + 'id': 3, + 'name': 'execute_tool add_numbers', + 'message': 'running tool: add_numbers', + } + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 5, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + { + 'id': 6, + 'name': 'execute_tool add_numbers', + 'message': 'running tool: add_numbers', + } + ], + }, + ], + } + ] + ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', @@ -1097,6 +1225,31 @@ async def add_numbers(x: int, y: int) -> int: } ) else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + { + 'id': 3, + 'name': 'execute_tool add_numbers', + 'message': 'running tool: add_numbers', + } + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + ], + } + ] + ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', @@ -1122,6 +1275,35 @@ async def add_numbers(x: int, y: int) -> int: ) else: if tool_error: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 5, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + {'id': 6, 'name': 'running tool', 'message': 'running tool: add_numbers'} + ], + }, + ], + } + ] + ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', @@ -1151,6 +1333,27 @@ async def add_numbers(x: int, y: int) -> int: } ) else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + ], + } + ] + ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', @@ -1177,6 +1380,74 @@ async def add_numbers(x: int, y: int) -> int: else: # Version 2 and 3 have identical snapshots when include_content=False if tool_error: + if version == 3: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + { + 'id': 3, + 'name': 'execute_tool add_numbers', + 'message': 'running tool: add_numbers', + } + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 5, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + { + 'id': 6, + 'name': 'execute_tool add_numbers', + 'message': 'running tool: add_numbers', + } + ], + }, + ], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 5, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + {'id': 6, 'name': 'running tool', 'message': 'running tool: add_numbers'} + ], + }, + ], + } + ] + ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', @@ -1190,6 +1461,54 @@ async def add_numbers(x: int, y: int) -> int: } ) else: + if version == 3: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + { + 'id': 3, + 'name': 'execute_tool add_numbers', + 'message': 'running tool: add_numbers', + } + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + ], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + { + 'id': 2, + 'name': 'running tools', + 'message': 'running 1 tool', + 'children': [ + {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} + ], + }, + {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + ], + } + ] + ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', @@ -1213,17 +1532,18 @@ def get_weather_info(city: str) -> WeatherInfo: @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') -@pytest.mark.parametrize('include_content', [True, False]) +@pytest.mark.parametrize('include_content, version', [(True, 2), (False, 2), (True, 3), (False, 3)]) def test_output_type_function_logfire_attributes( get_logfire_summary: Callable[[], LogfireSummary], include_content: bool, + version: Literal[2, 3], ) -> None: def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: assert info.output_tools is not None args_json = '{"city": "Mexico City"}' return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)]) - instrumentation_settings = InstrumentationSettings(include_content=include_content) + instrumentation_settings = InstrumentationSettings(include_content=include_content, version=version) my_agent = Agent(model=FunctionModel(call_tool), instrument=instrumentation_settings) result = my_agent.run_sync('Mexico City', output_type=get_weather_info) @@ -1231,36 +1551,99 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: summary = get_logfire_summary() + if version == 2: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, + { + 'id': 2, + 'name': 'running output function', + 'message': 'running output function: final_result', + }, + ], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, + { + 'id': 2, + 'name': 'execute_tool final_result', + 'message': 'running output function: final_result', + }, + ], + } + ] + ) + # Find the output function span attributes - [output_function_attributes] = [ + output_function_attributes = next( attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'final_result' - ] + ) if include_content: - assert output_function_attributes == snapshot( - { - 'gen_ai.tool.name': 'final_result', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"city":"Mexico City"}', - 'logfire.msg': 'running output function: final_result', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'tool_response': '{"temperature": 28.7, "description": "sunny"}', - } - ) + if version == 2: + assert output_function_attributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"city":"Mexico City"}', + 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'tool_response': '{"temperature": 28.7, "description": "sunny"}', + } + ) + else: + assert output_function_attributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'logfire.msg': 'running output function: final_result', + 'gen_ai.tool.call.id': IsStr(), + 'gen_ai.tool.call.arguments': '{"city":"Mexico City"}', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.call.arguments': {'type': 'object'}, + 'gen_ai.tool.call.result': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'gen_ai.tool.call.result': '{"temperature": 28.7, "description": "sunny"}', + } + ) else: + # Version 2 and 3 have the same output function attributes here assert output_function_attributes == snapshot( { 'gen_ai.tool.name': 'final_result', @@ -1275,10 +1658,11 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') -@pytest.mark.parametrize('include_content', [True, False]) +@pytest.mark.parametrize('include_content,version', [(True, 2), (False, 2), (True, 3), (False, 3)]) def test_output_type_function_with_run_context_logfire_attributes( get_logfire_summary: Callable[[], LogfireSummary], include_content: bool, + version: Literal[2, 3], ) -> None: def get_weather_with_ctx(ctx: RunContext[None], city: str) -> WeatherInfo: assert ctx is not None @@ -1289,7 +1673,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: args_json = '{"city": "Mexico City"}' return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)]) - instrumentation_settings = InstrumentationSettings(include_content=include_content) + instrumentation_settings = InstrumentationSettings(include_content=include_content, version=version) my_agent = Agent(model=FunctionModel(call_tool), instrument=instrumentation_settings) result = my_agent.run_sync('Mexico City', output_type=get_weather_with_ctx) @@ -1297,35 +1681,97 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: summary = get_logfire_summary() + if version == 2: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'agent run', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, + { + 'id': 2, + 'name': 'running output function', + 'message': 'running output function: final_result', + }, + ], + } + ] + ) + else: + assert summary.traces == snapshot( + [ + { + 'id': 0, + 'name': 'invoke_agent my_agent', + 'message': 'my_agent run', + 'children': [ + {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, + { + 'id': 2, + 'name': 'execute_tool final_result', + 'message': 'running output function: final_result', + }, + ], + } + ] + ) + # Find the output function span attributes - [output_function_attributes] = [ + output_function_attributes = next( attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'final_result' - ] + ) if include_content: - assert output_function_attributes == snapshot( - { - 'gen_ai.tool.name': 'final_result', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"city":"Mexico City"}', - 'logfire.msg': 'running output function: final_result', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'tool_response': '{"temperature": 28.7, "description": "sunny"}', - } - ) + if version == 2: + assert output_function_attributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"city":"Mexico City"}', + 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'tool_response': '{"temperature": 28.7, "description": "sunny"}', + } + ) + else: + assert output_function_attributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'logfire.msg': 'running output function: final_result', + 'gen_ai.tool.call.id': IsStr(), + 'gen_ai.tool.call.arguments': '{"city":"Mexico City"}', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.call.arguments': {'type': 'object'}, + 'gen_ai.tool.call.result': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'gen_ai.tool.call.result': '{"temperature": 28.7, "description": "sunny"}', + } + ) else: assert output_function_attributes == snapshot( { From 7e3f652bdc502a075d3ac548e969477bd1357897 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 19:05:45 +1000 Subject: [PATCH 17/25] refactor: revert logfire test by removing version parameter --- tests/test_logfire.py | 413 +++++++----------------------------------- 1 file changed, 66 insertions(+), 347 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 5da4c1b8bc..7f3f1adc21 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1111,28 +1111,15 @@ async def test_feedback(capfire: CaptureLogfire) -> None: @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') -@pytest.mark.parametrize( - 'include_content,tool_error,version', - [ - (True, False, 2), - (True, True, 2), - (False, False, 2), - (False, True, 2), - (True, False, 3), - (True, True, 3), - (False, False, 3), - (False, True, 3), - ], -) +@pytest.mark.parametrize('include_content,tool_error', [(True, False), (True, True), (False, False), (False, True)]) def test_include_tool_args_span_attributes( get_logfire_summary: Callable[[], LogfireSummary], include_content: bool, tool_error: bool, - version: Literal[2, 3], ) -> None: """Test that tool arguments are included/excluded in span attributes based on instrumentation settings.""" - instrumentation_settings = InstrumentationSettings(include_content=include_content, version=version) + instrumentation_settings = InstrumentationSettings(include_content=include_content) test_model = TestModel(seed=42) my_agent = Agent(model=test_model, instrument=instrumentation_settings) @@ -1157,365 +1144,97 @@ async def add_numbers(x: int, y: int) -> int: ) if include_content: - if version == 3: - if tool_error: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'invoke_agent my_agent', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - { - 'id': 3, - 'name': 'execute_tool add_numbers', - 'message': 'running tool: add_numbers', - } - ], - }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 5, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - { - 'id': 6, - 'name': 'execute_tool add_numbers', - 'message': 'running tool: add_numbers', - } - ], + if tool_error: + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"x":42,"y":42}', + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, - ], - } - ] - ) - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'gen_ai.tool.call.arguments': '{"x":42,"y":42}', - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.call.arguments': {'type': 'object'}, - 'gen_ai.tool.call.result': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'gen_ai.tool.call.result': """\ + } + ) + ), + 'logfire.span_type': 'span', + 'tool_response': """\ Tool error Fix the errors and try again.\ """, - 'logfire.level_num': 17, - } - ) - else: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'invoke_agent my_agent', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - { - 'id': 3, - 'name': 'execute_tool add_numbers', - 'message': 'running tool: add_numbers', - } - ], - }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - ], - } - ] - ) - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'gen_ai.tool.call.arguments': '{"x":42,"y":42}', - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.call.arguments': {'type': 'object'}, - 'gen_ai.tool.call.result': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'gen_ai.tool.call.result': '84', - } - ) + 'logfire.level_num': 17, + } + ) else: - if tool_error: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'agent run', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} - ], - }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 5, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - {'id': 6, 'name': 'running tool', 'message': 'running tool: add_numbers'} - ], - }, - ], - } - ] - ) - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"x":42,"y":42}', - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'tool_response': """\ -Tool error - -Fix the errors and try again.\ -""", - 'logfire.level_num': 17, - } - ) - else: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'agent run', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} - ], + assert tool_attributes == snapshot( + { + 'gen_ai.tool.name': 'add_numbers', + 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"x":42,"y":42}', + 'tool_response': '84', + 'logfire.msg': 'running tool: add_numbers', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - ], - } - ] - ) - assert tool_attributes == snapshot( - { - 'gen_ai.tool.name': 'add_numbers', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"x":42,"y":42}', - 'tool_response': '84', - 'logfire.msg': 'running tool: add_numbers', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - } - ) + } + ) + ), + 'logfire.span_type': 'span', + } + ) else: - # Version 2 and 3 have identical snapshots when include_content=False if tool_error: - if version == 3: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'invoke_agent my_agent', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - { - 'id': 3, - 'name': 'execute_tool add_numbers', - 'message': 'running tool: add_numbers', - } - ], - }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 5, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - { - 'id': 6, - 'name': 'execute_tool add_numbers', - 'message': 'running tool: add_numbers', - } - ], - }, - ], - } - ] - ) - else: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'agent run', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} - ], - }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 5, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - {'id': 6, 'name': 'running tool', 'message': 'running tool: add_numbers'} - ], - }, - ], - } - ] - ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running tool: add_numbers', 'logfire.json_schema': IsJson( - snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) ), 'logfire.span_type': 'span', 'logfire.level_num': 17, } ) else: - if version == 3: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'invoke_agent my_agent', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - { - 'id': 3, - 'name': 'execute_tool add_numbers', - 'message': 'running tool: add_numbers', - } - ], - }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - ], - } - ] - ) - else: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'agent run', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat test', 'message': 'chat test'}, - { - 'id': 2, - 'name': 'running tools', - 'message': 'running 1 tool', - 'children': [ - {'id': 3, 'name': 'running tool', 'message': 'running tool: add_numbers'} - ], - }, - {'id': 4, 'name': 'chat test', 'message': 'chat test'}, - ], - } - ] - ) assert tool_attributes == snapshot( { 'gen_ai.tool.name': 'add_numbers', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running tool: add_numbers', 'logfire.json_schema': IsJson( - snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) ), 'logfire.span_type': 'span', } From dd4b03cce6cbb4613cdf45e3928e16f732ccb161 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 19:28:34 +1000 Subject: [PATCH 18/25] refactor: consolidate output function test cases and add v2/v3 instrumentation settings --- tests/test_logfire.py | 282 ++++++++++++++++++++---------------------- 1 file changed, 132 insertions(+), 150 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 7f3f1adc21..1d2128dac3 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1251,26 +1251,34 @@ def get_weather_info(city: str) -> WeatherInfo: @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') -@pytest.mark.parametrize('include_content, version', [(True, 2), (False, 2), (True, 3), (False, 3)]) -def test_output_type_function_logfire_attributes( +@pytest.mark.parametrize( + 'instrument', + [ + True, + False, + InstrumentationSettings(version=2), + InstrumentationSettings(version=3), + ], +) +def test_logfire_output_function_v2_v3( get_logfire_summary: Callable[[], LogfireSummary], - include_content: bool, - version: Literal[2, 3], + instrument: InstrumentationSettings | bool, ) -> None: def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: assert info.output_tools is not None args_json = '{"city": "Mexico City"}' return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)]) - instrumentation_settings = InstrumentationSettings(include_content=include_content, version=version) - my_agent = Agent(model=FunctionModel(call_tool), instrument=instrumentation_settings) - + my_agent = Agent(model=FunctionModel(call_tool), instrument=instrument) result = my_agent.run_sync('Mexico City', output_type=get_weather_info) assert result.output == WeatherInfo(temperature=28.7, description='sunny') summary = get_logfire_summary() - if version == 2: + if instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 2: + output_function_atttributes = next( + (attr for attr in summary.attributes.values() if attr.get('gen_ai.tool.name') == 'final_result'), + ) assert summary.traces == snapshot( [ { @@ -1288,7 +1296,34 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: } ] ) - else: + assert output_function_atttributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'logfire.msg': 'running output function: final_result', + 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"city":"Mexico City"}', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'tool_response': '{"temperature": 28.7, "description": "sunny"}', + } + ) + + elif isinstance(instrument, InstrumentationSettings) and instrument.version == 3: + output_function_atttributes = next( + (attr for attr in summary.attributes.values() if attr.get('gen_ai.tool.name') == 'final_result'), + ) assert summary.traces == snapshot( [ { @@ -1306,6 +1341,52 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: } ] ) + assert output_function_atttributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'logfire.msg': 'running output function: final_result', + 'gen_ai.tool.call.id': IsStr(), + 'gen_ai.tool.call.arguments': '{"city":"Mexico City"}', + 'logfire.json_schema': IsJson( + snapshot( + { + 'type': 'object', + 'properties': { + 'gen_ai.tool.call.arguments': {'type': 'object'}, + 'gen_ai.tool.call.result': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, + }, + } + ) + ), + 'logfire.span_type': 'span', + 'gen_ai.tool.call.result': '{"temperature": 28.7, "description": "sunny"}', + } + ) + else: + assert summary.traces == snapshot([]) + assert summary.attributes == snapshot({}) + + +@pytest.mark.skipif(not logfire_installed, reason='logfire not installed') +@pytest.mark.parametrize('include_content', [True, False]) +def test_output_type_function_logfire_attributes( + get_logfire_summary: Callable[[], LogfireSummary], + include_content: bool, +) -> None: + def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: + assert info.output_tools is not None + args_json = '{"city": "Mexico City"}' + return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)]) + + instrumentation_settings = InstrumentationSettings(include_content=include_content) + my_agent = Agent(model=FunctionModel(call_tool), instrument=instrumentation_settings) + + result = my_agent.run_sync('Mexico City', output_type=get_weather_info) + assert result.output == WeatherInfo(temperature=28.7, description='sunny') + + summary = get_logfire_summary() # Find the output function span attributes output_function_attributes = next( @@ -1313,75 +1394,43 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: ) if include_content: - if version == 2: - assert output_function_attributes == snapshot( - { - 'gen_ai.tool.name': 'final_result', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"city":"Mexico City"}', - 'logfire.msg': 'running output function: final_result', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'tool_response': '{"temperature": 28.7, "description": "sunny"}', - } - ) - else: - assert output_function_attributes == snapshot( - { - 'gen_ai.tool.name': 'final_result', - 'logfire.msg': 'running output function: final_result', - 'gen_ai.tool.call.id': IsStr(), - 'gen_ai.tool.call.arguments': '{"city":"Mexico City"}', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.call.arguments': {'type': 'object'}, - 'gen_ai.tool.call.result': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'gen_ai.tool.call.result': '{"temperature": 28.7, "description": "sunny"}', - } - ) - else: - # Version 2 and 3 have the same output function attributes here assert output_function_attributes == snapshot( { 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"city":"Mexico City"}', 'logfire.msg': 'running output function: final_result', 'logfire.json_schema': IsJson( - snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + }, + } + ) ), 'logfire.span_type': 'span', + 'tool_response': '{"temperature": 28.7, "description": "sunny"}', + } + ) + else: + assert output_function_attributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running output function: final_result', + 'logfire.span_type': 'span', } ) @pytest.mark.skipif(not logfire_installed, reason='logfire not installed') -@pytest.mark.parametrize('include_content,version', [(True, 2), (False, 2), (True, 3), (False, 3)]) +@pytest.mark.parametrize('include_content', [True, False]) def test_output_type_function_with_run_context_logfire_attributes( get_logfire_summary: Callable[[], LogfireSummary], include_content: bool, - version: Literal[2, 3], ) -> None: def get_weather_with_ctx(ctx: RunContext[None], city: str) -> WeatherInfo: assert ctx is not None @@ -1392,7 +1441,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: args_json = '{"city": "Mexico City"}' return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)]) - instrumentation_settings = InstrumentationSettings(include_content=include_content, version=version) + instrumentation_settings = InstrumentationSettings(include_content=include_content) my_agent = Agent(model=FunctionModel(call_tool), instrument=instrumentation_settings) result = my_agent.run_sync('Mexico City', output_type=get_weather_with_ctx) @@ -1400,107 +1449,40 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: summary = get_logfire_summary() - if version == 2: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'agent run', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, - { - 'id': 2, - 'name': 'running output function', - 'message': 'running output function: final_result', - }, - ], - } - ] - ) - else: - assert summary.traces == snapshot( - [ - { - 'id': 0, - 'name': 'invoke_agent my_agent', - 'message': 'my_agent run', - 'children': [ - {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, - { - 'id': 2, - 'name': 'execute_tool final_result', - 'message': 'running output function: final_result', - }, - ], - } - ] - ) - # Find the output function span attributes output_function_attributes = next( attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'final_result' ) if include_content: - if version == 2: - assert output_function_attributes == snapshot( - { - 'gen_ai.tool.name': 'final_result', - 'gen_ai.tool.call.id': IsStr(), - 'tool_arguments': '{"city":"Mexico City"}', - 'logfire.msg': 'running output function: final_result', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'tool_arguments': {'type': 'object'}, - 'tool_response': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'tool_response': '{"temperature": 28.7, "description": "sunny"}', - } - ) - else: - assert output_function_attributes == snapshot( - { - 'gen_ai.tool.name': 'final_result', - 'logfire.msg': 'running output function: final_result', - 'gen_ai.tool.call.id': IsStr(), - 'gen_ai.tool.call.arguments': '{"city":"Mexico City"}', - 'logfire.json_schema': IsJson( - snapshot( - { - 'type': 'object', - 'properties': { - 'gen_ai.tool.call.arguments': {'type': 'object'}, - 'gen_ai.tool.call.result': {'type': 'object'}, - 'gen_ai.tool.name': {}, - 'gen_ai.tool.call.id': {}, - }, - } - ) - ), - 'logfire.span_type': 'span', - 'gen_ai.tool.call.result': '{"temperature": 28.7, "description": "sunny"}', - } - ) - else: assert output_function_attributes == snapshot( { 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), + 'tool_arguments': '{"city":"Mexico City"}', 'logfire.msg': 'running output function: final_result', 'logfire.json_schema': IsJson( - snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}}) + snapshot( + { + 'type': 'object', + 'properties': { + 'tool_arguments': {'type': 'object'}, + 'tool_response': {'type': 'object'}, + }, + } + ) ), 'logfire.span_type': 'span', + 'tool_response': '{"temperature": 28.7, "description": "sunny"}', + } + ) + else: + assert output_function_attributes == snapshot( + { + 'gen_ai.tool.name': 'final_result', + 'gen_ai.tool.call.id': IsStr(), + 'logfire.msg': 'running output function: final_result', + 'logfire.span_type': 'span', } ) From c887a8e1099e7d636c7e16a4c9bc524590a80c9f Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 19:42:49 +1000 Subject: [PATCH 19/25] fix: add missing gen_ai tool fields and json schema to logfire attributes --- tests/test_logfire.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 1d2128dac3..3ef9751c31 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1407,6 +1407,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1421,6 +1423,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': '{"type": "object", "properties": {"gen_ai.tool.name": {}, "gen_ai.tool.call.id": {}}}', 'logfire.span_type': 'span', } ) @@ -1468,6 +1471,8 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'properties': { 'tool_arguments': {'type': 'object'}, 'tool_response': {'type': 'object'}, + 'gen_ai.tool.name': {}, + 'gen_ai.tool.call.id': {}, }, } ) @@ -1482,6 +1487,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'gen_ai.tool.name': 'final_result', 'gen_ai.tool.call.id': IsStr(), 'logfire.msg': 'running output function: final_result', + 'logfire.json_schema': '{"type": "object", "properties": {"gen_ai.tool.name": {}, "gen_ai.tool.call.id": {}}}', 'logfire.span_type': 'span', } ) From cb3163e31ddd1d631b212d4cf1cc0848611f3f54 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Mon, 29 Sep 2025 19:58:49 +1000 Subject: [PATCH 20/25] fix: remove redundant version check in test_instructions_with_structured_output --- tests/test_logfire.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 3ef9751c31..03aa363072 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -598,7 +598,7 @@ class MyOutput: ] ) ) - elif instrument.version in (2, 3): + else: if instrument.version == 2: assert summary.traces == snapshot( [ From 903410a47cb613ea784781600dc6ca0fde98f98d Mon Sep 17 00:00:00 2001 From: bitnahian Date: Wed, 1 Oct 2025 11:27:35 +1000 Subject: [PATCH 21/25] refactor: rename InstrumentationConfig class to InstrumentationNames for clarity --- pydantic_ai_slim/pydantic_ai/_instrumentation.py | 2 +- pydantic_ai_slim/pydantic_ai/_output.py | 4 ++-- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 4 ++-- pydantic_ai_slim/pydantic_ai/agent/__init__.py | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_instrumentation.py b/pydantic_ai_slim/pydantic_ai/_instrumentation.py index 4e56f8c91a..cba1eb8139 100644 --- a/pydantic_ai_slim/pydantic_ai/_instrumentation.py +++ b/pydantic_ai_slim/pydantic_ai/_instrumentation.py @@ -11,7 +11,7 @@ @dataclass(frozen=True) -class InstrumentationConfig: +class InstrumentationNames: """Configuration for instrumentation span names and attributes based on version.""" # Agent run span configuration diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index 22318b5758..25d455f0f8 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -11,7 +11,7 @@ from pydantic_core import SchemaValidator, to_json from typing_extensions import Self, TypedDict, TypeVar, assert_never -from pydantic_ai._instrumentation import InstrumentationConfig +from pydantic_ai._instrumentation import InstrumentationNames from . import _function_schema, _utils, messages as _messages from ._run_context import AgentDepsT, RunContext @@ -96,7 +96,7 @@ async def execute_traced_output_function( ToolRetryError: When wrap_validation_errors is True and a ModelRetry is caught ModelRetry: When wrap_validation_errors is False and a ModelRetry occurs """ - instrumentation_config = InstrumentationConfig.for_version(run_context.instrumentation_version) + instrumentation_config = InstrumentationNames.for_version(run_context.instrumentation_version) # Set up span attributes tool_name = run_context.tool_name or getattr(function_schema.function, '__name__', 'output_function') attributes = { diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index efbfca9337..e145b71467 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -12,7 +12,7 @@ from typing_extensions import assert_never from . import messages as _messages -from ._instrumentation import InstrumentationConfig +from ._instrumentation import InstrumentationNames from ._run_context import AgentDepsT, RunContext from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior from .messages import ToolCallPart @@ -210,7 +210,7 @@ async def _call_tool_traced( usage_limits: UsageLimits | None = None, ) -> Any: """See .""" - instrumentation_config = InstrumentationConfig.for_version(instrumentation_version) + instrumentation_config = InstrumentationNames.for_version(instrumentation_version) span_attributes = { 'gen_ai.tool.name': call.tool_name, diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 83fe87631a..e808dd515c 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -14,7 +14,7 @@ from pydantic.json_schema import GenerateJsonSchema from typing_extensions import Self, TypeVar, deprecated -from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION, InstrumentationConfig +from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION, InstrumentationNames from pydantic_graph import Graph from .. import ( @@ -645,7 +645,7 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None: ) agent_name = self.name or 'agent' - instrumentation_config = InstrumentationConfig.for_version( + instrumentation_config = InstrumentationNames.for_version( instrumentation_settings.version if instrumentation_settings else DEFAULT_INSTRUMENTATION_VERSION ) From 2aeb162cf124f48d406ebd551f4f0981ac7dbfae Mon Sep 17 00:00:00 2001 From: bitnahian Date: Wed, 1 Oct 2025 18:49:29 +1000 Subject: [PATCH 22/25] refactor: replace next() with list comprehension for output function attributes lookup --- tests/test_logfire.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index 03aa363072..b1cf1be318 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1276,9 +1276,11 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: summary = get_logfire_summary() if instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 2: - output_function_atttributes = next( - (attr for attr in summary.attributes.values() if attr.get('gen_ai.tool.name') == 'final_result'), - ) + [output_function_atttributes] = [ + attributes + for attributes in summary.attributes.values() + if attributes.get('gen_ai.tool.name') == 'final_result' + ] assert summary.traces == snapshot( [ { @@ -1321,9 +1323,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: ) elif isinstance(instrument, InstrumentationSettings) and instrument.version == 3: - output_function_atttributes = next( + [output_function_atttributes] = [ (attr for attr in summary.attributes.values() if attr.get('gen_ai.tool.name') == 'final_result'), - ) + ] assert summary.traces == snapshot( [ { @@ -1389,9 +1391,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: summary = get_logfire_summary() # Find the output function span attributes - output_function_attributes = next( + [output_function_attributes] = [ attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'final_result' - ) + ] if include_content: assert output_function_attributes == snapshot( @@ -1453,9 +1455,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: summary = get_logfire_summary() # Find the output function span attributes - output_function_attributes = next( + [output_function_attributes] = [ attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'final_result' - ) + ] if include_content: assert output_function_attributes == snapshot( From e68656409c81ad9a62211c87eca4d23a8e5e3466 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Wed, 1 Oct 2025 18:52:37 +1000 Subject: [PATCH 23/25] fix: same --- tests/test_logfire.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index b1cf1be318..da2944449e 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1324,7 +1324,9 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: elif isinstance(instrument, InstrumentationSettings) and instrument.version == 3: [output_function_atttributes] = [ - (attr for attr in summary.attributes.values() if attr.get('gen_ai.tool.name') == 'final_result'), + attributes + for attributes in summary.attributes.values() + if attributes.get('gen_ai.tool.name') == 'final_result' ] assert summary.traces == snapshot( [ From 1e04c9a6e3c15df3abd185a3ba6c5ea157185d65 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Wed, 1 Oct 2025 19:05:41 +1000 Subject: [PATCH 24/25] refactor: rename instrumentation_config variable to instrumentation_names --- pydantic_ai_slim/pydantic_ai/_output.py | 12 ++++++------ pydantic_ai_slim/pydantic_ai/_tool_manager.py | 14 +++++++------- pydantic_ai_slim/pydantic_ai/agent/__init__.py | 4 ++-- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index 25d455f0f8..f1c1cc960b 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -96,7 +96,7 @@ async def execute_traced_output_function( ToolRetryError: When wrap_validation_errors is True and a ModelRetry is caught ModelRetry: When wrap_validation_errors is False and a ModelRetry occurs """ - instrumentation_config = InstrumentationNames.for_version(run_context.instrumentation_version) + instrumentation_names = InstrumentationNames.for_version(run_context.instrumentation_version) # Set up span attributes tool_name = run_context.tool_name or getattr(function_schema.function, '__name__', 'output_function') attributes = { @@ -106,7 +106,7 @@ async def execute_traced_output_function( if run_context.tool_call_id: attributes['gen_ai.tool.call.id'] = run_context.tool_call_id if run_context.trace_include_content: - attributes[instrumentation_config.tool_arguments_attr] = to_json(args).decode() + attributes[instrumentation_names.tool_arguments_attr] = to_json(args).decode() attributes['logfire.json_schema'] = json.dumps( { @@ -114,8 +114,8 @@ async def execute_traced_output_function( 'properties': { **( { - instrumentation_config.tool_arguments_attr: {'type': 'object'}, - instrumentation_config.tool_result_attr: {'type': 'object'}, + instrumentation_names.tool_arguments_attr: {'type': 'object'}, + instrumentation_names.tool_result_attr: {'type': 'object'}, } if run_context.trace_include_content else {} @@ -127,7 +127,7 @@ async def execute_traced_output_function( ) with run_context.tracer.start_as_current_span( - instrumentation_config.get_output_tool_span_name(tool_name), attributes=attributes + instrumentation_names.get_output_tool_span_name(tool_name), attributes=attributes ) as span: try: output = await function_schema.call(args, run_context) @@ -148,7 +148,7 @@ async def execute_traced_output_function( from .models.instrumented import InstrumentedModel span.set_attribute( - instrumentation_config.tool_result_attr, + instrumentation_names.tool_result_attr, output if isinstance(output, str) else json.dumps(InstrumentedModel.serialize_any(output)), ) diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index e145b71467..5cf66b00dd 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -210,13 +210,13 @@ async def _call_tool_traced( usage_limits: UsageLimits | None = None, ) -> Any: """See .""" - instrumentation_config = InstrumentationNames.for_version(instrumentation_version) + instrumentation_names = InstrumentationNames.for_version(instrumentation_version) span_attributes = { 'gen_ai.tool.name': call.tool_name, # NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai 'gen_ai.tool.call.id': call.tool_call_id, - **({instrumentation_config.tool_arguments_attr: call.args_as_json_str()} if include_content else {}), + **({instrumentation_names.tool_arguments_attr: call.args_as_json_str()} if include_content else {}), 'logfire.msg': f'running tool: {call.tool_name}', # add the JSON schema so these attributes are formatted nicely in Logfire 'logfire.json_schema': json.dumps( @@ -225,8 +225,8 @@ async def _call_tool_traced( 'properties': { **( { - instrumentation_config.tool_arguments_attr: {'type': 'object'}, - instrumentation_config.tool_result_attr: {'type': 'object'}, + instrumentation_names.tool_arguments_attr: {'type': 'object'}, + instrumentation_names.tool_result_attr: {'type': 'object'}, } if include_content else {} @@ -238,7 +238,7 @@ async def _call_tool_traced( ), } with tracer.start_as_current_span( - instrumentation_config.get_tool_span_name(call.tool_name), + instrumentation_names.get_tool_span_name(call.tool_name), attributes=span_attributes, ) as span: try: @@ -246,12 +246,12 @@ async def _call_tool_traced( except ToolRetryError as e: part = e.tool_retry if include_content and span.is_recording(): - span.set_attribute(instrumentation_config.tool_result_attr, part.model_response()) + span.set_attribute(instrumentation_names.tool_result_attr, part.model_response()) raise e if include_content and span.is_recording(): span.set_attribute( - instrumentation_config.tool_result_attr, + instrumentation_names.tool_result_attr, tool_result if isinstance(tool_result, str) else _messages.tool_return_ta.dump_json(tool_result).decode(), diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index e808dd515c..33f88f2abd 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -645,12 +645,12 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None: ) agent_name = self.name or 'agent' - instrumentation_config = InstrumentationNames.for_version( + instrumentation_names = InstrumentationNames.for_version( instrumentation_settings.version if instrumentation_settings else DEFAULT_INSTRUMENTATION_VERSION ) run_span = tracer.start_span( - instrumentation_config.get_agent_run_span_name(agent_name), + instrumentation_names.get_agent_run_span_name(agent_name), attributes={ 'model_name': model_used.model_name if model_used else 'no-model', 'agent_name': agent_name, From fb7ef44b5dfa2cef107aa8fc2f7345b436552645 Mon Sep 17 00:00:00 2001 From: bitnahian Date: Wed, 1 Oct 2025 20:26:10 +1000 Subject: [PATCH 25/25] fix: correct typo in output_function_attributes variable name --- tests/test_logfire.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index d687c1db53..091687f5a1 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -1275,7 +1275,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: summary = get_logfire_summary() if instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 2: - [output_function_atttributes] = [ + [output_function_attributes] = [ attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'final_result' @@ -1297,7 +1297,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: } ] ) - assert output_function_atttributes == snapshot( + assert output_function_attributes == snapshot( { 'gen_ai.tool.name': 'final_result', 'logfire.msg': 'running output function: final_result', @@ -1322,7 +1322,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: ) elif isinstance(instrument, InstrumentationSettings) and instrument.version == 3: - [output_function_atttributes] = [ + [output_function_attributes] = [ attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'final_result' @@ -1344,7 +1344,7 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: } ] ) - assert output_function_atttributes == snapshot( + assert output_function_attributes == snapshot( { 'gen_ai.tool.name': 'final_result', 'logfire.msg': 'running output function: final_result',