mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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:
parent
9343aeefca
commit
769dfa91be
2 changed files with 91 additions and 1 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue