Fix unsupported thinking_budget=0 for Gemini Pro

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -922,6 +922,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
     def _is_thinking_budget_zero(thinking_budget: Optional[int]) -> bool:
         return thinking_budget is not None and thinking_budget == 0
 
+    @staticmethod
+    def _model_supports_thinking_budget_zero(model: Optional[str]) -> bool:
+        """Only Flash-family Gemini 2.x models accept thinkingBudget=0 to disable thinking.
+        Pro models reject it with a 400 INVALID_ARGUMENT error."""
+        if model is None:
+            return False
+        return "flash" in model.lower()
+
     @staticmethod
     def _validate_thinking_config_conflicts(
         optional_params: Dict,
@@ -1005,7 +1013,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
             ):
                 params["includeThoughts"] = True
             if thinking_budget is not None and isinstance(thinking_budget, int):
-                params["thinkingBudget"] = thinking_budget
+                if (
+                    thinking_budget > 0
+                    or VertexGeminiConfig._model_supports_thinking_budget_zero(model)
+                ):
+                    params["thinkingBudget"] = thinking_budget
 
         return params
 

tests/llm_translation/test_gemini.py

@@ -1610,3 +1610,97 @@ async def test_gemini_openai_web_search_tool_to_google_search():
     print("response: ", response.model_dump_json(indent=4))
     assert hasattr(response, "vertex_ai_grounding_metadata")
     assert getattr(response, "vertex_ai_grounding_metadata") is not None
+
+
+@pytest.mark.asyncio
+async def test_gemini_flash_to_pro_fallback_thinking_budget_zero():
+    """
+    Regression test: when gemini-2.5-flash is rate-limited and falls back to
+    gemini-2.5-pro, thinking={"type": "enabled", "budget_tokens": 0} must NOT
+    produce thinkingBudget=0 in the Pro request body. Pro rejects that with a
+    400 INVALID_ARGUMENT ("The model does not support setting thinking_budget to 0").
+
+    Fix: _map_thinking_param skips thinkingBudget=0 for non-Flash models.
+    Flash-family models support disabling thinking via thinkingBudget=0; Pro does not.
+    """
+    from unittest.mock import patch
+    from litellm import Router
+    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+        VertexGeminiConfig,
+    )
+
+    # Translation-layer check: Pro must not receive thinkingBudget=0
+    pro_config = VertexGeminiConfig._map_thinking_param(
+        thinking_param={"type": "enabled", "budget_tokens": 0},
+        model="gemini-2.5-pro",
+    )
+    assert "thinkingBudget" not in pro_config, (
+        "thinkingBudget=0 must not be sent to Pro — it causes a 400 INVALID_ARGUMENT"
+    )
+
+    # Flash must still receive thinkingBudget=0 (it supports disabling thinking)
+    flash_config = VertexGeminiConfig._map_thinking_param(
+        thinking_param={"type": "enabled", "budget_tokens": 0},
+        model="gemini-2.5-flash",
+    )
+    assert flash_config.get("thinkingBudget") == 0
+
+    # Router-layer check: fallback from flash (429) to pro succeeds without 400
+    router = Router(
+        model_list=[
+            {
+                "model_name": "gemini-2.5-flash",
+                "litellm_params": {
+                    "model": "gemini/gemini-2.5-flash",
+                    "api_key": "fake-key",
+                },
+            },
+            {
+                "model_name": "gemini-2.5-pro",
+                "litellm_params": {
+                    "model": "gemini/gemini-2.5-pro",
+                    "api_key": "fake-key",
+                },
+            },
+        ],
+        fallbacks=[{"gemini-2.5-flash": ["gemini-2.5-pro"]}],
+        num_retries=0,
+    )
+
+    call_count = 0
+    pro_call_kwargs: dict = {}
+
+    async def mock_acompletion(**kwargs):
+        nonlocal call_count, pro_call_kwargs
+        call_count += 1
+        model = kwargs.get("model", "")
+        if "flash" in model:
+            raise litellm.RateLimitError(
+                message='vertex_aiException - {"error": {"code": 429, "message": "Resource exhausted.", "status": "RESOURCE_EXHAUSTED"}}',
+                llm_provider="vertex_ai",
+                model="gemini/gemini-2.5-flash",
+            )
+        pro_call_kwargs = kwargs.copy()
+        return litellm.ModelResponse(
+            id="chatcmpl-pro-mock",
+            choices=[
+                litellm.Choices(
+                    finish_reason="stop",
+                    index=0,
+                    message=litellm.Message(role="assistant", content="Hello!"),
+                )
+            ],
+            model="gemini-2.5-pro",
+        )
+
+    with patch("litellm.acompletion", side_effect=mock_acompletion):
+        response = await router.acompletion(
+            model="gemini-2.5-flash",
+            messages=[{"role": "user", "content": "Hello"}],
+            thinking={"type": "enabled", "budget_tokens": 0},
+        )
+
+    assert call_count == 2, f"Expected 2 calls (flash then pro fallback), got {call_count}"
+    assert response.choices[0].message.content == "Hello!"
+    # The thinking param reaches pro but thinkingBudget=0 is not in the translated body
+    assert pro_call_kwargs.get("thinking") == {"type": "enabled", "budget_tokens": 0}