From d80fc6fe3828c7d940e032146b1a982d6332e272 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Tue, 17 Jun 2025 16:59:57 +0000 Subject: [PATCH 1/2] Respect ModelSettings.timeout in GoogleModel --- pydantic_ai_slim/pydantic_ai/models/google.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index 897b5b0de1..8fd898d462 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/pydantic_ai_slim/pydantic_ai/models/google.py @@ -52,6 +52,7 @@ FunctionDeclarationDict, GenerateContentConfigDict, GenerateContentResponse, + HttpOptionsDict, Part, PartDict, SafetySettingDict, @@ -252,8 +253,17 @@ async def _generate_content( tool_config = self._get_tool_config(model_request_parameters, tools) system_instruction, contents = await self._map_messages(messages) + http_options: HttpOptionsDict = { + 'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()} + } + if timeout := model_settings.get('timeout'): + if isinstance(timeout, float): + http_options['timeout'] = int(1000 * timeout) + else: + raise ValueError('Google does not support setting ModelSettings.timeout to a httpx.Timeout') + config = GenerateContentConfigDict( - http_options={'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}}, + http_options=http_options, system_instruction=system_instruction, temperature=model_settings.get('temperature'), top_p=model_settings.get('top_p'), From 4e4924d97be6a95586f2f6c112cae8b165ad40f7 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Tue, 17 Jun 2025 17:22:14 +0000 Subject: [PATCH 2/2] Add test, handle int --- pydantic_ai_slim/pydantic_ai/models/google.py | 8 +-- .../test_google/test_google_timeout.yaml | 64 +++++++++++++++++++ tests/models/test_google.py | 13 +++- 3 files changed, 80 insertions(+), 5 deletions(-) create mode 100644 tests/models/cassettes/test_google/test_google_timeout.yaml diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index 8fd898d462..3f695fca15 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/pydantic_ai_slim/pydantic_ai/models/google.py @@ -10,9 +10,8 @@ from typing_extensions import assert_never -from pydantic_ai.providers import Provider - from .. import UnexpectedModelBehavior, _utils, usage +from ..exceptions import UserError from ..messages import ( BinaryContent, FileUrl, @@ -30,6 +29,7 @@ VideoUrl, ) from ..profiles import ModelProfileSpec +from ..providers import Provider from ..settings import ModelSettings from ..tools import ToolDefinition from . import ( @@ -257,10 +257,10 @@ async def _generate_content( 'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()} } if timeout := model_settings.get('timeout'): - if isinstance(timeout, float): + if isinstance(timeout, (int, float)): http_options['timeout'] = int(1000 * timeout) else: - raise ValueError('Google does not support setting ModelSettings.timeout to a httpx.Timeout') + raise UserError('Google does not support setting ModelSettings.timeout to a httpx.Timeout') config = GenerateContentConfigDict( http_options=http_options, diff --git a/tests/models/cassettes/test_google/test_google_timeout.yaml b/tests/models/cassettes/test_google/test_google_timeout.yaml new file mode 100644 index 0000000000..0f1e060315 --- /dev/null +++ b/tests/models/cassettes/test_google/test_google_timeout.yaml @@ -0,0 +1,64 @@ +interactions: +- request: + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '87' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + method: POST + parsed_body: + contents: + - parts: + - text: Hello! + role: user + generationConfig: {} + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent + response: + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '687' + content-type: + - application/json; charset=UTF-8 + server-timing: + - gfet4t7; dur=304 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + parsed_body: + candidates: + - avgLogprobs: -0.0009780019860376012 + content: + parts: + - text: | + Hello there! How can I help you today? + role: model + finishReason: STOP + modelVersion: gemini-1.5-flash + responseId: KKRRaJyZNJiFm9IPzMfNiAk + usageMetadata: + candidatesTokenCount: 11 + candidatesTokensDetails: + - modality: TEXT + tokenCount: 11 + promptTokenCount: 2 + promptTokensDetails: + - modality: TEXT + tokenCount: 2 + totalTokenCount: 13 + status: + code: 200 + message: OK +version: 1 diff --git a/tests/models/test_google.py b/tests/models/test_google.py index 2efa2c3d5a..e2a8ae5cee 100644 --- a/tests/models/test_google.py +++ b/tests/models/test_google.py @@ -6,7 +6,7 @@ from typing import Any, Union import pytest -from httpx import Request +from httpx import Request, Timeout from inline_snapshot import Is, snapshot from pytest_mock import MockerFixture from typing_extensions import TypedDict @@ -806,3 +806,14 @@ async def bar() -> str: ), ] ) + + +async def test_google_timeout(allow_model_requests: None, google_provider: GoogleProvider): + model = GoogleModel('gemini-1.5-flash', provider=google_provider) + agent = Agent(model=model) + + result = await agent.run('Hello!', model_settings={'timeout': 10}) + assert result.output == snapshot('Hello there! How can I help you today?\n') + + with pytest.raises(UserError, match='Google does not support setting ModelSettings.timeout to a httpx.Timeout'): + await agent.run('Hello!', model_settings={'timeout': Timeout(10)})