From 2e35f07e94cca006978c592b4dcae48a3feb17d1 Mon Sep 17 00:00:00 2001
From: Pranav Simha <pranav0521@gmail.com>
Date: Wed, 2 Apr 2025 07:50:44 -0700
Subject: [PATCH] Add support for max_completion_tokens to the Cohere chat
 transformation config (#9701)

---
 litellm/llms/cohere/chat/transformation.py    |  8 +++-
 .../llms/cohere/chat/test_transformation.py   | 45 +++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 tests/litellm/llms/cohere/chat/test_transformation.py

diff --git a/litellm/llms/cohere/chat/transformation.py b/litellm/llms/cohere/chat/transformation.py
index fbaedca8f63..70677214a74 100644
--- a/litellm/llms/cohere/chat/transformation.py
+++ b/litellm/llms/cohere/chat/transformation.py
@@ -54,7 +54,8 @@ class CohereChatConfig(BaseConfig):
         search_queries_only (bool, optional): When true, the response will only contain a list of generated search queries.
         documents (List[Dict[str, str]], optional): A list of relevant documents that the model can cite.
         temperature (float, optional): A non-negative float that tunes the degree of randomness in generation.
-        max_tokens (int, optional): The maximum number of tokens the model will generate as part of the response.
+        max_tokens [DEPRECATED - use max_completion_tokens] (int, optional): The maximum number of tokens the model will generate as part of the response.
+        max_completion_tokens (int, optional): The maximum number of tokens the model will generate as part of the response.
         k (int, optional): Ensures only the top k most likely tokens are considered for generation at each step.
         p (float, optional): Ensures that only the most likely tokens, with total probability mass of p, are considered for generation.
         frequency_penalty (float, optional): Used to reduce repetitiveness of generated tokens.
@@ -75,6 +76,7 @@ class CohereChatConfig(BaseConfig):
     documents: Optional[list] = None
     temperature: Optional[int] = None
     max_tokens: Optional[int] = None
+    max_completion_tokens: Optional[int] = None
     k: Optional[int] = None
     p: Optional[int] = None
     frequency_penalty: Optional[int] = None
@@ -96,6 +98,7 @@ class CohereChatConfig(BaseConfig):
         documents: Optional[list] = None,
         temperature: Optional[int] = None,
         max_tokens: Optional[int] = None,
+        max_completion_tokens: Optional[int] = None,
         k: Optional[int] = None,
         p: Optional[int] = None,
         frequency_penalty: Optional[int] = None,
@@ -131,6 +134,7 @@ class CohereChatConfig(BaseConfig):
             "stream",
             "temperature",
             "max_tokens",
+            "max_completion_tokens",
             "top_p",
             "frequency_penalty",
             "presence_penalty",
@@ -156,6 +160,8 @@ class CohereChatConfig(BaseConfig):
                 optional_params["temperature"] = value
             if param == "max_tokens":
                 optional_params["max_tokens"] = value
+            if param == "max_completion_tokens":
+                optional_params["max_tokens"] = value
             if param == "n":
                 optional_params["num_generations"] = value
             if param == "top_p":
diff --git a/tests/litellm/llms/cohere/chat/test_transformation.py b/tests/litellm/llms/cohere/chat/test_transformation.py
new file mode 100644
index 00000000000..30079abd434
--- /dev/null
+++ b/tests/litellm/llms/cohere/chat/test_transformation.py
@@ -0,0 +1,45 @@
+import os
+import sys
+from unittest.mock import MagicMock
+
+
+sys.path.insert(
+    0, os.path.abspath("../../../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.llms.cohere.chat.transformation import CohereChatConfig
+
+
+class TestCohereTransform:
+    def setup_method(self):
+        self.config = CohereChatConfig()
+        self.model = "command-r-plus-latest"
+        self.logging_obj = MagicMock()
+
+    def test_map_cohere_params(self):
+        """Test that parameters are correctly mapped"""
+        test_params = {"temperature": 0.7, "max_tokens": 200, "max_completion_tokens": 256}
+
+        result = self.config.map_openai_params(
+            non_default_params=test_params,
+            optional_params={},
+            model=self.model,
+            drop_params=False,
+        )
+
+        # The function should properly map max_completion_tokens to max_tokens and override max_tokens
+        assert result == {"temperature": 0.7, "max_tokens": 256}
+
+    def test_cohere_max_tokens_backward_compat(self):
+        """Test that parameters are correctly mapped"""
+        test_params = {"temperature": 0.7, "max_tokens": 200,}
+
+        result = self.config.map_openai_params(
+            non_default_params=test_params,
+            optional_params={},
+            model=self.model,
+            drop_params=False,
+        )
+
+        # The function should properly map max_tokens if max_completion_tokens is not provided
+        assert result == {"temperature": 0.7, "max_tokens": 200}