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

pydantic_ai/messages.py

@@ -1,20 +1,18 @@
 from __future__ import annotations as _annotations
 
-import json
 from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Annotated, Any, Literal, Union
 
 import pydantic
 import pydantic_core
-from pydantic import TypeAdapter
+from typing_extensions import Self
 
-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.
@@ -22,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],
@@ -36,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):
@@ -74,8 +79,11 @@ class ToolReturn:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
 
+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:
@@ -96,36 +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:
-            description = f'{len(self.content)} validation errors: {json.dumps(self.content, 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 ModelTextResponse:
+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 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
@@ -145,26 +171,31 @@ 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."""
+
+    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_json(cls, tool_name: str, args_json: str, tool_id: str | None = None) -> ToolCall:
-        return cls(tool_name, ArgsJson(args_json), tool_id)
+    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)
 
     @classmethod
-    def from_dict(cls, tool_name: str, args_dict: dict[str, Any]) -> ToolCall:
-        return cls(tool_name, ArgsDict(args_dict))
+    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 has_content(self) -> bool:
         if isinstance(self.args, ArgsDict):
@@ -173,28 +204,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

@@ -7,20 +7,21 @@ specific LLM being used.
 from __future__ import annotations as _annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Iterable, Iterator, Mapping, Sequence
+from collections.abc import AsyncIterator, Iterable, Iterator
 from contextlib import asynccontextmanager, contextmanager
 from datetime import datetime
 from functools import cache
-from typing import TYPE_CHECKING, Literal, Protocol, Union
+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 .._utils import ObjectJsonSchema
     from ..result import Cost
+    from ..tools import ToolDefinition
 
 
 KnownModelName = Literal[
@@ -31,6 +32,7 @@ KnownModelName = Literal[
     'openai:o1-preview',
     'openai:o1-mini',
     '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 +49,35 @@ 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',
+    'ollama:llama3',
+    'ollama:llama3.1',
+    'ollama:llama3.2',
+    'ollama:llama3.2-vision',
+    'ollama:llama3.3',
+    'ollama:mistral',
+    'ollama:mistral-nemo',
+    'ollama:mixtral',
+    'ollama:phi3',
+    'ollama:qwq',
+    'ollama:qwen',
+    '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].
@@ -63,11 +92,12 @@ class Model(ABC):
     @abstractmethod
     async def agent_model(
         self,
-        function_tools: Mapping[str, AbstractToolDefinition],
+        *,
+        function_tools: list[ToolDefinition],
         allow_text_result: bool,
-        result_tools: Sequence[AbstractToolDefinition] | None,
+        result_tools: list[ToolDefinition],
     ) -> AgentModel:
-        """Create an agent model.
+        """Create an agent model, this is called for each step of an agent run.
 
         This is async in case slow/async config checks need to be performed that can't be done in `__init__`.
 
@@ -87,15 +117,19 @@ class Model(ABC):
 
 
 class AgentModel(ABC):
-    """Model configured for a specific agent."""
+    """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, Cost]:
         """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
@@ -160,10 +194,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.
@@ -238,7 +272,7 @@ def infer_model(model: Model | KnownModelName) -> Model:
     elif model.startswith('openai:'):
         from .openai import OpenAIModel
 
-        return OpenAIModel(model[7:])  # pyright: ignore[reportArgumentType]
+        return OpenAIModel(model[7:])
     elif model.startswith('gemini'):
         from .gemini import GeminiModel
 
@@ -252,37 +286,20 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .vertexai import VertexAIModel
 
         return VertexAIModel(model[9:])  # pyright: ignore[reportArgumentType]
-    else:
-        raise UserError(f'Unknown model: {model}')
-
-
-class AbstractToolDefinition(Protocol):
-    """Abstract definition of a function/tool.
-
-    This is used for both function tools and result tools.
-    """
+    elif model.startswith('mistral:'):
+        from .mistral import MistralModel
 
-    @property
-    def name(self) -> str:
-        """The name of the tool."""
-        ...
-
-    @property
-    def description(self) -> str:
-        """The description of the tool."""
-        ...
-
-    @property
-    def json_schema(self) -> ObjectJsonSchema:
-        """The JSON schema for the tool's arguments."""
-        ...
+        return MistralModel(model[8:])
+    elif model.startswith('ollama:'):
+        from .ollama import OllamaModel
 
-    @property
-    def outer_typed_dict_key(self) -> str | None:
-        """The key in the outer [TypedDict] that wraps a result tool.
+        return OllamaModel(model[7:])
+    elif model.startswith('claude'):
+        from .anthropic import AnthropicModel
 
-        This will only be set for result tools which don't have an `object` JSON schema.
-        """
+        return AnthropicModel(model)
+    else:
+        raise UserError(f'Unknown model: {model}')
 
 
 @cache