diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index 41b1569bc4a..cddcccc4ef5 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -1,5 +1,5 @@ import time -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -13,7 +13,7 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAnnotation, ChatCompletionAnnotationURLCitation, - ChatCompletionToolCallChunk, + ChatCompletionAssistantToolCall, ) from litellm.types.utils import ModelResponse, Usage @@ -169,15 +169,38 @@ class CohereV2ChatConfig(OpenAIGPTConfig): def transform_request( self, model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: - """ - Cohere v2 chat api is in openai format, so we can use the openai transform request function to transform the request. - """ - data: Final = super().transform_request(model, messages, optional_params, litellm_params, headers) + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + ) -> dict[str, object]: + data: Final[dict[str, object]] + data = super().transform_request( # pyright: ignore[reportUnknownMemberType] # untyped parent dictionaries + model, list(messages), dict(optional_params), dict(litellm_params), dict(headers) + ) + messages_to_send: Final[Sequence[AllMessageValues] | None] + messages_to_send = data.get("messages") # pyright: ignore[reportAssignmentType] # parent erases message type + if messages_to_send is not None: + litellm.remove_index_from_tool_calls(messages=list(messages_to_send)) + + return data + + async def async_transform_request( + self, + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + ) -> dict[str, object]: + data: Final[dict[str, object]] + data = await super().async_transform_request( # pyright: ignore[reportUnknownMemberType] # untyped parent + model, list(messages), dict(optional_params), dict(litellm_params), dict(headers) + ) + messages_to_send: Final[Sequence[AllMessageValues] | None] + messages_to_send = data.get("messages") # pyright: ignore[reportAssignmentType] # parent erases message type + if messages_to_send is not None: + litellm.remove_index_from_tool_calls(messages=list(messages_to_send)) return data @@ -225,11 +248,10 @@ class CohereV2ChatConfig(OpenAIGPTConfig): cohere_tools_response: Final = cohere_v2_chat_response["message"].get("tool_calls", []) if cohere_tools_response is not None and cohere_tools_response != []: # convert cohere_tools_response to OpenAI response format - tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] - for index, tool in enumerate(cohere_tools_response): - tool_call: ChatCompletionToolCallChunk = { + tool_calls: Final[list[ChatCompletionAssistantToolCall]] = [] + for tool in cohere_tools_response: + tool_call: ChatCompletionAssistantToolCall = { **tool, - "index": index, } tool_calls.append(tool_call) _message: Final = litellm.Message( diff --git a/tests/unit/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py index 087a5254ae0..222b2814b53 100644 --- a/tests/unit/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/unit/llms/cohere/chat/test_cohere_transformation.py @@ -1,7 +1,7 @@ import json from collections.abc import Callable, Iterator from typing import Final -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest @@ -11,7 +11,8 @@ import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.cohere.chat.transformation import CohereChatConfig from litellm.llms.cohere.chat.v2_transformation import CohereV2ChatConfig -from unittest.mock import AsyncMock, patch +from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantToolCall, ChatCompletionToolCallChunk +from litellm.types.utils import ModelResponse COHERE_V1_CHAT_URL: Final = "https://api.cohere.ai/v1/chat" COHERE_V2_CHAT_URL: Final = "https://api.cohere.com/v2/chat" @@ -141,6 +142,70 @@ class TestCohereV2Transform: assert optional_params["max_tokens"] == 256 + def test_v2_response_preserves_tool_call_fields_without_synthetic_index(self) -> None: + expected: Final[ChatCompletionAssistantToolCall] = { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "Paris"}'}, + } + response: Final = httpx.Response( + status_code=200, + json={ + "id": "res-1", + "message": {"role": "assistant", "tool_calls": [expected]}, + "usage": {"tokens": {"input_tokens": 10, "output_tokens": 20}}, + }, + ) + result: Final = self.config.transform_response( + model=self.model, + raw_response=response, + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].message.model_dump(exclude_none=True)["tool_calls"] == [expected] + + @pytest.mark.parametrize("use_async", [False, True]) + @pytest.mark.asyncio + async def test_v2_request_preserves_tool_call_fields_without_stream_index(self, use_async: bool) -> None: + indexed_call: Final[ChatCompletionToolCallChunk] = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "Paris"}'}, + } + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [indexed_call]}, + {"role": "tool", "tool_call_id": "call_1", "content": "18C, sunny"}, + ] + result: Final = ( + await self.config.async_transform_request(self.model, messages, {}, {}, {}) + if use_async + else self.config.transform_request(self.model, messages, {}, {}, {}) + ) + + assert result["messages"] == [ + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"location": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "18C, sunny"}, + ] + @pytest.mark.asyncio async def test_cohere_request_body_with_allowed_params(): @@ -613,9 +678,7 @@ async def test_cohere_v2_annotations(respx_mock: respx.MockRouter) -> None: def test_cohere_v2_parameter_mapping(respx_mock: respx.MockRouter) -> None: messages: Final = [{"role": "user", "content": "Generate a creative story."}] route: Final = respx_mock.post(COHERE_V2_CHAT_URL).mock( - return_value=httpx.Response( - 200, json=_cohere_v2_response({"content": [{"type": "text", "text": "A story."}]}) - ) + return_value=httpx.Response(200, json=_cohere_v2_response({"content": [{"type": "text", "text": "A story."}]})) ) litellm.completion(