From 443bee300294796081a201bc1da57833ed77a64b Mon Sep 17 00:00:00 2001 From: cohml <62400541+cohml@users.noreply.github.com> Date: Tue, 23 Jun 2026 11:04:30 -0400 Subject: [PATCH] fix(databricks): translate reasoning_effort to thinking for Gemini 2.5 models. --- .../llms/databricks/chat/transformation.py | 25 +++- .../test_databricks_chat_transformation.py | 116 ++++++++++++++++++ 2 files changed, 139 insertions(+), 2 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index ba8c312ea51..a79dfe8a59b 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -264,6 +264,22 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): "thinking", ] + @staticmethod + def _databricks_model_uses_anthropic_thinking_param(model: str) -> bool: + """ + Per Databricks docs, Claude and Gemini 2.5 endpoints accept the + Anthropic-style `thinking={"type":"enabled","budget_tokens":N}` payload + and do NOT accept OpenAI's top-level `reasoning_effort`. Gemini 3+ and + GPT-5/GPT-OSS accept `reasoning_effort` natively and need no + translation. + """ + model_lower = model.lower() + if "claude" in model_lower: + return True + if "gemini-2" in model_lower: + return True + return False + def convert_anthropic_tool_to_databricks_tool( self, tool: Optional[AllAnthropicToolsValues] ) -> Optional[DatabricksTool]: @@ -367,19 +383,24 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): "response_format", None ) # unsupported for claude models - if json_schema -> convert to tool call - if "reasoning_effort" in non_default_params and "claude" in model: + if ( + "reasoning_effort" in non_default_params + and self._databricks_model_uses_anthropic_thinking_param(model) + ): reasoning_effort_value = non_default_params.get("reasoning_effort") mapped_thinking = AnthropicConfig._map_reasoning_effort( reasoning_effort=reasoning_effort_value, model=model, llm_provider="databricks", ) + is_claude = "claude" in model.lower() if mapped_thinking is None: optional_params.pop("thinking", None) optional_params.pop("output_config", None) else: optional_params["thinking"] = mapped_thinking - if AnthropicConfig._is_adaptive_thinking_model(model): + # output_config + adaptive thinking is an Anthropic-only feature. + if is_claude and AnthropicConfig._is_adaptive_thinking_model(model): mapped_effort: Optional[str] = None if isinstance(reasoning_effort_value, str): mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort_value) 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 cfdb76a97f4..e1a126298d7 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 @@ -8,6 +8,11 @@ from fastapi.testclient import TestClient sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path from unittest.mock import MagicMock, patch +from litellm.constants import ( + DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, +) from litellm.llms.databricks.chat.transformation import ( DatabricksChatResponseIterator, DatabricksConfig, @@ -416,3 +421,114 @@ def test_transform_request_keeps_parallel_tool_calls_for_claude(): )["messages"] assert len([m for m in result if m.get("role") == "assistant"]) == 1 + + +# --------------------------------------------------------------------------- +# reasoning_effort translation +# +# Databricks foundation-model endpoints take reasoning controls via different +# payload shapes depending on the underlying model family: +# +# Claude: Anthropic-style `thinking={"type":"enabled","budget_tokens":N}` +# Gemini 2.5: Same Anthropic-style `thinking` payload as Claude +# (per docs.databricks.com/.../query-reason-models) +# Gemini 3+: Native OpenAI-style top-level `reasoning_effort` +# GPT-5/GPT-OSS: Native OpenAI-style top-level `reasoning_effort` +# +# LiteLLM should translate `reasoning_effort` into the right shape for the +# first two families and pass it through unchanged for the latter two. +# --------------------------------------------------------------------------- + + +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} + non_default.update(extra_non_default) + return DatabricksConfig().map_openai_params( + non_default_params=non_default, + optional_params={}, + model=model, + drop_params=False, + ) + + +def test_claude_translates_reasoning_effort_to_thinking(): + """Regression: Claude path must still translate to Anthropic-style thinking.""" + params = _map_reasoning_effort("databricks-claude-3-7-sonnet", "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_5_low_translates_to_thinking_budget(): + 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_5_medium_translates_to_thinking_budget(): + params = _map_reasoning_effort("databricks-gemini-2-5-flash", "medium") + assert params.get("thinking") == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + } + assert "reasoning_effort" not in params + + +def test_gemini_2_5_high_translates_to_thinking_budget(): + params = _map_reasoning_effort("databricks-gemini-2-5-flash", "high") + assert params.get("thinking") == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + } + assert "reasoning_effort" not in params + + +def test_gemini_2_5_pro_translates_to_thinking_budget(): + """Cover the gemini-2-5-pro endpoint too, not just flash.""" + params = _map_reasoning_effort("databricks-gemini-2-5-pro", "high") + assert params.get("thinking") == { + "type": "enabled", + "budget_tokens": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, + } + assert "reasoning_effort" not in params + + +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") + assert "thinking" not in params + assert "reasoning_effort" not in params + + +def test_gemini_3_passes_reasoning_effort_through(): + """Databricks-Gemini-3+ accepts reasoning_effort natively — do not translate.""" + params = _map_reasoning_effort("databricks-gemini-3-1-pro", "low") + assert params.get("reasoning_effort") == "low" + assert "thinking" not in params + + +def test_gpt_5_passes_reasoning_effort_through(): + """Databricks-GPT-5 family accepts reasoning_effort natively.""" + params = _map_reasoning_effort("databricks-gpt-5-1", "low") + assert params.get("reasoning_effort") == "low" + assert "thinking" not in params + + +def test_gpt_oss_passes_reasoning_effort_through(): + """Databricks-GPT-OSS accepts reasoning_effort natively.""" + params = _map_reasoning_effort("databricks-gpt-oss-120b", "high") + assert params.get("reasoning_effort") == "high" + assert "thinking" not in params