Remove the AgentModel class (#800)

Author: dmontagu

docs/api/models/base.md

@@ -4,8 +4,8 @@
     options:
       members:
         - KnownModelName
+        - ModelRequestParameters
         - Model
-        - AgentModel
         - AbstractToolDefinition
         - StreamedResponse
         - ALLOW_MODEL_REQUESTS

docs/api/models/test.md

@@ -17,7 +17,7 @@ async def test_my_agent():
     with my_agent.override(model=m):
         result = await my_agent.run('Testing my agent...')
         assert result.data == 'success (no tool calls)'
-    assert m.agent_model_function_tools == []
+    assert m.last_model_request_parameters.function_tools == []
 ```
 
 See [Unit testing with `TestModel`](../../testing-evals.md#unit-testing-with-testmodel) for detailed documentation.

docs/api/models/vertexai.md

@@ -2,8 +2,8 @@
 
 Custom interface to the `*-aiplatform.googleapis.com` API for Gemini models.
 
-This model uses [`GeminiAgentModel`][pydantic_ai.models.gemini.GeminiAgentModel] with just the URL and auth method
-changed from [`GeminiModel`][pydantic_ai.models.gemini.GeminiModel], it relies on the VertexAI
+This model inherits from [`GeminiModel`][pydantic_ai.models.gemini.GeminiModel] with just the URL and auth method
+changed, it relies on the VertexAI
 [`generateContent`](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/generateContent)
 and
 [`streamGenerateContent`](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.endpoints/streamGenerateContent)

docs/models.md

@@ -628,9 +628,8 @@ agent = Agent(model)
 
 To implement support for models not already supported, you will need to subclass the [`Model`][pydantic_ai.models.Model] abstract base class.
 
-This in turn will require you to implement the following other abstract base classes:
+For streaming, you'll also need to implement the following abstract base class:
 
-* [`AgentModel`][pydantic_ai.models.AgentModel]
 * [`StreamedResponse`][pydantic_ai.models.StreamedResponse]
 
 The best place to start is to review the source code for existing implementations, e.g. [`OpenAIModel`](https://github.com/pydantic/pydantic-ai/blob/main/pydantic_ai_slim/pydantic_ai/models/openai.py).

docs/tools.md

@@ -292,7 +292,7 @@ The return type of tool can be anything which Pydantic can serialize to JSON as
 
 If a tool has a single parameter that can be represented as an object in JSON schema (e.g. dataclass, TypedDict, pydantic model), the schema for the tool is simplified to be just that object.
 
-Here's an example, we use [`TestModel.agent_model_function_tools`][pydantic_ai.models.test.TestModel.agent_model_function_tools] to inspect the tool schema that would be passed to the model.
+Here's an example where we use [`TestModel.last_model_request_parameters`][pydantic_ai.models.test.TestModel.last_model_request_parameters] to inspect the tool schema that would be passed to the model.
 
 ```python {title="single_parameter_tool.py"}
 from pydantic import BaseModel
@@ -320,7 +320,7 @@ test_model = TestModel()
 result = agent.run_sync('hello', model=test_model)
 print(result.data)
 #> {"foobar":"x=0 y='a' z=3.14"}
-print(test_model.agent_model_function_tools)
+print(test_model.last_model_request_parameters.function_tools)
 """
 [
     ToolDefinition(
@@ -425,7 +425,7 @@ agent = Agent(test_model, tools=[greet_tool], deps_type=Literal['human', 'machin
 result = agent.run_sync('testing...', deps='human')
 print(result.data)
 #> {"greet":"hello a"}
-print(test_model.agent_model_function_tools)
+print(test_model.last_model_request_parameters.function_tools)
 """
 [
     ToolDefinition(

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -204,9 +204,9 @@ async def run(
         return StreamModelRequestNode[DepsT, NodeRunEndT](request=await self._get_first_message(ctx))
 
 
-async def _prepare_model(
+async def _prepare_request_parameters(
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
-) -> models.AgentModel:
+) -> models.ModelRequestParameters:
     """Build tools and create an agent model."""
     function_tool_defs: list[ToolDefinition] = []
 
@@ -220,7 +220,7 @@ async def add_tool(tool: Tool[DepsT]) -> None:
     await asyncio.gather(*map(add_tool, ctx.deps.function_tools.values()))
 
     result_schema = ctx.deps.result_schema
-    return await run_context.model.agent_model(
+    return models.ModelRequestParameters(
         function_tools=function_tool_defs,
         allow_text_result=_allow_text_result(result_schema),
         result_tools=result_schema.tool_defs() if result_schema is not None else [],
@@ -245,13 +245,15 @@ async def run(
         # Increment run_step
         ctx.state.run_step += 1
 
-        with _logfire.span('preparing model and tools {run_step=}', run_step=ctx.state.run_step):
-            agent_model = await _prepare_model(ctx)
+        with _logfire.span('preparing model request params {run_step=}', run_step=ctx.state.run_step):
+            model_request_parameters = await _prepare_request_parameters(ctx)
 
         # Actually make the model request
         model_settings = merge_model_settings(ctx.deps.model_settings, None)
         with _logfire.span('model request') as span:
-            model_response, request_usage = await agent_model.request(ctx.state.message_history, model_settings)
+            model_response, request_usage = await ctx.deps.model.request(
+                ctx.state.message_history, model_settings, model_request_parameters
+            )
             span.set_attribute('response', model_response)
             span.set_attribute('usage', request_usage)
 
@@ -405,12 +407,14 @@ async def run_to_result(
         ctx.state.run_step += 1
 
         with _logfire.span('preparing model and tools {run_step=}', run_step=ctx.state.run_step):
-            agent_model = await _prepare_model(ctx)
+            model_request_parameters = await _prepare_request_parameters(ctx)
 
         # Actually make the model request
         model_settings = merge_model_settings(ctx.deps.model_settings, None)
         with _logfire.span('model request {run_step=}', run_step=ctx.state.run_step) as model_req_span:
-            async with agent_model.request_stream(ctx.state.message_history, model_settings) as streamed_response:
+            async with ctx.deps.model.request_stream(
+                ctx.state.message_history, model_settings, model_request_parameters
+            ) as streamed_response:
                 ctx.state.usage.requests += 1
                 model_req_span.set_attribute('response_type', streamed_response.__class__.__name__)
                 # We want to end the "model request" span here, but we can't exit the context manager

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -161,49 +161,38 @@
 """
 
 
+@dataclass
+class ModelRequestParameters:
+    """Configuration for an agent's request to a model, specifically related to tools and result handling."""
+
+    function_tools: list[ToolDefinition]
+    allow_text_result: bool
+    result_tools: list[ToolDefinition]
+
+
 class Model(ABC):
     """Abstract class for a model."""
 
-    @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()
-
     @abstractmethod
     def name(self) -> str:
         raise NotImplementedError()
 
-
-class AgentModel(ABC):
-    """Model configured for each step of an Agent run."""
-
     @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
@@ -274,7 +263,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.

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -28,8 +28,8 @@
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
-    AgentModel,
     Model,
+    ModelRequestParameters,
     StreamedResponse,
     cached_async_http_client,
     check_allow_model_requests,
@@ -134,81 +134,70 @@ def __init__(
         else:
             self.client = AsyncAnthropic(api_key=api_key, http_client=cached_async_http_client())
 
-    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 AnthropicAgentModel(
-            self.client,
-            self.model_name,
-            allow_text_result,
-            tools,
-        )
-
     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
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
     ) -> tuple[ModelResponse, usage.Usage]:
-        response = await self._messages_create(messages, False, cast(AnthropicModelSettings, model_settings or {}))
+        check_allow_model_requests()
+        response = await self._messages_create(
+            messages, False, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
+        )
         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'}
@@ -223,7 +212,7 @@ async def _messages_create(
             system=system_prompt or NOT_GIVEN,
             messages=anthropic_messages,
             model=self.model_name,
-            tools=self.tools or NOT_GIVEN,
+            tools=tools or NOT_GIVEN,
             tool_choice=tool_choice or NOT_GIVEN,
             stream=stream,
             temperature=model_settings.get('temperature', NOT_GIVEN),
@@ -260,8 +249,13 @@ async def _process_streamed_response(self, response: AsyncStream[RawMessageStrea
         timestamp = datetime.now(tz=timezone.utc)
         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 +304,28 @@ def _map_message(messages: list[ModelMessage]) -> tuple[str, list[MessageParam]]
                         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: