Skip to content

Commit 7ef0b9b

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 78fb707 commit 7ef0b9b

File tree

2 files changed

+139
-4
lines changed

2 files changed

+139
-4
lines changed

pydantic_ai_slim/pydantic_ai/models/google.py

Lines changed: 43 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,
@@ -50,7 +50,7 @@
5050
)
5151

5252
try:
53-
from google.genai import Client
53+
from google.genai import Client, errors
5454
from google.genai.types import (
5555
BlobDict,
5656
CodeExecutionResult,
@@ -357,6 +357,37 @@ def _get_tool_config(
357357
else:
358358
return None
359359

360+
def _handle_google_error(self, error: Exception) -> ModelHTTPError:
361+
"""Helper method to convert Google API errors to ModelHTTPError."""
362+
if isinstance(error, errors.APIError):
363+
return ModelHTTPError(
364+
status_code=getattr(error, 'code', 500),
365+
model_name=self._model_name,
366+
body=error.details,
367+
)
368+
369+
error_mappings = {
370+
errors.UnknownFunctionCallArgumentError: (400, 'the function call argument cannot be converted to the parameter annotation.', 'BAD_REQUEST'),
371+
errors.UnsupportedFunctionError: (404, 'the function is not supported.', 'NOT_FOUND'),
372+
errors.FunctionInvocationError: (400, 'the function cannot be invoked with the given arguments.', 'BAD_REQUEST'),
373+
errors.UnknownApiResponseError: (422, 'the response from the API cannot be parsed as JSON.', 'UNPROCESSABLE_CONTENT')
374+
}
375+
376+
if error.__class__ in error_mappings:
377+
code, message, status = error_mappings[error.__class__]
378+
return ModelHTTPError(
379+
status_code=code,
380+
model_name=self._model_name,
381+
body={'error': {'code': code, 'message': message, 'status': status}}
382+
)
383+
384+
# Handle unknown errors as 500 Internal Server Error
385+
return ModelHTTPError(
386+
status_code=500,
387+
model_name=self._model_name,
388+
body={'error': {'code': 500, 'message': str(error), 'status': 'INTERNAL_ERROR'}}
389+
)
390+
360391
@overload
361392
async def _generate_content(
362393
self,
@@ -384,7 +415,16 @@ async def _generate_content(
384415
) -> GenerateContentResponse | Awaitable[AsyncIterator[GenerateContentResponse]]:
385416
contents, config = await self._build_content_and_config(messages, model_settings, model_request_parameters)
386417
func = self.client.aio.models.generate_content_stream if stream else self.client.aio.models.generate_content
387-
return await func(model=self._model_name, contents=contents, config=config) # type: ignore
418+
try:
419+
return await func(model=self._model_name, contents=contents, config=config) # type: ignore
420+
except (
421+
errors.APIError,
422+
errors.UnknownFunctionCallArgumentError,
423+
errors.UnsupportedFunctionError,
424+
errors.FunctionInvocationError,
425+
errors.UnknownApiResponseError
426+
) as e:
427+
raise self._handle_google_error(e) from e
388428

389429
async def _build_content_and_config(
390430
self,

tests/models/test_google.py

Lines changed: 96 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from httpx import Timeout
1010
from inline_snapshot import Is, snapshot
1111
from pydantic import BaseModel
12+
from pytest_mock import MockerFixture
1213
from typing_extensions import TypedDict
1314

1415
from pydantic_ai import (
@@ -41,7 +42,7 @@
4142
)
4243
from pydantic_ai.agent import Agent
4344
from pydantic_ai.builtin_tools import CodeExecutionTool, ImageGenerationTool, UrlContextTool, WebSearchTool
44-
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError
45+
from pydantic_ai.exceptions import ModelHTTPError, ModelRetry, UnexpectedModelBehavior, UserError
4546
from pydantic_ai.messages import (
4647
BuiltinToolCallEvent, # pyright: ignore[reportDeprecated]
4748
BuiltinToolResultEvent, # pyright: ignore[reportDeprecated]
@@ -54,6 +55,7 @@
5455
from ..parts_from_messages import part_types_from_messages
5556

5657
with try_import() as imports_successful:
58+
from google.genai import errors
5759
from google.genai.types import (
5860
GenerateContentResponse,
5961
GenerateContentResponseUsageMetadata,
@@ -2929,3 +2931,96 @@ async def test_google_vertexai_image_generation(allow_model_requests: None, vert
29292931
identifier='f3edd8',
29302932
)
29312933
)
2934+
2935+
2936+
# API 에러 테스트 데이터
2937+
@pytest.mark.parametrize(
2938+
'error_class,error_response,expected_status',
2939+
[(
2940+
errors.ServerError,
2941+
{
2942+
'error': {
2943+
'code': 503,
2944+
'message': 'The service is currently unavailable.',
2945+
'status': 'UNAVAILABLE'
2946+
}
2947+
},
2948+
503
2949+
),
2950+
(
2951+
errors.ClientError,
2952+
{
2953+
'error': {
2954+
'code': 400,
2955+
'message': 'Invalid request parameters',
2956+
'status': 'INVALID_ARGUMENT'
2957+
}
2958+
},
2959+
400
2960+
),
2961+
(
2962+
errors.ClientError,
2963+
{
2964+
'error': {
2965+
'code': 429,
2966+
'message': 'Rate limit exceeded',
2967+
'status': 'RESOURCE_EXHAUSTED'
2968+
}
2969+
},
2970+
429
2971+
),
2972+
])
2973+
async def test_google_api_errors_are_handled(
2974+
allow_model_requests: None,
2975+
google_provider: GoogleProvider,
2976+
mocker: MockerFixture,
2977+
error_class: type[errors.APIError],
2978+
error_response: dict,
2979+
expected_status: int,
2980+
):
2981+
model = GoogleModel('gemini-1.5-flash', provider=google_provider)
2982+
mocked_error = error_class(expected_status, error_response)
2983+
mocker.patch.object(
2984+
model.client.aio.models,
2985+
'generate_content',
2986+
side_effect=mocked_error
2987+
)
2988+
2989+
agent = Agent(model=model)
2990+
2991+
with pytest.raises(ModelHTTPError) as exc_info:
2992+
await agent.run('This prompt will trigger the mocked error.')
2993+
2994+
assert exc_info.value.status_code == expected_status
2995+
assert error_response['error']['message'] in str(exc_info.value.body)
2996+
2997+
2998+
@pytest.mark.parametrize(
2999+
'error_class,expected_status',
3000+
[
3001+
(errors.UnknownFunctionCallArgumentError, 400),
3002+
(errors.UnsupportedFunctionError, 404),
3003+
(errors.FunctionInvocationError, 400),
3004+
(errors.UnknownApiResponseError, 422),
3005+
])
3006+
async def test_google_specific_errors_are_handled(
3007+
allow_model_requests: None,
3008+
google_provider: GoogleProvider,
3009+
mocker: MockerFixture,
3010+
error_class: type[errors.APIError],
3011+
expected_status: int,
3012+
):
3013+
3014+
model = GoogleModel('gemini-1.5-flash', provider=google_provider)
3015+
mocked_error = error_class
3016+
mocker.patch.object(
3017+
model.client.aio.models,
3018+
'generate_content',
3019+
side_effect=mocked_error
3020+
)
3021+
3022+
agent = Agent(model=model)
3023+
3024+
with pytest.raises(ModelHTTPError) as exc_info:
3025+
await agent.run('This prompt will trigger the mocked error.')
3026+
assert exc_info.value.status_code == expected_status

0 commit comments

Comments
 (0)