Added Missing allow request Check (#697)

Author: YanSte

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -99,6 +99,7 @@ async def agent_model(
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ) -> GeminiAgentModel:
+        check_allow_model_requests()
         return GeminiAgentModel(
             http_client=self.http_client,
             model_name=self.model_name,
@@ -151,7 +152,6 @@ def __init__(
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ):
-        check_allow_model_requests()
         tools = [_function_from_abstract_tool(t) for t in function_tools]
         if result_tools:
             tools += [_function_from_abstract_tool(t) for t in result_tools]

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -36,6 +36,7 @@
     Model,
     StreamedResponse,
     cached_async_http_client,
+    check_allow_model_requests,
 )
 
 try:
@@ -130,6 +131,7 @@ async def agent_model(
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
         """Create an agent model, this is called for each step of an agent run from Pydantic AI call."""
+        check_allow_model_requests()
         return MistralAgentModel(
             self.client,
             self.model_name,

pydantic_ai_slim/pydantic_ai/models/ollama.py

@@ -10,6 +10,7 @@
     AgentModel,
     Model,
     cached_async_http_client,
+    check_allow_model_requests,
 )
 
 try:
@@ -110,6 +111,7 @@ async def agent_model(
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
+        check_allow_model_requests()
         return await self.openai_model.agent_model(
             function_tools=function_tools,
             allow_text_result=allow_text_result,

pydantic_ai_slim/pydantic_ai/models/vertexai.py

@@ -10,7 +10,7 @@
 from .._utils import run_in_executor
 from ..exceptions import UserError
 from ..tools import ToolDefinition
-from . import Model, cached_async_http_client
+from . import Model, cached_async_http_client, check_allow_model_requests
 from .gemini import GeminiAgentModel, GeminiModelName
 
 try:
@@ -114,6 +114,7 @@ async def agent_model(
         allow_text_result: bool,
         result_tools: list[ToolDefinition],
     ) -> GeminiAgentModel:
+        check_allow_model_requests()
         url, auth = await self.ainit()
         return GeminiAgentModel(
             http_client=self.http_client,

tests/models/test_vertexai.py

@@ -92,7 +92,7 @@ async def test_init_right_project_id(tmp_path: Path, allow_model_requests: None)
     assert model.auth is not None
 
 
-async def test_init_service_account_wrong_project_id(tmp_path: Path):
+async def test_init_service_account_wrong_project_id(tmp_path: Path, allow_model_requests: None):
     service_account_path = tmp_path / 'service_account.json'
     save_service_account(service_account_path, 'my-project-id')
 
@@ -105,7 +105,7 @@ async def test_init_service_account_wrong_project_id(tmp_path: Path):
     )
 
 
-async def test_init_env_wrong_project_id(mocker: MockerFixture):
+async def test_init_env_wrong_project_id(mocker: MockerFixture, allow_model_requests: None):
     mocker.patch('pydantic_ai.models.vertexai.google.auth.default', return_value=(NoOpCredentials(), 'my-project-id'))
     model = VertexAIModel('gemini-1.5-flash', project_id='different')
 
@@ -116,7 +116,7 @@ async def test_init_env_wrong_project_id(mocker: MockerFixture):
     )
 
 
-async def test_init_env_no_project_id(mocker: MockerFixture):
+async def test_init_env_no_project_id(mocker: MockerFixture, allow_model_requests: None):
     mocker.patch(
         'pydantic_ai.models.vertexai.google.auth.default',
         return_value=(NoOpCredentials(), None),