From d7b793c054a184ee9d7784559d959ba124ffe0d5 Mon Sep 17 00:00:00 2001 From: cohml <62400541+cohml@users.noreply.github.com> Date: Tue, 23 Jun 2026 12:02:20 -0400 Subject: [PATCH] fix(databricks,base_llm): narrow gemini-2.5 match and guard thinking-tokens helper. --- litellm/llms/base_llm/chat/transformation.py | 11 +++++-- .../llms/databricks/chat/transformation.py | 5 +++- .../test_databricks_chat_transformation.py | 30 +++++++++++++------ 3 files changed, 34 insertions(+), 12 deletions(-) diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index ab901a467e8..693ef38bd29 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -14,7 +14,6 @@ from typing import ( Tuple, Type, Union, - cast, ) import httpx @@ -124,7 +123,15 @@ class BaseConfig(ABC): if is_thinking_enabled and ( "max_tokens" not in non_default_params and "max_completion_tokens" not in non_default_params ): - thinking_token_budget = cast(dict, optional_params["thinking"]).get("budget_tokens", None) + # `is_thinking_enabled` is True when `reasoning_effort` is set OR + # when `thinking` is set. Providers that pass `reasoning_effort` + # through natively (e.g. Databricks-Gemini-3+, Databricks-GPT-5) + # never populate an Anthropic-style `thinking` block, so guard + # against that case here. + thinking = optional_params.get("thinking") + if not isinstance(thinking, dict): + return + thinking_token_budget = thinking.get("budget_tokens", None) if thinking_token_budget is not None: optional_params["max_tokens"] = thinking_token_budget + DEFAULT_MAX_TOKENS diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index a79dfe8a59b..e6e63392e9c 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -276,7 +276,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): model_lower = model.lower() if "claude" in model_lower: return True - if "gemini-2" in model_lower: + # Match gemini-2-5 / gemini-2.5 only — not the broader 2.x range, which + # could catch hypothetical future 2.0/2.6/etc variants that may use a + # different reasoning contract. + if "gemini-2-5" in model_lower or "gemini-2.5" in model_lower: return True return False diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index e1a126298d7..02936e740f7 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -441,15 +441,8 @@ def test_transform_request_keeps_parallel_tool_calls_for_claude(): def _map_reasoning_effort(model: str, reasoning_effort, **extra_non_default): - """Run map_openai_params with reasoning_effort + optional extras. - - `max_tokens` is included by default to mirror real client behavior. Without - it the base-class `update_optional_params_with_thinking_tokens` helper - KeyErrors on pure pass-through models (a pre-existing issue orthogonal to - this fix — `is_thinking_enabled` returns True whenever `reasoning_effort` is - set, but the helper then assumes `optional_params["thinking"]` exists). - """ - non_default = {"reasoning_effort": reasoning_effort, "max_tokens": 1024} + """Run map_openai_params with reasoning_effort + optional extras.""" + non_default = {"reasoning_effort": reasoning_effort} non_default.update(extra_non_default) return DatabricksConfig().map_openai_params( non_default_params=non_default, @@ -506,6 +499,25 @@ def test_gemini_2_5_pro_translates_to_thinking_budget(): assert "reasoning_effort" not in params +def test_gemini_2_5_with_dot_notation_translates(): + """A user passing the upstream Google-style `gemini-2.5-...` form should + still trigger the Anthropic-thinking translation, not pass through.""" + params = _map_reasoning_effort("databricks-gemini-2.5-flash", "low") + assert params.get("thinking") == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, + } + assert "reasoning_effort" not in params + + +def test_gemini_2_0_does_not_match(): + """Guard against over-matching: `gemini-2-0` (hypothetical or future) is + NOT a Gemini 2.5 endpoint and must not get the thinking translation.""" + params = _map_reasoning_effort("databricks-gemini-2-0-flash", "low") + assert "thinking" not in params + assert params.get("reasoning_effort") == "low" + + def test_gemini_2_5_none_drops_thinking_and_reasoning_effort(): """`reasoning_effort='none'` mirrors the Claude behavior: no thinking emitted.""" params = _map_reasoning_effort("databricks-gemini-2-5-flash", "none")