Replace model request span with InstrumentedModel (#1012)

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -286,21 +286,18 @@ async def _stream(
         assert not self._did_stream, 'stream() should only be called once per node'
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
-        with _logfire.span('model request', run_step=ctx.state.run_step) as span:
-            async with ctx.deps.model.request_stream(
-                ctx.state.message_history, model_settings, model_request_parameters
-            ) as streamed_response:
-                self._did_stream = True
-                ctx.state.usage.incr(_usage.Usage(), requests=1)
-                yield streamed_response
-                # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
-                # otherwise usage won't be properly counted:
-                async for _ in streamed_response:
-                    pass
-            model_response = streamed_response.get()
-            request_usage = streamed_response.usage()
-            span.set_attribute('response', model_response)
-            span.set_attribute('usage', request_usage)
+        async with ctx.deps.model.request_stream(
+            ctx.state.message_history, model_settings, model_request_parameters
+        ) as streamed_response:
+            self._did_stream = True
+            ctx.state.usage.incr(_usage.Usage(), requests=1)
+            yield streamed_response
+            # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
+            # otherwise usage won't be properly counted:
+            async for _ in streamed_response:
+                pass
+        model_response = streamed_response.get()
+        request_usage = streamed_response.usage()
 
         self._finish_handling(ctx, model_response, request_usage)
         assert self._result is not None  # this should be set by the previous line
@@ -312,13 +309,10 @@ async def _make_request(
             return self._result
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
-        with _logfire.span('model request', run_step=ctx.state.run_step) as span:
-            model_response, request_usage = await ctx.deps.model.request(
-                ctx.state.message_history, model_settings, model_request_parameters
-            )
-            ctx.state.usage.incr(_usage.Usage(), requests=1)
-            span.set_attribute('response', model_response)
-            span.set_attribute('usage', request_usage)
+        model_response, request_usage = await ctx.deps.model.request(
+            ctx.state.message_history, model_settings, model_request_parameters
+        )
+        ctx.state.usage.incr(_usage.Usage(), requests=1)
 
         return self._finish_handling(ctx, model_response, request_usage)
 

pydantic_ai_slim/pydantic_ai/agent.py

@@ -25,6 +25,7 @@
     result,
     usage as _usage,
 )
+from .models.instrumented import InstrumentedModel
 from .result import FinalResult, ResultDataT, StreamedRunResult
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
@@ -1115,6 +1116,9 @@ def _get_model(self, model: models.Model | models.KnownModelName | None) -> mode
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
+        if not isinstance(model_, InstrumentedModel):
+            model_ = InstrumentedModel(model_)
+
         return model_
 
     def _get_deps(self: Agent[T, ResultDataT], deps: T) -> T:

pydantic_ai_slim/pydantic_ai/messages.py

@@ -8,6 +8,7 @@
 
 import pydantic
 import pydantic_core
+from opentelemetry._events import Event
 from typing_extensions import TypeAlias
 
 from ._utils import now_utc as _now_utc
@@ -33,6 +34,9 @@ class SystemPromptPart:
     part_kind: Literal['system-prompt'] = 'system-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 
+    def otel_event(self) -> Event:
+        return Event('gen_ai.system.message', body={'content': self.content, 'role': 'system'})
+
 
 @dataclass
 class AudioUrl:
@@ -138,6 +142,14 @@ class UserPromptPart:
     part_kind: Literal['user-prompt'] = 'user-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 
+    def otel_event(self) -> Event:
+        if isinstance(self.content, str):
+            content = self.content
+        else:
+            # TODO figure out what to record for images and audio
+            content = [part if isinstance(part, str) else {'kind': part.kind} for part in self.content]
+        return Event('gen_ai.user.message', body={'content': content, 'role': 'user'})
+
 
 tool_return_ta: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter(Any, config=pydantic.ConfigDict(defer_build=True))
 
@@ -176,6 +188,9 @@ def model_response_object(self) -> dict[str, Any]:
         else:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
+    def otel_event(self) -> Event:
+        return Event('gen_ai.tool.message', body={'content': self.content, 'role': 'tool', 'id': self.tool_call_id})
+
 
 error_details_ta = pydantic.TypeAdapter(list[pydantic_core.ErrorDetails], config=pydantic.ConfigDict(defer_build=True))
 
@@ -224,6 +239,14 @@ def model_response(self) -> str:
             description = f'{len(self.content)} validation errors: {json_errors.decode()}'
         return f'{description}\n\nFix the errors and try again.'
 
+    def otel_event(self) -> Event:
+        if self.tool_name is None:
+            return Event('gen_ai.user.message', body={'content': self.model_response(), 'role': 'user'})
+        else:
+            return Event(
+                'gen_ai.tool.message', body={'content': self.model_response(), 'role': 'tool', 'id': self.tool_call_id}
+            )
+
 
 ModelRequestPart = Annotated[
     Union[SystemPromptPart, UserPromptPart, ToolReturnPart, RetryPromptPart], pydantic.Discriminator('part_kind')
@@ -329,6 +352,36 @@ class ModelResponse:
     kind: Literal['response'] = 'response'
     """Message type identifier, this is available on all parts as a discriminator."""
 
+    def otel_events(self) -> list[Event]:
+        """Return OpenTelemetry events for the response."""
+        result: list[Event] = []
+
+        def new_event_body():
+            new_body: dict[str, Any] = {'role': 'assistant'}
+            ev = Event('gen_ai.assistant.message', body=new_body)
+            result.append(ev)
+            return new_body
+
+        body = new_event_body()
+        for part in self.parts:
+            if isinstance(part, ToolCallPart):
+                body.setdefault('tool_calls', []).append(
+                    {
+                        'id': part.tool_call_id,
+                        'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
+                        'function': {
+                            'name': part.tool_name,
+                            'arguments': part.args,
+                        },
+                    }
+                )
+            elif isinstance(part, TextPart):
+                if body.get('content'):
+                    body = new_event_body()
+                body['content'] = part.content
+
+        return result
+
 
 ModelMessage = Annotated[Union[ModelRequest, ModelResponse], pydantic.Discriminator('kind')]
 """Any message sent to or returned by a model."""

pydantic_ai_slim/pydantic_ai/models/instrumented.py

@@ -1,28 +1,21 @@
 from __future__ import annotations
 
 import json
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Mapping
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
-from functools import partial
 from typing import Any, Callable, Literal
 
 import logfire_api
 from opentelemetry._events import Event, EventLogger, EventLoggerProvider, get_event_logger_provider
-from opentelemetry.trace import Tracer, TracerProvider, get_tracer_provider
+from opentelemetry.trace import Span, Tracer, TracerProvider, get_tracer_provider
 from opentelemetry.util.types import AttributeValue
+from pydantic import TypeAdapter
 
 from ..messages import (
     ModelMessage,
     ModelRequest,
-    ModelRequestPart,
     ModelResponse,
-    RetryPromptPart,
-    SystemPromptPart,
-    TextPart,
-    ToolCallPart,
-    ToolReturnPart,
-    UserPromptPart,
 )
 from ..settings import ModelSettings
 from ..usage import Usage
@@ -48,6 +41,8 @@
     'frequency_penalty',
 )
 
+ANY_ADAPTER = TypeAdapter[Any](Any)
+
 
 @dataclass
 class InstrumentedModel(WrapperModel):
@@ -115,7 +110,7 @@ async def request_stream(
                     finish(response_stream.get(), response_stream.usage())
 
     @contextmanager
-    def _instrument(  # noqa: C901
+    def _instrument(
         self,
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
@@ -141,35 +136,24 @@ def _instrument(  # noqa: C901
                 if isinstance(value := model_settings.get(key), (float, int)):
                     attributes[f'gen_ai.request.{key}'] = value
 
-        events_list = []
-        emit_event = partial(self._emit_event, system, events_list)
-
         with self.tracer.start_as_current_span(span_name, attributes=attributes) as span:
-            if span.is_recording():
-                for message in messages:
-                    if isinstance(message, ModelRequest):
-                        for part in message.parts:
-                            event_name, body = _request_part_body(part)
-                            if event_name:
-                                emit_event(event_name, body)
-                    elif isinstance(message, ModelResponse):
-                        for body in _response_bodies(message):
-                            emit_event('gen_ai.assistant.message', body)
 
             def finish(response: ModelResponse, usage: Usage):
                 if not span.is_recording():
                     return
 
-                for response_body in _response_bodies(response):
-                    if response_body:
-                        emit_event(
+                events = self.messages_to_otel_events(messages)
+                for event in self.messages_to_otel_events([response]):
+                    events.append(
+                        Event(
                             'gen_ai.choice',
-                            {
+                            body={
                                 # TODO finish_reason
                                 'index': 0,
-                                'message': response_body,
+                                'message': event.body,
                             },
                         )
+                    )
                 span.set_attributes(
                     {
                         # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
@@ -178,67 +162,56 @@ def finish(response: ModelResponse, usage: Usage):
                         **usage.opentelemetry_attributes(),
                     }
                 )
-                if events_list:
-                    attr_name = 'events'
-                    span.set_attributes(
-                        {
-                            attr_name: json.dumps(events_list),
-                            'logfire.json_schema': json.dumps(
-                                {
-                                    'type': 'object',
-                                    'properties': {attr_name: {'type': 'array'}},
-                                }
-                            ),
-                        }
-                    )
+                self._emit_events(system, span, events)
 
             yield finish
 
-    def _emit_event(
-        self, system: str, events_list: list[dict[str, Any]], event_name: str, body: dict[str, Any]
-    ) -> None:
-        attributes = {'gen_ai.system': system}
+    def _emit_events(self, system: str, span: Span, events: list[Event]) -> None:
+        for event in events:
+            event.attributes = {'gen_ai.system': system, **(event.attributes or {})}
         if self.event_mode == 'logs':
-            self.event_logger.emit(Event(event_name, body=body, attributes=attributes))
-        else:
-            events_list.append({'event.name': event_name, **body, **attributes})
-
-
-def _request_part_body(part: ModelRequestPart) -> tuple[str, dict[str, Any]]:
-    if isinstance(part, SystemPromptPart):
-        return 'gen_ai.system.message', {'content': part.content, 'role': 'system'}
-    elif isinstance(part, UserPromptPart):
-        return 'gen_ai.user.message', {'content': part.content, 'role': 'user'}
-    elif isinstance(part, ToolReturnPart):
-        return 'gen_ai.tool.message', {'content': part.content, 'role': 'tool', 'id': part.tool_call_id}
-    elif isinstance(part, RetryPromptPart):
-        if part.tool_name is None:
-            return 'gen_ai.user.message', {'content': part.model_response(), 'role': 'user'}
+            for event in events:
+                self.event_logger.emit(event)
         else:
-            return 'gen_ai.tool.message', {'content': part.model_response(), 'role': 'tool', 'id': part.tool_call_id}
-    else:
-        return '', {}
-
-
-def _response_bodies(message: ModelResponse) -> list[dict[str, Any]]:
-    body: dict[str, Any] = {'role': 'assistant'}
-    result = [body]
-    for part in message.parts:
-        if isinstance(part, ToolCallPart):
-            body.setdefault('tool_calls', []).append(
+            attr_name = 'events'
+            span.set_attributes(
                 {
-                    'id': part.tool_call_id,
-                    'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
-                    'function': {
-                        'name': part.tool_name,
-                        'arguments': part.args,
-                    },
+                    attr_name: json.dumps([self.event_to_dict(event) for event in events]),
+                    'logfire.json_schema': json.dumps(
+                        {
+                            'type': 'object',
+                            'properties': {attr_name: {'type': 'array'}},
+                        }
+                    ),
                 }
             )
-        elif isinstance(part, TextPart):
-            if body.get('content'):
-                body = {'role': 'assistant'}
-                result.append(body)
-            body['content'] = part.content
 
-    return result
+    @staticmethod
+    def event_to_dict(event: Event) -> dict[str, Any]:
+        if not event.body:
+            body = {}
+        elif isinstance(event.body, Mapping):
+            body = event.body  # type: ignore
+        else:
+            body = {'body': event.body}
+        return {**body, **(event.attributes or {})}
+
+    @staticmethod
+    def messages_to_otel_events(messages: list[ModelMessage]) -> list[Event]:
+        result: list[Event] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                for part in message.parts:
+                    if hasattr(part, 'otel_event'):
+                        result.append(part.otel_event())
+            elif isinstance(message, ModelResponse):
+                result.extend(message.otel_events())
+        for event in result:
+            try:
+                event.body = ANY_ADAPTER.dump_python(event.body, mode='json')
+            except Exception:
+                try:
+                    event.body = str(event.body)
+                except Exception:
+                    event.body = 'Unable to serialize event body'
+        return result

pydantic_ai_slim/pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "pydantic>=2.10",
     "pydantic-graph==0.0.30",
     "exceptiongroup; python_version < '3.11'",
+    "opentelemetry-api>=1.28.0",
 ]
 
 [project.optional-dependencies]