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

pydantic_ai/models/__init__.py

@@ -42,7 +42,7 @@ from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
 from ..profiles._json_schema import JsonSchemaTransformer
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from ..usage import Usage
+from ..usage import RequestUsage
 
 KnownModelName = TypeAliasType(
     'KnownModelName',
@@ -418,7 +418,7 @@ class Model(ABC):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
-    ) -> Usage:
+    ) -> RequestUsage:
         """Make a request to the model for counting tokens."""
         # This method is not required, but you need to implement it if you want to support `UsageLimits.count_tokens_before_request`.
         raise NotImplementedError(f'Token counting ahead of the request is not supported by {self.__class__.__name__}')
@@ -480,7 +480,7 @@ class Model(ABC):
     @property
     @abstractmethod
     def system(self) -> str:
-        """The system / model provider, ex: openai.
+        """The model provider, ex: openai.
 
         Use to populate the `gen_ai.system` OpenTelemetry semantic convention attribute,
         so should use well-known values listed in
@@ -547,7 +547,7 @@ class StreamedResponse(ABC):
 
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
     _event_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
-    _usage: Usage = field(default_factory=Usage, init=False)
+    _usage: RequestUsage = field(default_factory=RequestUsage, init=False)
 
     def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
         """Stream the response as an async iterable of [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s.
@@ -600,7 +600,7 @@ class StreamedResponse(ABC):
             usage=self.usage(),
         )
 
-    def usage(self) -> Usage:
+    def usage(self) -> RequestUsage:
         """Get the usage of the response so far. This will not be the final usage until the stream is exhausted."""
         return self._usage
 

pydantic_ai/models/anthropic.py

@@ -8,14 +8,6 @@ from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import Any, Literal, Union, cast, overload
 
-from anthropic.types.beta import (
-    BetaCitationsDelta,
-    BetaCodeExecutionToolResultBlock,
-    BetaCodeExecutionToolResultBlockParam,
-    BetaInputJSONDelta,
-    BetaServerToolUseBlockParam,
-    BetaWebSearchToolResultBlockParam,
-)
 from typing_extensions import assert_never
 
 from pydantic_ai.builtin_tools import CodeExecutionTool, WebSearchTool
@@ -47,24 +39,21 @@ 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
 
 try:
     from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropic, AsyncStream
     from anthropic.types.beta import (
         BetaBase64PDFBlockParam,
         BetaBase64PDFSourceParam,
+        BetaCitationsDelta,
         BetaCodeExecutionTool20250522Param,
+        BetaCodeExecutionToolResultBlock,
+        BetaCodeExecutionToolResultBlockParam,
         BetaContentBlock,
         BetaContentBlockParam,
         BetaImageBlockParam,
+        BetaInputJSONDelta,
         BetaMessage,
         BetaMessageParam,
         BetaMetadataParam,
@@ -78,6 +67,7 @@ try:
         BetaRawMessageStreamEvent,
         BetaRedactedThinkingBlock,
         BetaServerToolUseBlock,
+        BetaServerToolUseBlockParam,
         BetaSignatureDelta,
         BetaTextBlock,
         BetaTextBlockParam,
@@ -94,6 +84,7 @@ try:
         BetaToolUseBlockParam,
         BetaWebSearchTool20250305Param,
         BetaWebSearchToolResultBlock,
+        BetaWebSearchToolResultBlockParam,
     )
     from anthropic.types.beta.beta_web_search_tool_20250305_param import UserLocation
     from anthropic.types.model_param import ModelParam
@@ -146,7 +137,7 @@ class AnthropicModel(Model):
     client: AsyncAnthropic = field(repr=False)
 
     _model_name: AnthropicModelName = field(repr=False)
-    _system: str = field(default='anthropic', repr=False)
+    _provider: Provider[AsyncAnthropic] = field(repr=False)
 
     def __init__(
         self,
@@ -170,6 +161,7 @@ class AnthropicModel(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)
@@ -178,6 +170,16 @@ class AnthropicModel(Model):
     def base_url(self) -> str:
         return str(self.client.base_url)
 
+    @property
+    def model_name(self) -> AnthropicModelName:
+        """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],
@@ -189,7 +191,6 @@ class AnthropicModel(Model):
             messages, False, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
         )
         model_response = self._process_response(response)
-        model_response.usage.requests = 1
         return model_response
 
     @asynccontextmanager
@@ -207,16 +208,6 @@ class AnthropicModel(Model):
         async with response:
             yield await self._process_streamed_response(response, model_request_parameters)
 
-    @property
-    def model_name(self) -> AnthropicModelName:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider."""
-        return self._system
-
     @overload
     async def _messages_create(
         self,
@@ -246,7 +237,9 @@ class AnthropicModel(Model):
     ) -> BetaMessage | AsyncStream[BetaRawMessageStreamEvent]:
         # standalone function to make it easier to override
         tools = self._get_tools(model_request_parameters)
-        tools += self._get_builtin_tools(model_request_parameters)
+        builtin_tools, tool_headers = self._get_builtin_tools(model_request_parameters)
+        tools += builtin_tools
+
         tool_choice: BetaToolChoiceParam | None
 
         if not tools:
@@ -264,8 +257,10 @@ class AnthropicModel(Model):
 
         try:
             extra_headers = model_settings.get('extra_headers', {})
+            for k, v in tool_headers.items():
+                extra_headers.setdefault(k, v)
             extra_headers.setdefault('User-Agent', get_user_agent())
-            extra_headers.setdefault('anthropic-beta', 'code-execution-2025-05-22')
+
             return await self.client.beta.messages.create(
                 max_tokens=model_settings.get('max_tokens', 4096),
                 system=system_prompt or NOT_GIVEN,
@@ -330,7 +325,9 @@ class AnthropicModel(Model):
                     )
                 )
 
-        return ModelResponse(items, usage=_map_usage(response), model_name=response.model, vendor_id=response.id)
+        return ModelResponse(
+            items, usage=_map_usage(response), model_name=response.model, provider_request_id=response.id
+        )
 
     async def _process_streamed_response(
         self, response: AsyncStream[BetaRawMessageStreamEvent], model_request_parameters: ModelRequestParameters
@@ -352,8 +349,11 @@ class AnthropicModel(Model):
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[BetaToolParam]:
         return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
-    def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -> list[BetaToolUnionParam]:
+    def _get_builtin_tools(
+        self, model_request_parameters: ModelRequestParameters
+    ) -> tuple[list[BetaToolUnionParam], dict[str, str]]:
         tools: list[BetaToolUnionParam] = []
+        extra_headers: dict[str, str] = {}
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):
                 user_location = UserLocation(type='approximate', **tool.user_location) if tool.user_location else None
@@ -361,18 +361,20 @@ class AnthropicModel(Model):
                     BetaWebSearchTool20250305Param(
                         name='web_search',
                         type='web_search_20250305',
+                        max_uses=tool.max_uses,
                         allowed_domains=tool.allowed_domains,
                         blocked_domains=tool.blocked_domains,
                         user_location=user_location,
                     )
                 )
             elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
+                extra_headers['anthropic-beta'] = 'code-execution-2025-05-22'
                 tools.append(BetaCodeExecutionTool20250522Param(name='code_execution', type='code_execution_20250522'))
             else:  # pragma: no cover
                 raise UserError(
                     f'`{tool.__class__.__name__}` is not supported by `AnthropicModel`. If it should be, please file an issue.'
                 )
-        return tools
+        return tools, extra_headers
 
     async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[BetaMessageParam]]:  # noqa: C901
         """Just maps a `pydantic_ai.Message` to a `anthropic.types.MessageParam`."""
@@ -528,7 +530,7 @@ class AnthropicModel(Model):
         }
 
 
-def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.Usage:
+def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.RequestUsage:
     if isinstance(message, BetaMessage):
         response_usage = message.usage
     elif isinstance(message, BetaRawMessageStartEvent):
@@ -541,7 +543,7 @@ def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.Usage:
         # - RawContentBlockStartEvent
         # - RawContentBlockDeltaEvent
         # - RawContentBlockStopEvent
-        return usage.Usage()
+        return usage.RequestUsage()
 
     # Store all integer-typed usage values in the details, except 'output_tokens' which is represented exactly by
     # `response_tokens`
@@ -552,17 +554,16 @@ def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.Usage:
     # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence using `get`
     # Tokens are only counted once between input_tokens, cache_creation_input_tokens, and cache_read_input_tokens
     # This approach maintains request_tokens as the count of all input tokens, with cached counts as details
-    request_tokens = (
-        details.get('input_tokens', 0)
-        + details.get('cache_creation_input_tokens', 0)
-        + details.get('cache_read_input_tokens', 0)
-    )
-
-    return usage.Usage(
-        request_tokens=request_tokens or None,
-        response_tokens=response_usage.output_tokens,
-        total_tokens=request_tokens + response_usage.output_tokens,
-        details=details or None,
+    cache_write_tokens = details.get('cache_creation_input_tokens', 0)
+    cache_read_tokens = details.get('cache_read_input_tokens', 0)
+    request_tokens = details.get('input_tokens', 0) + cache_write_tokens + cache_read_tokens
+
+    return usage.RequestUsage(
+        input_tokens=request_tokens,
+        cache_read_tokens=cache_read_tokens,
+        cache_write_tokens=cache_write_tokens,
+        output_tokens=response_usage.output_tokens,
+        details=details,
     )
 
 

pydantic_ai/models/bedrock.py

@@ -190,17 +190,7 @@ class BedrockConverseModel(Model):
     client: BedrockRuntimeClient
 
     _model_name: BedrockModelName = field(repr=False)
-    _system: str = field(default='bedrock', repr=False)
-
-    @property
-    def model_name(self) -> str:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider, ex: openai."""
-        return self._system
+    _provider: Provider[BaseClient] = field(repr=False)
 
     def __init__(
         self,
@@ -226,10 +216,25 @@ class BedrockConverseModel(Model):
 
         if isinstance(provider, str):
             provider = infer_provider(provider)
+        self._provider = provider
         self.client = cast('BedrockRuntimeClient', provider.client)
 
         super().__init__(settings=settings, profile=profile or provider.model_profile)
 
+    @property
+    def base_url(self) -> str:
+        return str(self.client.meta.endpoint_url)
+
+    @property
+    def model_name(self) -> str:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The model provider."""
+        return self._provider.name
+
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolTypeDef]:
         return [self._map_tool_definition(r) for r in model_request_parameters.tool_defs.values()]
 
@@ -245,10 +250,6 @@ class BedrockConverseModel(Model):
 
         return {'toolSpec': tool_spec}
 
-    @property
-    def base_url(self) -> str:
-        return str(self.client.meta.endpoint_url)
-
     async def request(
         self,
         messages: list[ModelMessage],
@@ -258,7 +259,6 @@ class BedrockConverseModel(Model):
         settings = cast(BedrockModelSettings, model_settings or {})
         response = await self._messages_create(messages, False, settings, model_request_parameters)
         model_response = await self._process_response(response)
-        model_response.usage.requests = 1
         return model_response
 
     @asynccontextmanager
@@ -299,13 +299,12 @@ class BedrockConverseModel(Model):
                             tool_call_id=tool_use['toolUseId'],
                         ),
                     )
-        u = usage.Usage(
-            request_tokens=response['usage']['inputTokens'],
-            response_tokens=response['usage']['outputTokens'],
-            total_tokens=response['usage']['totalTokens'],
+        u = usage.RequestUsage(
+            input_tokens=response['usage']['inputTokens'],
+            output_tokens=response['usage']['outputTokens'],
         )
         vendor_id = response.get('ResponseMetadata', {}).get('RequestId', None)
-        return ModelResponse(items, usage=u, model_name=self.model_name, vendor_id=vendor_id)
+        return ModelResponse(items, usage=u, model_name=self.model_name, provider_request_id=vendor_id)
 
     @overload
     async def _messages_create(
@@ -648,7 +647,7 @@ class BedrockStreamedResponse(StreamedResponse):
                         )
                 if 'text' in delta:
                     maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
-                    if maybe_event is not None:
+                    if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
                 if 'toolUse' in delta:
                     tool_use = delta['toolUse']
@@ -670,11 +669,10 @@ class BedrockStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
-    def _map_usage(self, metadata: ConverseStreamMetadataEventTypeDef) -> usage.Usage:
-        return usage.Usage(
-            request_tokens=metadata['usage']['inputTokens'],
-            response_tokens=metadata['usage']['outputTokens'],
-            total_tokens=metadata['usage']['totalTokens'],
+    def _map_usage(self, metadata: ConverseStreamMetadataEventTypeDef) -> usage.RequestUsage:
+        return usage.RequestUsage(
+            input_tokens=metadata['usage']['inputTokens'],
+            output_tokens=metadata['usage']['outputTokens'],
         )
 
 

pydantic_ai/models/cohere.py

@@ -30,11 +30,7 @@ from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import (
-    Model,
-    ModelRequestParameters,
-    check_allow_model_requests,
-)
+from . import Model, ModelRequestParameters, check_allow_model_requests
 
 try:
     from cohere import (
@@ -106,7 +102,7 @@ class CohereModel(Model):
     client: AsyncClientV2 = field(repr=False)
 
     _model_name: CohereModelName = field(repr=False)
-    _system: str = field(default='cohere', repr=False)
+    _provider: Provider[AsyncClientV2] = field(repr=False)
 
     def __init__(
         self,
@@ -131,6 +127,7 @@ class CohereModel(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)
@@ -140,6 +137,16 @@ class CohereModel(Model):
         client_wrapper = self.client._client_wrapper  # type: ignore
         return str(client_wrapper.get_base_url())
 
+    @property
+    def model_name(self) -> CohereModelName:
+        """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],
@@ -149,19 +156,8 @@ class CohereModel(Model):
         check_allow_model_requests()
         response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}), model_request_parameters)
         model_response = self._process_response(response)
-        model_response.usage.requests = 1
         return model_response
 
-    @property
-    def model_name(self) -> CohereModelName:
-        """The model name."""
-        return self._model_name
-
-    @property
-    def system(self) -> str:
-        """The system / model provider."""
-        return self._system
-
     async def _chat(
         self,
         messages: list[ModelMessage],
@@ -301,10 +297,10 @@ class CohereModel(Model):
                 assert_never(part)
 
 
-def _map_usage(response: V2ChatResponse) -> usage.Usage:
+def _map_usage(response: V2ChatResponse) -> usage.RequestUsage:
     u = response.usage
     if u is None:
-        return usage.Usage()
+        return usage.RequestUsage()
     else:
         details: dict[str, int] = {}
         if u.billed_units is not None:
@@ -317,11 +313,10 @@ def _map_usage(response: V2ChatResponse) -> usage.Usage:
             if u.billed_units.classifications:  # pragma: no cover
                 details['classifications'] = int(u.billed_units.classifications)
 
-        request_tokens = int(u.tokens.input_tokens) if u.tokens and u.tokens.input_tokens else None
-        response_tokens = int(u.tokens.output_tokens) if u.tokens and u.tokens.output_tokens else None
-        return usage.Usage(
-            request_tokens=request_tokens,
-            response_tokens=response_tokens,
-            total_tokens=(request_tokens or 0) + (response_tokens or 0),
+        request_tokens = int(u.tokens.input_tokens) if u.tokens and u.tokens.input_tokens else 0
+        response_tokens = int(u.tokens.output_tokens) if u.tokens and u.tokens.output_tokens else 0
+        return usage.RequestUsage(
+            input_tokens=request_tokens,
+            output_tokens=response_tokens,
             details=details,
         )

pydantic_ai/models/fallback.py

@@ -33,8 +33,8 @@ class FallbackModel(Model):
 
     def __init__(
         self,
-        default_model: Model | KnownModelName,
-        *fallback_models: Model | KnownModelName,
+        default_model: Model | KnownModelName | str,
+        *fallback_models: Model | KnownModelName | str,
         fallback_on: Callable[[Exception], bool] | tuple[type[Exception], ...] = (ModelHTTPError,),
     ):
         """Initialize a fallback model instance.
@@ -52,6 +52,19 @@ class FallbackModel(Model):
         else:
             self._fallback_on = fallback_on
 
+    @property
+    def model_name(self) -> str:
+        """The model name."""
+        return f'fallback:{",".join(model.model_name for model in self.models)}'
+
+    @property
+    def system(self) -> str:
+        return f'fallback:{",".join(model.system for model in self.models)}'
+
+    @property
+    def base_url(self) -> str | None:
+        return self.models[0].base_url
+
     async def request(
         self,
         messages: list[ModelMessage],
@@ -121,19 +134,6 @@ class FallbackModel(Model):
                 if attributes.get('gen_ai.request.model') == self.model_name:  # pragma: no branch
                     span.set_attributes(InstrumentedModel.model_attributes(model))
 
-    @property
-    def model_name(self) -> str:
-        """The model name."""
-        return f'fallback:{",".join(model.model_name for model in self.models)}'
-
-    @property
-    def system(self) -> str:
-        return f'fallback:{",".join(model.system for model in self.models)}'
-
-    @property
-    def base_url(self) -> str | None:
-        return self.models[0].base_url
-
 
 def _default_fallback_condition_factory(exceptions: tuple[type[Exception], ...]) -> Callable[[Exception], bool]:
     """Create a default fallback condition for the given exceptions."""

pydantic_ai/models/function.py

@@ -138,7 +138,6 @@ class FunctionModel(Model):
         # Add usage data if not already present
         if not response.usage.has_values():  # pragma: no branch
             response.usage = _estimate_usage(chain(messages, [response]))
-            response.usage.requests = 1
         return response
 
     @asynccontextmanager
@@ -270,7 +269,7 @@ class FunctionStreamedResponse(StreamedResponse):
         async for item in self._iter:
             if isinstance(item, str):
                 response_tokens = _estimate_string_tokens(item)
-                self._usage += usage.Usage(response_tokens=response_tokens, total_tokens=response_tokens)
+                self._usage += usage.RequestUsage(output_tokens=response_tokens)
                 maybe_event = self._parts_manager.handle_text_delta(vendor_part_id='content', content=item)
                 if maybe_event is not None:  # pragma: no branch
                     yield maybe_event
@@ -279,7 +278,7 @@ class FunctionStreamedResponse(StreamedResponse):
                     if isinstance(delta, DeltaThinkingPart):
                         if delta.content:  # pragma: no branch
                             response_tokens = _estimate_string_tokens(delta.content)
-                            self._usage += usage.Usage(response_tokens=response_tokens, total_tokens=response_tokens)
+                            self._usage += usage.RequestUsage(output_tokens=response_tokens)
                         yield self._parts_manager.handle_thinking_delta(
                             vendor_part_id=dtc_index,
                             content=delta.content,
@@ -288,7 +287,7 @@ class FunctionStreamedResponse(StreamedResponse):
                     elif isinstance(delta, DeltaToolCall):
                         if delta.json_args:
                             response_tokens = _estimate_string_tokens(delta.json_args)
-                            self._usage += usage.Usage(response_tokens=response_tokens, total_tokens=response_tokens)
+                            self._usage += usage.RequestUsage(output_tokens=response_tokens)
                         maybe_event = self._parts_manager.handle_tool_call_delta(
                             vendor_part_id=dtc_index,
                             tool_name=delta.name,
@@ -311,7 +310,7 @@ class FunctionStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.Usage:
+def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.RequestUsage:
     """Very rough guesstimate of the token usage associated with a series of messages.
 
     This is designed to be used solely to give plausible numbers for testing!
@@ -349,10 +348,9 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.Usage:
                     assert_never(part)
         else:
             assert_never(message)
-    return usage.Usage(
-        request_tokens=request_tokens,
-        response_tokens=response_tokens,
-        total_tokens=request_tokens + response_tokens,
+    return usage.RequestUsage(
+        input_tokens=request_tokens,
+        output_tokens=response_tokens,
     )