From 22e86cde2afd4718436b0c6d17f050e7ee73bbc4 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 16 Dec 2025 17:52:12 +0530 Subject: [PATCH] Add support for gemini 3 flash via v1/messages endpoint --- .../vertex_and_google_ai_studio_gemini.py | 36 +++- tests/llm_translation/test_gemini.py | 172 ++++++++++++++++++ 2 files changed, 201 insertions(+), 7 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 667bf55d4e..b13afa2bfd 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -775,17 +775,38 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _map_thinking_param( thinking_param: AnthropicThinkingParam, + model: Optional[str] = None, ) -> GeminiThinkingConfig: thinking_enabled = thinking_param.get("type") == "enabled" thinking_budget = thinking_param.get("budget_tokens") params: GeminiThinkingConfig = {} - if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero( - thinking_budget - ): - params["includeThoughts"] = True - if thinking_budget is not None and isinstance(thinking_budget, int): - params["thinkingBudget"] = thinking_budget + + # For Gemini 3+ models, use thinkingLevel instead of thinkingBudget + if model and VertexGeminiConfig._is_gemini_3_or_newer(model): + if thinking_enabled: + if thinking_budget is None or thinking_budget == 0: + params["includeThoughts"] = False + else: + params["includeThoughts"] = True + if thinking_budget >= 10000: + is_fiercefalcon = "fiercefalcon" in model.lower() or "gemini-3-flash" in model.lower() + params["thinkingLevel"] = "minimal" if is_fiercefalcon else "low" + else: + is_fiercefalcon = "fiercefalcon" in model.lower() or "gemini-3-flash" in model.lower() + params["thinkingLevel"] = "minimal" if is_fiercefalcon else "low" + else: + # Thinking disabled + params["includeThoughts"] = False + else: + # For older Gemini models, use thinkingBudget + if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero( + thinking_budget + ): + params["includeThoughts"] = True + if thinking_budget is not None and isinstance(thinking_budget, int): + params["thinkingBudget"] = thinking_budget + return params def map_response_modalities(self, value: list) -> list: @@ -962,7 +983,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params[ "thinkingConfig" ] = VertexGeminiConfig._map_thinking_param( - cast(AnthropicThinkingParam, value) + cast(AnthropicThinkingParam, value), + model=model, ) elif param == "modalities" and isinstance(value, list): response_modalities = self.map_response_modalities(value) diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index dbbf0d31f1..d9ec7ff915 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -1229,3 +1229,175 @@ def test_gemini_function_args_preserve_unicode(): assert parsed_args["recipient"] == "José" assert "\\u" not in arguments_str assert "José" in arguments_str + + +def test_anthropic_thinking_param_to_gemini_3_thinkingLevel(): + """ + Test that Anthropic thinking parameters are correctly transformed to Gemini 3 thinkingLevel + instead of thinkingBudget. + + For Gemini 3+ models (gemini-3-flash, gemini-3-pro, fiercefalcon): + - Should use thinkingLevel instead of thinkingBudget + - budget_tokens should map to thinkingLevel + + Related issue: https://github.com/BerriAI/litellm/issues/XXXX + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.anthropic import AnthropicThinkingParam + + # Test 1: Anthropic thinking enabled with budget_tokens for Gemini 3 model + thinking_param: AnthropicThinkingParam = { + "type": "enabled", + "budget_tokens": 10000, + } + + result = VertexGeminiConfig._map_thinking_param( + thinking_param=thinking_param, + model="gemini-3-flash", + ) + + # For Gemini 3, should use thinkingLevel, not thinkingBudget + assert "thinkingLevel" in result, "Should have thinkingLevel for Gemini 3" + assert "thinkingBudget" not in result, "Should NOT have thinkingBudget for Gemini 3" + assert result["includeThoughts"] is True + assert result["thinkingLevel"] in ["minimal", "low"], "thinkingLevel should be 'minimal' or 'low'" + + # Test 2: Anthropic thinking disabled for Gemini 3 + thinking_param_disabled: AnthropicThinkingParam = { + "type": "disabled", + "budget_tokens": None, + } + + result_disabled = VertexGeminiConfig._map_thinking_param( + thinking_param=thinking_param_disabled, + model="gemini-3-pro-preview", + ) + + assert result_disabled.get("includeThoughts") is False + assert "thinkingLevel" not in result_disabled or result_disabled.get("thinkingLevel") is None + + # Test 3: Budget tokens = 0 for Gemini 3 + thinking_param_zero: AnthropicThinkingParam = { + "type": "enabled", + "budget_tokens": 0, + } + + result_zero = VertexGeminiConfig._map_thinking_param( + thinking_param=thinking_param_zero, + model="gemini-3-flash", + ) + + assert result_zero["includeThoughts"] is False + assert "thinkingLevel" not in result_zero or result_zero.get("thinkingLevel") is None + + # Test 4: Fiercefalcon model (Gemini 3 Flash checkpoint) should use thinkingLevel + result_fiercefalcon = VertexGeminiConfig._map_thinking_param( + thinking_param=thinking_param, + model="fiercefalcon", + ) + + assert "thinkingLevel" in result_fiercefalcon, "Should have thinkingLevel for fiercefalcon" + assert "thinkingBudget" not in result_fiercefalcon, "Should NOT have thinkingBudget for fiercefalcon" + assert result_fiercefalcon["includeThoughts"] is True + + +def test_anthropic_thinking_param_to_gemini_2_thinkingBudget(): + """ + Test that Anthropic thinking parameters are correctly transformed to Gemini 2 thinkingBudget + (not thinkingLevel). + + For Gemini 2.x models (gemini-2.5-flash, gemini-2.0-flash): + - Should continue using thinkingBudget + - thinkingLevel should NOT be used + + Related issue: https://github.com/BerriAI/litellm/issues/XXXX + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.anthropic import AnthropicThinkingParam + + # Test 1: Anthropic thinking enabled with budget_tokens for Gemini 2 model + thinking_param: AnthropicThinkingParam = { + "type": "enabled", + "budget_tokens": 10000, + } + + result = VertexGeminiConfig._map_thinking_param( + thinking_param=thinking_param, + model="gemini-2.5-flash", + ) + + # For Gemini 2, should use thinkingBudget, not thinkingLevel + assert "thinkingBudget" in result, "Should have thinkingBudget for Gemini 2" + assert "thinkingLevel" not in result, "Should NOT have thinkingLevel for Gemini 2" + assert result["includeThoughts"] is True + assert result["thinkingBudget"] == 10000 + + # Test 2: Anthropic thinking enabled for gemini-2.0-flash model + result_gemini2 = VertexGeminiConfig._map_thinking_param( + thinking_param=thinking_param, + model="gemini-2.0-flash-thinking-exp-01-21", + ) + + assert "thinkingBudget" in result_gemini2, "Should have thinkingBudget for Gemini 2" + assert "thinkingLevel" not in result_gemini2, "Should NOT have thinkingLevel for Gemini 2" + assert result_gemini2["includeThoughts"] is True + assert result_gemini2["thinkingBudget"] == 10000 + + +def test_anthropic_thinking_param_via_map_openai_params(): + """ + Test that the thinking parameter is correctly transformed through the full map_openai_params flow + for Gemini 3 models, resulting in thinkingConfig with thinkingLevel. + + This tests the full integration from Anthropic API format to Gemini format. + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.types.llms.anthropic import AnthropicThinkingParam + + config = VertexGeminiConfig() + + # Test with Gemini 3 model + non_default_params = { + "thinking": { + "type": "enabled", + "budget_tokens": 10000, + } + } + optional_params: dict = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gemini-3-flash", + drop_params=False, + ) + + # Check that thinkingConfig was created with thinkingLevel + assert "thinkingConfig" in result, "Should have thinkingConfig in optional_params" + thinking_config = result["thinkingConfig"] + assert "thinkingLevel" in thinking_config, "Should have thinkingLevel for Gemini 3" + assert "thinkingBudget" not in thinking_config, "Should NOT have thinkingBudget for Gemini 3" + assert thinking_config["includeThoughts"] is True + + # Test with Gemini 2 model + optional_params_2 = {} + result_2 = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params_2, + model="gemini-2.5-flash", + drop_params=False, + ) + + # Check that thinkingConfig was created with thinkingBudget + assert "thinkingConfig" in result_2, "Should have thinkingConfig in optional_params" + thinking_config_2 = result_2["thinkingConfig"] + assert "thinkingBudget" in thinking_config_2, "Should have thinkingBudget for Gemini 2" + assert "thinkingLevel" not in thinking_config_2, "Should NOT have thinkingLevel for Gemini 2" + assert thinking_config_2["includeThoughts"] is True + assert thinking_config_2["thinkingBudget"] == 10000