Fix instrumentation of FallbackModel (#1076)

alexmojaki · web-flow · commit f96849c19df3 · 2025-03-13T12:18:53.000+02:00
diff --git a/pydantic_ai_slim/pydantic_ai/models/fallback.py b/pydantic_ai_slim/pydantic_ai/models/fallback.py
@@ -61,7 +61,9 @@ async def request(
 
         for model in self.models:
             try:
-                return await model.request(messages, model_settings, model_request_parameters)
+                response, usage = await model.request(messages, model_settings, model_request_parameters)
+                response.model_used = model  # type: ignore
+                return response, usage
             except Exception as exc:
                 if self._fallback_on(exc):
                     exceptions.append(exc)
diff --git a/pydantic_ai_slim/pydantic_ai/models/instrumented.py b/pydantic_ai_slim/pydantic_ai/models/instrumented.py
@@ -88,6 +88,10 @@ def __init__(
         self.event_mode = event_mode
 
 
+GEN_AI_SYSTEM_ATTRIBUTE = 'gen_ai.system'
+GEN_AI_REQUEST_MODEL_ATTRIBUTE = 'gen_ai.request.model'
+
+
 @dataclass
 class InstrumentedModel(WrapperModel):
     """Model which is instrumented with OpenTelemetry."""
@@ -138,27 +142,14 @@ def _instrument(
         model_settings: ModelSettings | None,
     ) -> Iterator[Callable[[ModelResponse, Usage], None]]:
         operation = 'chat'
-        model_name = self.model_name
-        span_name = f'{operation} {model_name}'
-        system = getattr(self.wrapped, 'system', '') or self.wrapped.__class__.__name__.removesuffix('Model').lower()
-        system = {'google-gla': 'gemini', 'google-vertex': 'vertex_ai', 'mistral': 'mistral_ai'}.get(system, system)
+        span_name = f'{operation} {self.model_name}'
         # TODO Missing attributes:
         #  - error.type: unclear if we should do something here or just always rely on span exceptions
         #  - gen_ai.request.stop_sequences/top_k: model_settings doesn't include these
         attributes: dict[str, AttributeValue] = {
             'gen_ai.operation.name': operation,
-            'gen_ai.system': system,
-            'gen_ai.request.model': model_name,
+            **self.model_attributes(self.wrapped),
         }
-        if base_url := self.wrapped.base_url:
-            try:
-                parsed = urlparse(base_url)
-                if parsed.hostname:
-                    attributes['server.address'] = parsed.hostname
-                if parsed.port:
-                    attributes['server.port'] = parsed.port
-            except Exception:  # pragma: no cover
-                pass
 
         if model_settings:
             for key in MODEL_SETTING_ATTRIBUTES:
@@ -183,21 +174,26 @@ def finish(response: ModelResponse, usage: Usage):
                             },
                         )
                     )
-                span.set_attributes(
-                    {
-                        # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
-                        #  https://github.com/pydantic/pydantic-ai/issues/886
-                        'gen_ai.response.model': response.model_name or model_name,
-                        **usage.opentelemetry_attributes(),
+                new_attributes: dict[str, AttributeValue] = usage.opentelemetry_attributes()  # type: ignore
+                if model_used := getattr(response, 'model_used', None):
+                    # FallbackModel sets model_used on the response so that we can report the attributes
+                    # of the model that was actually used.
+                    new_attributes.update(self.model_attributes(model_used))
+                    attributes.update(new_attributes)
+                request_model = attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE]
+                new_attributes['gen_ai.response.model'] = response.model_name or request_model
+                span.set_attributes(new_attributes)
+                span.update_name(f'{operation} {request_model}')
+                for event in events:
+                    event.attributes = {
+                        GEN_AI_SYSTEM_ATTRIBUTE: attributes[GEN_AI_SYSTEM_ATTRIBUTE],
+                        **(event.attributes or {}),
                     }
-                )
-                self._emit_events(system, span, events)
+                self._emit_events(span, events)
 
             yield finish
 
-    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 {})}
+    def _emit_events(self, span: Span, events: list[Event]) -> None:
         if self.options.event_mode == 'logs':
             for event in events:
                 self.options.event_logger.emit(event)
@@ -215,6 +211,27 @@ def _emit_events(self, system: str, span: Span, events: list[Event]) -> None:
                 }
             )
 
+    @staticmethod
+    def model_attributes(model: Model):
+        system = getattr(model, 'system', '') or model.__class__.__name__.removesuffix('Model').lower()
+        system = {'google-gla': 'gemini', 'google-vertex': 'vertex_ai', 'mistral': 'mistral_ai'}.get(system, system)
+        attributes: dict[str, AttributeValue] = {
+            GEN_AI_SYSTEM_ATTRIBUTE: system,
+            GEN_AI_REQUEST_MODEL_ATTRIBUTE: model.model_name,
+        }
+        if base_url := model.base_url:
+            try:
+                parsed = urlparse(base_url)
+            except Exception:  # pragma: no cover
+                pass
+            else:
+                if parsed.hostname:
+                    attributes['server.address'] = parsed.hostname
+                if parsed.port:
+                    attributes['server.port'] = parsed.port
+
+        return attributes
+
     @staticmethod
     def event_to_dict(event: Event) -> dict[str, Any]:
         if not event.body:
diff --git a/pydantic_ai_slim/pydantic_ai/usage.py b/pydantic_ai_slim/pydantic_ai/usage.py
@@ -64,7 +64,7 @@ def opentelemetry_attributes(self) -> dict[str, int]:
         }
         for key, value in (self.details or {}).items():
             result[f'gen_ai.usage.details.{key}'] = value
-        return {k: v for k, v in result.items() if v is not None}
+        return {k: v for k, v in result.items() if v}
 
 
 @dataclass
diff --git a/tests/models/test_fallback.py b/tests/models/test_fallback.py
@@ -1,22 +1,29 @@
+from __future__ import annotations
+
 import sys
 from collections.abc import AsyncIterator
 from datetime import timezone
 
 import pytest
+from dirty_equals import IsJson
 from inline_snapshot import snapshot
 
 from pydantic_ai import Agent, ModelHTTPError
 from pydantic_ai.messages import ModelMessage, ModelRequest, ModelResponse, TextPart, UserPromptPart
 from pydantic_ai.models.fallback import FallbackModel
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 
-from ..conftest import IsNow
+from ..conftest import IsNow, try_import
 
 if sys.version_info < (3, 11):
     from exceptiongroup import ExceptionGroup as ExceptionGroup
 else:
     ExceptionGroup = ExceptionGroup
 
+with try_import() as logfire_imports_successful:
+    from logfire.testing import CaptureLogfire
+
+
 pytestmark = pytest.mark.anyio
 
 
@@ -86,6 +93,100 @@ def test_first_failed() -> None:
     )
 
 
+@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
+def test_first_failed_instrumented(capfire: CaptureLogfire) -> None:
+    fallback_model = FallbackModel(failure_model, success_model)
+    agent = Agent(model=fallback_model, instrument=True)
+    result = agent.run_sync('hello')
+    assert result.data == snapshot('success')
+    assert result.all_messages() == snapshot(
+        [
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content='hello',
+                        timestamp=IsNow(tz=timezone.utc),
+                    )
+                ]
+            ),
+            ModelResponse(
+                parts=[TextPart(content='success')],
+                model_name='function:success_response:',
+                timestamp=IsNow(tz=timezone.utc),
+            ),
+        ]
+    )
+    assert capfire.exporter.exported_spans_as_dict() == snapshot(
+        [
+            {
+                'name': 'preparing model request params',
+                'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
+                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'start_time': 2000000000,
+                'end_time': 3000000000,
+                'attributes': {
+                    'run_step': 1,
+                    'logfire.span_type': 'span',
+                    'logfire.msg': 'preparing model request params',
+                },
+            },
+            {
+                'name': 'chat function:success_response:',
+                'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
+                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'start_time': 4000000000,
+                'end_time': 5000000000,
+                'attributes': {
+                    'gen_ai.operation.name': 'chat',
+                    'logfire.span_type': 'span',
+                    'logfire.msg': 'chat FallBackModel[function:failure_response:, function:success_response:]',
+                    'gen_ai.usage.input_tokens': 51,
+                    'gen_ai.usage.output_tokens': 1,
+                    'gen_ai.system': 'function',
+                    'gen_ai.request.model': 'function:success_response:',
+                    'gen_ai.response.model': 'function:success_response:',
+                    'events': IsJson(
+                        [
+                            {
+                                'content': 'hello',
+                                'role': 'user',
+                                'gen_ai.system': 'function',
+                                'gen_ai.message.index': 0,
+                                'event.name': 'gen_ai.user.message',
+                            },
+                            {
+                                'index': 0,
+                                'message': {'role': 'assistant', 'content': 'success'},
+                                'gen_ai.system': 'function',
+                                'event.name': 'gen_ai.choice',
+                            },
+                        ]
+                    ),
+                    'logfire.json_schema': '{"type": "object", "properties": {"events": {"type": "array"}}}',
+                },
+            },
+            {
+                'name': 'agent run',
+                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'parent': None,
+                'start_time': 1000000000,
+                'end_time': 6000000000,
+                'attributes': {
+                    'model_name': 'FallBackModel[function:failure_response:, function:success_response:]',
+                    'agent_name': 'agent',
+                    'logfire.msg': 'agent run',
+                    'logfire.span_type': 'span',
+                    'gen_ai.usage.input_tokens': 51,
+                    'gen_ai.usage.output_tokens': 1,
+                    'all_messages_events': '[{"content": "hello", "role": "user", "gen_ai.message.index": 0, "event.name": "gen_ai.user.message"}, {"role": "assistant", "content": "success", "gen_ai.message.index": 1, "event.name": "gen_ai.assistant.message"}]',
+                    'final_result': 'success',
+                    'logfire.json_schema': '{"type": "object", "properties": {"all_messages_events": {"type": "array"}, "final_result": {"type": "object"}}}',
+                },
+            },
+        ]
+    )
+
+
 def test_all_failed() -> None:
     fallback_model = FallbackModel(failure_model, failure_model)
     agent = Agent(model=fallback_model)

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ def opentelemetry_attributes(self) -> dict[str, int]:`
`64`	`64`	`}`
`65`	`65`	`for key, value in (self.details or {}).items():`
`66`	`66`	`result[f'gen_ai.usage.details.{key}'] = value`
`67`		`- return {k: v for k, v in result.items() if v is not None}`
	`67`	`+ return {k: v for k, v in result.items() if v}`
`68`	`68`
`69`	`69`
`70`	`70`	`@dataclass`