Skip to content

Commit df333bd

Browse files
feat: Otel instrumentation version 3 (#3021)
Co-authored-by: Alex Hall <[email protected]>
1 parent 3164cb2 commit df333bd

File tree

9 files changed

+479
-49
lines changed

9 files changed

+479
-49
lines changed

pydantic_ai_slim/pydantic_ai/_agent_graph.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from typing_extensions import TypeVar, assert_never
1717

1818
from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx # type: ignore
19+
from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION
1920
from pydantic_ai._tool_manager import ToolManager
2021
from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, is_async_callable, run_in_executor
2122
from pydantic_ai.builtin_tools import AbstractBuiltinTool
@@ -704,6 +705,9 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
704705
tracer=ctx.deps.tracer,
705706
trace_include_content=ctx.deps.instrumentation_settings is not None
706707
and ctx.deps.instrumentation_settings.include_content,
708+
instrumentation_version=ctx.deps.instrumentation_settings.version
709+
if ctx.deps.instrumentation_settings
710+
else DEFAULT_INSTRUMENTATION_VERSION,
707711
run_step=ctx.state.run_step,
708712
tool_call_approved=ctx.state.run_step == 0,
709713
)
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
from __future__ import annotations
2+
3+
from dataclasses import dataclass
4+
from typing import TYPE_CHECKING
5+
6+
if TYPE_CHECKING:
7+
from typing_extensions import Self
8+
9+
DEFAULT_INSTRUMENTATION_VERSION = 2
10+
"""Default instrumentation version for `InstrumentationSettings`."""
11+
12+
13+
@dataclass(frozen=True)
14+
class InstrumentationNames:
15+
"""Configuration for instrumentation span names and attributes based on version."""
16+
17+
# Agent run span configuration
18+
agent_run_span_name: str
19+
agent_name_attr: str
20+
21+
# Tool execution span configuration
22+
tool_span_name: str
23+
tool_arguments_attr: str
24+
tool_result_attr: str
25+
26+
# Output Tool execution span configuration
27+
output_tool_span_name: str
28+
29+
@classmethod
30+
def for_version(cls, version: int) -> Self:
31+
"""Create instrumentation configuration for a specific version.
32+
33+
Args:
34+
version: The instrumentation version (1, 2, or 3+)
35+
36+
Returns:
37+
InstrumentationConfig instance with version-appropriate settings
38+
"""
39+
if version <= 2:
40+
return cls(
41+
agent_run_span_name='agent run',
42+
agent_name_attr='agent_name',
43+
tool_span_name='running tool',
44+
tool_arguments_attr='tool_arguments',
45+
tool_result_attr='tool_response',
46+
output_tool_span_name='running output function',
47+
)
48+
else:
49+
return cls(
50+
agent_run_span_name='invoke_agent',
51+
agent_name_attr='gen_ai.agent.name',
52+
tool_span_name='execute_tool', # Will be formatted with tool name
53+
tool_arguments_attr='gen_ai.tool.call.arguments',
54+
tool_result_attr='gen_ai.tool.call.result',
55+
output_tool_span_name='execute_tool',
56+
)
57+
58+
def get_agent_run_span_name(self, agent_name: str) -> str:
59+
"""Get the formatted agent span name.
60+
61+
Args:
62+
agent_name: Name of the agent being executed
63+
64+
Returns:
65+
Formatted span name
66+
"""
67+
if self.agent_run_span_name == 'invoke_agent':
68+
return f'invoke_agent {agent_name}'
69+
return self.agent_run_span_name
70+
71+
def get_tool_span_name(self, tool_name: str) -> str:
72+
"""Get the formatted tool span name.
73+
74+
Args:
75+
tool_name: Name of the tool being executed
76+
77+
Returns:
78+
Formatted span name
79+
"""
80+
if self.tool_span_name == 'execute_tool':
81+
return f'execute_tool {tool_name}'
82+
return self.tool_span_name
83+
84+
def get_output_tool_span_name(self, tool_name: str) -> str:
85+
"""Get the formatted output tool span name.
86+
87+
Args:
88+
tool_name: Name of the tool being executed
89+
90+
Returns:
91+
Formatted span name
92+
"""
93+
if self.output_tool_span_name == 'execute_tool':
94+
return f'execute_tool {tool_name}'
95+
return self.output_tool_span_name

pydantic_ai_slim/pydantic_ai/_output.py

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
from pydantic_core import SchemaValidator, to_json
1212
from typing_extensions import Self, TypedDict, TypeVar, assert_never
1313

14+
from pydantic_ai._instrumentation import InstrumentationNames
15+
1416
from . import _function_schema, _utils, messages as _messages
1517
from ._run_context import AgentDepsT, RunContext
1618
from .exceptions import ModelRetry, ToolRetryError, UserError
@@ -95,6 +97,7 @@ async def execute_traced_output_function(
9597
ToolRetryError: When wrap_validation_errors is True and a ModelRetry is caught
9698
ModelRetry: When wrap_validation_errors is False and a ModelRetry occurs
9799
"""
100+
instrumentation_names = InstrumentationNames.for_version(run_context.instrumentation_version)
98101
# Set up span attributes
99102
tool_name = run_context.tool_name or getattr(function_schema.function, '__name__', 'output_function')
100103
attributes = {
@@ -104,18 +107,29 @@ async def execute_traced_output_function(
104107
if run_context.tool_call_id:
105108
attributes['gen_ai.tool.call.id'] = run_context.tool_call_id
106109
if run_context.trace_include_content:
107-
attributes['tool_arguments'] = to_json(args).decode()
108-
attributes['logfire.json_schema'] = json.dumps(
109-
{
110-
'type': 'object',
111-
'properties': {
112-
'tool_arguments': {'type': 'object'},
113-
'tool_response': {'type': 'object'},
114-
},
115-
}
116-
)
110+
attributes[instrumentation_names.tool_arguments_attr] = to_json(args).decode()
111+
112+
attributes['logfire.json_schema'] = json.dumps(
113+
{
114+
'type': 'object',
115+
'properties': {
116+
**(
117+
{
118+
instrumentation_names.tool_arguments_attr: {'type': 'object'},
119+
instrumentation_names.tool_result_attr: {'type': 'object'},
120+
}
121+
if run_context.trace_include_content
122+
else {}
123+
),
124+
'gen_ai.tool.name': {},
125+
**({'gen_ai.tool.call.id': {}} if run_context.tool_call_id else {}),
126+
},
127+
}
128+
)
117129

118-
with run_context.tracer.start_as_current_span('running output function', attributes=attributes) as span:
130+
with run_context.tracer.start_as_current_span(
131+
instrumentation_names.get_output_tool_span_name(tool_name), attributes=attributes
132+
) as span:
119133
try:
120134
output = await function_schema.call(args, run_context)
121135
except ModelRetry as r:
@@ -135,7 +149,7 @@ async def execute_traced_output_function(
135149
from .models.instrumented import InstrumentedModel
136150

137151
span.set_attribute(
138-
'tool_response',
152+
instrumentation_names.tool_result_attr,
139153
output if isinstance(output, str) else json.dumps(InstrumentedModel.serialize_any(output)),
140154
)
141155

pydantic_ai_slim/pydantic_ai/_run_context.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
from opentelemetry.trace import NoOpTracer, Tracer
99
from typing_extensions import TypeVar
1010

11+
from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION
12+
1113
from . import _utils, messages as _messages
1214

1315
if TYPE_CHECKING:
@@ -36,6 +38,8 @@ class RunContext(Generic[AgentDepsT]):
3638
"""The tracer to use for tracing the run."""
3739
trace_include_content: bool = False
3840
"""Whether to include the content of the messages in the trace."""
41+
instrumentation_version: int = DEFAULT_INSTRUMENTATION_VERSION
42+
"""Instrumentation settings version, if instrumentation is enabled."""
3943
retries: dict[str, int] = field(default_factory=dict)
4044
"""Number of retries for each tool so far."""
4145
tool_call_id: str | None = None

pydantic_ai_slim/pydantic_ai/_tool_manager.py

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from typing_extensions import assert_never
1313

1414
from . import messages as _messages
15+
from ._instrumentation import InstrumentationNames
1516
from ._run_context import AgentDepsT, RunContext
1617
from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
1718
from .messages import ToolCallPart
@@ -115,6 +116,7 @@ async def handle_call(
115116
wrap_validation_errors,
116117
self.ctx.tracer,
117118
self.ctx.trace_include_content,
119+
self.ctx.instrumentation_version,
118120
usage_limits,
119121
)
120122

@@ -203,15 +205,18 @@ async def _call_tool_traced(
203205
allow_partial: bool,
204206
wrap_validation_errors: bool,
205207
tracer: Tracer,
206-
include_content: bool = False,
208+
include_content: bool,
209+
instrumentation_version: int,
207210
usage_limits: UsageLimits | None = None,
208211
) -> Any:
209212
"""See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
213+
instrumentation_names = InstrumentationNames.for_version(instrumentation_version)
214+
210215
span_attributes = {
211216
'gen_ai.tool.name': call.tool_name,
212217
# NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai
213218
'gen_ai.tool.call.id': call.tool_call_id,
214-
**({'tool_arguments': call.args_as_json_str()} if include_content else {}),
219+
**({instrumentation_names.tool_arguments_attr: call.args_as_json_str()} if include_content else {}),
215220
'logfire.msg': f'running tool: {call.tool_name}',
216221
# add the JSON schema so these attributes are formatted nicely in Logfire
217222
'logfire.json_schema': json.dumps(
@@ -220,8 +225,8 @@ async def _call_tool_traced(
220225
'properties': {
221226
**(
222227
{
223-
'tool_arguments': {'type': 'object'},
224-
'tool_response': {'type': 'object'},
228+
instrumentation_names.tool_arguments_attr: {'type': 'object'},
229+
instrumentation_names.tool_result_attr: {'type': 'object'},
225230
}
226231
if include_content
227232
else {}
@@ -232,18 +237,21 @@ async def _call_tool_traced(
232237
}
233238
),
234239
}
235-
with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
240+
with tracer.start_as_current_span(
241+
instrumentation_names.get_tool_span_name(call.tool_name),
242+
attributes=span_attributes,
243+
) as span:
236244
try:
237245
tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits)
238246
except ToolRetryError as e:
239247
part = e.tool_retry
240248
if include_content and span.is_recording():
241-
span.set_attribute('tool_response', part.model_response())
249+
span.set_attribute(instrumentation_names.tool_result_attr, part.model_response())
242250
raise e
243251

244252
if include_content and span.is_recording():
245253
span.set_attribute(
246-
'tool_response',
254+
instrumentation_names.tool_result_attr,
247255
tool_result
248256
if isinstance(tool_result, str)
249257
else _messages.tool_return_ta.dump_json(tool_result).decode(),

pydantic_ai_slim/pydantic_ai/agent/__init__.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from pydantic.json_schema import GenerateJsonSchema
1515
from typing_extensions import Self, TypeVar, deprecated
1616

17+
from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION, InstrumentationNames
1718
from pydantic_graph import Graph
1819

1920
from .. import (
@@ -644,11 +645,16 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
644645
)
645646

646647
agent_name = self.name or 'agent'
648+
instrumentation_names = InstrumentationNames.for_version(
649+
instrumentation_settings.version if instrumentation_settings else DEFAULT_INSTRUMENTATION_VERSION
650+
)
651+
647652
run_span = tracer.start_span(
648-
'agent run',
653+
instrumentation_names.get_agent_run_span_name(agent_name),
649654
attributes={
650655
'model_name': model_used.model_name if model_used else 'no-model',
651656
'agent_name': agent_name,
657+
'gen_ai.agent.name': agent_name,
652658
'logfire.msg': f'{agent_name} run',
653659
},
654660
)

pydantic_ai_slim/pydantic_ai/models/instrumented.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
from opentelemetry.util.types import AttributeValue
2222
from pydantic import TypeAdapter
2323

24+
from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION
25+
2426
from .. import _otel_messages
2527
from .._run_context import RunContext
2628
from ..messages import (
@@ -90,7 +92,7 @@ class InstrumentationSettings:
9092
event_mode: Literal['attributes', 'logs'] = 'attributes'
9193
include_binary_content: bool = True
9294
include_content: bool = True
93-
version: Literal[1, 2] = 1
95+
version: Literal[1, 2, 3] = DEFAULT_INSTRUMENTATION_VERSION
9496

9597
def __init__(
9698
self,
@@ -99,7 +101,7 @@ def __init__(
99101
meter_provider: MeterProvider | None = None,
100102
include_binary_content: bool = True,
101103
include_content: bool = True,
102-
version: Literal[1, 2] = 2,
104+
version: Literal[1, 2, 3] = DEFAULT_INSTRUMENTATION_VERSION,
103105
event_mode: Literal['attributes', 'logs'] = 'attributes',
104106
event_logger_provider: EventLoggerProvider | None = None,
105107
):

tests/models/test_fallback.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,7 @@ def test_first_failed_instrumented(capfire: CaptureLogfire) -> None:
169169
'attributes': {
170170
'model_name': 'fallback:function:failure_response:,function:success_response:',
171171
'agent_name': 'agent',
172+
'gen_ai.agent.name': 'agent',
172173
'logfire.msg': 'agent run',
173174
'logfire.span_type': 'span',
174175
'gen_ai.usage.input_tokens': 51,
@@ -268,6 +269,7 @@ async def test_first_failed_instrumented_stream(capfire: CaptureLogfire) -> None
268269
'attributes': {
269270
'model_name': 'fallback:function::failure_response_stream,function::success_response_stream',
270271
'agent_name': 'agent',
272+
'gen_ai.agent.name': 'agent',
271273
'logfire.msg': 'agent run',
272274
'logfire.span_type': 'span',
273275
'gen_ai.usage.input_tokens': 50,
@@ -375,6 +377,7 @@ def test_all_failed_instrumented(capfire: CaptureLogfire) -> None:
375377
'attributes': {
376378
'model_name': 'fallback:function:failure_response:,function:failure_response:',
377379
'agent_name': 'agent',
380+
'gen_ai.agent.name': 'agent',
378381
'logfire.msg': 'agent run',
379382
'logfire.span_type': 'span',
380383
'pydantic_ai.all_messages': [{'role': 'user', 'parts': [{'type': 'text', 'content': 'hello'}]}],

0 commit comments

Comments
 (0)