Merge pull request #40909 from BerriAI/litellm_databricks_reasoning_effort_thinking

fix(databricks): translate reasoning_effort to thinking for Gemini 2.5
This commit is contained in:
Mateo Wang 2026-09-12 17:46:44 -07:00 committed by GitHub
commit 9d984371fd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 180 additions and 4 deletions

View file

@ -1084,9 +1084,7 @@ jobs:
name: Run tests
command: |
mkdir -p test-results
TEST_FILES=$(printf "%s\n%s\n" \
"$(circleci tests glob "tests/ocr_tests/**/test_*.py")" \
"tests/test_litellm/ocr/test_rust_bridge.py")
TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \

View file

@ -272,6 +272,15 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
"thinking",
]
@staticmethod
def _uses_anthropic_thinking_param(model: str) -> bool:
from litellm.utils import supports_anthropic_thinking_payload
normalized: Final = model.lower().replace(".", "-")
return "claude" in normalized or supports_anthropic_thinking_payload(
model=normalized, custom_llm_provider="databricks"
)
def convert_anthropic_tool_to_databricks_tool(self, tool: AllAnthropicToolsValues | None) -> DatabricksTool | None:
if tool is None:
return None
@ -377,7 +386,7 @@ 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._uses_anthropic_thinking_param(model):
reasoning_effort_value: Final = non_default_params.get("reasoning_effort")
mapped_thinking: Final = AnthropicConfig._map_reasoning_effort(
reasoning_effort=reasoning_effort_value,

View file

@ -17599,6 +17599,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-fable-5": {
@ -17624,6 +17625,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": false,
@ -17653,6 +17655,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -17678,6 +17681,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 4096
},
@ -17701,6 +17705,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17724,6 +17729,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17747,6 +17753,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"supports_output_config": true,
"prompt_cache_min_tokens": 4096
@ -17772,6 +17779,7 @@
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 4096
},
@ -17797,6 +17805,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -17824,6 +17833,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -17851,6 +17861,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -17877,6 +17888,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4-1": {
@ -17899,6 +17911,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4-5": {
@ -17921,6 +17934,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17945,6 +17959,7 @@
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17971,6 +17986,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -18049,6 +18065,7 @@
"output_dbu_cost_per_token": 3.5714e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_anthropic_thinking_payload": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
},
@ -18069,6 +18086,7 @@
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_anthropic_thinking_payload": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
},

View file

@ -169,6 +169,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None]
supports_output_config: bool | None
supports_image_size: bool | None
supports_anthropic_thinking_payload: ReadOnly[bool | None]
supported_audio_formats: ReadOnly[Sequence[Literal["mp3", "wav"]] | None]
vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None]
bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None

View file

@ -2850,6 +2850,12 @@ def supports_reasoning(model: str, custom_llm_provider: str | None = None) -> bo
return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key="supports_reasoning")
def supports_anthropic_thinking_payload(model: str, custom_llm_provider: str | None = None) -> bool:
return _supports_factory(
model=model, custom_llm_provider=custom_llm_provider, key="supports_anthropic_thinking_payload"
)
def supports_none_reasoning_effort(model: str, custom_llm_provider: str | None = None) -> bool:
"""
Check if the given model accepts reasoning effort "none" and return a boolean value.
@ -5996,6 +6002,7 @@ def _get_model_info_helper(
thinking_always_on=_model_info.get("thinking_always_on", None),
supports_tool_search=_model_info.get("supports_tool_search", None),
supports_mid_conversation_system=_model_info.get("supports_mid_conversation_system", None),
supports_anthropic_thinking_payload=_model_info.get("supports_anthropic_thinking_payload", None),
supports_none_reasoning_effort=_model_info.get("supports_none_reasoning_effort", None),
supports_minimal_reasoning_effort=_model_info.get("supports_minimal_reasoning_effort", None),
supports_low_reasoning_effort=_model_info.get("supports_low_reasoning_effort", None),

View file

@ -17599,6 +17599,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-fable-5": {
@ -17624,6 +17625,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": false,
@ -17653,6 +17655,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true,
@ -17678,6 +17681,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 4096
},
@ -17701,6 +17705,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17724,6 +17729,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17747,6 +17753,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"supports_output_config": true,
"prompt_cache_min_tokens": 4096
@ -17772,6 +17779,7 @@
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 4096
},
@ -17797,6 +17805,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -17824,6 +17833,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -17851,6 +17861,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -17877,6 +17888,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4-1": {
@ -17899,6 +17911,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4-5": {
@ -17921,6 +17934,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17945,6 +17959,7 @@
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_tool_choice": true,
"prompt_cache_min_tokens": 1024
},
@ -17971,6 +17986,7 @@
"supports_mid_conversation_system": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_anthropic_thinking_payload": true,
"supports_sampling_params": false,
"supports_tool_choice": true,
"supports_vision": true
@ -18049,6 +18065,7 @@
"output_dbu_cost_per_token": 3.5714e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_anthropic_thinking_payload": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
},
@ -18069,6 +18086,7 @@
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_anthropic_thinking_payload": true,
"supports_prompt_caching": true,
"supports_tool_choice": true
},

View file

@ -682,6 +682,9 @@
"supports_adaptive_thinking": {
"type": "boolean"
},
"supports_anthropic_thinking_payload": {
"type": "boolean"
},
"supports_assistant_prefill": {
"type": "boolean"
},

View file

@ -5,6 +5,12 @@ from fastapi.testclient import TestClient
from unittest.mock import MagicMock, patch
import litellm
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,
@ -12,6 +18,12 @@ from litellm.llms.databricks.chat.transformation import (
)
@pytest.fixture()
def _use_local_model_cost_map(monkeypatch):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
def test_transform_choices():
config = DatabricksConfig()
databricks_choices = [
@ -532,6 +544,115 @@ def test_map_openai_params_upgrades_legacy_thinking_on_adaptive_only_claude(
assert mapped.get("output_config") == expected_output_config
def _map_reasoning_effort(model: str, reasoning_effort: str):
return DatabricksConfig().map_openai_params(
non_default_params={"reasoning_effort": reasoning_effort},
optional_params={},
model=model,
drop_params=False,
)
def test_claude_translates_reasoning_effort_to_thinking(_use_local_model_cost_map):
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_adaptive_claude_translates_reasoning_effort_to_output_config(_use_local_model_cost_map):
params = _map_reasoning_effort("databricks-claude-opus-4-7", "high")
assert params.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert params.get("output_config") == {"effort": "high"}
assert "reasoning_effort" not in params
def test_unmapped_claude_endpoint_still_translates(_use_local_model_cost_map):
params = _map_reasoning_effort("my-claude-serving-endpoint", "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(_use_local_model_cost_map):
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(_use_local_model_cost_map):
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(_use_local_model_cost_map):
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(_use_local_model_cost_map):
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_with_dot_notation_translates(_use_local_model_cost_map):
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(_use_local_model_cost_map):
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(_use_local_model_cost_map):
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(_use_local_model_cost_map):
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(_use_local_model_cost_map):
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(_use_local_model_cost_map):
params = _map_reasoning_effort("databricks-gpt-oss-120b", "high")
assert params.get("reasoning_effort") == "high"
assert "thinking" not in params
def _streaming_chunk(usage=None, choices=None):
base = {
"id": "chatcmpl-test",

View file

@ -1150,6 +1150,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"supports_adaptive_thinking": {"type": "boolean"},
"supports_anthropic_thinking_payload": {"type": "boolean"},
"supports_legacy_thinking": {"type": "boolean"},
"thinking_always_on": {"type": "boolean"},
"supports_mid_conversation_system": {"type": "boolean"},