encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

chore: Enhance the messaging with possible fixes for RESOURCE_EXHAUSTED errors from Gemini

Browse Source 
Co-authored-by: Ankur Sharma <ankusharma@google.com>
PiperOrigin-RevId: 833538475
This commit is contained in:
Ankur Sharma
2025-11-17 16:15:56 -08:00
committed by Copybara-Service
parent 5ac5129fb0
commit b2c45f8d91
2 changed files with 136 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/models/test_google_llm.py
+57
View File
@@ -26,11 +26,14 @@ from google.adk.models.google_llm import _AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NA
from google.adk.models.google_llm import _AGENT_ENGINE_TELEMETRY_TAG
from google.adk.models.google_llm import _build_function_declaration_log
from google.adk.models.google_llm import _build_request_log
from google.adk.models.google_llm import _RESOURCE_EXHAUSTED_POSSIBLE_FIX_MESSAGE
from google.adk.models.google_llm import _ResourceExhaustedError
from google.adk.models.google_llm import Gemini
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.utils.variant_utils import GoogleLLMVariant
from google.genai import types
from google.genai.errors import ClientError
from google.genai.types import Content
from google.genai.types import Part
import pytest
@@ -386,6 +389,60 @@ async def test_generate_content_async_stream_preserves_thinking_and_text_parts(
mock_client.aio.models.generate_content_stream.assert_called_once()
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_generate_content_async_resource_exhausted_error(
stream, gemini_llm, llm_request
):
with mock.patch.object(gemini_llm, "api_client") as mock_client:
err = ClientError(code=429, response_json={})
err.code = 429
if stream:
mock_client.aio.models.generate_content_stream.side_effect = err
else:
mock_client.aio.models.generate_content.side_effect = err
with pytest.raises(_ResourceExhaustedError) as excinfo:
responses = []
async for resp in gemini_llm.generate_content_async(
llm_request, stream=stream
):
responses.append(resp)
assert _RESOURCE_EXHAUSTED_POSSIBLE_FIX_MESSAGE in str(excinfo.value)
assert excinfo.value.code == 429
if stream:
mock_client.aio.models.generate_content_stream.assert_called_once()
else:
mock_client.aio.models.generate_content.assert_called_once()
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_generate_content_async_other_client_error(
stream, gemini_llm, llm_request
):
with mock.patch.object(gemini_llm, "api_client") as mock_client:
err = ClientError(code=500, response_json={})
err.code = 500
if stream:
mock_client.aio.models.generate_content_stream.side_effect = err
else:
mock_client.aio.models.generate_content.side_effect = err
with pytest.raises(ClientError) as excinfo:
responses = []
async for resp in gemini_llm.generate_content_async(
llm_request, stream=stream
):
responses.append(resp)
assert excinfo.value.code == 500
assert not isinstance(excinfo.value, _ResourceExhaustedError)
if stream:
mock_client.aio.models.generate_content_stream.assert_called_once()
else:
mock_client.aio.models.generate_content.assert_called_once()
@pytest.mark.asyncio
async def test_connect(gemini_llm, llm_request):
# Create a mock connection