Handle errors in cost calculation in InstrumentedModel (#2834)

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/models/instrumented.py

@@ -420,16 +420,20 @@ def _record_metrics():
                         return
 
                     self.instrumentation_settings.handle_messages(messages, response, system, span)
-                    try:
-                        cost_attributes = {'operation.cost': float(response.cost().total_price)}
-                    except LookupError:
-                        cost_attributes = {}
 
                     attributes_to_set = {
                         **response.usage.opentelemetry_attributes(),
                         'gen_ai.response.model': response_model,
-                        **cost_attributes,
                     }
+                    try:
+                        attributes_to_set['operation.cost'] = float(response.cost().total_price)
+                    except LookupError:
+                        # The cost of this provider/model is unknown, which is common.
+                        pass
+                    except Exception as e:
+                        warnings.warn(
+                            f'Failed to get cost from response: {type(e).__name__}: {e}', CostCalculationFailedWarning
+                        )
                     if response.provider_response_id is not None:
                         attributes_to_set['gen_ai.response.id'] = response.provider_response_id
                     span.set_attributes(attributes_to_set)
@@ -480,3 +484,7 @@ def serialize_any(value: Any) -> str:
                 return str(value)
             except Exception as e:
                 return f'Unable to serialize: {e}'
+
+
+class CostCalculationFailedWarning(Warning):
+    """Warning raised when cost calculation fails."""

pydantic_ai_slim/pyproject.toml

@@ -60,7 +60,7 @@ dependencies = [
     "exceptiongroup; python_version < '3.11'",
     "opentelemetry-api>=1.28.0",
     "typing-inspection>=0.4.0",
-    "genai-prices>=0.0.22",
+    "genai-prices>=0.0.23",
 ]
 
 [tool.hatch.metadata.hooks.uv-dynamic-versioning.optional-dependencies]

tests/models/test_anthropic.py

@@ -252,7 +252,7 @@ async def test_async_request_prompt_caching(allow_model_requests: None):
     )
     last_message = result.all_messages()[-1]
     assert isinstance(last_message, ModelResponse)
-    assert last_message.cost().total_price == snapshot(Decimal('0.00003488'))
+    assert last_message.cost().total_price == snapshot(Decimal('0.00002688'))
 
 
 async def test_async_request_text_response(allow_model_requests: None):

tests/models/test_instrumented.py

@@ -7,6 +7,7 @@
 
 import pytest
 from inline_snapshot import snapshot
+from inline_snapshot.extra import warns
 from logfire_api import DEFAULT_LOGFIRE_INSTANCE
 from opentelemetry._events import NoOpEventLoggerProvider
 from opentelemetry.trace import NoOpTracerProvider
@@ -1278,3 +1279,73 @@ def test_deprecated_event_mode_warning():
     assert settings.event_mode == 'logs'
     assert settings.version == 1
     assert InstrumentationSettings().version == 2
+
+
+async def test_response_cost_error(capfire: CaptureLogfire, monkeypatch: pytest.MonkeyPatch):
+    model = InstrumentedModel(MyModel())
+
+    messages: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart('user_prompt')])]
+    monkeypatch.setattr(ModelResponse, 'cost', None)
+
+    with warns(
+        snapshot(
+            [
+                "CostCalculationFailedWarning: Failed to get cost from response: TypeError: 'NoneType' object is not callable"
+            ]
+        )
+    ):
+        await model.request(messages, model_settings=ModelSettings(), model_request_parameters=ModelRequestParameters())
+
+    assert capfire.exporter.exported_spans_as_dict(parse_json_attributes=True) == snapshot(
+        [
+            {
+                'name': 'chat gpt-4o',
+                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'parent': None,
+                'start_time': 1000000000,
+                'end_time': 2000000000,
+                'attributes': {
+                    'gen_ai.operation.name': 'chat',
+                    'gen_ai.system': 'openai',
+                    'gen_ai.request.model': 'gpt-4o',
+                    'server.address': 'example.com',
+                    'server.port': 8000,
+                    'model_request_parameters': {
+                        'function_tools': [],
+                        'builtin_tools': [],
+                        'output_mode': 'text',
+                        'output_object': None,
+                        'output_tools': [],
+                        'allow_text_output': True,
+                    },
+                    'logfire.span_type': 'span',
+                    'logfire.msg': 'chat gpt-4o',
+                    'gen_ai.input.messages': [{'role': 'user', 'parts': [{'type': 'text', 'content': 'user_prompt'}]}],
+                    'gen_ai.output.messages': [
+                        {
+                            'role': 'assistant',
+                            'parts': [
+                                {'type': 'text', 'content': 'text1'},
+                                {'type': 'tool_call', 'id': 'tool_call_1', 'name': 'tool1', 'arguments': 'args1'},
+                                {'type': 'tool_call', 'id': 'tool_call_2', 'name': 'tool2', 'arguments': {'args2': 3}},
+                                {'type': 'text', 'content': 'text2'},
+                            ],
+                            'finish_reason': 'stop',
+                        }
+                    ],
+                    'logfire.json_schema': {
+                        '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': 100,
+                    'gen_ai.usage.output_tokens': 200,
+                    'gen_ai.response.model': 'gpt-4o-2024-11-20',
+                    'gen_ai.response.id': 'response_id',
+                },
+            }
+        ]
+    )