revert agent_model() name change

Author: samuelcolvin

pydantic_ai_slim/pydantic_ai/agent.py

@@ -669,7 +669,7 @@ async def add_tool(key: str, tool: Tool[AgentDeps]) -> None:
 
         await asyncio.gather(*(add_tool(k, t) for k, t in self._function_tools.items()))
 
-        return await model.prepare(
+        return await model.agent_model(
             function_tools=function_tools,
             allow_text_result=self._allow_text_result,
             result_tools=self._result_schema.tool_defs() if self._result_schema is not None else None,

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -61,7 +61,7 @@ class Model(ABC):
     """Abstract class for a model."""
 
     @abstractmethod
-    async def prepare(
+    async def agent_model(
         self,
         *,
         function_tools: dict[str, ToolDefinition],

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -51,7 +51,7 @@ def __init__(self, function: FunctionDef | None = None, *, stream_function: Stre
         self.function = function
         self.stream_function = stream_function
 
-    async def prepare(
+    async def agent_model(
         self,
         *,
         function_tools: dict[str, ToolDefinition],

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -88,7 +88,7 @@ def __init__(
         self.http_client = http_client or cached_async_http_client()
         self.url = url_template.format(model=model_name)
 
-    async def prepare(
+    async def agent_model(
         self,
         *,
         function_tools: dict[str, ToolDefinition],

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -107,7 +107,7 @@ def __init__(
         else:
             self.client = AsyncGroq(api_key=api_key, http_client=cached_async_http_client())
 
-    async def prepare(
+    async def agent_model(
         self,
         *,
         function_tools: dict[str, ToolDefinition],

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -87,7 +87,7 @@ def __init__(
         else:
             self.client = AsyncOpenAI(api_key=api_key, http_client=cached_async_http_client())
 
-    async def prepare(
+    async def agent_model(
         self,
         *,
         function_tools: dict[str, ToolDefinition],

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -59,9 +59,9 @@ class TestModel(Model):
     agent_model_function_tools: dict[str, ToolDefinition] | None = field(default=None, init=False)
     agent_model_allow_text_result: bool | None = field(default=None, init=False)
     agent_model_result_tools: dict[str, ToolDefinition] | None = field(default=None, init=False)
-    agent_model: TestAgentModel | None = field(default=None, init=False)
+    agent_model_cache: TestAgentModel | None = field(default=None, init=False)
 
-    async def prepare(
+    async def agent_model(
         self,
         *,
         function_tools: dict[str, ToolDefinition],
@@ -71,9 +71,9 @@ async def prepare(
         self.agent_model_function_tools = function_tools
         self.agent_model_allow_text_result = allow_text_result
         self.agent_model_result_tools = result_tools
-        if self.agent_model is None:
-            self.agent_model = self._build_agent(function_tools, allow_text_result, result_tools)
-        return self.agent_model
+        if self.agent_model_cache is None:
+            self.agent_model_cache = self._build_agent(function_tools, allow_text_result, result_tools)
+        return self.agent_model_cache
 
     def _build_agent(
         self,

pydantic_ai_slim/pydantic_ai/models/vertexai.py

@@ -107,7 +107,7 @@ def __init__(
         self.auth = None
         self.url = None
 
-    async def prepare(
+    async def agent_model(
         self,
         *,
         function_tools: dict[str, ToolDefinition],

tests/models/test_gemini.py

@@ -77,7 +77,7 @@ def test_api_key_empty(env: TestEnv):
 
 async def test_agent_model_simple(allow_model_requests: None):
     m = GeminiModel('gemini-1.5-flash', api_key='via-arg')
-    agent_model = await m.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+    agent_model = await m.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
     assert isinstance(agent_model.http_client, httpx.AsyncClient)
     assert agent_model.model_name == 'gemini-1.5-flash'
     assert isinstance(agent_model.auth, ApiKeyAuth)
@@ -110,7 +110,9 @@ async def test_agent_model_tools(allow_model_requests: None):
         'This is the tool for the final Result',
         {'type': 'object', 'title': 'Result', 'properties': {'spam': {'type': 'number'}}, 'required': ['spam']},
     )
-    agent_model = await m.prepare(function_tools=tools, allow_text_result=True, result_tools={'result': result_tool})
+    agent_model = await m.agent_model(
+        function_tools=tools, allow_text_result=True, result_tools={'result': result_tool}
+    )
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -149,7 +151,7 @@ async def test_require_response_tool(allow_model_requests: None):
         'This is the tool for the final Result',
         {'type': 'object', 'title': 'Result', 'properties': {'spam': {'type': 'number'}}},
     )
-    agent_model = await m.prepare(function_tools={}, allow_text_result=False, result_tools={'result': result_tool})
+    agent_model = await m.agent_model(function_tools={}, allow_text_result=False, result_tools={'result': result_tool})
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -206,7 +208,7 @@ class Locations(BaseModel):
         'This is the tool for the final Result',
         json_schema,
     )
-    agent_model = await m.prepare(function_tools={}, allow_text_result=True, result_tools={'result': result_tool})
+    agent_model = await m.agent_model(function_tools={}, allow_text_result=True, result_tools={'result': result_tool})
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -252,7 +254,7 @@ class Locations(BaseModel):
         'This is the tool for the final Result',
         json_schema,
     )
-    agent_model = await m.prepare(function_tools={}, allow_text_result=True, result_tools={'result': result_tool})
+    agent_model = await m.agent_model(function_tools={}, allow_text_result=True, result_tools={'result': result_tool})
     assert agent_model.tools == snapshot(
         _GeminiTools(
             function_declarations=[
@@ -316,7 +318,7 @@ class Location(BaseModel):
         json_schema,
     )
     with pytest.raises(UserError, match=r'Recursive `\$ref`s in JSON Schema are not supported by Gemini'):
-        await m.prepare(function_tools={}, allow_text_result=True, result_tools={'result': result_tool})
+        await m.agent_model(function_tools={}, allow_text_result=True, result_tools={'result': result_tool})
 
 
 @dataclass

tests/models/test_vertexai.py

@@ -33,7 +33,7 @@ async def test_init_service_account(tmp_path: Path, allow_model_requests: None):
     assert model.url is None
     assert model.auth is None
 
-    await model.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+    await model.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
 
     assert model.url == snapshot(
         'https://us-central1-aiplatform.googleapis.com/v1/projects/my-project-id/locations/us-central1/'
@@ -58,7 +58,7 @@ async def test_init_env(mocker: MockerFixture, allow_model_requests: None):
 
     assert patch.call_count == 0
 
-    await model.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+    await model.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
 
     assert patch.call_count == 1
 
@@ -69,7 +69,7 @@ async def test_init_env(mocker: MockerFixture, allow_model_requests: None):
     assert model.auth is not None
     assert model.name() == snapshot('vertexai:gemini-1.5-flash')
 
-    await model.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+    await model.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
     assert model.url is not None
     assert model.auth is not None
     assert patch.call_count == 1
@@ -83,7 +83,7 @@ async def test_init_right_project_id(tmp_path: Path, allow_model_requests: None)
     assert model.url is None
     assert model.auth is None
 
-    await model.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+    await model.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
 
     assert model.url == snapshot(
         'https://us-central1-aiplatform.googleapis.com/v1/projects/my-project-id/locations/us-central1/'
@@ -99,7 +99,7 @@ async def test_init_service_account_wrong_project_id(tmp_path: Path):
     model = VertexAIModel('gemini-1.5-flash', service_account_file=service_account_path, project_id='different')
 
     with pytest.raises(UserError) as exc_info:
-        await model.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+        await model.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
     assert str(exc_info.value) == snapshot(
         "The project_id you provided does not match the one from service account file: 'different' != 'my-project-id'"
     )
@@ -110,7 +110,7 @@ async def test_init_env_wrong_project_id(mocker: MockerFixture):
     model = VertexAIModel('gemini-1.5-flash', project_id='different')
 
     with pytest.raises(UserError) as exc_info:
-        await model.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+        await model.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
     assert str(exc_info.value) == snapshot(
         "The project_id you provided does not match the one from `google.auth.default()`: 'different' != 'my-project-id'"
     )
@@ -124,7 +124,7 @@ async def test_init_env_no_project_id(mocker: MockerFixture):
     model = VertexAIModel('gemini-1.5-flash')
 
     with pytest.raises(UserError) as exc_info:
-        await model.prepare(function_tools={}, allow_text_result=True, result_tools=None)
+        await model.agent_model(function_tools={}, allow_text_result=True, result_tools=None)
     assert str(exc_info.value) == snapshot('No project_id provided and none found in `google.auth.default()`')