From e64036f8100ac4a5fb308edefe9e8e423a9ffe4e Mon Sep 17 00:00:00 2001 From: Felipe Musse Date: Tue, 28 Apr 2026 14:48:15 +0100 Subject: [PATCH] Fix test --- tests/llm_translation/test_gemini.py | 88 ++++++++----------- .../llms/gemini/test_gemini_thinking_param.py | 33 ++----- 2 files changed, 43 insertions(+), 78 deletions(-) diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index c4362f46866..d4c791ffbfb 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -454,36 +454,27 @@ 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 + import os + from litellm.types.utils import CallTypes from litellm.utils import return_raw_request import json - # 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}, - }, - ) + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + 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"]) @@ -1636,40 +1627,31 @@ async def test_gemini_flash_to_pro_fallback_thinking_budget_zero(): Fix: _map_thinking_param skips thinkingBudget=0 for non-Flash models. Flash-family models support disabling thinking via thinkingBudget=0; Pro does not. """ + import os from unittest.mock import patch from litellm import Router from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) - # 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, - } - } + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") - 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" - ) + # 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 + # 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 06b52ab9d94..72497b4891a 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_thinking_param.py +++ b/tests/test_litellm/llms/gemini/test_gemini_thinking_param.py @@ -1,4 +1,4 @@ -from unittest.mock import patch +import os import pytest @@ -7,30 +7,13 @@ 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 +def local_model_cost_map(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + yield + del os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] class TestMapThinkingParamBudgetZero: @@ -50,7 +33,7 @@ class TestMapThinkingParamBudgetZero: ) assert "thinkingBudget" not in result - def test_flash_budget_zero_emits_thinking_budget(self, flash_model_cost_flags): + def test_flash_budget_zero_emits_thinking_budget(self, local_model_cost_map): result = VertexGeminiConfig._map_thinking_param( thinking_param={"type": "enabled", "budget_tokens": 0}, model="gemini-2.5-flash", @@ -73,7 +56,7 @@ class TestMapThinkingParamBudgetZero: assert result.get("thinkingBudget") == 1024 assert result.get("includeThoughts") is True - def test_model_supports_thinking_budget_zero_flash(self, flash_model_cost_flags): + def test_model_supports_thinking_budget_zero_flash(self, local_model_cost_map): 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