From fc68752d80ed28346ffa16e52e46247dcb44b254 Mon Sep 17 00:00:00 2001 From: Amit-kr26 Date: Wed, 25 Mar 2026 11:48:47 +0530 Subject: [PATCH] fix(cohere): address review suggestions - Guard data["messages"] write with 'if messages in data' to avoid injecting an empty list when parent transform omits the key - Add hasattr(tc, "model_dump") guard before calling model_dump on non-dict tool calls to avoid AttributeError on unknown types - Assert type and function fields are preserved after Pydantic model_dump --- litellm/llms/cohere/chat/v2_transformation.py | 6 +++++- .../llms/cohere/chat/test_cohere_transformation.py | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index bbf2787244b..cd01416e55c 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -179,8 +179,10 @@ class CohereV2ChatConfig(OpenAIGPTConfig): # 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 + if "messages" not in data: + return data sanitized: List[AllMessageValues] = [] - for message in data.get("messages", []): + for message in data["messages"]: if hasattr(message, "model_dump"): message = message.model_dump(exclude_unset=True) if isinstance(message, dict): @@ -190,6 +192,8 @@ class CohereV2ChatConfig(OpenAIGPTConfig): {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_unset=True).items() if k != "index"} + if hasattr(tc, "model_dump") + else tc for tc in message["tool_calls"] ] message = {**message, "tool_calls": cleaned_tool_calls} 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 78960f41090..d727e352d97 100644 --- a/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/test_litellm/llms/cohere/chat/test_cohere_transformation.py @@ -133,6 +133,8 @@ class TestCohereV2Transform: assistant_msg = result["messages"][1] assert "index" not in assistant_msg["tool_calls"][0] assert assistant_msg["tool_calls"][0]["id"] == "call_abc" + assert assistant_msg["tool_calls"][0]["type"] == "function" + assert assistant_msg["tool_calls"][0]["function"]["name"] == "get_time" def test_preserves_messages_without_offending_fields(self): """Messages that don't have index or name are passed through unchanged."""