From 47d06d9fdd5973ea2e72daec07b1a38bae24b2bd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 11:56:02 -0700 Subject: [PATCH 1/2] test(unified_google_tests): use the Vertex global endpoint and retry 429s with backoff The google_generate_content_endpoint_testing job went red on main when us-central1 ran out of shared gemini-2.5-flash-lite capacity for a few hours. The suite's proxy config now sends the Vertex deployment to the global endpoint and retries rate limit errors 5 times with exponential backoff, and a regression test pins that the config rides out 3 consecutive 429s --- .../google_genai_proxy_test_config.yaml | 5 ++ .../test_google_genai_proxy_test_config.py | 67 +++++++++++++++++++ 2 files changed, 72 insertions(+) create mode 100644 tests/unified_google_tests/test_google_genai_proxy_test_config.py diff --git a/tests/unified_google_tests/google_genai_proxy_test_config.yaml b/tests/unified_google_tests/google_genai_proxy_test_config.yaml index 9913c05d434..64a83ef3d81 100644 --- a/tests/unified_google_tests/google_genai_proxy_test_config.yaml +++ b/tests/unified_google_tests/google_genai_proxy_test_config.yaml @@ -7,6 +7,11 @@ model_list: - model_name: vertex-gemini-2.5-flash-lite litellm_params: model: vertex_ai/gemini-2.5-flash-lite + vertex_location: global + +router_settings: + retry_policy: + RateLimitErrorRetries: 5 general_settings: master_key: sk-1234 diff --git a/tests/unified_google_tests/test_google_genai_proxy_test_config.py b/tests/unified_google_tests/test_google_genai_proxy_test_config.py new file mode 100644 index 00000000000..d84eefb406b --- /dev/null +++ b/tests/unified_google_tests/test_google_genai_proxy_test_config.py @@ -0,0 +1,67 @@ +import time +from pathlib import Path +from typing import Final, ReadOnly, TypedDict + +import httpx +import pytest +import respx +import yaml +from pydantic import TypeAdapter + +import litellm +from litellm import Router + +CONFIG_PATH: Final = Path(__file__).parent / "google_genai_proxy_test_config.yaml" +GEMINI_HOST: Final = "generativelanguage.googleapis.com" +GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/gemini-2.5-flash-lite:generateContent" +RESOURCE_EXHAUSTED: Final = { + "error": {"code": 429, "message": "Resource exhausted. Please try again later.", "status": "RESOURCE_EXHAUSTED"} +} +PONG: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "pong"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 1, "totalTokenCount": 9}, +} +CONSECUTIVE_RATE_LIMITS: Final = 3 +MINIMUM_BACKOFF_SECONDS: Final = 0.5 + 1.0 + 2.0 + + +class _Deployment(TypedDict): + model_name: ReadOnly[str] + litellm_params: ReadOnly[dict[str, str]] + + +class _ProxyConfig(TypedDict): + model_list: ReadOnly[list[_Deployment]] + router_settings: ReadOnly[dict[str, dict[str, int]]] + + +def _router_from_ci_proxy_config() -> Router: + config: Final = TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text())) + gemini_deployments: Final = [ + {"model_name": deployment["model_name"], "litellm_params": {**deployment["litellm_params"], "api_key": "test"}} + for deployment in config["model_list"] + if deployment["model_name"] == "gemini-2.5-flash-lite" + ] + return Router(model_list=gemini_deployments, retry_policy=config["router_settings"]["retry_policy"]) + + +@pytest.mark.asyncio +async def test_ci_proxy_config_rides_out_consecutive_429s_with_backoff( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + route: Final = respx_mock.post(host=GEMINI_HOST, path=GEMINI_GENERATE_CONTENT_PATH).mock( + side_effect=[httpx.Response(429, json=RESOURCE_EXHAUSTED)] * CONSECUTIVE_RATE_LIMITS + + [httpx.Response(200, json=PONG)] + ) + started: Final = time.monotonic() + response: Final = await _router_from_ci_proxy_config().agenerate_content( + model="gemini-2.5-flash-lite", + contents=[{"role": "user", "parts": [{"text": "Reply with only the single word: pong"}]}], + ) + elapsed: Final = time.monotonic() - started + + assert response.model_dump()["candidates"][0]["content"]["parts"][0]["text"] == "pong" + assert route.call_count == CONSECUTIVE_RATE_LIMITS + 1 + assert elapsed >= MINIMUM_BACKOFF_SECONDS From a7870a902a281e61a4842dbc4bd079fd621a21cf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 12:15:34 -0700 Subject: [PATCH 2/2] test(unified_google_tests): import ReadOnly from typing_extensions and cover the Vertex global endpoint The first commit imported ReadOnly from typing, which only exists on Python 3.13 and up. CircleCI runs this suite on 3.12, so the module failed at import and the job stopped at collection before any of its tests ran. ReadOnly and TypedDict now come from typing_extensions, like the rest of the repo A new test resolves the Vertex deployment's location from the suite's config with VERTEXAI_LOCATION set to a region, and fails if the vertex_location line is removed The expected minimum backoff is now derived from litellm's INITIAL_RETRY_DELAY and MAX_RETRY_DELAY, so the test holds when those are overridden through the environment --- .../test_google_genai_proxy_test_config.py | 50 +++++++++++++++---- 1 file changed, 40 insertions(+), 10 deletions(-) diff --git a/tests/unified_google_tests/test_google_genai_proxy_test_config.py b/tests/unified_google_tests/test_google_genai_proxy_test_config.py index d84eefb406b..694ec336bac 100644 --- a/tests/unified_google_tests/test_google_genai_proxy_test_config.py +++ b/tests/unified_google_tests/test_google_genai_proxy_test_config.py @@ -1,19 +1,26 @@ import time from pathlib import Path -from typing import Final, ReadOnly, TypedDict +from typing import Final import httpx import pytest import respx import yaml from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import Router +from litellm.constants import INITIAL_RETRY_DELAY, MAX_RETRY_DELAY +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase CONFIG_PATH: Final = Path(__file__).parent / "google_genai_proxy_test_config.yaml" +GEMINI_DEPLOYMENT: Final = "gemini-2.5-flash-lite" +VERTEX_DEPLOYMENT: Final = "vertex-gemini-2.5-flash-lite" GEMINI_HOST: Final = "generativelanguage.googleapis.com" GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/gemini-2.5-flash-lite:generateContent" +VERTEX_GLOBAL_BASE_URL: Final = "https://aiplatform.googleapis.com" RESOURCE_EXHAUSTED: Final = { "error": {"code": 429, "message": "Resource exhausted. Please try again later.", "status": "RESOURCE_EXHAUSTED"} } @@ -22,7 +29,9 @@ PONG: Final = { "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 1, "totalTokenCount": 9}, } CONSECUTIVE_RATE_LIMITS: Final = 3 -MINIMUM_BACKOFF_SECONDS: Final = 0.5 + 1.0 + 2.0 +MINIMUM_BACKOFF_SECONDS: Final = sum( + min(INITIAL_RETRY_DELAY * 2**attempt, MAX_RETRY_DELAY) for attempt in range(CONSECUTIVE_RATE_LIMITS) +) class _Deployment(TypedDict): @@ -35,14 +44,35 @@ class _ProxyConfig(TypedDict): router_settings: ReadOnly[dict[str, dict[str, int]]] +def _ci_proxy_config() -> _ProxyConfig: + return TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text())) + + +def _litellm_params(config: _ProxyConfig, model_name: str) -> dict[str, str]: + return next( + deployment["litellm_params"] for deployment in config["model_list"] if deployment["model_name"] == model_name + ) + + def _router_from_ci_proxy_config() -> Router: - config: Final = TypeAdapter(_ProxyConfig).validate_python(yaml.safe_load(CONFIG_PATH.read_text())) - gemini_deployments: Final = [ - {"model_name": deployment["model_name"], "litellm_params": {**deployment["litellm_params"], "api_key": "test"}} - for deployment in config["model_list"] - if deployment["model_name"] == "gemini-2.5-flash-lite" - ] - return Router(model_list=gemini_deployments, retry_policy=config["router_settings"]["retry_policy"]) + config: Final = _ci_proxy_config() + return Router( + model_list=[ + { + "model_name": GEMINI_DEPLOYMENT, + "litellm_params": {**_litellm_params(config, GEMINI_DEPLOYMENT), "api_key": "test"}, + } + ], + retry_policy=config["router_settings"]["retry_policy"], + ) + + +def test_ci_proxy_config_sends_vertex_calls_to_the_global_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5") + location: Final = VertexBase.safe_get_vertex_ai_location(_litellm_params(_ci_proxy_config(), VERTEX_DEPLOYMENT)) + + assert location == "global" + assert get_vertex_base_url(location) == VERTEX_GLOBAL_BASE_URL @pytest.mark.asyncio @@ -57,7 +87,7 @@ async def test_ci_proxy_config_rides_out_consecutive_429s_with_backoff( ) started: Final = time.monotonic() response: Final = await _router_from_ci_proxy_config().agenerate_content( - model="gemini-2.5-flash-lite", + model=GEMINI_DEPLOYMENT, contents=[{"role": "user", "parts": [{"text": "Reply with only the single word: pong"}]}], ) elapsed: Final = time.monotonic() - started