From 092876d377948c11fcc12e2056fd4f7fd02c695d Mon Sep 17 00:00:00 2001 From: balazss Date: Mon, 4 May 2026 13:38:16 -0700 Subject: [PATCH] fix(gemini): address remaining greptile review comments - google_code_assist: log unrecognized optional_params keys for visibility - gemini/authenticator: document timeout overshoot in OAuth loopback loop - gemini/fallback_handler: take primary_call as Callable so async sync-setup errors are caught by the fallback try/except (symmetric with sync path) - gemini/google_genai: pre-fetch OAuth via asyncio.to_thread in async get_auth_token_and_url to avoid blocking the event loop - tests/google_code_assist: add missing __init__.py for parallel pytest discovery - tests/fallback_handler: pass deferred coroutine factory to async fallback wrapper Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/llms/gemini/authenticator.py | 1 + litellm/llms/gemini/fallback_handler.py | 4 ++-- litellm/llms/gemini/google_genai/transformation.py | 11 +++++++++++ litellm/llms/google_code_assist/transformation.py | 10 ++++++++++ litellm/main.py | 2 +- .../test_litellm/llms/gemini/test_fallback_handler.py | 2 +- .../test_litellm/llms/google_code_assist/__init__.py | 0 7 files changed, 26 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/llms/google_code_assist/__init__.py diff --git a/litellm/llms/gemini/authenticator.py b/litellm/llms/gemini/authenticator.py index 9e317d4fe20..89b44a09601 100644 --- a/litellm/llms/gemini/authenticator.py +++ b/litellm/llms/gemini/authenticator.py @@ -257,6 +257,7 @@ class GeminiAuthenticator: deadline = time.monotonic() + loopback_timeout_s server.timeout = 1.0 try: + # Note: deadline is checked after handle_request() returns, so the loop can overshoot by up to server.timeout (1s) — acceptable for a one-shot CLI login. while auth_code is None and error is None: server.handle_request() if time.monotonic() >= deadline: diff --git a/litellm/llms/gemini/fallback_handler.py b/litellm/llms/gemini/fallback_handler.py index 548cf5d5918..cd1cf75d5bc 100644 --- a/litellm/llms/gemini/fallback_handler.py +++ b/litellm/llms/gemini/fallback_handler.py @@ -8,7 +8,7 @@ _google_code_assist_chat = get_google_code_assist_chat() async def run_gemini_acompletion_with_code_assist_fallback( - primary_call: Awaitable[Any], + primary_call: Callable[[], Awaitable[Any]], fallback_kwargs: Dict[str, Any], auto_fallback_to_google_code_assist: bool = False, ) -> Any: @@ -17,7 +17,7 @@ async def run_gemini_acompletion_with_code_assist_fallback( OAuth scope is insufficient. """ try: - return await primary_call + return await primary_call() except Exception as e: if not auto_fallback_to_google_code_assist: raise e diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index 7c4c7dba626..7b5dafea8b6 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -2,6 +2,7 @@ Transformation for Calling Google models in their native format. """ +import asyncio from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast import httpx @@ -11,6 +12,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) +from litellm.llms.gemini.common_utils import get_gemini_oauth_token from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.types.router import GenericLiteLLMParams @@ -200,6 +202,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): stream: bool, api_base: Optional[str], litellm_params: dict, + gemini_auth_data: Optional[dict] = None, ) -> Tuple[dict, str]: """ Build final headers and API URL from auth components. @@ -217,6 +220,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): custom_llm_provider=self.custom_llm_provider, api_base=api_base, should_use_v1beta1_features=True, + gemini_auth_data=gemini_auth_data, ) headers = self.validate_environment( @@ -291,6 +295,12 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): custom_llm_provider=self.custom_llm_provider, ) + gemini_auth_data: Optional[dict] = None + if self.custom_llm_provider == "gemini": + gemini_api_key = self._get_google_ai_studio_api_key(dict(litellm_params)) + if gemini_api_key is None: + gemini_auth_data = await asyncio.to_thread(get_gemini_oauth_token) + return self._build_final_headers_and_url( model=model, auth_header=_auth_header, @@ -300,6 +310,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): stream=stream, api_base=api_base, litellm_params=litellm_params, + gemini_auth_data=gemini_auth_data, ) def transform_generate_content_request( diff --git a/litellm/llms/google_code_assist/transformation.py b/litellm/llms/google_code_assist/transformation.py index 441af4e9988..6de5b608ea0 100644 --- a/litellm/llms/google_code_assist/transformation.py +++ b/litellm/llms/google_code_assist/transformation.py @@ -108,6 +108,16 @@ class GoogleCodeAssistConfig(VertexGeminiConfig): ) generation_config[key] = value + # Log unmapped/unrecognized caller keys so users have visibility (do not auto-forward — + # they may not be valid Gemini fields). + for key in optional_params: + if key not in base_params and key not in generation_config: + verbose_logger.debug( + "google_code_assist: dropping unrecognized optional_param '%s' " + "(not produced by VertexGeminiConfig.map_openai_params)", + key, + ) + vertex_request = { "contents": contents, "session_id": litellm_params.get("session_id", str(uuid.uuid4())), diff --git a/litellm/main.py b/litellm/main.py index 6b9c4f776d8..0734613f50c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3447,7 +3447,7 @@ def completion( # type: ignore # noqa: PLR0915 new_params = safe_deep_copy(optional_params or {}) if acompletion is True: response = run_gemini_acompletion_with_code_assist_fallback( - primary_call=vertex_chat_completion.completion( # type: ignore + primary_call=lambda: vertex_chat_completion.completion( # type: ignore model=model, messages=messages, model_response=model_response, diff --git a/tests/test_litellm/llms/gemini/test_fallback_handler.py b/tests/test_litellm/llms/gemini/test_fallback_handler.py index d70e9984ba9..74e2f144627 100644 --- a/tests/test_litellm/llms/gemini/test_fallback_handler.py +++ b/tests/test_litellm/llms/gemini/test_fallback_handler.py @@ -66,7 +66,7 @@ async def test_run_gemini_acompletion_with_code_assist_fallback_enabled(): ): mock_acompletion.return_value = "fallback-ok" result = await run_gemini_acompletion_with_code_assist_fallback( - primary_call=_raise_scope_error(), + primary_call=_raise_scope_error, fallback_kwargs={}, auto_fallback_to_google_code_assist=True, ) diff --git a/tests/test_litellm/llms/google_code_assist/__init__.py b/tests/test_litellm/llms/google_code_assist/__init__.py new file mode 100644 index 00000000000..e69de29bb2d