mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
Merge PR #32786 (databricks gemini reasoning_effort) into internal_staging, resolving conflicts
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
f5e178634c
7 changed files with 191 additions and 4 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -268,6 +268,24 @@ 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.
|
||||
"""
|
||||
from litellm.utils import _supports_factory
|
||||
|
||||
normalized = model.lower().replace(".", "-")
|
||||
return _supports_factory(
|
||||
model=normalized,
|
||||
custom_llm_provider="databricks",
|
||||
key="supports_anthropic_thinking_payload",
|
||||
)
|
||||
|
||||
def convert_anthropic_tool_to_databricks_tool(
|
||||
self, tool: Optional[AllAnthropicToolsValues]
|
||||
) -> Optional[DatabricksTool]:
|
||||
|
|
@ -371,7 +389,10 @@ 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,
|
||||
|
|
@ -379,12 +400,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
custom_llm_provider="databricks",
|
||||
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, "databricks"):
|
||||
# output_config + adaptive thinking is an Anthropic-only feature.
|
||||
if is_claude and AnthropicConfig._is_adaptive_thinking_model(model, "databricks"):
|
||||
mapped_effort: Optional[str] = None
|
||||
if isinstance(reasoning_effort_value, str):
|
||||
mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort_value)
|
||||
|
|
|
|||
|
|
@ -13339,6 +13339,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-haiku-4-5": {
|
||||
|
|
@ -13358,6 +13359,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-opus-4": {
|
||||
|
|
@ -13377,6 +13379,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-opus-4-1": {
|
||||
|
|
@ -13396,6 +13399,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-opus-4-5": {
|
||||
|
|
@ -13415,6 +13419,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_output_config": true
|
||||
},
|
||||
|
|
@ -13435,6 +13440,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-sonnet-4-1": {
|
||||
|
|
@ -13454,6 +13460,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-sonnet-4-5": {
|
||||
|
|
@ -13473,6 +13480,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-gemini-2-5-flash": {
|
||||
|
|
@ -13490,6 +13498,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_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-gemini-2-5-pro": {
|
||||
|
|
@ -13507,6 +13516,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_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-gemma-3-12b": {
|
||||
|
|
|
|||
|
|
@ -152,6 +152,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
|
|||
supports_max_reasoning_effort: Optional[bool]
|
||||
supports_output_config: Optional[bool]
|
||||
supports_image_size: Optional[bool]
|
||||
supports_anthropic_thinking_payload: Optional[bool]
|
||||
bedrock_output_config_effort_ceiling: Optional[Literal["low", "medium", "high", "max", "xhigh"]]
|
||||
bedrock_converse_supports_strict_tools: Optional[bool]
|
||||
|
||||
|
|
|
|||
|
|
@ -5513,6 +5513,7 @@ def _get_model_info_helper(
|
|||
supports_reasoning=_model_info.get("supports_reasoning", None),
|
||||
supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", 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),
|
||||
|
|
|
|||
|
|
@ -13339,6 +13339,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-haiku-4-5": {
|
||||
|
|
@ -13358,6 +13359,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-opus-4": {
|
||||
|
|
@ -13377,6 +13379,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-opus-4-1": {
|
||||
|
|
@ -13396,6 +13399,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-opus-4-5": {
|
||||
|
|
@ -13415,6 +13419,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_output_config": true
|
||||
},
|
||||
|
|
@ -13435,6 +13440,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-sonnet-4-1": {
|
||||
|
|
@ -13454,6 +13460,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-claude-sonnet-4-5": {
|
||||
|
|
@ -13473,6 +13480,7 @@
|
|||
"supports_assistant_prefill": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_anthropic_thinking_payload": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-gemini-2-5-flash": {
|
||||
|
|
@ -13490,6 +13498,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_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-gemini-2-5-pro": {
|
||||
|
|
@ -13507,6 +13516,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_tool_choice": true
|
||||
},
|
||||
"databricks/databricks-gemma-3-12b": {
|
||||
|
|
|
|||
|
|
@ -8,6 +8,12 @@ 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
|
||||
|
||||
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,
|
||||
|
|
@ -15,6 +21,12 @@ from litellm.llms.databricks.chat.transformation import (
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
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 = [
|
||||
|
|
@ -423,3 +435,126 @@ def test_databricks_config_probes_capabilities_under_databricks_namespace():
|
|||
without this override they probed the ``anthropic`` cost-map namespace and
|
||||
ignored the exact ``databricks/databricks-claude-*`` entries."""
|
||||
assert DatabricksConfig().custom_llm_provider == "databricks"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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."""
|
||||
non_default = {"reasoning_effort": reasoning_effort}
|
||||
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_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")
|
||||
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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue