mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-16 22:17:21 +00:00
Fix: respect max_completion_tokens in thinking feat
This commit is contained in:
@@ -132,10 +132,10 @@ class BaseConfig(ABC):
|
||||
|
||||
Checks 'non_default_params' for 'thinking' and 'max_tokens'
|
||||
|
||||
if 'thinking' is enabled and 'max_tokens' is not specified, set 'max_tokens' to the thinking token budget + DEFAULT_MAX_TOKENS
|
||||
if 'thinking' is enabled and 'max_tokens' or 'max_completion_tokens' is not specified, set 'max_tokens' to the thinking token budget + DEFAULT_MAX_TOKENS
|
||||
"""
|
||||
is_thinking_enabled = self.is_thinking_enabled(optional_params)
|
||||
if is_thinking_enabled and "max_tokens" not in non_default_params:
|
||||
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
|
||||
)
|
||||
|
||||
@@ -2735,3 +2735,75 @@ def test_is_nova_lite_2_model():
|
||||
assert config._is_nova_lite_2_model("anthropic.claude-3-5-sonnet-20240620-v1:0") is False
|
||||
assert config._is_nova_lite_2_model("meta.llama3-70b-instruct-v1:0") is False
|
||||
assert config._is_nova_lite_2_model("mistral.mistral-7b-instruct-v0:2") is False
|
||||
|
||||
|
||||
def test_thinking_with_max_completion_tokens():
|
||||
"""Test that thinking respects max_completion_tokens parameter."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
# Test case 1: max_completion_tokens is specified - should NOT set maxTokens automatically
|
||||
non_default_params_with_max_completion = {
|
||||
"thinking": {"type": "enabled", "budget_tokens": 5000},
|
||||
"max_completion_tokens": 10000,
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
result = config.map_openai_params(
|
||||
non_default_params=non_default_params_with_max_completion,
|
||||
optional_params=optional_params,
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
# Should have maxTokens set to max_completion_tokens value
|
||||
assert "maxTokens" in result
|
||||
assert result["maxTokens"] == 10000
|
||||
# Should have thinking config
|
||||
assert "thinking" in result
|
||||
assert result["thinking"]["type"] == "enabled"
|
||||
assert result["thinking"]["budget_tokens"] == 5000
|
||||
|
||||
# Test case 2: max_tokens is specified - should NOT set maxTokens automatically
|
||||
non_default_params_with_max_tokens = {
|
||||
"thinking": {"type": "enabled", "budget_tokens": 5000},
|
||||
"max_tokens": 8000,
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
result = config.map_openai_params(
|
||||
non_default_params=non_default_params_with_max_tokens,
|
||||
optional_params=optional_params,
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
# Should have maxTokens set to max_tokens value
|
||||
assert "maxTokens" in result
|
||||
assert result["maxTokens"] == 8000
|
||||
# Should have thinking config
|
||||
assert "thinking" in result
|
||||
assert result["thinking"]["type"] == "enabled"
|
||||
assert result["thinking"]["budget_tokens"] == 5000
|
||||
|
||||
# Test case 3: Neither max_tokens nor max_completion_tokens specified - should set maxTokens automatically
|
||||
from litellm.constants import DEFAULT_MAX_TOKENS
|
||||
|
||||
non_default_params_without_max = {
|
||||
"thinking": {"type": "enabled", "budget_tokens": 5000},
|
||||
}
|
||||
optional_params = {}
|
||||
|
||||
result = config.map_openai_params(
|
||||
non_default_params=non_default_params_without_max,
|
||||
optional_params=optional_params,
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
# Should have maxTokens set to budget_tokens + DEFAULT_MAX_TOKENS
|
||||
assert "maxTokens" in result
|
||||
assert result["maxTokens"] == 5000 + DEFAULT_MAX_TOKENS
|
||||
# Should have thinking config
|
||||
assert "thinking" in result
|
||||
assert result["thinking"]["type"] == "enabled"
|
||||
assert result["thinking"]["budget_tokens"] == 5000
|
||||
|
||||
Reference in New Issue
Block a user