pydantic-ai-slim 0.7.2__py3-none-any.whl → 0.7.4__py3-none-any.whl

pydantic_ai/models/google.py

@@ -13,7 +13,7 @@ from typing_extensions import assert_never
 from .. import UnexpectedModelBehavior, _utils, usage
 from .._output import OutputObjectDefinition
 from .._run_context import RunContext
-from ..builtin_tools import CodeExecutionTool, WebSearchTool
+from ..builtin_tools import CodeExecutionTool, UrlContextTool, WebSearchTool
 from ..exceptions import UserError
 from ..messages import (
     BinaryContent,
@@ -72,6 +72,7 @@ try:
         ToolConfigDict,
         ToolDict,
         ToolListUnionDict,
+        UrlContextDict,
     )
 
     from ..providers.google import GoogleProvider
@@ -144,7 +145,6 @@ class GoogleModel(Model):
     _model_name: GoogleModelName = field(repr=False)
     _provider: Provider[Client] = field(repr=False)
     _url: str | None = field(repr=False)
-    _system: str = field(default='google', repr=False)
 
     def __init__(
         self,
@@ -168,9 +168,7 @@ class GoogleModel(Model):
 
         if isinstance(provider, str):
             provider = GoogleProvider(vertexai=provider == 'google-vertex')
-
         self._provider = provider
-        self._system = provider.name
         self.client = provider.client
 
         super().__init__(settings=settings, profile=profile or provider.model_profile)
@@ -179,6 +177,16 @@ class GoogleModel(Model):
     def base_url(self) -> str:
         return self._provider.base_url
 
+    @property
+    def model_name(self) -> GoogleModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The model provider."""
+        return self._provider.name
+
     async def request(
         self,
         messages: list[ModelMessage],
@@ -195,7 +203,7 @@ class GoogleModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
-    ) -> usage.Usage:
+    ) -> usage.RequestUsage:
         check_allow_model_requests()
         model_settings = cast(GoogleModelSettings, model_settings or {})
         contents, generation_config = await self._build_content_and_config(
@@ -209,9 +217,9 @@ class GoogleModel(Model):
         config = CountTokensConfigDict(
             http_options=generation_config.get('http_options'),
         )
-        if self.system != 'google-gla':
+        if self._provider.name != 'google-gla':
             # The fields are not supported by the Gemini API per https://github.com/googleapis/python-genai/blob/7e4ec284dc6e521949626f3ed54028163ef9121d/google/genai/models.py#L1195-L1214
-            config.update(
+            config.update(  # pragma: lax no cover
                 system_instruction=generation_config.get('system_instruction'),
                 tools=cast(list[ToolDict], generation_config.get('tools')),
                 # Annoyingly, GenerationConfigDict has fewer fields than GenerateContentConfigDict, and no extra fields are allowed.
@@ -238,9 +246,8 @@ class GoogleModel(Model):
             raise UnexpectedModelBehavior(  # pragma: no cover
                 'Total tokens missing from Gemini response', str(response)
             )
-        return usage.Usage(
-            request_tokens=response.total_tokens,
-            total_tokens=response.total_tokens,
+        return usage.RequestUsage(
+            input_tokens=response.total_tokens,
         )
 
     @asynccontextmanager
@@ -256,16 +263,6 @@ class GoogleModel(Model):
         response = await self._generate_content(messages, True, model_settings, model_request_parameters)
         yield await self._process_streamed_response(response, model_request_parameters)  # type: ignore
 
-    @property
-    def model_name(self) -> GoogleModelName:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider."""
-        return self._system
-
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolDict] | None:
         tools: list[ToolDict] = [
             ToolDict(function_declarations=[_function_declaration_from_tool(t)])
@@ -274,6 +271,8 @@ class GoogleModel(Model):
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):
                 tools.append(ToolDict(google_search=GoogleSearchDict()))
+            elif isinstance(tool, UrlContextTool):
+                tools.append(ToolDict(url_context=UrlContextDict()))
             elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
                 tools.append(ToolDict(code_execution=ToolCodeExecutionDict()))
             else:  # pragma: no cover
@@ -378,23 +377,27 @@ class GoogleModel(Model):
     def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
         if not response.candidates or len(response.candidates) != 1:
             raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')  # pragma: no cover
-        if response.candidates[0].content is None or response.candidates[0].content.parts is None:
-            if response.candidates[0].finish_reason == 'SAFETY':
+        candidate = response.candidates[0]
+        if candidate.content is None or candidate.content.parts is None:
+            if candidate.finish_reason == 'SAFETY':
                 raise UnexpectedModelBehavior('Safety settings triggered', str(response))
             else:
                 raise UnexpectedModelBehavior(
                     'Content field missing from Gemini response', str(response)
                 )  # pragma: no cover
-        parts = response.candidates[0].content.parts or []
+        parts = candidate.content.parts or []
         vendor_id = response.response_id or None
         vendor_details: dict[str, Any] | None = None
-        finish_reason = response.candidates[0].finish_reason
+        finish_reason = candidate.finish_reason
         if finish_reason:  # pragma: no branch
             vendor_details = {'finish_reason': finish_reason.value}
         usage = _metadata_as_usage(response)
-        usage.requests = 1
         return _process_response_from_parts(
-            parts, response.model_version or self._model_name, usage, vendor_id=vendor_id, vendor_details=vendor_details
+            parts,
+            response.model_version or self._model_name,
+            usage,
+            vendor_id=vendor_id,
+            vendor_details=vendor_details,
         )
 
     async def _process_streamed_response(
@@ -527,10 +530,13 @@ class GeminiStreamedResponse(StreamedResponse):
 
             assert chunk.candidates is not None
             candidate = chunk.candidates[0]
-            if candidate.content is None:
-                raise UnexpectedModelBehavior('Streamed response has no content field')  # pragma: no cover
-            assert candidate.content.parts is not None
-            for part in candidate.content.parts:
+            if candidate.content is None or candidate.content.parts is None:
+                if candidate.finish_reason == 'SAFETY':  # pragma: no cover
+                    raise UnexpectedModelBehavior('Safety settings triggered', str(chunk))
+                else:  # pragma: no cover
+                    raise UnexpectedModelBehavior('Content field missing from streaming Gemini response', str(chunk))
+            parts = candidate.content.parts or []
+            for part in parts:
                 if part.text is not None:
                     if part.thought:
                         yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)
@@ -590,7 +596,7 @@ def _content_model_response(m: ModelResponse) -> ContentDict:
 def _process_response_from_parts(
     parts: list[Part],
     model_name: GoogleModelName,
-    usage: usage.Usage,
+    usage: usage.RequestUsage,
     vendor_id: str | None,
     vendor_details: dict[str, Any] | None = None,
 ) -> ModelResponse:
@@ -627,7 +633,7 @@ 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, vendor_id=vendor_id, vendor_details=vendor_details
+        parts=items, model_name=model_name, usage=usage, provider_request_id=vendor_id, provider_details=vendor_details
     )
 
 
@@ -647,31 +653,51 @@ def _tool_config(function_names: list[str]) -> ToolConfigDict:
     return ToolConfigDict(function_calling_config=function_calling_config)
 
 
-def _metadata_as_usage(response: GenerateContentResponse) -> usage.Usage:
+def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
     metadata = response.usage_metadata
     if metadata is None:
-        return usage.Usage()  # pragma: no cover
-    metadata = metadata.model_dump(exclude_defaults=True)
-
+        return usage.RequestUsage()
     details: dict[str, int] = {}
-    if cached_content_token_count := metadata.get('cached_content_token_count'):
-        details['cached_content_tokens'] = cached_content_token_count  # pragma: no cover
+    if cached_content_token_count := metadata.cached_content_token_count:
+        details['cached_content_tokens'] = cached_content_token_count
 
-    if thoughts_token_count := metadata.get('thoughts_token_count'):
+    if thoughts_token_count := metadata.thoughts_token_count:
         details['thoughts_tokens'] = thoughts_token_count
 
-    if tool_use_prompt_token_count := metadata.get('tool_use_prompt_token_count'):
+    if tool_use_prompt_token_count := metadata.tool_use_prompt_token_count:
         details['tool_use_prompt_tokens'] = tool_use_prompt_token_count
 
-    for key, metadata_details in metadata.items():
-        if key.endswith('_details') and metadata_details:
-            suffix = key.removesuffix('_details')
-            for detail in metadata_details:
-                details[f'{detail["modality"].lower()}_{suffix}'] = detail.get('token_count', 0)
-
-    return usage.Usage(
-        request_tokens=metadata.get('prompt_token_count', 0),
-        response_tokens=metadata.get('candidates_token_count', 0),
-        total_tokens=metadata.get('total_token_count', 0),
+    input_audio_tokens = 0
+    output_audio_tokens = 0
+    cache_audio_read_tokens = 0
+    for prefix, metadata_details in [
+        ('prompt', metadata.prompt_tokens_details),
+        ('cache', metadata.cache_tokens_details),
+        ('candidates', metadata.candidates_tokens_details),
+        ('tool_use_prompt', metadata.tool_use_prompt_tokens_details),
+    ]:
+        assert getattr(metadata, f'{prefix}_tokens_details') is metadata_details
+        if not metadata_details:
+            continue
+        for detail in metadata_details:
+            if not detail.modality or not detail.token_count:  # pragma: no cover
+                continue
+            details[f'{detail.modality.lower()}_{prefix}_tokens'] = detail.token_count
+            if detail.modality != 'AUDIO':
+                continue
+            if metadata_details is metadata.prompt_tokens_details:
+                input_audio_tokens = detail.token_count
+            elif metadata_details is metadata.candidates_tokens_details:
+                output_audio_tokens = detail.token_count
+            elif metadata_details is metadata.cache_tokens_details:  # pragma: no branch
+                cache_audio_read_tokens = detail.token_count
+
+    return usage.RequestUsage(
+        input_tokens=metadata.prompt_token_count or 0,
+        output_tokens=metadata.candidates_token_count or 0,
+        cache_read_tokens=cached_content_token_count or 0,
+        input_audio_tokens=input_audio_tokens,
+        output_audio_tokens=output_audio_tokens,
+        cache_audio_read_tokens=cache_audio_read_tokens,
         details=details,
     )

pydantic_ai/models/groq.py

@@ -118,7 +118,7 @@ class GroqModel(Model):
     client: AsyncGroq = field(repr=False)
 
     _model_name: GroqModelName = field(repr=False)
-    _system: str = field(default='groq', repr=False)
+    _provider: Provider[AsyncGroq] = field(repr=False)
 
     def __init__(
         self,
@@ -143,6 +143,7 @@ class GroqModel(Model):
 
         if isinstance(provider, str):
             provider = infer_provider(provider)
+        self._provider = provider
         self.client = provider.client
 
         super().__init__(settings=settings, profile=profile or provider.model_profile)
@@ -151,6 +152,16 @@ class GroqModel(Model):
     def base_url(self) -> str:
         return str(self.client.base_url)
 
+    @property
+    def model_name(self) -> GroqModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The model provider."""
+        return self._provider.name
+
     async def request(
         self,
         messages: list[ModelMessage],
@@ -162,7 +173,6 @@ class GroqModel(Model):
             messages, False, cast(GroqModelSettings, model_settings or {}), model_request_parameters
         )
         model_response = self._process_response(response)
-        model_response.usage.requests = 1
         return model_response
 
     @asynccontextmanager
@@ -180,16 +190,6 @@ class GroqModel(Model):
         async with response:
             yield await self._process_streamed_response(response, model_request_parameters)
 
-    @property
-    def model_name(self) -> GroqModelName:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider."""
-        return self._system
-
     @overload
     async def _completions_create(
         self,
@@ -285,7 +285,11 @@ class GroqModel(Model):
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id))
         return ModelResponse(
-            items, usage=_map_usage(response), model_name=response.model, timestamp=timestamp, vendor_id=response.id
+            items,
+            usage=_map_usage(response),
+            model_name=response.model,
+            timestamp=timestamp,
+            provider_request_id=response.id,
         )
 
     async def _process_streamed_response(
@@ -484,7 +488,7 @@ class GroqStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _map_usage(completion: chat.ChatCompletionChunk | chat.ChatCompletion) -> usage.Usage:
+def _map_usage(completion: chat.ChatCompletionChunk | chat.ChatCompletion) -> usage.RequestUsage:
     response_usage = None
     if isinstance(completion, chat.ChatCompletion):
         response_usage = completion.usage
@@ -492,10 +496,9 @@ def _map_usage(completion: chat.ChatCompletionChunk | chat.ChatCompletion) -> us
         response_usage = completion.x_groq.usage
 
     if response_usage is None:
-        return usage.Usage()
+        return usage.RequestUsage()
 
-    return usage.Usage(
-        request_tokens=response_usage.prompt_tokens,
-        response_tokens=response_usage.completion_tokens,
-        total_tokens=response_usage.total_tokens,
+    return usage.RequestUsage(
+        input_tokens=response_usage.prompt_tokens,
+        output_tokens=response_usage.completion_tokens,
     )

pydantic_ai/models/huggingface.py

@@ -114,7 +114,7 @@ class HuggingFaceModel(Model):
     client: AsyncInferenceClient = field(repr=False)
 
     _model_name: str = field(repr=False)
-    _system: str = field(default='huggingface', repr=False)
+    _provider: Provider[AsyncInferenceClient] = field(repr=False)
 
     def __init__(
         self,
@@ -134,13 +134,23 @@ class HuggingFaceModel(Model):
             settings: Model-specific settings that will be used as defaults for this model.
         """
         self._model_name = model_name
-        self._provider = provider
         if isinstance(provider, str):
             provider = infer_provider(provider)
+        self._provider = provider
         self.client = provider.client
 
         super().__init__(settings=settings, profile=profile or provider.model_profile)
 
+    @property
+    def model_name(self) -> HuggingFaceModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The system / model provider."""
+        return self._provider.name
+
     async def request(
         self,
         messages: list[ModelMessage],
@@ -152,7 +162,6 @@ class HuggingFaceModel(Model):
             messages, False, cast(HuggingFaceModelSettings, model_settings or {}), model_request_parameters
         )
         model_response = self._process_response(response)
-        model_response.usage.requests = 1
         return model_response
 
     @asynccontextmanager
@@ -169,16 +178,6 @@ class HuggingFaceModel(Model):
         )
         yield await self._process_streamed_response(response, model_request_parameters)
 
-    @property
-    def model_name(self) -> HuggingFaceModelName:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider."""
-        return self._system
-
     @overload
     async def _completions_create(
         self,
@@ -272,7 +271,7 @@ class HuggingFaceModel(Model):
             usage=_map_usage(response),
             model_name=response.model,
             timestamp=timestamp,
-            vendor_id=response.id,
+            provider_request_id=response.id,
         )
 
     async def _process_streamed_response(
@@ -481,14 +480,12 @@ class HuggingFaceStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _map_usage(response: ChatCompletionOutput | ChatCompletionStreamOutput) -> usage.Usage:
+def _map_usage(response: ChatCompletionOutput | ChatCompletionStreamOutput) -> usage.RequestUsage:
     response_usage = response.usage
     if response_usage is None:
-        return usage.Usage()
+        return usage.RequestUsage()
 
-    return usage.Usage(
-        request_tokens=response_usage.prompt_tokens,
-        response_tokens=response_usage.completion_tokens,
-        total_tokens=response_usage.total_tokens,
-        details=None,
+    return usage.RequestUsage(
+        input_tokens=response_usage.prompt_tokens,
+        output_tokens=response_usage.completion_tokens,
     )

pydantic_ai/models/instrumented.py

@@ -280,14 +280,14 @@ class InstrumentedModel(WrapperModel):
                             'gen_ai.request.model': request_model,
                             'gen_ai.response.model': response_model,
                         }
-                        if response.usage.request_tokens:  # pragma: no branch
+                        if response.usage.input_tokens:  # pragma: no branch
                             self.instrumentation_settings.tokens_histogram.record(
-                                response.usage.request_tokens,
+                                response.usage.input_tokens,
                                 {**metric_attributes, 'gen_ai.token.type': 'input'},
                             )
-                        if response.usage.response_tokens:  # pragma: no branch
+                        if response.usage.output_tokens:  # pragma: no branch
                             self.instrumentation_settings.tokens_histogram.record(
-                                response.usage.response_tokens,
+                                response.usage.output_tokens,
                                 {**metric_attributes, 'gen_ai.token.type': 'output'},
                             )
 

pydantic_ai/models/mcp_sampling.py

@@ -5,7 +5,7 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
 
-from .. import _mcp, exceptions, usage
+from .. import _mcp, exceptions
 from .._run_context import RunContext
 from ..messages import ModelMessage, ModelResponse
 from ..settings import ModelSettings
@@ -63,7 +63,6 @@ class MCPSamplingModel(Model):
         if result.role == 'assistant':
             return ModelResponse(
                 parts=[_mcp.map_from_sampling_content(result.content)],
-                usage=usage.Usage(requests=1),
                 model_name=result.model,
             )
         else:

pydantic_ai/models/mistral.py

@@ -40,7 +40,7 @@ from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from ..usage import Usage
+from ..usage import RequestUsage
 from . import (
     Model,
     ModelRequestParameters,
@@ -120,7 +120,7 @@ class MistralModel(Model):
     json_mode_schema_prompt: str = """Answer in JSON Object, respect the format:\n```\n{schema}\n```\n"""
 
     _model_name: MistralModelName = field(repr=False)
-    _system: str = field(default='mistral_ai', repr=False)
+    _provider: Provider[Mistral] = field(repr=False)
 
     def __init__(
         self,
@@ -147,13 +147,24 @@ class MistralModel(Model):
 
         if isinstance(provider, str):
             provider = infer_provider(provider)
+        self._provider = provider
         self.client = provider.client
 
         super().__init__(settings=settings, profile=profile or provider.model_profile)
 
     @property
     def base_url(self) -> str:
-        return self.client.sdk_configuration.get_server_details()[0]
+        return self._provider.base_url
+
+    @property
+    def model_name(self) -> MistralModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The model provider."""
+        return self._provider.name
 
     async def request(
         self,
@@ -167,7 +178,6 @@ class MistralModel(Model):
             messages, cast(MistralModelSettings, model_settings or {}), model_request_parameters
         )
         model_response = self._process_response(response)
-        model_response.usage.requests = 1
         return model_response
 
     @asynccontextmanager
@@ -186,16 +196,6 @@ class MistralModel(Model):
         async with response:
             yield await self._process_streamed_response(response, model_request_parameters)
 
-    @property
-    def model_name(self) -> MistralModelName:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider."""
-        return self._system
-
     async def _completions_create(
         self,
         messages: list[ModelMessage],
@@ -348,7 +348,11 @@ class MistralModel(Model):
                 parts.append(tool)
 
         return ModelResponse(
-            parts, usage=_map_usage(response), model_name=response.model, timestamp=timestamp, vendor_id=response.id
+            parts,
+            usage=_map_usage(response),
+            model_name=response.model,
+            timestamp=timestamp,
+            provider_request_id=response.id,
         )
 
     async def _process_streamed_response(
@@ -699,17 +703,15 @@ SIMPLE_JSON_TYPE_MAPPING = {
 }
 
 
-def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk) -> Usage:
+def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk) -> RequestUsage:
     """Maps a Mistral Completion Chunk or Chat Completion Response to a Usage."""
     if response.usage:
-        return Usage(
-            request_tokens=response.usage.prompt_tokens,
-            response_tokens=response.usage.completion_tokens,
-            total_tokens=response.usage.total_tokens,
-            details=None,
+        return RequestUsage(
+            input_tokens=response.usage.prompt_tokens,
+            output_tokens=response.usage.completion_tokens,
         )
     else:
-        return Usage()  # pragma: no cover
+        return RequestUsage()  # pragma: no cover
 
 
 def _map_content(content: MistralOptionalNullable[MistralContent]) -> str | None: