Fix test

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(

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