pydantic-ai-slim 0.0.22__py3-none-any.whl → 0.0.24__py3-none-any.whl

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,38 @@ class CohereModel(Model):
         else:
             self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)  # type: ignore
 
-    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, result.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 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."""
+        response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}), model_request_parameters)
+        return self._process_response(response), _map_usage(response)
 
-    client: AsyncClientV2
-    model_name: CohereModelName
-    allow_text_result: bool
-    tools: list[ToolV2]
+    @property
+    def model_name(self) -> CohereModelName:
+        """The model name."""
+        return self._model_name
 
-    async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, result.Usage]:
-        response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}))
-        return self._process_response(response), _map_usage(response)
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
 
     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 +183,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 +196,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 +208,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 +262,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:

pydantic_ai/models/function.py

@@ -4,7 +4,7 @@ import inspect
 import re
 from collections.abc import AsyncIterator, Awaitable, Iterable
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, field
 from datetime import datetime
 from itertools import chain
 from typing import Callable, Union
@@ -27,7 +27,7 @@ from ..messages import (
 )
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
-from . import AgentModel, Model, StreamedResponse
+from . import Model, ModelRequestParameters, StreamedResponse
 
 
 @dataclass(init=False)
@@ -40,6 +40,9 @@ class FunctionModel(Model):
     function: FunctionDef | None = None
     stream_function: StreamFunctionDef | None = None
 
+    _model_name: str = field(repr=False)
+    _system: str | None = field(default=None, repr=False)
+
     @overload
     def __init__(self, function: FunctionDef) -> None: ...
 
@@ -63,23 +66,70 @@ class FunctionModel(Model):
         self.function = function
         self.stream_function = stream_function
 
-    async def agent_model(
+        function_name = self.function.__name__ if self.function is not None else ''
+        stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
+        self._model_name = f'function:{function_name}:{stream_function_name}'
+
+    async def request(
         self,
-        *,
-        function_tools: list[ToolDefinition],
-        allow_text_result: bool,
-        result_tools: list[ToolDefinition],
-    ) -> AgentModel:
-        return FunctionAgentModel(
-            self.function,
-            self.stream_function,
-            AgentInfo(function_tools, allow_text_result, result_tools, None),
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, usage.Usage]:
+        agent_info = AgentInfo(
+            model_request_parameters.function_tools,
+            model_request_parameters.allow_text_result,
+            model_request_parameters.result_tools,
+            model_settings,
         )
 
-    def name(self) -> str:
-        function_name = self.function.__name__ if self.function is not None else ''
-        stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
-        return f'function:{function_name}:{stream_function_name}'
+        assert self.function is not None, 'FunctionModel must receive a `function` to support non-streamed requests'
+
+        if inspect.iscoroutinefunction(self.function):
+            response = await self.function(messages, agent_info)
+        else:
+            response_ = await _utils.run_in_executor(self.function, messages, agent_info)
+            assert isinstance(response_, ModelResponse), response_
+            response = response_
+        response.model_name = f'function:{self.function.__name__}'
+        # TODO is `messages` right here? Should it just be new messages?
+        return response, _estimate_usage(chain(messages, [response]))
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        agent_info = AgentInfo(
+            model_request_parameters.function_tools,
+            model_request_parameters.allow_text_result,
+            model_request_parameters.result_tools,
+            model_settings,
+        )
+
+        assert (
+            self.stream_function is not None
+        ), 'FunctionModel must receive a `stream_function` to support streamed requests'
+
+        response_stream = PeekableAsyncStream(self.stream_function(messages, agent_info))
+
+        first = await response_stream.peek()
+        if isinstance(first, _utils.Unset):
+            raise ValueError('Stream function must return at least one item')
+
+        yield FunctionStreamedResponse(_model_name=f'function:{self.stream_function.__name__}', _iter=response_stream)
+
+    @property
+    def model_name(self) -> str:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
 
 
 @dataclass(frozen=True)
@@ -119,9 +169,11 @@ class DeltaToolCall:
 DeltaToolCalls: TypeAlias = dict[int, DeltaToolCall]
 """A mapping of tool call IDs to incremental changes."""
 
+# TODO: Change the signature to Callable[[list[ModelMessage], ModelSettings, ModelRequestParameters], ...]
 FunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], Union[ModelResponse, Awaitable[ModelResponse]]]
 """A function used to generate a non-streamed response."""
 
+# TODO: Change signature as indicated above
 StreamFunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], AsyncIterator[Union[str, DeltaToolCalls]]]
 """A function used to generate a streamed response.
 
@@ -132,54 +184,11 @@ E.g. you need to yield all text or all `DeltaToolCalls`, not mix them.
 """
 
 
-@dataclass
-class FunctionAgentModel(AgentModel):
-    """Implementation of `AgentModel` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""
-
-    function: FunctionDef | None
-    stream_function: StreamFunctionDef | None
-    agent_info: AgentInfo
-
-    async def request(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> tuple[ModelResponse, usage.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'
-        model_name = f'function:{self.function.__name__}'
-
-        if inspect.iscoroutinefunction(self.function):
-            response = await self.function(messages, agent_info)
-        else:
-            response_ = await _utils.run_in_executor(self.function, messages, agent_info)
-            assert isinstance(response_, ModelResponse), response_
-            response = response_
-        response.model_name = model_name
-        # TODO is `messages` right here? Should it just be new messages?
-        return response, _estimate_usage(chain(messages, [response]))
-
-    @asynccontextmanager
-    async def request_stream(
-        self, messages: list[ModelMessage], model_settings: ModelSettings | None
-    ) -> AsyncIterator[StreamedResponse]:
-        assert (
-            self.stream_function is not None
-        ), 'FunctionModel must receive a `stream_function` to support streamed requests'
-        model_name = f'function:{self.stream_function.__name__}'
-
-        response_stream = PeekableAsyncStream(self.stream_function(messages, self.agent_info))
-
-        first = await response_stream.peek()
-        if isinstance(first, _utils.Unset):
-            raise ValueError('Stream function must return at least one item')
-
-        yield FunctionStreamedResponse(_model_name=model_name, _iter=response_stream)
-
-
 @dataclass
 class FunctionStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""
 
+    _model_name: str
     _iter: AsyncIterator[str | DeltaToolCalls]
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
@@ -207,7 +216,14 @@ class FunctionStreamedResponse(StreamedResponse):
                     if maybe_event is not None:
                         yield maybe_event
 
+    @property
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
+
+    @property
     def timestamp(self) -> datetime:
+        """Get the timestamp of the response."""
         return self._timestamp