From ae45a9bb7cffd4f8c7fe657073a04f241969cae8 Mon Sep 17 00:00:00 2001 From: timothy-jeong Date: Sat, 11 Oct 2025 17:58:50 +0900 Subject: [PATCH 1/2] 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 --- pydantic_ai_slim/pydantic_ai/models/google.py | 58 ++++++++++++++- tests/models/test_google.py | 73 ++++++++++++++++++- 2 files changed, 127 insertions(+), 4 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index ad13e298f3..a04bc7a82b 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/pydantic_ai_slim/pydantic_ai/models/google.py @@ -14,7 +14,7 @@ from .._output import OutputObjectDefinition from .._run_context import RunContext from ..builtin_tools import CodeExecutionTool, ImageGenerationTool, UrlContextTool, WebSearchTool -from ..exceptions import UserError +from ..exceptions import ModelHTTPError, UserError from ..messages import ( BinaryContent, BuiltinToolCallPart, @@ -50,7 +50,7 @@ ) try: - from google.genai import Client + from google.genai import Client, errors from google.genai.types import ( BlobDict, CodeExecutionResult, @@ -357,6 +357,49 @@ def _get_tool_config( else: return None + def _handle_google_error(self, error: Exception) -> ModelHTTPError: + """Helper method to convert Google API errors to ModelHTTPError.""" + if isinstance(error, errors.APIError): + return ModelHTTPError( + status_code=getattr(error, 'code', 500), + model_name=self._model_name, + body=error.details, + ) + + error_mappings = { + errors.UnknownFunctionCallArgumentError: ( + 400, + 'the function call argument cannot be converted to the parameter annotation.', + 'BAD_REQUEST', + ), + errors.UnsupportedFunctionError: (404, 'the function is not supported.', 'NOT_FOUND'), + errors.FunctionInvocationError: ( + 400, + 'the function cannot be invoked with the given arguments.', + 'BAD_REQUEST', + ), + errors.UnknownApiResponseError: ( + 422, + 'the response from the API cannot be parsed as JSON.', + 'UNPROCESSABLE_CONTENT', + ), + } + + if error.__class__ in error_mappings: + code, message, status = error_mappings[error.__class__] + return ModelHTTPError( + status_code=code, + model_name=self._model_name, + body={'error': {'code': code, 'message': message, 'status': status}}, + ) + + # Handle unknown errors as 500 Internal Server Error + return ModelHTTPError( + status_code=500, + model_name=self._model_name, + body={'error': {'code': 500, 'message': str(error), 'status': 'INTERNAL_ERROR'}}, + ) + @overload async def _generate_content( self, @@ -384,7 +427,16 @@ async def _generate_content( ) -> GenerateContentResponse | Awaitable[AsyncIterator[GenerateContentResponse]]: contents, config = await self._build_content_and_config(messages, model_settings, model_request_parameters) func = self.client.aio.models.generate_content_stream if stream else self.client.aio.models.generate_content - return await func(model=self._model_name, contents=contents, config=config) # type: ignore + try: + return await func(model=self._model_name, contents=contents, config=config) # type: ignore + except ( + errors.APIError, + errors.UnknownFunctionCallArgumentError, + errors.UnsupportedFunctionError, + errors.FunctionInvocationError, + errors.UnknownApiResponseError, + ) as e: + raise self._handle_google_error(e) from e async def _build_content_and_config( self, diff --git a/tests/models/test_google.py b/tests/models/test_google.py index ce461e2810..a721e02b57 100644 --- a/tests/models/test_google.py +++ b/tests/models/test_google.py @@ -9,6 +9,7 @@ from httpx import Timeout from inline_snapshot import Is, snapshot from pydantic import BaseModel +from pytest_mock import MockerFixture from typing_extensions import TypedDict from pydantic_ai import ( @@ -41,7 +42,7 @@ ) from pydantic_ai.agent import Agent from pydantic_ai.builtin_tools import CodeExecutionTool, ImageGenerationTool, UrlContextTool, WebSearchTool -from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError +from pydantic_ai.exceptions import ModelHTTPError, ModelRetry, UnexpectedModelBehavior, UserError from pydantic_ai.messages import ( BuiltinToolCallEvent, # pyright: ignore[reportDeprecated] BuiltinToolResultEvent, # pyright: ignore[reportDeprecated] @@ -54,6 +55,7 @@ from ..parts_from_messages import part_types_from_messages with try_import() as imports_successful: + from google.genai import errors from google.genai.types import ( GenerateContentResponse, GenerateContentResponseUsageMetadata, @@ -2929,3 +2931,72 @@ async def test_google_vertexai_image_generation(allow_model_requests: None, vert identifier='f3edd8', ) ) + + +# API 에러 테스트 데이터 +@pytest.mark.parametrize( + 'error_class,error_response,expected_status', + [ + ( + errors.ServerError, + {'error': {'code': 503, 'message': 'The service is currently unavailable.', 'status': 'UNAVAILABLE'}}, + 503, + ), + ( + errors.ClientError, + {'error': {'code': 400, 'message': 'Invalid request parameters', 'status': 'INVALID_ARGUMENT'}}, + 400, + ), + ( + errors.ClientError, + {'error': {'code': 429, 'message': 'Rate limit exceeded', 'status': 'RESOURCE_EXHAUSTED'}}, + 429, + ), + ], +) +async def test_google_api_errors_are_handled( + allow_model_requests: None, + google_provider: GoogleProvider, + mocker: MockerFixture, + error_class: type[errors.APIError], + error_response: dict[str, Any], + expected_status: int, +): + model = GoogleModel('gemini-1.5-flash', provider=google_provider) + mocked_error = error_class(expected_status, error_response) + mocker.patch.object(model.client.aio.models, 'generate_content', side_effect=mocked_error) + + agent = Agent(model=model) + + with pytest.raises(ModelHTTPError) as exc_info: + await agent.run('This prompt will trigger the mocked error.') + + assert exc_info.value.status_code == expected_status + assert error_response['error']['message'] in str(exc_info.value.body) + + +@pytest.mark.parametrize( + 'error_class,expected_status', + [ + (errors.UnknownFunctionCallArgumentError, 400), + (errors.UnsupportedFunctionError, 404), + (errors.FunctionInvocationError, 400), + (errors.UnknownApiResponseError, 422), + ], +) +async def test_google_specific_errors_are_handled( + allow_model_requests: None, + google_provider: GoogleProvider, + mocker: MockerFixture, + error_class: type[errors.APIError], + expected_status: int, +): + model = GoogleModel('gemini-1.5-flash', provider=google_provider) + mocked_error = error_class + mocker.patch.object(model.client.aio.models, 'generate_content', side_effect=mocked_error) + + agent = Agent(model=model) + + with pytest.raises(ModelHTTPError) as exc_info: + await agent.run('This prompt will trigger the mocked error.') + assert exc_info.value.status_code == expected_status From eecaa4a3e47ab81d67b0f740b6d124f2e79629b2 Mon Sep 17 00:00:00 2001 From: timothy-jeong Date: Fri, 17 Oct 2025 15:51:08 +0900 Subject: [PATCH 2/2] Refactor: simplify `ModelHTTPError` handling - Align with other model implementations (openai, groq) by removing unnecessary error transformation logic Resolves: #3088 --- pydantic_ai_slim/pydantic_ai/models/google.py | 55 +------------- tests/models/test_google.py | 73 +------------------ 2 files changed, 5 insertions(+), 123 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index a04bc7a82b..8a103f2765 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/pydantic_ai_slim/pydantic_ai/models/google.py @@ -357,49 +357,6 @@ def _get_tool_config( else: return None - def _handle_google_error(self, error: Exception) -> ModelHTTPError: - """Helper method to convert Google API errors to ModelHTTPError.""" - if isinstance(error, errors.APIError): - return ModelHTTPError( - status_code=getattr(error, 'code', 500), - model_name=self._model_name, - body=error.details, - ) - - error_mappings = { - errors.UnknownFunctionCallArgumentError: ( - 400, - 'the function call argument cannot be converted to the parameter annotation.', - 'BAD_REQUEST', - ), - errors.UnsupportedFunctionError: (404, 'the function is not supported.', 'NOT_FOUND'), - errors.FunctionInvocationError: ( - 400, - 'the function cannot be invoked with the given arguments.', - 'BAD_REQUEST', - ), - errors.UnknownApiResponseError: ( - 422, - 'the response from the API cannot be parsed as JSON.', - 'UNPROCESSABLE_CONTENT', - ), - } - - if error.__class__ in error_mappings: - code, message, status = error_mappings[error.__class__] - return ModelHTTPError( - status_code=code, - model_name=self._model_name, - body={'error': {'code': code, 'message': message, 'status': status}}, - ) - - # Handle unknown errors as 500 Internal Server Error - return ModelHTTPError( - status_code=500, - model_name=self._model_name, - body={'error': {'code': 500, 'message': str(error), 'status': 'INTERNAL_ERROR'}}, - ) - @overload async def _generate_content( self, @@ -429,14 +386,10 @@ async def _generate_content( func = self.client.aio.models.generate_content_stream if stream else self.client.aio.models.generate_content try: return await func(model=self._model_name, contents=contents, config=config) # type: ignore - except ( - errors.APIError, - errors.UnknownFunctionCallArgumentError, - errors.UnsupportedFunctionError, - errors.FunctionInvocationError, - errors.UnknownApiResponseError, - ) as e: - raise self._handle_google_error(e) from e + except errors.APIError as e: + if (status_code := e.code) >= 400: + raise ModelHTTPError(status_code=status_code, model_name=self._model_name, body=e.details) from e + raise # pragma: lax no cover async def _build_content_and_config( self, diff --git a/tests/models/test_google.py b/tests/models/test_google.py index a721e02b57..ce461e2810 100644 --- a/tests/models/test_google.py +++ b/tests/models/test_google.py @@ -9,7 +9,6 @@ from httpx import Timeout from inline_snapshot import Is, snapshot from pydantic import BaseModel -from pytest_mock import MockerFixture from typing_extensions import TypedDict from pydantic_ai import ( @@ -42,7 +41,7 @@ ) from pydantic_ai.agent import Agent from pydantic_ai.builtin_tools import CodeExecutionTool, ImageGenerationTool, UrlContextTool, WebSearchTool -from pydantic_ai.exceptions import ModelHTTPError, ModelRetry, UnexpectedModelBehavior, UserError +from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError from pydantic_ai.messages import ( BuiltinToolCallEvent, # pyright: ignore[reportDeprecated] BuiltinToolResultEvent, # pyright: ignore[reportDeprecated] @@ -55,7 +54,6 @@ from ..parts_from_messages import part_types_from_messages with try_import() as imports_successful: - from google.genai import errors from google.genai.types import ( GenerateContentResponse, GenerateContentResponseUsageMetadata, @@ -2931,72 +2929,3 @@ async def test_google_vertexai_image_generation(allow_model_requests: None, vert identifier='f3edd8', ) ) - - -# API 에러 테스트 데이터 -@pytest.mark.parametrize( - 'error_class,error_response,expected_status', - [ - ( - errors.ServerError, - {'error': {'code': 503, 'message': 'The service is currently unavailable.', 'status': 'UNAVAILABLE'}}, - 503, - ), - ( - errors.ClientError, - {'error': {'code': 400, 'message': 'Invalid request parameters', 'status': 'INVALID_ARGUMENT'}}, - 400, - ), - ( - errors.ClientError, - {'error': {'code': 429, 'message': 'Rate limit exceeded', 'status': 'RESOURCE_EXHAUSTED'}}, - 429, - ), - ], -) -async def test_google_api_errors_are_handled( - allow_model_requests: None, - google_provider: GoogleProvider, - mocker: MockerFixture, - error_class: type[errors.APIError], - error_response: dict[str, Any], - expected_status: int, -): - model = GoogleModel('gemini-1.5-flash', provider=google_provider) - mocked_error = error_class(expected_status, error_response) - mocker.patch.object(model.client.aio.models, 'generate_content', side_effect=mocked_error) - - agent = Agent(model=model) - - with pytest.raises(ModelHTTPError) as exc_info: - await agent.run('This prompt will trigger the mocked error.') - - assert exc_info.value.status_code == expected_status - assert error_response['error']['message'] in str(exc_info.value.body) - - -@pytest.mark.parametrize( - 'error_class,expected_status', - [ - (errors.UnknownFunctionCallArgumentError, 400), - (errors.UnsupportedFunctionError, 404), - (errors.FunctionInvocationError, 400), - (errors.UnknownApiResponseError, 422), - ], -) -async def test_google_specific_errors_are_handled( - allow_model_requests: None, - google_provider: GoogleProvider, - mocker: MockerFixture, - error_class: type[errors.APIError], - expected_status: int, -): - model = GoogleModel('gemini-1.5-flash', provider=google_provider) - mocked_error = error_class - mocker.patch.object(model.client.aio.models, 'generate_content', side_effect=mocked_error) - - agent = Agent(model=model) - - with pytest.raises(ModelHTTPError) as exc_info: - await agent.run('This prompt will trigger the mocked error.') - assert exc_info.value.status_code == expected_status