fix(cohere): strip index from tool_calls and name from tool results in v2 API

Cohere v2 API rejects two fields that are valid in OpenAI format:
- `index` in assistant tool_calls (added by LiteLLM when building responses)
- `name` in tool result messages (commonly added by users following OpenAI patterns)

Override transform_request in CohereV2ChatConfig to sanitize messages
before sending, handling both dict and Pydantic model message objects.

Fixes #24031
This commit is contained in:
Amit-kr26 2026-03-24 16:10:54 +05:30
parent 9343aeefca
commit 769dfa91be
2 changed files with 91 additions and 1 deletions

View file

@ -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(

View file

@ -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