From 508e4da40e9db56fdb650df6760b352c17624f89 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Mon, 12 Jan 2026 12:44:46 +0530
Subject: [PATCH] Fix: respect max_completion_tokens in thinking feat

---
 litellm/llms/base_llm/chat/transformation.py  |  4 +-
 .../chat/test_converse_transformation.py      | 72 +++++++++++++++++++
 2 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py
index b592c23846..41a1797ceb 100644
--- a/litellm/llms/base_llm/chat/transformation.py
+++ b/litellm/llms/base_llm/chat/transformation.py
@@ -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
             )
diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
index e603f94ab8..692866f855 100644
--- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
+++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
@@ -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