pydantic-ai-slim 0.0.21__py3-none-any.whl → 0.0.23__py3-none-any.whl

pydantic_ai/models/__init__.py

@@ -52,11 +52,15 @@ KnownModelName = Literal[
     'google-gla:gemini-1.5-flash-8b',
     'google-gla:gemini-1.5-pro',
     'google-gla:gemini-2.0-flash-exp',
+    'google-gla:gemini-2.0-flash-thinking-exp-01-21',
+    'google-gla:gemini-exp-1206',
     'google-vertex:gemini-1.0-pro',
     'google-vertex:gemini-1.5-flash',
     'google-vertex:gemini-1.5-flash-8b',
     'google-vertex:gemini-1.5-pro',
     'google-vertex:gemini-2.0-flash-exp',
+    'google-vertex:gemini-2.0-flash-thinking-exp-01-21',
+    'google-vertex:gemini-exp-1206',
     'gpt-3.5-turbo',
     'gpt-3.5-turbo-0125',
     'gpt-3.5-turbo-0301',
@@ -108,6 +112,8 @@ KnownModelName = Literal[
     'o1-mini-2024-09-12',
     'o1-preview',
     'o1-preview-2024-09-12',
+    'o3-mini',
+    'o3-mini-2025-01-31',
     'openai:chatgpt-4o-latest',
     'openai:gpt-3.5-turbo',
     'openai:gpt-3.5-turbo-0125',
@@ -145,6 +151,8 @@ KnownModelName = Literal[
     'openai:o1-mini-2024-09-12',
     'openai:o1-preview',
     'openai:o1-preview-2024-09-12',
+    'openai:o3-mini',
+    'openai:o3-mini-2025-01-31',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -153,49 +161,37 @@ KnownModelName = Literal[
 """
 
 
-class Model(ABC):
-    """Abstract class for a model."""
+@dataclass
+class ModelRequestParameters:
+    """Configuration for an agent's request to a model, specifically related to tools and result handling."""
 
-    @abstractmethod
-    async def agent_model(
-        self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> AgentModel:
-        """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__`.
-
-        Args:
-            function_tools: The tools available to the agent.
-            allow_text_result: Whether a plain text final response/result is permitted.
-            result_tools: Tool definitions for the final result tool(s), if any.
-
-        Returns:
-            An agent model.
-        """
-        raise NotImplementedError()
+    function_tools: list[ToolDefinition]
+    allow_text_result: bool
+    result_tools: list[ToolDefinition]
 
-    @abstractmethod
-    def name(self) -> str:
-        raise NotImplementedError()
 
+class Model(ABC):
+    """Abstract class for a model."""
 
-class AgentModel(ABC):
-    """Model configured for each step of an Agent run."""
+    _model_name: str
+    _system: str | None
 
     @abstractmethod
     async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> tuple[ModelResponse, Usage]:
         """Make a request to the model."""
         raise NotImplementedError()
 
     @asynccontextmanager
     async def request_stream(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> AsyncIterator[StreamedResponse]:
         """Make a request to the model and return a streaming response."""
         # This method is not required, but you need to implement it if you want to support streamed responses
@@ -204,6 +200,16 @@ class AgentModel(ABC):
         # noinspection PyUnreachableCode
         yield  # pragma: no cover
 
+    @property
+    def model_name(self) -> str:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider, ex: openai."""
+        return self._system
+
 
 @dataclass
 class StreamedResponse(ABC):
@@ -266,7 +272,7 @@ def check_allow_model_requests() -> None:
     """Check if model requests are allowed.
 
     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].
+    [`Model.request`][pydantic_ai.models.Model.request] and [`Model.request_stream`][pydantic_ai.models.Model.request_stream].
 
     Raises:
         RuntimeError: If model requests are not allowed.
@@ -307,33 +313,33 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .openai import OpenAIModel
 
         return OpenAIModel(model[7:])
-    elif model.startswith(('gpt', 'o1')):
+    elif model.startswith(('gpt', 'o1', 'o3')):
         from .openai import OpenAIModel
 
         return OpenAIModel(model)
     elif model.startswith('google-gla'):
         from .gemini import GeminiModel
 
-        return GeminiModel(model[11:])  # pyright: ignore[reportArgumentType]
+        return GeminiModel(model[11:])
     # backwards compatibility with old model names (ex, gemini-1.5-flash -> google-gla:gemini-1.5-flash)
     elif model.startswith('gemini'):
         from .gemini import GeminiModel
 
         # noinspection PyTypeChecker
-        return GeminiModel(model)  # pyright: ignore[reportArgumentType]
+        return GeminiModel(model)
     elif model.startswith('groq:'):
         from .groq import GroqModel
 
-        return GroqModel(model[5:])  # pyright: ignore[reportArgumentType]
+        return GroqModel(model[5:])
     elif model.startswith('google-vertex'):
         from .vertexai import VertexAIModel
 
-        return VertexAIModel(model[14:])  # pyright: ignore[reportArgumentType]
+        return VertexAIModel(model[14:])
     # backwards compatibility with old model names (ex, vertexai:gemini-1.5-flash -> google-vertex:gemini-1.5-flash)
     elif model.startswith('vertexai:'):
         from .vertexai import VertexAIModel
 
-        return VertexAIModel(model[9:])  # pyright: ignore[reportArgumentType]
+        return VertexAIModel(model[9:])
     elif model.startswith('mistral:'):
         from .mistral import MistralModel
 

pydantic_ai/models/anthropic.py

@@ -28,8 +28,8 @@ from ..messages import (
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
-    AgentModel,
     Model,
+    ModelRequestParameters,
     StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
@@ -68,14 +68,14 @@ LatestAnthropicModelNames = Literal[
     'claude-3-5-sonnet-latest',
     'claude-3-opus-latest',
 ]
-"""Latest named Anthropic models."""
+"""Latest Anthropic models."""
 
 AnthropicModelName = Union[str, LatestAnthropicModelNames]
 """Possible Anthropic model names.
 
 Since Anthropic supports a variety of date-stamped models, we explicitly list the latest models but
 allow any name in the type hints.
-Since [the Anthropic docs](https://docs.anthropic.com/en/docs/about-claude/models) for a full list.
+See [the Anthropic docs](https://docs.anthropic.com/en/docs/about-claude/models) for a full list.
 """
 
 
@@ -101,9 +101,11 @@ class AnthropicModel(Model):
         We anticipate adding support for streaming responses in a near-term future release.
     """
 
-    model_name: AnthropicModelName
     client: AsyncAnthropic = field(repr=False)
 
+    _model_name: AnthropicModelName = field(repr=False)
+    _system: str | None = field(default='anthropic', repr=False)
+
     def __init__(
         self,
         model_name: AnthropicModelName,
@@ -124,7 +126,7 @@ class AnthropicModel(Model):
                 client to use, if provided, `api_key` and `http_client` must be `None`.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
         """
-        self.model_name = model_name
+        self._model_name = model_name
         if anthropic_client is not None:
             assert http_client is None, 'Cannot provide both `anthropic_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `anthropic_client` and `api_key`'
@@ -134,81 +136,67 @@ class AnthropicModel(Model):
         else:
             self.client = AsyncAnthropic(api_key=api_key, http_client=cached_async_http_client())
 
-    async def agent_model(
+    async def request(
         self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> AgentModel:
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, usage.Usage]:
         check_allow_model_requests()
-        tools = [self._map_tool_definition(r) for r in function_tools]
-        if result_tools:
-            tools += [self._map_tool_definition(r) for r in result_tools]
-        return AnthropicAgentModel(
-            self.client,
-            self.model_name,
-            allow_text_result,
-            tools,
+        response = await self._messages_create(
+            messages, False, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
         )
-
-    def name(self) -> str:
-        return f'anthropic:{self.model_name}'
-
-    @staticmethod
-    def _map_tool_definition(f: ToolDefinition) -> ToolParam:
-        return {
-            'name': f.name,
-            'description': f.description,
-            'input_schema': f.parameters_json_schema,
-        }
-
-
-@dataclass
-class AnthropicAgentModel(AgentModel):
-    """Implementation of `AgentModel` for Anthropic models."""
-
-    client: AsyncAnthropic
-    model_name: AnthropicModelName
-    allow_text_result: bool
-    tools: list[ToolParam]
-
-    async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, usage.Usage]:
-        response = await self._messages_create(messages, False, cast(AnthropicModelSettings, model_settings or {}))
         return self._process_response(response), _map_usage(response)
 
     @asynccontextmanager
     async def request_stream(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> AsyncIterator[StreamedResponse]:
-        response = await self._messages_create(messages, True, cast(AnthropicModelSettings, model_settings or {}))
+        check_allow_model_requests()
+        response = await self._messages_create(
+            messages, True, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
+        )
         async with response:
             yield await self._process_streamed_response(response)
 
     @overload
     async def _messages_create(
-        self, messages: list[ModelMessage], stream: Literal[True], model_settings: AnthropicModelSettings
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[True],
+        model_settings: AnthropicModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> AsyncStream[RawMessageStreamEvent]:
         pass
 
     @overload
     async def _messages_create(
-        self, messages: list[ModelMessage], stream: Literal[False], model_settings: AnthropicModelSettings
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[False],
+        model_settings: AnthropicModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> AnthropicMessage:
         pass
 
     async def _messages_create(
-        self, messages: list[ModelMessage], stream: bool, model_settings: AnthropicModelSettings
+        self,
+        messages: list[ModelMessage],
+        stream: bool,
+        model_settings: AnthropicModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> AnthropicMessage | AsyncStream[RawMessageStreamEvent]:
         # standalone function to make it easier to override
+        tools = self._get_tools(model_request_parameters)
         tool_choice: ToolChoiceParam | None
 
-        if not self.tools:
+        if not tools:
             tool_choice = None
         else:
-            if not self.allow_text_result:
+            if not model_request_parameters.allow_text_result:
                 tool_choice = {'type': 'any'}
             else:
                 tool_choice = {'type': 'auto'}
@@ -222,8 +210,8 @@ class AnthropicAgentModel(AgentModel):
             max_tokens=model_settings.get('max_tokens', 1024),
             system=system_prompt or NOT_GIVEN,
             messages=anthropic_messages,
-            model=self.model_name,
-            tools=self.tools or NOT_GIVEN,
+            model=self._model_name,
+            tools=tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
             temperature=model_settings.get('temperature', NOT_GIVEN),
@@ -248,7 +236,7 @@ class AnthropicAgentModel(AgentModel):
                     )
                 )
 
-        return ModelResponse(items, model_name=self.model_name)
+        return ModelResponse(items, model_name=self._model_name)
 
     async def _process_streamed_response(self, response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse:
         peekable_response = _utils.PeekableAsyncStream(response)
@@ -258,10 +246,17 @@ class AnthropicAgentModel(AgentModel):
 
         # Since Anthropic doesn't provide a timestamp in the message, we'll use the current time
         timestamp = datetime.now(tz=timezone.utc)
-        return AnthropicStreamedResponse(_model_name=self.model_name, _response=peekable_response, _timestamp=timestamp)
+        return AnthropicStreamedResponse(
+            _model_name=self._model_name, _response=peekable_response, _timestamp=timestamp
+        )
 
-    @staticmethod
-    def _map_message(messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]:
+    def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolParam]:
+        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
+        if model_request_parameters.result_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        return tools
+
+    def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]:
         """Just maps a `pydantic_ai.Message` to a `anthropic.types.MessageParam`."""
         system_prompt: str = ''
         anthropic_messages: list[MessageParam] = []
@@ -310,20 +305,28 @@ class AnthropicAgentModel(AgentModel):
                         content.append(TextBlockParam(text=item.content, type='text'))
                     else:
                         assert isinstance(item, ToolCallPart)
-                        content.append(_map_tool_call(item))
+                        content.append(self._map_tool_call(item))
                 anthropic_messages.append(MessageParam(role='assistant', content=content))
             else:
                 assert_never(m)
         return system_prompt, anthropic_messages
 
+    @staticmethod
+    def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
+        return ToolUseBlockParam(
+            id=_guard_tool_call_id(t=t, model_source='Anthropic'),
+            type='tool_use',
+            name=t.tool_name,
+            input=t.args_as_dict(),
+        )
 
-def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
-    return ToolUseBlockParam(
-        id=_guard_tool_call_id(t=t, model_source='Anthropic'),
-        type='tool_use',
-        name=t.tool_name,
-        input=t.args_as_dict(),
-    )
+    @staticmethod
+    def _map_tool_definition(f: ToolDefinition) -> ToolParam:
+        return {
+            'name': f.name,
+            'description': f.description,
+            'input_schema': f.parameters_json_schema,
+        }
 
 
 def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage:

pydantic_ai/models/cohere.py

@@ -26,8 +26,8 @@ from ..messages import (
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
-    AgentModel,
     Model,
+    ModelRequestParameters,
     check_allow_model_requests,
 )
 
@@ -52,7 +52,7 @@ except ImportError as _import_error:
         "you can use the `cohere` optional group — `pip install 'pydantic-ai-slim[cohere]'`"
     ) from _import_error
 
-NamedCohereModels = Literal[
+LatestCohereModelNames = Literal[
     'c4ai-aya-expanse-32b',
     'c4ai-aya-expanse-8b',
     'command',
@@ -67,9 +67,15 @@ NamedCohereModels = Literal[
     'command-r-plus-08-2024',
     'command-r7b-12-2024',
 ]
-"""Latest / most popular named Cohere models."""
+"""Latest Cohere models."""
 
-CohereModelName = Union[NamedCohereModels, str]
+CohereModelName = Union[str, LatestCohereModelNames]
+"""Possible Cohere model names.
+
+Since Cohere supports a variety of date-stamped models, we explicitly list the latest models but
+allow any name in the type hints.
+See [Cohere's docs](https://docs.cohere.com/v2/docs/models) for a list of all available models.
+"""
 
 
 class CohereModelSettings(ModelSettings):
@@ -88,9 +94,11 @@ class CohereModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    model_name: CohereModelName
     client: AsyncClientV2 = field(repr=False)
 
+    _model_name: CohereModelName = field(repr=False)
+    _system: str | None = field(default='cohere', repr=False)
+
     def __init__(
         self,
         model_name: CohereModelName,
@@ -110,7 +118,7 @@ class CohereModel(Model):
                 `api_key` and `http_client` must be `None`.
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
         """
-        self.model_name: CohereModelName = model_name
+        self._model_name: CohereModelName = model_name
         if cohere_client is not None:
             assert http_client is None, 'Cannot provide both `cohere_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `cohere_client` and `api_key`'
@@ -118,64 +126,28 @@ class CohereModel(Model):
         else:
             self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)  # type: ignore
 
-    async def agent_model(
-        self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> AgentModel:
-        check_allow_model_requests()
-        tools = [self._map_tool_definition(r) for r in function_tools]
-        if result_tools:
-            tools += [self._map_tool_definition(r) for r in result_tools]
-        return CohereAgentModel(
-            self.client,
-            self.model_name,
-            allow_text_result,
-            tools,
-        )
-
-    def name(self) -> str:
-        return f'cohere:{self.model_name}'
-
-    @staticmethod
-    def _map_tool_definition(f: ToolDefinition) -> ToolV2:
-        return ToolV2(
-            type='function',
-            function=ToolV2Function(
-                name=f.name,
-                description=f.description,
-                parameters=f.parameters_json_schema,
-            ),
-        )
-
-
-@dataclass
-class CohereAgentModel(AgentModel):
-    """Implementation of `AgentModel` for Cohere models."""
-
-    client: AsyncClientV2
-    model_name: CohereModelName
-    allow_text_result: bool
-    tools: list[ToolV2]
-
     async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> tuple[ModelResponse, result.Usage]:
-        response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}))
+        check_allow_model_requests()
+        response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}), model_request_parameters)
         return self._process_response(response), _map_usage(response)
 
     async def _chat(
         self,
         messages: list[ModelMessage],
         model_settings: CohereModelSettings,
+        model_request_parameters: ModelRequestParameters,
     ) -> ChatResponse:
+        tools = self._get_tools(model_request_parameters)
         cohere_messages = list(chain(*(self._map_message(m) for m in messages)))
         return await self.client.chat(
-            model=self.model_name,
+            model=self._model_name,
             messages=cohere_messages,
-            tools=self.tools or OMIT,
+            tools=tools or OMIT,
             max_tokens=model_settings.get('max_tokens', OMIT),
             temperature=model_settings.get('temperature', OMIT),
             p=model_settings.get('top_p', OMIT),
@@ -201,13 +173,12 @@ class CohereAgentModel(AgentModel):
                         tool_call_id=c.id,
                     )
                 )
-        return ModelResponse(parts=parts, model_name=self.model_name)
+        return ModelResponse(parts=parts, model_name=self._model_name)
 
-    @classmethod
-    def _map_message(cls, message: ModelMessage) -> Iterable[ChatMessageV2]:
+    def _map_message(self, message: ModelMessage) -> Iterable[ChatMessageV2]:
         """Just maps a `pydantic_ai.Message` to a `cohere.ChatMessageV2`."""
         if isinstance(message, ModelRequest):
-            yield from cls._map_user_message(message)
+            yield from self._map_user_message(message)
         elif isinstance(message, ModelResponse):
             texts: list[str] = []
             tool_calls: list[ToolCallV2] = []
@@ -215,7 +186,7 @@ class CohereAgentModel(AgentModel):
                 if isinstance(item, TextPart):
                     texts.append(item.content)
                 elif isinstance(item, ToolCallPart):
-                    tool_calls.append(_map_tool_call(item))
+                    tool_calls.append(self._map_tool_call(item))
                 else:
                     assert_never(item)
             message_param = AssistantChatMessageV2(role='assistant')
@@ -227,6 +198,34 @@ class CohereAgentModel(AgentModel):
         else:
             assert_never(message)
 
+    def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolV2]:
+        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
+        if model_request_parameters.result_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        return tools
+
+    @staticmethod
+    def _map_tool_call(t: ToolCallPart) -> ToolCallV2:
+        return ToolCallV2(
+            id=_guard_tool_call_id(t=t, model_source='Cohere'),
+            type='function',
+            function=ToolCallV2Function(
+                name=t.tool_name,
+                arguments=t.args_as_json_str(),
+            ),
+        )
+
+    @staticmethod
+    def _map_tool_definition(f: ToolDefinition) -> ToolV2:
+        return ToolV2(
+            type='function',
+            function=ToolV2Function(
+                name=f.name,
+                description=f.description,
+                parameters=f.parameters_json_schema,
+            ),
+        )
+
     @classmethod
     def _map_user_message(cls, message: ModelRequest) -> Iterable[ChatMessageV2]:
         for part in message.parts:
@@ -253,17 +252,6 @@ class CohereAgentModel(AgentModel):
                 assert_never(part)
 
 
-def _map_tool_call(t: ToolCallPart) -> ToolCallV2:
-    return ToolCallV2(
-        id=_guard_tool_call_id(t=t, model_source='Cohere'),
-        type='function',
-        function=ToolCallV2Function(
-            name=t.tool_name,
-            arguments=t.args_as_json_str(),
-        ),
-    )
-
-
 def _map_usage(response: ChatResponse) -> result.Usage:
     usage = response.usage
     if usage is None: