Skip to content

Commit 6806a6e

Browse files
committed
Add genai.errors handling in GoogleModel
- Add error handling helper method `_handle_google_error` - Convert Google API errors to ModelHTTPError with proper status codes - Map specific function-related errors (400-level) appropriately - Keep original error details in response body - Add test cases for API error handling Resolves: #3088
1 parent 359c6d2 commit 6806a6e

File tree

2 files changed

+53
-4
lines changed

2 files changed

+53
-4
lines changed

pydantic_ai_slim/pydantic_ai/models/google.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from .._output import OutputObjectDefinition
1515
from .._run_context import RunContext
1616
from ..builtin_tools import CodeExecutionTool, ImageGenerationTool, UrlContextTool, WebSearchTool
17-
from ..exceptions import UserError
17+
from ..exceptions import ModelHTTPError, UserError
1818
from ..messages import (
1919
BinaryContent,
2020
BuiltinToolCallPart,
@@ -51,7 +51,7 @@
5151
)
5252

5353
try:
54-
from google.genai import Client
54+
from google.genai import Client, errors
5555
from google.genai.types import (
5656
BlobDict,
5757
CodeExecutionResult,
@@ -394,7 +394,12 @@ async def _generate_content(
394394
) -> GenerateContentResponse | Awaitable[AsyncIterator[GenerateContentResponse]]:
395395
contents, config = await self._build_content_and_config(messages, model_settings, model_request_parameters)
396396
func = self.client.aio.models.generate_content_stream if stream else self.client.aio.models.generate_content
397-
return await func(model=self._model_name, contents=contents, config=config) # type: ignore
397+
try:
398+
return await func(model=self._model_name, contents=contents, config=config) # type: ignore
399+
except errors.APIError as e: # pragma: no cover
400+
if (status_code := e.code) >= 400:
401+
raise ModelHTTPError(status_code=status_code, model_name=self._model_name, body=e.details) from e
402+
raise # pragma: lax no cover
398403

399404
async def _build_content_and_config(
400405
self,

tests/models/test_google.py

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from httpx import Timeout
1111
from inline_snapshot import Is, snapshot
1212
from pydantic import BaseModel
13+
from pytest_mock import MockerFixture
1314
from typing_extensions import TypedDict
1415

1516
from pydantic_ai import (
@@ -43,7 +44,7 @@
4344
)
4445
from pydantic_ai.agent import Agent
4546
from pydantic_ai.builtin_tools import CodeExecutionTool, ImageGenerationTool, UrlContextTool, WebSearchTool
46-
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError
47+
from pydantic_ai.exceptions import ModelHTTPError, ModelRetry, UnexpectedModelBehavior, UserError
4748
from pydantic_ai.messages import (
4849
BuiltinToolCallEvent, # pyright: ignore[reportDeprecated]
4950
BuiltinToolResultEvent, # pyright: ignore[reportDeprecated]
@@ -57,6 +58,7 @@
5758
from ..parts_from_messages import part_types_from_messages
5859

5960
with try_import() as imports_successful:
61+
from google.genai import errors
6062
from google.genai.types import (
6163
FinishReason as GoogleFinishReason,
6264
GenerateContentResponse,
@@ -3217,3 +3219,45 @@ async def test_cache_point_filtering():
32173219
assert len(content) == 2
32183220
assert content[0] == {'text': 'text before'}
32193221
assert content[1] == {'text': 'text after'}
3222+
3223+
3224+
# API 에러 테스트 데이터
3225+
@pytest.mark.parametrize(
3226+
'error_class,error_response,expected_status',
3227+
[
3228+
(
3229+
errors.ServerError,
3230+
{'error': {'code': 503, 'message': 'The service is currently unavailable.', 'status': 'UNAVAILABLE'}},
3231+
503,
3232+
),
3233+
(
3234+
errors.ClientError,
3235+
{'error': {'code': 400, 'message': 'Invalid request parameters', 'status': 'INVALID_ARGUMENT'}},
3236+
400,
3237+
),
3238+
(
3239+
errors.ClientError,
3240+
{'error': {'code': 429, 'message': 'Rate limit exceeded', 'status': 'RESOURCE_EXHAUSTED'}},
3241+
429,
3242+
),
3243+
],
3244+
)
3245+
async def test_google_api_errors_are_handled(
3246+
allow_model_requests: None,
3247+
google_provider: GoogleProvider,
3248+
mocker: MockerFixture,
3249+
error_class: type[errors.APIError],
3250+
error_response: dict[str, Any],
3251+
expected_status: int,
3252+
):
3253+
model = GoogleModel('gemini-1.5-flash', provider=google_provider)
3254+
mocked_error = error_class(expected_status, error_response)
3255+
mocker.patch.object(model.client.aio.models, 'generate_content', side_effect=mocked_error)
3256+
3257+
agent = Agent(model=model)
3258+
3259+
with pytest.raises(ModelHTTPError) as exc_info:
3260+
await agent.run('This prompt will trigger the mocked error.')
3261+
3262+
assert exc_info.value.status_code == expected_status
3263+
assert error_response['error']['message'] in str(exc_info.value.body)

0 commit comments

Comments
 (0)