pydantic-ai-slim 0.0.13__py3-none-any.whl → 0.0.15__py3-none-any.whl

pydantic_ai/exceptions.py

@@ -2,7 +2,7 @@ from __future__ import annotations as _annotations
 
 import json
 
-__all__ = 'ModelRetry', 'UserError', 'UnexpectedModelBehavior'
+__all__ = 'ModelRetry', 'UserError', 'AgentRunError', 'UnexpectedModelBehavior', 'UsageLimitExceeded'
 
 
 class ModelRetry(Exception):
@@ -30,7 +30,25 @@ class UserError(RuntimeError):
         super().__init__(message)
 
 
-class UnexpectedModelBehavior(RuntimeError):
+class AgentRunError(RuntimeError):
+    """Base class for errors occurring during an agent run."""
+
+    message: str
+    """The error message."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class UsageLimitExceeded(AgentRunError):
+    """Error raised when a Model's usage exceeds the specified limits."""
+
+
+class UnexpectedModelBehavior(AgentRunError):
     """Error caused by unexpected Model behavior, e.g. an unexpected response code."""
 
     message: str

pydantic_ai/messages.py

@@ -2,11 +2,11 @@ from __future__ import annotations as _annotations
 
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Annotated, Any, Literal, Union
+from typing import Annotated, Any, Literal, Union, cast
 
 import pydantic
 import pydantic_core
-from typing_extensions import Self
+from typing_extensions import Self, assert_never
 
 from ._utils import now_utc as _now_utc
 
@@ -190,12 +190,34 @@ class ToolCallPart:
     """Part type identifier, this is available on all parts as a discriminator."""
 
     @classmethod
-    def from_json(cls, tool_name: str, args_json: str, tool_call_id: str | None = None) -> Self:
-        return cls(tool_name, ArgsJson(args_json), tool_call_id)
+    def from_raw_args(cls, tool_name: str, args: str | dict[str, Any], tool_call_id: str | None = None) -> Self:
+        """Create a `ToolCallPart` from raw arguments."""
+        if isinstance(args, str):
+            return cls(tool_name, ArgsJson(args), tool_call_id)
+        elif isinstance(args, dict):
+            return cls(tool_name, ArgsDict(args), tool_call_id)
+        else:
+            assert_never(args)
 
-    @classmethod
-    def from_dict(cls, tool_name: str, args_dict: dict[str, Any], tool_call_id: str | None = None) -> Self:
-        return cls(tool_name, ArgsDict(args_dict), tool_call_id)
+    def args_as_dict(self) -> dict[str, Any]:
+        """Return the arguments as a Python dictionary.
+
+        This is just for convenience with models that require dicts as input.
+        """
+        if isinstance(self.args, ArgsDict):
+            return self.args.args_dict
+        args = pydantic_core.from_json(self.args.args_json)
+        assert isinstance(args, dict), 'args should be a dict'
+        return cast(dict[str, Any], args)
+
+    def args_as_json_str(self) -> str:
+        """Return the arguments as a JSON string.
+
+        This is just for convenience with models that require JSON strings as input.
+        """
+        if isinstance(self.args, ArgsJson):
+            return self.args.args_json
+        return pydantic_core.to_json(self.args.args_dict).decode()
 
     def has_content(self) -> bool:
         if isinstance(self.args, ArgsDict):

pydantic_ai/models/__init__.py

@@ -20,7 +20,7 @@ from ..messages import ModelMessage, ModelResponse
 from ..settings import ModelSettings
 
 if TYPE_CHECKING:
-    from ..result import Cost
+    from ..result import Usage
     from ..tools import ToolDefinition
 
 
@@ -31,6 +31,7 @@ KnownModelName = Literal[
     'openai:gpt-4',
     'openai:o1-preview',
     'openai:o1-mini',
+    'openai:o1',
     'openai:gpt-3.5-turbo',
     'groq:llama-3.3-70b-versatile',
     'groq:llama-3.1-70b-versatile',
@@ -122,7 +123,7 @@ class AgentModel(ABC):
     @abstractmethod
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, Cost]:
+    ) -> tuple[ModelResponse, Usage]:
         """Make a request to the model."""
         raise NotImplementedError()
 
@@ -164,10 +165,10 @@ class StreamTextResponse(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def cost(self) -> Cost:
-        """Return the cost of the request.
+    def usage(self) -> Usage:
+        """Return the usage of the request.
 
-        NOTE: this won't return the ful cost until the stream is finished.
+        NOTE: this won't return the full usage until the stream is finished.
         """
         raise NotImplementedError()
 
@@ -205,10 +206,10 @@ class StreamStructuredResponse(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def cost(self) -> Cost:
-        """Get the cost of the request.
+    def usage(self) -> Usage:
+        """Get the usage of the request.
 
-        NOTE: this won't return the full cost until the stream is finished.
+        NOTE: this won't return the full usage until the stream is finished.
         """
         raise NotImplementedError()
 
@@ -235,7 +236,7 @@ The testing models [`TestModel`][pydantic_ai.models.test.TestModel] and
 def check_allow_model_requests() -> None:
     """Check if model requests are allowed.
 
-    If you're defining your own models that have cost or latency associated with their use, you should call this in
+    If you're defining your own models that have costs or latency associated with their use, you should call this in
     [`Model.agent_model`][pydantic_ai.models.Model.agent_model].
 
     Raises:

pydantic_ai/models/anthropic.py

@@ -158,9 +158,9 @@ class AnthropicAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Cost]:
+    ) -> tuple[ModelResponse, result.Usage]:
         response = await self._messages_create(messages, False, model_settings)
-        return self._process_response(response), _map_cost(response)
+        return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
@@ -220,7 +220,7 @@ class AnthropicAgentModel(AgentModel):
             else:
                 assert isinstance(item, ToolUseBlock), 'unexpected item type'
                 items.append(
-                    ToolCallPart.from_dict(
+                    ToolCallPart.from_raw_args(
                         item.name,
                         cast(dict[str, Any], item.input),
                         item.id,
@@ -282,11 +282,11 @@ class AnthropicAgentModel(AgentModel):
                                 MessageParam(
                                     role='user',
                                     content=[
-                                        ToolUseBlockParam(
-                                            id=_guard_tool_call_id(t=part, model_source='Anthropic'),
-                                            input=part.model_response(),
-                                            name=part.tool_name,
-                                            type='tool_use',
+                                        ToolResultBlockParam(
+                                            tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'),
+                                            type='tool_result',
+                                            content=part.model_response(),
+                                            is_error=True,
                                         ),
                                     ],
                                 )
@@ -311,11 +311,11 @@ def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
         id=_guard_tool_call_id(t=t, model_source='Anthropic'),
         type='tool_use',
         name=t.tool_name,
-        input=t.args.args_dict,
+        input=t.args_as_dict(),
     )
 
 
-def _map_cost(message: AnthropicMessage | RawMessageStreamEvent) -> result.Cost:
+def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> result.Usage:
     if isinstance(message, AnthropicMessage):
         usage = message.usage
     else:
@@ -332,11 +332,11 @@ def _map_cost(message: AnthropicMessage | RawMessageStreamEvent) -> result.Cost:
             usage = None
 
     if usage is None:
-        return result.Cost()
+        return result.Usage()
 
     request_tokens = getattr(usage, 'input_tokens', None)
 
-    return result.Cost(
+    return result.Usage(
         # Usage coming from the RawMessageDeltaEvent doesn't have input token data, hence this getattr
         request_tokens=request_tokens,
         response_tokens=usage.output_tokens,

pydantic_ai/models/function.py

@@ -9,12 +9,10 @@ from datetime import datetime
 from itertools import chain
 from typing import Callable, Union, cast
 
-import pydantic_core
 from typing_extensions import TypeAlias, assert_never, overload
 
 from .. import _utils, result
 from ..messages import (
-    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -144,7 +142,7 @@ class FunctionAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Cost]:
+    ) -> tuple[ModelResponse, result.Usage]:
         agent_info = replace(self.agent_info, model_settings=model_settings)
 
         assert self.function is not None, 'FunctionModel must receive a `function` to support non-streamed requests'
@@ -155,7 +153,7 @@ class FunctionAgentModel(AgentModel):
             assert isinstance(response_, ModelResponse), response_
             response = response_
         # TODO is `messages` right here? Should it just be new messages?
-        return response, _estimate_cost(chain(messages, [response]))
+        return response, _estimate_usage(chain(messages, [response]))
 
     @asynccontextmanager
     async def request_stream(
@@ -198,8 +196,8 @@ class FunctionStreamTextResponse(StreamTextResponse):
         yield from self._buffer
         self._buffer.clear()
 
-    def cost(self) -> result.Cost:
-        return result.Cost()
+    def usage(self) -> result.Usage:
+        return result.Usage()
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -232,19 +230,19 @@ class FunctionStreamStructuredResponse(StreamStructuredResponse):
         calls: list[ModelResponsePart] = []
         for c in self._delta_tool_calls.values():
             if c.name is not None and c.json_args is not None:
-                calls.append(ToolCallPart.from_json(c.name, c.json_args))
+                calls.append(ToolCallPart.from_raw_args(c.name, c.json_args))
 
         return ModelResponse(calls, timestamp=self._timestamp)
 
-    def cost(self) -> result.Cost:
-        return result.Cost()
+    def usage(self) -> result.Usage:
+        return _estimate_usage([self.get()])
 
     def timestamp(self) -> datetime:
         return self._timestamp
 
 
-def _estimate_cost(messages: Iterable[ModelMessage]) -> result.Cost:
-    """Very rough guesstimate of the number of tokens associate with a series of messages.
+def _estimate_usage(messages: Iterable[ModelMessage]) -> result.Usage:
+    """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!
     """
@@ -255,32 +253,28 @@ def _estimate_cost(messages: Iterable[ModelMessage]) -> result.Cost:
         if isinstance(message, ModelRequest):
             for part in message.parts:
                 if isinstance(part, (SystemPromptPart, UserPromptPart)):
-                    request_tokens += _string_cost(part.content)
+                    request_tokens += _estimate_string_usage(part.content)
                 elif isinstance(part, ToolReturnPart):
-                    request_tokens += _string_cost(part.model_response_str())
+                    request_tokens += _estimate_string_usage(part.model_response_str())
                 elif isinstance(part, RetryPromptPart):
-                    request_tokens += _string_cost(part.model_response())
+                    request_tokens += _estimate_string_usage(part.model_response())
                 else:
                     assert_never(part)
         elif isinstance(message, ModelResponse):
             for part in message.parts:
                 if isinstance(part, TextPart):
-                    response_tokens += _string_cost(part.content)
+                    response_tokens += _estimate_string_usage(part.content)
                 elif isinstance(part, ToolCallPart):
                     call = part
-                    if isinstance(call.args, ArgsJson):
-                        args_str = call.args.args_json
-                    else:
-                        args_str = pydantic_core.to_json(call.args.args_dict).decode()
-                    response_tokens += 1 + _string_cost(args_str)
+                    response_tokens += 1 + _estimate_string_usage(call.args_as_json_str())
                 else:
                     assert_never(part)
         else:
             assert_never(message)
-    return result.Cost(
+    return result.Usage(
         request_tokens=request_tokens, response_tokens=response_tokens, total_tokens=request_tokens + response_tokens
     )
 
 
-def _string_cost(content: str) -> int:
+def _estimate_string_usage(content: str) -> int:
     return len(re.split(r'[\s",.:]+', content))

pydantic_ai/models/gemini.py

@@ -16,7 +16,6 @@ from typing_extensions import NotRequired, TypedDict, TypeGuard, assert_never
 
 from .. import UnexpectedModelBehavior, _utils, exceptions, result
 from ..messages import (
-    ArgsDict,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -172,10 +171,10 @@ class GeminiAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Cost]:
+    ) -> tuple[ModelResponse, result.Usage]:
         async with self._make_request(messages, False, model_settings) as http_response:
             response = _gemini_response_ta.validate_json(await http_response.aread())
-        return self._process_response(response), _metadata_as_cost(response)
+        return self._process_response(response), _metadata_as_usage(response)
 
     @asynccontextmanager
     async def request_stream(
@@ -301,7 +300,7 @@ class GeminiStreamTextResponse(StreamTextResponse):
     _stream: AsyncIterator[bytes]
     _position: int = 0
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
-    _cost: result.Cost = field(default_factory=result.Cost, init=False)
+    _usage: result.Usage = field(default_factory=result.Usage, init=False)
 
     async def __anext__(self) -> None:
         chunk = await self._stream.__anext__()
@@ -321,7 +320,7 @@ class GeminiStreamTextResponse(StreamTextResponse):
                 new_items, experimental_allow_partial='trailing-strings'
             )
         for r in new_responses:
-            self._cost += _metadata_as_cost(r)
+            self._usage += _metadata_as_usage(r)
             parts = r['candidates'][0]['content']['parts']
             if _all_text_parts(parts):
                 for part in parts:
@@ -331,8 +330,8 @@ class GeminiStreamTextResponse(StreamTextResponse):
                     'Streamed response with unexpected content, expected all parts to be text'
                 )
 
-    def cost(self) -> result.Cost:
-        return self._cost
+    def usage(self) -> result.Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -345,7 +344,7 @@ class GeminiStreamStructuredResponse(StreamStructuredResponse):
     _content: bytearray
     _stream: AsyncIterator[bytes]
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
-    _cost: result.Cost = field(default_factory=result.Cost, init=False)
+    _usage: result.Usage = field(default_factory=result.Usage, init=False)
 
     async def __anext__(self) -> None:
         chunk = await self._stream.__anext__()
@@ -365,15 +364,15 @@ class GeminiStreamStructuredResponse(StreamStructuredResponse):
             experimental_allow_partial='off' if final else 'trailing-strings',
         )
         combined_parts: list[_GeminiPartUnion] = []
-        self._cost = result.Cost()
+        self._usage = result.Usage()
         for r in responses:
-            self._cost += _metadata_as_cost(r)
+            self._usage += _metadata_as_usage(r)
             candidate = r['candidates'][0]
             combined_parts.extend(candidate['content']['parts'])
         return _process_response_from_parts(combined_parts, timestamp=self._timestamp)
 
-    def cost(self) -> result.Cost:
-        return self._cost
+    def usage(self) -> result.Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -460,8 +459,7 @@ class _GeminiFunctionCallPart(TypedDict):
 
 
 def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart:
-    assert isinstance(tool.args, ArgsDict), f'Expected ArgsObject, got {tool.args}'
-    return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args.args_dict))
+    return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args_as_dict()))
 
 
 def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: datetime | None = None) -> ModelResponse:
@@ -470,7 +468,7 @@ def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: d
         if 'text' in part:
             items.append(TextPart(part['text']))
         elif 'function_call' in part:
-            items.append(ToolCallPart.from_dict(part['function_call']['name'], part['function_call']['args']))
+            items.append(ToolCallPart.from_raw_args(part['function_call']['name'], part['function_call']['args']))
         elif 'function_response' in part:
             raise exceptions.UnexpectedModelBehavior(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
@@ -640,14 +638,14 @@ class _GeminiUsageMetaData(TypedDict, total=False):
     cached_content_token_count: NotRequired[Annotated[int, pydantic.Field(alias='cachedContentTokenCount')]]
 
 
-def _metadata_as_cost(response: _GeminiResponse) -> result.Cost:
+def _metadata_as_usage(response: _GeminiResponse) -> result.Usage:
     metadata = response.get('usage_metadata')
     if metadata is None:
-        return result.Cost()
+        return result.Usage()
     details: dict[str, int] = {}
     if cached_content_token_count := metadata.get('cached_content_token_count'):
         details['cached_content_token_count'] = cached_content_token_count
-    return result.Cost(
+    return result.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),

pydantic_ai/models/groq.py

@@ -13,7 +13,6 @@ from typing_extensions import assert_never
 from .. import UnexpectedModelBehavior, _utils, result
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
-    ArgsJson,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -25,7 +24,7 @@ from ..messages import (
     ToolReturnPart,
     UserPromptPart,
 )
-from ..result import Cost
+from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
@@ -158,9 +157,9 @@ class GroqAgentModel(AgentModel):
 
     async def request(
         self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Cost]:
+    ) -> tuple[ModelResponse, result.Usage]:
         response = await self._completions_create(messages, False, model_settings)
-        return self._process_response(response), _map_cost(response)
+        return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
@@ -221,14 +220,14 @@ class GroqAgentModel(AgentModel):
             items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
-                items.append(ToolCallPart.from_json(c.function.name, c.function.arguments, c.id))
+                items.append(ToolCallPart.from_raw_args(c.function.name, c.function.arguments, c.id))
         return ModelResponse(items, timestamp=timestamp)
 
     @staticmethod
     async def _process_streamed_response(response: AsyncStream[ChatCompletionChunk]) -> EitherStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         timestamp: datetime | None = None
-        start_cost = Cost()
+        start_usage = Usage()
         # the first chunk may contain enough information so we iterate until we get either `tool_calls` or `content`
         while True:
             try:
@@ -236,19 +235,19 @@ class GroqAgentModel(AgentModel):
             except StopAsyncIteration as e:
                 raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') from e
             timestamp = timestamp or datetime.fromtimestamp(chunk.created, tz=timezone.utc)
-            start_cost += _map_cost(chunk)
+            start_usage += _map_usage(chunk)
 
             if chunk.choices:
                 delta = chunk.choices[0].delta
 
                 if delta.content is not None:
-                    return GroqStreamTextResponse(delta.content, response, timestamp, start_cost)
+                    return GroqStreamTextResponse(delta.content, response, timestamp, start_usage)
                 elif delta.tool_calls is not None:
                     return GroqStreamStructuredResponse(
                         response,
                         {c.index: c for c in delta.tool_calls},
                         timestamp,
-                        start_cost,
+                        start_usage,
                     )
 
     @classmethod
@@ -308,7 +307,7 @@ class GroqStreamTextResponse(StreamTextResponse):
     _first: str | None
     _response: AsyncStream[ChatCompletionChunk]
     _timestamp: datetime
-    _cost: result.Cost
+    _usage: result.Usage
     _buffer: list[str] = field(default_factory=list, init=False)
 
     async def __anext__(self) -> None:
@@ -318,7 +317,7 @@ class GroqStreamTextResponse(StreamTextResponse):
             return None
 
         chunk = await self._response.__anext__()
-        self._cost = _map_cost(chunk)
+        self._usage = _map_usage(chunk)
 
         try:
             choice = chunk.choices[0]
@@ -335,8 +334,8 @@ class GroqStreamTextResponse(StreamTextResponse):
         yield from self._buffer
         self._buffer.clear()
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
@@ -349,11 +348,11 @@ class GroqStreamStructuredResponse(StreamStructuredResponse):
     _response: AsyncStream[ChatCompletionChunk]
     _delta_tool_calls: dict[int, ChoiceDeltaToolCall]
     _timestamp: datetime
-    _cost: result.Cost
+    _usage: result.Usage
 
     async def __anext__(self) -> None:
         chunk = await self._response.__anext__()
-        self._cost = _map_cost(chunk)
+        self._usage = _map_usage(chunk)
 
         try:
             choice = chunk.choices[0]
@@ -380,27 +379,26 @@ class GroqStreamStructuredResponse(StreamStructuredResponse):
         for c in self._delta_tool_calls.values():
             if f := c.function:
                 if f.name is not None and f.arguments is not None:
-                    items.append(ToolCallPart.from_json(f.name, f.arguments, c.id))
+                    items.append(ToolCallPart.from_raw_args(f.name, f.arguments, c.id))
 
         return ModelResponse(items, timestamp=self._timestamp)
 
-    def cost(self) -> Cost:
-        return self._cost
+    def usage(self) -> Usage:
+        return self._usage
 
     def timestamp(self) -> datetime:
         return self._timestamp
 
 
 def _map_tool_call(t: ToolCallPart) -> chat.ChatCompletionMessageToolCallParam:
-    assert isinstance(t.args, ArgsJson), f'Expected ArgsJson, got {t.args}'
     return chat.ChatCompletionMessageToolCallParam(
         id=_guard_tool_call_id(t=t, model_source='Groq'),
         type='function',
-        function={'name': t.tool_name, 'arguments': t.args.args_json},
+        function={'name': t.tool_name, 'arguments': t.args_as_json_str()},
     )
 
 
-def _map_cost(completion: ChatCompletionChunk | ChatCompletion) -> result.Cost:
+def _map_usage(completion: ChatCompletionChunk | ChatCompletion) -> result.Usage:
     usage = None
     if isinstance(completion, ChatCompletion):
         usage = completion.usage
@@ -408,9 +406,9 @@ def _map_cost(completion: ChatCompletionChunk | ChatCompletion) -> result.Cost:
         usage = completion.x_groq.usage
 
     if usage is None:
-        return result.Cost()
+        return result.Usage()
 
-    return result.Cost(
+    return result.Usage(
         request_tokens=usage.prompt_tokens,
         response_tokens=usage.completion_tokens,
         total_tokens=usage.total_tokens,