pydantic-ai-slim 0.0.12__py3-none-any.whl → 0.0.14__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,18 +2,17 @@ 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 pydantic import TypeAdapter
+from typing_extensions import Self, assert_never
 
-from . import _pydantic
 from ._utils import now_utc as _now_utc
 
 
 @dataclass
-class SystemPrompt:
+class SystemPromptPart:
     """A system prompt, generally written by the application developer.
 
     This gives the model context and guidance on how to respond.
@@ -21,12 +20,13 @@ class SystemPrompt:
 
     content: str
     """The content of the prompt."""
-    role: Literal['system'] = 'system'
-    """Message type identifier, this type is available on all message as a discriminator."""
+
+    part_kind: Literal['system-prompt'] = 'system-prompt'
+    """Part type identifier, this is available on all parts as a discriminator."""
 
 
 @dataclass
-class UserPrompt:
+class UserPromptPart:
     """A user prompt, generally written by the end user.
 
     Content comes from the `user_prompt` parameter of [`Agent.run`][pydantic_ai.Agent.run],
@@ -35,29 +35,35 @@ class UserPrompt:
 
     content: str
     """The content of the prompt."""
+
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp of the prompt."""
-    role: Literal['user'] = 'user'
-    """Message type identifier, this type is available on all message as a discriminator."""
+
+    part_kind: Literal['user-prompt'] = 'user-prompt'
+    """Part type identifier, this is available on all parts as a discriminator."""
 
 
-tool_return_ta: TypeAdapter[Any] = TypeAdapter(Any)
+tool_return_ta: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter(Any, config=pydantic.ConfigDict(defer_build=True))
 
 
 @dataclass
-class ToolReturn:
+class ToolReturnPart:
     """A tool return message, this encodes the result of running a tool."""
 
     tool_name: str
     """The name of the "tool" was called."""
+
     content: Any
     """The return value."""
-    tool_id: str | None = None
-    """Optional tool identifier, this is used by some models including OpenAI."""
+
+    tool_call_id: str | None = None
+    """Optional tool call identifier, this is used by some models including OpenAI."""
+
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp, when the tool returned."""
-    role: Literal['tool-return'] = 'tool-return'
-    """Message type identifier, this type is available on all message as a discriminator."""
+
+    part_kind: Literal['tool-return'] = 'tool-return'
+    """Part type identifier, this is available on all parts as a discriminator."""
 
     def model_response_str(self) -> str:
         if isinstance(self.content, str):
@@ -73,11 +79,11 @@ class ToolReturn:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
 
-ErrorDetailsTa = _pydantic.LazyTypeAdapter(list[pydantic_core.ErrorDetails])
+error_details_ta = pydantic.TypeAdapter(list[pydantic_core.ErrorDetails], config=pydantic.ConfigDict(defer_build=True))
 
 
 @dataclass
-class RetryPrompt:
+class RetryPromptPart:
     """A message back to a model asking it to try again.
 
     This can be sent for a number of reasons:
@@ -98,37 +104,54 @@ class RetryPrompt:
     If the retry was triggered by a [`ValidationError`][pydantic_core.ValidationError], this will be a list of
     error details.
     """
+
     tool_name: str | None = None
     """The name of the tool that was called, if any."""
-    tool_id: str | None = None
-    """The tool identifier, if any."""
+
+    tool_call_id: str | None = None
+    """Optional tool call identifier, this is used by some models including OpenAI."""
+
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp, when the retry was triggered."""
-    role: Literal['retry-prompt'] = 'retry-prompt'
-    """Message type identifier, this type is available on all message as a discriminator."""
+
+    part_kind: Literal['retry-prompt'] = 'retry-prompt'
+    """Part type identifier, this is available on all parts as a discriminator."""
 
     def model_response(self) -> str:
         if isinstance(self.content, str):
             description = self.content
         else:
-            json_errors = ErrorDetailsTa.dump_json(self.content, exclude={'__all__': {'ctx'}}, indent=2)
+            json_errors = error_details_ta.dump_json(self.content, exclude={'__all__': {'ctx'}}, indent=2)
             description = f'{len(self.content)} validation errors: {json_errors.decode()}'
         return f'{description}\n\nFix the errors and try again.'
 
 
+ModelRequestPart = Annotated[
+    Union[SystemPromptPart, UserPromptPart, ToolReturnPart, RetryPromptPart], pydantic.Discriminator('part_kind')
+]
+"""A message part sent by PydanticAI to a model."""
+
+
+@dataclass
+class ModelRequest:
+    """A request generated by PydanticAI and sent to a model, e.g. a message from the PydanticAI app to the model."""
+
+    parts: list[ModelRequestPart]
+    """The parts of the user message."""
+
+    kind: Literal['request'] = 'request'
+    """Message type identifier, this is available on all parts as a discriminator."""
+
+
 @dataclass
-class ModelTextResponse:
+class TextPart:
     """A plain text response from a model."""
 
     content: str
     """The text content of the response."""
-    timestamp: datetime = field(default_factory=_now_utc)
-    """The timestamp of the response.
 
-    If the model provides a timestamp in the response (as OpenAI does) that will be used.
-    """
-    role: Literal['model-text-response'] = 'model-text-response'
-    """Message type identifier, this type is available on all message as a discriminator."""
+    part_kind: Literal['text'] = 'text'
+    """Part type identifier, this is available on all parts as a discriminator."""
 
 
 @dataclass
@@ -148,26 +171,53 @@ class ArgsDict:
 
 
 @dataclass
-class ToolCall:
-    """Either a tool call from the agent."""
+class ToolCallPart:
+    """A tool call from a model."""
 
     tool_name: str
     """The name of the tool to call."""
+
     args: ArgsJson | ArgsDict
     """The arguments to pass to the tool.
 
     Either as JSON or a Python dictionary depending on how data was returned.
     """
-    tool_id: str | None = None
-    """Optional tool identifier, this is used by some models including OpenAI."""
 
-    @classmethod
-    def from_json(cls, tool_name: str, args_json: str, tool_id: str | None = None) -> ToolCall:
-        return cls(tool_name, ArgsJson(args_json), tool_id)
+    tool_call_id: str | None = None
+    """Optional tool call identifier, this is used by some models including OpenAI."""
+
+    part_kind: Literal['tool-call'] = 'tool-call'
+    """Part type identifier, this is available on all parts as a discriminator."""
 
     @classmethod
-    def from_dict(cls, tool_name: str, args_dict: dict[str, Any]) -> ToolCall:
-        return cls(tool_name, ArgsDict(args_dict))
+    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)
+
+    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):
@@ -176,28 +226,39 @@ class ToolCall:
             return bool(self.args.args_json)
 
 
+ModelResponsePart = Annotated[Union[TextPart, ToolCallPart], pydantic.Discriminator('part_kind')]
+"""A message part returned by a model."""
+
+
 @dataclass
-class ModelStructuredResponse:
-    """A structured response from a model.
+class ModelResponse:
+    """A response from a model, e.g. a message from the model to the PydanticAI app."""
 
-    This is used either to call a tool or to return a structured response from an agent run.
-    """
+    parts: list[ModelResponsePart]
+    """The parts of the model message."""
 
-    calls: list[ToolCall]
-    """The tool calls being made."""
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp of the response.
 
     If the model provides a timestamp in the response (as OpenAI does) that will be used.
     """
-    role: Literal['model-structured-response'] = 'model-structured-response'
-    """Message type identifier, this type is available on all message as a discriminator."""
+
+    kind: Literal['response'] = 'response'
+    """Message type identifier, this is available on all parts as a discriminator."""
+
+    @classmethod
+    def from_text(cls, content: str, timestamp: datetime | None = None) -> Self:
+        return cls([TextPart(content)], timestamp=timestamp or _now_utc())
+
+    @classmethod
+    def from_tool_call(cls, tool_call: ToolCallPart) -> Self:
+        return cls([tool_call])
 
 
-ModelAnyResponse = Union[ModelTextResponse, ModelStructuredResponse]
-"""Any response from a model."""
-Message = Union[SystemPrompt, UserPrompt, ToolReturn, RetryPrompt, ModelTextResponse, ModelStructuredResponse]
+ModelMessage = Union[ModelRequest, ModelResponse]
 """Any message send to or returned by a model."""
 
-MessagesTypeAdapter = _pydantic.LazyTypeAdapter(list[Annotated[Message, pydantic.Field(discriminator='role')]])
+ModelMessagesTypeAdapter = pydantic.TypeAdapter(
+    list[Annotated[ModelMessage, pydantic.Discriminator('kind')]], config=pydantic.ConfigDict(defer_build=True)
+)
 """Pydantic [`TypeAdapter`][pydantic.type_adapter.TypeAdapter] for (de)serializing messages."""

pydantic_ai/models/__init__.py

@@ -16,10 +16,11 @@ from typing import TYPE_CHECKING, Literal, Union
 import httpx
 
 from ..exceptions import UserError
-from ..messages import Message, ModelAnyResponse, ModelStructuredResponse
+from ..messages import ModelMessage, ModelResponse
+from ..settings import ModelSettings
 
 if TYPE_CHECKING:
-    from ..result import Cost
+    from ..result import Usage
     from ..tools import ToolDefinition
 
 
@@ -30,7 +31,9 @@ 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',
     'groq:llama3-groq-70b-8192-tool-use-preview',
     'groq:llama3-groq-8b-8192-tool-use-preview',
@@ -47,8 +50,15 @@ KnownModelName = Literal[
     'groq:gemma-7b-it',
     'gemini-1.5-flash',
     'gemini-1.5-pro',
+    'gemini-2.0-flash-exp',
     'vertexai:gemini-1.5-flash',
     'vertexai:gemini-1.5-pro',
+    # since mistral models are supported by other providers (e.g. ollama), and some of their models (e.g. "codestral")
+    # don't start with "mistral", we add the "mistral:" prefix to all to be explicit
+    'mistral:mistral-small-latest',
+    'mistral:mistral-large-latest',
+    'mistral:codestral-latest',
+    'mistral:mistral-moderation-latest',
     'ollama:codellama',
     'ollama:gemma',
     'ollama:gemma2',
@@ -66,6 +76,9 @@ KnownModelName = Literal[
     'ollama:qwen2',
     'ollama:qwen2.5',
     'ollama:starcoder2',
+    'claude-3-5-haiku-latest',
+    'claude-3-5-sonnet-latest',
+    'claude-3-opus-latest',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -108,12 +121,16 @@ class AgentModel(ABC):
     """Model configured for each step of an Agent run."""
 
     @abstractmethod
-    async def request(self, messages: list[Message]) -> tuple[ModelAnyResponse, Cost]:
+    async def request(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> tuple[ModelResponse, Usage]:
         """Make a request to the model."""
         raise NotImplementedError()
 
     @asynccontextmanager
-    async def request_stream(self, messages: list[Message]) -> AsyncIterator[EitherStreamedResponse]:
+    async def request_stream(
+        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+    ) -> AsyncIterator[EitherStreamedResponse]:
         """Make a request to the model and return a streaming response."""
         raise NotImplementedError(f'Streamed requests not supported by this {self.__class__.__name__}')
         # yield is required to make this a generator for type checking
@@ -148,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()
 
@@ -178,10 +195,10 @@ class StreamStructuredResponse(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def get(self, *, final: bool = False) -> ModelStructuredResponse:
-        """Get the `ModelStructuredResponse` at this point.
+    def get(self, *, final: bool = False) -> ModelResponse:
+        """Get the `ModelResponse` at this point.
 
-        The `ModelStructuredResponse` may or may not be complete, depending on whether the stream is finished.
+        The `ModelResponse` may or may not be complete, depending on whether the stream is finished.
 
         Args:
             final: If True, this is the final call, after iteration is complete, the response should be fully validated.
@@ -189,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()
 
@@ -219,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:
@@ -270,10 +287,18 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .vertexai import VertexAIModel
 
         return VertexAIModel(model[9:])  # pyright: ignore[reportArgumentType]
+    elif model.startswith('mistral:'):
+        from .mistral import MistralModel
+
+        return MistralModel(model[8:])
     elif model.startswith('ollama:'):
         from .ollama import OllamaModel
 
         return OllamaModel(model[7:])
+    elif model.startswith('claude'):
+        from .anthropic import AnthropicModel
+
+        return AnthropicModel(model)
     else:
         raise UserError(f'Unknown model: {model}')