Add span_attributes to InstrumentationSettings.

This allows users to add extra attributes to the agent's span,
either as a direct dict, or a callable that takes the RunContext.
The attributes are added/computed after the agent finishes.

Author: sstanovnik

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -15,6 +15,7 @@
 from typing_extensions import Self, TypeVar, deprecated
 
 from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION, InstrumentationNames
+from pydantic_graph import GraphRunContext
 
 from .. import (
     _agent_graph,
@@ -667,21 +668,22 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
         finally:
             try:
                 if instrumentation_settings and run_span.is_recording():
-                    run_span.set_attributes(
-                        self._run_span_end_attributes(
-                            instrumentation_settings, usage, state.message_history, graph_deps.new_message_index
-                        )
-                    )
+                    graph_ctx = GraphRunContext(state=state, deps=graph_deps)
+                    span_attributes = self._run_span_end_attributes(instrumentation_settings, graph_ctx)
+                    run_span.set_attributes(span_attributes)
             finally:
                 run_span.end()
 
     def _run_span_end_attributes(
         self,
         settings: InstrumentationSettings,
-        usage: _usage.RunUsage,
-        message_history: list[_messages.ModelMessage],
-        new_message_index: int,
+        graph_ctx: GraphRunContext[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, OutputDataT]],
     ):
+        run_ctx = _agent_graph.build_run_context(graph_ctx)
+        usage = run_ctx.usage
+        message_history = run_ctx.messages
+        new_message_index = graph_ctx.deps.new_message_index
+
         if settings.version == 1:
             attrs = {
                 'all_messages_events': json.dumps(
@@ -714,9 +716,20 @@ def _run_span_end_attributes(
             ):
                 attrs['pydantic_ai.variable_instructions'] = True
 
+        extra_attributes: dict[str, str] = {}
+        span_attrs_setting = settings.span_attributes
+        if span_attrs_setting:
+            if callable(span_attrs_setting):
+                resolved_attrs = span_attrs_setting(run_ctx)
+                if resolved_attrs:
+                    extra_attributes = dict(resolved_attrs)
+            else:
+                extra_attributes = dict(span_attrs_setting)
+
         return {
             **usage.opentelemetry_attributes(),
             **attrs,
+            **extra_attributes,
             'logfire.json_schema': json.dumps(
                 {
                     'type': 'object',

pydantic_ai_slim/pydantic_ai/models/instrumented.py

@@ -93,6 +93,7 @@ class InstrumentationSettings:
     include_binary_content: bool = True
     include_content: bool = True
     version: Literal[1, 2, 3] = DEFAULT_INSTRUMENTATION_VERSION
+    span_attributes: dict[str, str] | Callable[[RunContext[Any]], dict[str, str]] = field(default_factory=dict)
 
     def __init__(
         self,
@@ -104,6 +105,7 @@ def __init__(
         version: Literal[1, 2, 3] = DEFAULT_INSTRUMENTATION_VERSION,
         event_mode: Literal['attributes', 'logs'] = 'attributes',
         event_logger_provider: EventLoggerProvider | None = None,
+        span_attributes: dict[str, str] | Callable[[RunContext[Any]], dict[str, str]] | None = None,
     ):
         """Create instrumentation options.
 
@@ -132,6 +134,9 @@ def __init__(
                 If not provided, the global event logger provider is used.
                 Calling `logfire.configure()` sets the global event logger provider, so most users don't need this.
                 This is only used if `event_mode='logs'` and `version=1`.
+            span_attributes: Additional attributes to add to the agent's span.
+                Can either be a dict of attributes, or a callable that accepts the RunContext to compute them.
+                Attributes are added or computed after the agent finishes running.
         """
         from pydantic_ai import __version__
 
@@ -145,6 +150,7 @@ def __init__(
         self.event_mode = event_mode
         self.include_binary_content = include_binary_content
         self.include_content = include_content
+        self.span_attributes = span_attributes or dict()
 
         if event_mode == 'logs' and version != 1:
             warnings.warn(

tests/test_logfire.py

@@ -1331,6 +1331,44 @@ async def add_numbers(x: int, y: int) -> int:
             )
 
 
+def _callable_span_attributes(ctx: RunContext[Any]) -> dict[str, str]:
+    """Helper callable that reflects the prompt in a span attribute."""
+    return {'callable_attr': str(ctx.prompt)}
+
+
+@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
+@pytest.mark.parametrize(
+    ('span_attrs', 'expected'),
+    [
+        pytest.param({'custom_attr': 'value'}, {'custom_attr': 'value'}, id='dict'),
+        pytest.param(_callable_span_attributes, {'callable_attr': 'span attributes prompt'}, id='callable'),
+    ],
+)
+def test_agent_span_attributes_extra(
+    get_logfire_summary: Callable[[], LogfireSummary],
+    span_attrs: dict[str, str] | Callable[[RunContext[Any]], dict[str, str]],
+    expected: dict[str, str],
+) -> None:
+    instrumentation_settings = InstrumentationSettings(span_attributes=span_attrs)
+    agent = Agent(
+        model=TestModel(),
+        instrument=instrumentation_settings,
+        name='span_attributes_agent',
+    )
+
+    agent.run_sync('span attributes prompt')
+
+    summary = get_logfire_summary()
+    agent_span_attributes = next(
+        attributes
+        for attributes in summary.attributes.values()
+        if attributes.get('gen_ai.agent.name') == 'span_attributes_agent'
+    )
+
+    for key, value in expected.items():
+        assert agent_span_attributes[key] == value
+
+
 class WeatherInfo(BaseModel):
     temperature: float
     description: str