pydantic-ai-slim 0.7.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

pydantic_ai/models/gemini.py

@@ -40,14 +40,7 @@ from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import (
-    Model,
-    ModelRequestParameters,
-    StreamedResponse,
-    check_allow_model_requests,
-    download_item,
-    get_user_agent,
-)
+from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
 
 LatestGeminiModelNames = Literal[
     'gemini-2.0-flash',
@@ -108,10 +101,9 @@ class GeminiModel(Model):
     client: httpx.AsyncClient = field(repr=False)
 
     _model_name: GeminiModelName = field(repr=False)
-    _provider: Literal['google-gla', 'google-vertex'] | Provider[httpx.AsyncClient] | None = field(repr=False)
+    _provider: Provider[httpx.AsyncClient] = field(repr=False)
     _auth: AuthProtocol | None = field(repr=False)
     _url: str | None = field(repr=False)
-    _system: str = field(default='gemini', repr=False)
 
     def __init__(
         self,
@@ -132,11 +124,10 @@ class GeminiModel(Model):
             settings: Default model settings for this model instance.
         """
         self._model_name = model_name
-        self._provider = provider
 
         if isinstance(provider, str):
             provider = infer_provider(provider)
-        self._system = provider.name
+        self._provider = provider
         self.client = provider.client
         self._url = str(self.client.base_url)
 
@@ -147,6 +138,16 @@ class GeminiModel(Model):
         assert self._url is not None, 'URL not initialized'  # pragma: no cover
         return self._url  # pragma: no cover
 
+    @property
+    def model_name(self) -> GeminiModelName:
+        """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],
@@ -175,16 +176,6 @@ class GeminiModel(Model):
         ) as http_response:
             yield await self._process_streamed_response(http_response, model_request_parameters)
 
-    @property
-    def model_name(self) -> GeminiModelName:
-        """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) -> _GeminiTools | None:
         tools = [_function_from_abstract_tool(t) for t in model_request_parameters.tool_defs.values()]
         return _GeminiTools(function_declarations=tools) if tools else None
@@ -237,7 +228,7 @@ class GeminiModel(Model):
             request_data['safetySettings'] = gemini_safety_settings
 
         if gemini_labels := model_settings.get('gemini_labels'):
-            if self._system == 'google-vertex':
+            if self._provider.name == 'google-vertex':
                 request_data['labels'] = gemini_labels  # pragma: lax no cover
 
         headers = {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
@@ -278,7 +269,6 @@ class GeminiModel(Model):
         if finish_reason:
             vendor_details = {'finish_reason': finish_reason}
         usage = _metadata_as_usage(response)
-        usage.requests = 1
         return _process_response_from_parts(
             parts,
             response.get('model_version', self._model_name),
@@ -673,7 +663,7 @@ def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart
 def _process_response_from_parts(
     parts: Sequence[_GeminiPartUnion],
     model_name: GeminiModelName,
-    usage: usage.Usage,
+    usage: usage.RequestUsage,
     vendor_id: str | None,
     vendor_details: dict[str, Any] | None = None,
 ) -> ModelResponse:
@@ -693,7 +683,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, usage=usage, model_name=model_name, vendor_id=vendor_id, vendor_details=vendor_details
+        parts=items, usage=usage, model_name=model_name, provider_request_id=vendor_id, provider_details=vendor_details
     )
 
 
@@ -859,31 +849,45 @@ class _GeminiUsageMetaData(TypedDict, total=False):
     ]
 
 
-def _metadata_as_usage(response: _GeminiResponse) -> usage.Usage:
+def _metadata_as_usage(response: _GeminiResponse) -> usage.RequestUsage:
     metadata = response.get('usage_metadata')
     if metadata is None:
-        return usage.Usage()  # pragma: no cover
+        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.get('cached_content_token_count', 0):
+        details['cached_content_tokens'] = cached_content_token_count
 
-    if thoughts_token_count := metadata.get('thoughts_token_count'):
+    if thoughts_token_count := metadata.get('thoughts_token_count', 0):
         details['thoughts_tokens'] = thoughts_token_count
 
-    if tool_use_prompt_token_count := metadata.get('tool_use_prompt_token_count'):
-        details['tool_use_prompt_tokens'] = tool_use_prompt_token_count  # pragma: no cover
+    if tool_use_prompt_token_count := metadata.get('tool_use_prompt_token_count', 0):
+        details['tool_use_prompt_tokens'] = tool_use_prompt_token_count
 
+    input_audio_tokens = 0
+    output_audio_tokens = 0
+    cache_audio_read_tokens = 0
     for key, metadata_details in metadata.items():
         if key.endswith('_details') and metadata_details:
             metadata_details = cast(list[_GeminiModalityTokenCount], 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),
+                modality = detail['modality']
+                details[f'{modality.lower()}_{suffix}'] = value = detail.get('token_count', 0)
+                if value and modality == 'AUDIO':
+                    if key == 'prompt_tokens_details':
+                        input_audio_tokens = value
+                    elif key == 'candidates_tokens_details':
+                        output_audio_tokens = value
+                    elif key == 'cache_tokens_details':  # pragma: no branch
+                        cache_audio_read_tokens = value
+
+    return usage.RequestUsage(
+        input_tokens=metadata.get('prompt_token_count', 0),
+        output_tokens=metadata.get('candidates_token_count', 0),
+        cache_read_tokens=cached_content_token_count,
+        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/google.py

@@ -144,7 +144,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 +167,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 +176,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 +202,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,7 +216,7 @@ 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(
                 system_instruction=generation_config.get('system_instruction'),
@@ -238,9 +245,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 +262,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)])
@@ -392,9 +388,12 @@ class GoogleModel(Model):
         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(
@@ -590,7 +589,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 +626,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 +646,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(
@@ -457,6 +461,7 @@ class GroqStreamedResponse(StreamedResponse):
                     vendor_part_id='content',
                     content=content,
                     thinking_tags=self._model_profile.thinking_tags,
+                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
                 )
                 if maybe_event is not None:  # pragma: no branch
                     yield maybe_event
@@ -483,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
@@ -491,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

@@ -35,7 +35,7 @@ from ..messages import (
     UserPromptPart,
     VideoUrl,
 )
-from ..profiles import ModelProfile
+from ..profiles import ModelProfile, ModelProfileSpec
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
@@ -114,13 +114,15 @@ 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,
         model_name: str,
         *,
         provider: Literal['huggingface'] | Provider[AsyncInferenceClient] = 'huggingface',
+        profile: ModelProfileSpec | None = None,
+        settings: ModelSettings | None = None,
     ):
         """Initialize a Hugging Face model.
 
@@ -128,13 +130,27 @@ class HuggingFaceModel(Model):
             model_name: The name of the Model to use. You can browse available models [here](https://huggingface.co/models?pipeline_tag=text-generation&inference_provider=all&sort=trending).
             provider: The provider to use for Hugging Face Inference Providers. Can be either the string 'huggingface' or an
                 instance of `Provider[AsyncInferenceClient]`. If not provided, the other parameters will be used.
+            profile: The model profile to use. Defaults to a profile picked by the provider based on the model name.
+            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],
@@ -146,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
@@ -163,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,
@@ -266,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(
@@ -444,11 +449,12 @@ class HuggingFaceStreamedResponse(StreamedResponse):
 
             # Handle the text part of the response
             content = choice.delta.content
-            if content:
+            if content is not None:
                 maybe_event = self._parts_manager.handle_text_delta(
                     vendor_part_id='content',
                     content=content,
                     thinking_tags=self._model_profile.thinking_tags,
+                    ignore_leading_whitespace=self._model_profile.ignore_streamed_leading_whitespace,
                 )
                 if maybe_event is not None:  # pragma: no branch
                     yield maybe_event
@@ -474,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: