Add price() method to ModelResponse (#2584)

Co-authored-by: Alex Hall <[email protected]>

Author: Kludex

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py

@@ -55,6 +55,10 @@ def usage(self) -> RequestUsage:
     def model_name(self) -> str:
         return self.response.model_name or ''  # pragma: no cover
 
+    @property
+    def provider_name(self) -> str:
+        return self.response.provider_name or ''  # pragma: no cover
+
     @property
     def timestamp(self) -> datetime:
         return self.response.timestamp  # pragma: no cover

pydantic_ai_slim/pydantic_ai/messages.py

@@ -10,6 +10,7 @@
 
 import pydantic
 import pydantic_core
+from genai_prices import calc_price, types as genai_types
 from opentelemetry._events import Event  # pyright: ignore[reportPrivateImportUsage]
 from typing_extensions import TypeAlias, deprecated
 
@@ -848,6 +849,9 @@ class ModelResponse:
     kind: Literal['response'] = 'response'
     """Message type identifier, this is available on all parts as a discriminator."""
 
+    provider_name: str | None = None
+    """The name of the LLM provider that generated the response."""
+
     provider_details: dict[str, Any] | None = field(default=None)
     """Additional provider-specific details in a serializable format.
 
@@ -858,6 +862,19 @@ class ModelResponse:
     provider_request_id: str | None = None
     """request ID as specified by the model provider. This can be used to track the specific request to the model."""
 
+    def price(self) -> genai_types.PriceCalculation:
+        """Calculate the price of the usage.
+
+        Uses [`genai-prices`](https://github.com/pydantic/genai-prices).
+        """
+        assert self.model_name, 'Model name is required to calculate price'
+        return calc_price(
+            self.usage,
+            self.model_name,
+            provider_id=self.provider_name,
+            genai_request_timestamp=self.timestamp,
+        )
+
     def otel_events(self, settings: InstrumentationSettings) -> list[Event]:
         """Return OpenTelemetry events for the response."""
         result: list[Event] = []

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -598,6 +598,7 @@ def get(self) -> ModelResponse:
             model_name=self.model_name,
             timestamp=self.timestamp,
             usage=self.usage(),
+            provider_name=self.provider_name,
         )
 
     def usage(self) -> RequestUsage:
@@ -610,6 +611,12 @@ def model_name(self) -> str:
         """Get the model name of the response."""
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def provider_name(self) -> str | None:
+        """Get the provider name."""
+        raise NotImplementedError()
+
     @property
     @abstractmethod
     def timestamp(self) -> datetime:

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -326,7 +326,11 @@ def _process_response(self, response: BetaMessage) -> ModelResponse:
                 )
 
         return ModelResponse(
-            items, usage=_map_usage(response), model_name=response.model, provider_request_id=response.id
+            items,
+            usage=_map_usage(response),
+            model_name=response.model,
+            provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -344,6 +348,7 @@ async def _process_streamed_response(
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=timestamp,
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[BetaToolParam]:
@@ -574,6 +579,7 @@ class AnthropicStreamedResponse(StreamedResponse):
     _model_name: AnthropicModelName
     _response: AsyncIterable[BetaRawMessageStreamEvent]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         current_block: BetaContentBlock | None = None
@@ -655,6 +661,11 @@ def model_name(self) -> AnthropicModelName:
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -240,10 +240,7 @@ def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[T
 
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolTypeDef:
-        tool_spec: ToolSpecificationTypeDef = {
-            'name': f.name,
-            'inputSchema': {'json': f.parameters_json_schema},
-        }
+        tool_spec: ToolSpecificationTypeDef = {'name': f.name, 'inputSchema': {'json': f.parameters_json_schema}}
 
         if f.description:  # pragma: no branch
             tool_spec['description'] = f.description
@@ -275,6 +272,7 @@ async def request_stream(
             model_request_parameters=model_request_parameters,
             _model_name=self.model_name,
             _event_stream=response,
+            _provider_name=self._provider.name,
         )
 
     async def _process_response(self, response: ConverseResponseTypeDef) -> ModelResponse:
@@ -304,7 +302,9 @@ async def _process_response(self, response: ConverseResponseTypeDef) -> ModelRes
             output_tokens=response['usage']['outputTokens'],
         )
         vendor_id = response.get('ResponseMetadata', {}).get('RequestId', None)
-        return ModelResponse(items, usage=u, model_name=self.model_name, provider_request_id=vendor_id)
+        return ModelResponse(
+            items, usage=u, model_name=self.model_name, provider_request_id=vendor_id, provider_name=self._provider.name
+        )
 
     @overload
     async def _messages_create(
@@ -594,6 +594,7 @@ class BedrockStreamedResponse(StreamedResponse):
 
     _model_name: BedrockModelName
     _event_stream: EventStream[ConverseStreamOutputTypeDef]
+    _provider_name: str
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
@@ -660,15 +661,20 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     if maybe_event:  # pragma: no branch
                         yield maybe_event
 
-    @property
-    def timestamp(self) -> datetime:
-        return self._timestamp
-
     @property
     def model_name(self) -> str:
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
+    @property
+    def timestamp(self) -> datetime:
+        return self._timestamp
+
     def _map_usage(self, metadata: ConverseStreamMetadataEventTypeDef) -> usage.RequestUsage:
         return usage.RequestUsage(
             input_tokens=metadata['usage']['inputTokens'],

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -205,7 +205,9 @@ def _process_response(self, response: V2ChatResponse) -> ModelResponse:
                         tool_call_id=c.id or _generate_tool_call_id(),
                     )
                 )
-        return ModelResponse(parts=parts, usage=_map_usage(response), model_name=self._model_name)
+        return ModelResponse(
+            parts=parts, usage=_map_usage(response), model_name=self._model_name, provider_name=self._provider.name
+        )
 
     def _map_messages(self, messages: list[ModelMessage]) -> list[ChatMessageV2]:
         """Just maps a `pydantic_ai.Message` to a `cohere.ChatMessageV2`."""

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -304,6 +304,11 @@ def model_name(self) -> str:
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> None:
+        """Get the provider name."""
+        return None
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -305,6 +305,7 @@ async def _process_streamed_response(
             _model_name=self._model_name,
             _content=content,
             _stream=aiter_bytes,
+            _provider_name=self._provider.name,
         )
 
     async def _message_to_gemini_content(
@@ -425,6 +426,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _model_name: GeminiModelName
     _content: bytearray
     _stream: AsyncIterator[bytes]
+    _provider_name: str
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
@@ -495,6 +497,11 @@ def model_name(self) -> GeminiModelName:
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -395,6 +395,7 @@ def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
         return _process_response_from_parts(
             parts,
             response.model_version or self._model_name,
+            self._provider.name,
             usage,
             vendor_id=vendor_id,
             vendor_details=vendor_details,
@@ -414,6 +415,7 @@ async def _process_streamed_response(
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=first_chunk.create_time or _utils.now_utc(),
+            _provider_name=self._provider.name,
         )
 
     async def _map_messages(self, messages: list[ModelMessage]) -> tuple[ContentDict | None, list[ContentUnionDict]]:
@@ -523,6 +525,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _model_name: GoogleModelName
     _response: AsyncIterator[GenerateContentResponse]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -564,6 +567,11 @@ def model_name(self) -> GoogleModelName:
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
@@ -599,6 +607,7 @@ def _content_model_response(m: ModelResponse) -> ContentDict:
 def _process_response_from_parts(
     parts: list[Part],
     model_name: GoogleModelName,
+    provider_name: str,
     usage: usage.RequestUsage,
     vendor_id: str | None,
     vendor_details: dict[str, Any] | None = None,
@@ -636,7 +645,12 @@ def _process_response_from_parts(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
             )
     return ModelResponse(
-        parts=items, model_name=model_name, usage=usage, provider_request_id=vendor_id, provider_details=vendor_details
+        parts=items,
+        model_name=model_name,
+        usage=usage,
+        provider_request_id=vendor_id,
+        provider_details=vendor_details,
+        provider_name=provider_name,
     )
 
 

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -290,6 +290,7 @@ def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
             model_name=response.model,
             timestamp=timestamp,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -309,6 +310,7 @@ async def _process_streamed_response(
             _model_name=self._model_name,
             _model_profile=self.profile,
             _timestamp=number_to_datetime(first_chunk.created),
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
@@ -444,6 +446,7 @@ class GroqStreamedResponse(StreamedResponse):
     _model_profile: ModelProfile
     _response: AsyncIterable[chat.ChatCompletionChunk]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -482,6 +485,11 @@ def model_name(self) -> GroqModelName:
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""