diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index 190491adfc7..c499dea2456 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -176,6 +176,27 @@ class CohereV2ChatConfig(OpenAIGPTConfig): model, messages, optional_params, litellm_params, headers ) + # Cohere v2 rejects fields that are valid in OpenAI but not in Cohere: + # 1. 'index' in assistant tool_calls + # 2. 'name' in tool result messages + sanitized: List[AllMessageValues] = [] + for message in data.get("messages", []): + if hasattr(message, "model_dump"): + message = message.model_dump(exclude_none=True) + if isinstance(message, dict): + role = message.get("role") + if role == "assistant" and message.get("tool_calls"): + cleaned_tool_calls = [ + {k: v for k, v in tc.items() if k != "index"} + if isinstance(tc, dict) + else {k: v for k, v in tc.model_dump(exclude_none=True).items() if k != "index"} + for tc in message["tool_calls"] + ] + message = {**message, "tool_calls": cleaned_tool_calls} + elif role == "tool": + message = {k: v for k, v in message.items() if k != "name"} + sanitized.append(message) # type: ignore + data["messages"] = sanitized return data def transform_response( diff --git a/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py b/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py index 4fe8f8a88a9..70826fe4a58 100644 --- a/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py @@ -1,12 +1,13 @@ import os import sys -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path from litellm.llms.cohere.chat.transformation import CohereChatConfig +from litellm.llms.cohere.chat.v2_transformation import CohereV2ChatConfig class TestCohereTransform: @@ -49,3 +50,71 @@ class TestCohereTransform: # The function should properly map max_tokens if max_completion_tokens is not provided assert result == {"temperature": 0.7, "max_tokens": 200} + + +class TestCohereV2Transform: + def setup_method(self): + self.config = CohereV2ChatConfig() + self.model = "command-r-08-2024" + + def _make_transform_request(self, messages): + with patch.object( + self.config.__class__.__bases__[0], + "transform_request", + return_value={"model": self.model, "messages": messages}, + ): + return self.config.transform_request( + model=self.model, + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + def test_strips_index_from_assistant_tool_calls(self): + """Cohere v2 rejects 'index' in tool_calls — it must be stripped before sending.""" + messages = [ + {"role": "user", "content": "What time is it?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "index": 0, + "id": "call_abc", + "type": "function", + "function": {"name": "get_time", "arguments": "{}"}, + } + ], + }, + ] + result = self._make_transform_request(messages) + assistant_msg = result["messages"][1] + assert "index" not in assistant_msg["tool_calls"][0] + assert assistant_msg["tool_calls"][0]["id"] == "call_abc" + + def test_strips_name_from_tool_result_messages(self): + """Cohere v2 rejects 'name' in tool result messages — it must be stripped.""" + messages = [ + {"role": "user", "content": "What time is it?"}, + { + "role": "tool", + "tool_call_id": "call_abc", + "name": "get_time", + "content": "12:00", + }, + ] + result = self._make_transform_request(messages) + tool_msg = result["messages"][1] + assert "name" not in tool_msg + assert tool_msg["tool_call_id"] == "call_abc" + assert tool_msg["content"] == "12:00" + + def test_preserves_messages_without_offending_fields(self): + """Messages that don't have index or name are passed through unchanged.""" + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + ] + result = self._make_transform_request(messages) + assert result["messages"] == messages