From 0338b36392129944bd8c6e2295abeba0eed97e1e Mon Sep 17 00:00:00 2001 From: Felipe Musse Date: Tue, 28 Apr 2026 14:23:06 +0100 Subject: [PATCH] Fix test --- litellm/utils.py | 3 + tests/llm_translation/test_gemini.py | 80 ++++++++++++------- .../llms/gemini/test_gemini_thinking_param.py | 32 +++++++- 3 files changed, 86 insertions(+), 29 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 1df751f915f..fd7595d2835 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5916,6 +5916,9 @@ def _get_model_info_helper( # noqa: PLR0915 supports_max_reasoning_effort=_model_info.get( "supports_max_reasoning_effort", None ), + supports_thinking_budget_zero=_model_info.get( + "supports_thinking_budget_zero", None + ), supports_computer_use=_model_info.get("supports_computer_use", None), search_context_cost_per_query=_model_info.get( "search_context_cost_per_query", None diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index 13ca1c87cbe..c4362f46866 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -454,23 +454,36 @@ def test_gemini_thinking(): def test_gemini_thinking_budget_0(): litellm._turn_on_debug() + from unittest.mock import patch from litellm.types.utils import Message, CallTypes from litellm.utils import return_raw_request import json - raw_request = return_raw_request( - endpoint=CallTypes.completion, - kwargs={ - "model": "gemini/gemini-2.5-flash", - "messages": [ - { - "role": "user", - "content": "Explain the concept of Occam's Razor and provide a simple, everyday example", - } - ], - "thinking": {"type": "enabled", "budget_tokens": 0}, - }, - ) + # Inject supports_thinking_budget_zero so the translation layer emits + # thinkingBudget=0 for Flash regardless of whether the remote JSON has the + # flag yet (CI fetches GitHub main which lags behind this PR). + # _map_thinking_param receives the stripped model name ("gemini-2.5-flash"), + # which resolves to key "gemini-2.5-flash" in model_cost (vertex_ai provider). + flash_cost_patch = { + "gemini-2.5-flash": { + **litellm.model_cost.get("gemini-2.5-flash", {}), + "supports_thinking_budget_zero": True, + } + } + with patch.dict(litellm.model_cost, flash_cost_patch): + raw_request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": "gemini/gemini-2.5-flash", + "messages": [ + { + "role": "user", + "content": "Explain the concept of Occam's Razor and provide a simple, everyday example", + } + ], + "thinking": {"type": "enabled", "budget_tokens": 0}, + }, + ) print(json.dumps(raw_request, indent=4, default=str)) assert "0" in json.dumps(raw_request["raw_request_body"]) @@ -1629,21 +1642,34 @@ async def test_gemini_flash_to_pro_fallback_thinking_budget_zero(): 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" - ) + # Patch model_cost to declare Flash support for thinkingBudget=0. + # In CI, litellm fetches the remote JSON from GitHub main which won't have + # this flag until after the PR merges, so we inject it here explicitly. + # _map_thinking_param receives the stripped model name ("gemini-2.5-flash"), + # which resolves to key "gemini-2.5-flash" in model_cost (vertex_ai provider). + flash_cost_patch = { + "gemini-2.5-flash": { + **litellm.model_cost.get("gemini-2.5-flash", {}), + "supports_thinking_budget_zero": True, + } + } - # 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 + with patch.dict(litellm.model_cost, flash_cost_patch): + # 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( diff --git a/tests/test_litellm/llms/gemini/test_gemini_thinking_param.py b/tests/test_litellm/llms/gemini/test_gemini_thinking_param.py index 97eae8a35a2..06b52ab9d94 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_thinking_param.py +++ b/tests/test_litellm/llms/gemini/test_gemini_thinking_param.py @@ -1,9 +1,37 @@ +from unittest.mock import patch + import pytest +import litellm from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) +# Keys that _get_model_info_helper resolves to for Flash models when the remote +# model_cost JSON is used (which won't have the flag until after this PR merges). +_FLASH_MODEL_COST_KEYS = [ + "gemini-2.5-flash", + "gemini/gemini-2.5-flash", + "gemini-2.5-flash-lite", + "gemini/gemini-2.5-flash-lite", +] + + +@pytest.fixture() +def flash_model_cost_flags(): + """Inject supports_thinking_budget_zero=True into model_cost for Flash models. + + In CI, litellm fetches the remote model_cost JSON from GitHub main, which + won't have this flag until after the PR merges. Patching model_cost here + makes tests independent of the remote file state. + """ + patch_entries = { + key: {**litellm.model_cost.get(key, {}), "supports_thinking_budget_zero": True} + for key in _FLASH_MODEL_COST_KEYS + } + with patch.dict(litellm.model_cost, patch_entries): + yield + class TestMapThinkingParamBudgetZero: """ @@ -22,7 +50,7 @@ class TestMapThinkingParamBudgetZero: ) assert "thinkingBudget" not in result - def test_flash_budget_zero_emits_thinking_budget(self): + def test_flash_budget_zero_emits_thinking_budget(self, flash_model_cost_flags): result = VertexGeminiConfig._map_thinking_param( thinking_param={"type": "enabled", "budget_tokens": 0}, model="gemini-2.5-flash", @@ -45,7 +73,7 @@ class TestMapThinkingParamBudgetZero: assert result.get("thinkingBudget") == 1024 assert result.get("includeThoughts") is True - def test_model_supports_thinking_budget_zero_flash(self): + def test_model_supports_thinking_budget_zero_flash(self, flash_model_cost_flags): assert VertexGeminiConfig._model_supports_thinking_budget_zero("gemini-2.5-flash") is True assert VertexGeminiConfig._model_supports_thinking_budget_zero("gemini/gemini-2.5-flash") is True assert VertexGeminiConfig._model_supports_thinking_budget_zero("gemini-2.5-flash-lite") is True