Add provider_url to ModelResponse and use it in cost() (#3648)

Co-authored-by: Alex Hall <alex.mojaki@gmail.com>

Author: mahiro72

pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_model.py

@@ -46,6 +46,10 @@ def model_name(self) -> str:
     def provider_name(self) -> str:
         return self.response.provider_name or ''  # pragma: no cover
 
+    @property
+    def provider_url(self) -> str | None:
+        return self.response.provider_url  # pragma: no cover
+
     @property
     def timestamp(self) -> datetime:
         return self.response.timestamp  # pragma: no cover

pydantic_ai_slim/pydantic_ai/durable_exec/prefect/_model.py

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

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py

@@ -60,6 +60,10 @@ def model_name(self) -> str:
     def provider_name(self) -> str:
         return self.response.provider_name or ''  # pragma: no cover
 
+    @property
+    def provider_url(self) -> str | None:
+        return self.response.provider_url  # pragma: no cover
+
     @property
     def timestamp(self) -> datetime:
         return self.response.timestamp  # pragma: no cover

pydantic_ai_slim/pydantic_ai/messages.py

@@ -1249,6 +1249,9 @@ class ModelResponse:
     provider_name: str | None = None
     """The name of the LLM provider that generated the response."""
 
+    provider_url: str | None = None
+    """The base URL of the LLM provider that generated the response."""
+
     provider_details: Annotated[
         dict[str, Any] | None,
         # `vendor_details` is deprecated, but we still want to support deserializing model responses stored in a DB before the name was changed
@@ -1337,6 +1340,17 @@ def cost(self) -> genai_types.PriceCalculation:
         Uses [`genai-prices`](https://github.com/pydantic/genai-prices).
         """
         assert self.model_name, 'Model name is required to calculate price'
+        # Try matching on provider_api_url first as this is more specific, then fall back to provider_id.
+        if self.provider_url:
+            try:
+                return calc_price(
+                    self.usage,
+                    self.model_name,
+                    provider_api_url=self.provider_url,
+                    genai_request_timestamp=self.timestamp,
+                )
+            except LookupError:
+                pass
         return calc_price(
             self.usage,
             self.model_name,

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -875,6 +875,7 @@ def get(self) -> ModelResponse:
             timestamp=self.timestamp,
             usage=self.usage(),
             provider_name=self.provider_name,
+            provider_url=self.provider_url,
             provider_response_id=self.provider_response_id,
             provider_details=self.provider_details,
             finish_reason=self.finish_reason,
@@ -897,6 +898,12 @@ def provider_name(self) -> str | None:
         """Get the provider name."""
         raise NotImplementedError()
 
+    @property
+    @abstractmethod
+    def provider_url(self) -> str | None:
+        """Get the provider base URL."""
+        raise NotImplementedError()
+
     @property
     @abstractmethod
     def timestamp(self) -> datetime:

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -561,6 +561,7 @@ def _process_response(self, response: BetaMessage) -> ModelResponse:
             model_name=response.model,
             provider_response_id=response.id,
             provider_name=self._provider.name,
+            provider_url=self._provider.base_url,
             finish_reason=finish_reason,
             provider_details=provider_details,
         )
@@ -1298,6 +1299,11 @@ def provider_name(self) -> str:
         """Get the provider name."""
         return self._provider_name
 
+    @property
+    def provider_url(self) -> str:
+        """Get the provider base URL."""
+        return self._provider_url
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -330,6 +330,7 @@ async def request_stream(
             _model_name=self.model_name,
             _event_stream=response['stream'],
             _provider_name=self._provider.name,
+            _provider_url=self.base_url,
             _provider_response_id=response.get('ResponseMetadata', {}).get('RequestId', None),
         )
 
@@ -381,6 +382,7 @@ async def _process_response(self, response: ConverseResponseTypeDef) -> ModelRes
             model_name=self.model_name,
             provider_response_id=response_id,
             provider_name=self._provider.name,
+            provider_url=self.base_url,
             finish_reason=finish_reason,
             provider_details=provider_details,
         )
@@ -700,6 +702,7 @@ class BedrockStreamedResponse(StreamedResponse):
     _model_name: BedrockModelName
     _event_stream: EventStream[ConverseStreamOutputTypeDef]
     _provider_name: str
+    _provider_url: str
     _timestamp: datetime = field(default_factory=_utils.now_utc)
     _provider_response_id: str | None = None
 
@@ -787,6 +790,11 @@ def provider_name(self) -> str:
         """Get the provider name."""
         return self._provider_name
 
+    @property
+    def provider_url(self) -> str:
+        """Get the provider base URL."""
+        return self._provider_url
+
     @property
     def timestamp(self) -> datetime:
         return self._timestamp

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -225,6 +225,7 @@ def _process_response(self, response: V2ChatResponse) -> ModelResponse:
             usage=_map_usage(response),
             model_name=self._model_name,
             provider_name=self._provider.name,
+            provider_url=self.base_url,
             finish_reason=finish_reason,
             provider_details=provider_details,
         )

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -347,6 +347,11 @@ def provider_name(self) -> None:
         """Get the provider name."""
         return None
 
+    @property
+    def provider_url(self) -> None:
+        """Get the provider base URL."""
+        return None
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -148,8 +148,8 @@ def __init__(
 
     @property
     def base_url(self) -> str:
-        assert self._url is not None, 'URL not initialized'  # pragma: no cover
-        return self._url  # pragma: no cover
+        assert self._url is not None, 'URL not initialized'
+        return self._url
 
     @property
     def model_name(self) -> GeminiModelName:
@@ -298,6 +298,7 @@ def _process_response(self, response: _GeminiResponse) -> ModelResponse:
             usage,
             vendor_id=vendor_id,
             vendor_details=vendor_details,
+            provider_url=self.base_url,
         )
 
     async def _process_streamed_response(
@@ -329,6 +330,7 @@ async def _process_streamed_response(
             _content=content,
             _stream=aiter_bytes,
             _provider_name=self._provider.name,
+            _provider_url=self.base_url,
         )
 
     async def _message_to_gemini_content(
@@ -453,6 +455,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _content: bytearray
     _stream: AsyncIterator[bytes]
     _provider_name: str
+    _provider_url: str
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
@@ -527,6 +530,11 @@ def provider_name(self) -> str:
         """Get the provider name."""
         return self._provider_name
 
+    @property
+    def provider_url(self) -> str:
+        """Get the provider base URL."""
+        return self._provider_url
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
@@ -713,6 +721,7 @@ def _process_response_from_parts(
     model_name: GeminiModelName,
     usage: usage.RequestUsage,
     vendor_id: str | None,
+    provider_url: str,
     vendor_details: dict[str, Any] | None = None,
 ) -> ModelResponse:
     items: list[ModelResponsePart] = []
@@ -731,7 +740,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, usage=usage, model_name=model_name, provider_response_id=vendor_id, provider_details=vendor_details
+        parts=items,
+        usage=usage,
+        model_name=model_name,
+        provider_response_id=vendor_id,
+        provider_details=vendor_details,
+        provider_url=provider_url,
     )