diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index c964e60fac7..0d633e9d60b 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -1,5 +1,5 @@ import json -from typing import Final, Literal +from typing import Final, Literal, cast from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -234,20 +234,21 @@ class CohereV2ModelResponseIterator: return content return "" - def _parse_tool_call_delta(self, chunk: dict) -> ChatCompletionToolCallChunk | None: - """Parse tool-call-delta chunks to extract tool calls.""" - delta: Final = chunk.get("delta", {}) - tool_calls: Final = delta.get("tool_calls", []) - if tool_calls: - return { - "id": tool_calls[0].get("id", ""), - "type": "function", - "function": { - "name": tool_calls[0].get("name", ""), - "arguments": tool_calls[0].get("arguments", ""), - }, - } - return None + def _parse_tool_call_delta(self, chunk: dict, index: int) -> ChatCompletionToolCallChunk | None: + """Parse tool-call-start and tool-call-delta chunks to extract tool calls.""" + delta: Final = cast("dict[str, dict[str, object]]", chunk.get("delta") or {}) + raw: Final = delta.get("message", {}).get("tool_calls") or {} + tool_call: Final = cast("dict[str, object]", (raw[0] if raw else {}) if isinstance(raw, list) else raw) + function: Final = cast("dict[str, str]", tool_call.get("function") or {}) + return { + "id": cast("str", tool_call.get("id")), + "type": "function", + "function": { + "name": function.get("name"), + "arguments": function.get("arguments", ""), + }, + "index": index, + } def _parse_tool_plan_delta(self, chunk: dict) -> dict | None: """Parse tool-plan-delta events to extract tool plan.""" @@ -301,7 +302,7 @@ class CohereV2ModelResponseIterator: v2 format: - Content: chunk.type == "content-delta" -> chunk.delta.message.content.text - - Tool calls: chunk.type == "tool-call-delta" -> chunk.delta.tool_calls + - Tool calls: chunk.type in ("tool-call-start", "tool-call-delta") -> chunk.delta.message.tool_calls - Tool plan: chunk.event == "tool-plan-delta" -> chunk.data.delta.message.tool_plan - Citations: chunk.event == "citation-start" -> chunk.data.delta.message.citations - Finish: chunk.event == "message-end" -> chunk.data.delta.finish_reason @@ -321,8 +322,8 @@ class CohereV2ModelResponseIterator: # Handle different chunk types if chunk_type == "content-delta": text = self._parse_content_delta(chunk) - elif chunk_type == "tool-call-delta": - tool_use = self._parse_tool_call_delta(chunk) + elif chunk_type in ("tool-call-start", "tool-call-delta"): + tool_use = self._parse_tool_call_delta(chunk, index) elif event_type == "tool-plan-delta": provider_specific_fields = self._parse_tool_plan_delta(chunk) elif event_type == "citation-start": diff --git a/tests/unit/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py index 8d39c301029..d81e5732af9 100644 --- a/tests/unit/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/unit/llms/cohere/chat/test_cohere_transformation.py @@ -653,9 +653,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( @@ -705,3 +703,68 @@ def test_cohere_v2_error_handling(respx_mock: respx.MockRouter) -> None: assert error.value.status_code == 400 assert error.value.llm_provider == "cohere" assert "invalid request: model 'invalid-model' not found" in error.value.message + + +def test_cohere_v2_stream_tool_calls_are_emitted_with_index() -> None: + events: Final = [ + { + "type": "tool-call-start", + "index": 0, + "delta": { + "message": { + "tool_calls": { + "id": "get_weather_a1", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, + } + } + }, + }, + { + "type": "tool-call-delta", + "index": 0, + "delta": {"message": {"tool_calls": {"function": {"arguments": '{"city": '}}}}, + }, + { + "type": "tool-call-delta", + "index": 0, + "delta": {"message": {"tool_calls": {"function": {"arguments": '"Paris"}'}}}}, + }, + {"type": "tool-call-end", "index": 0}, + { + "type": "tool-call-start", + "index": 1, + "delta": { + "message": { + "tool_calls": { + "id": "get_time_b2", + "type": "function", + "function": {"name": "get_time", "arguments": ""}, + } + } + }, + }, + { + "type": "tool-call-delta", + "index": 1, + "delta": {"message": {"tool_calls": {"function": {"arguments": "{}"}}}}, + }, + ] + iterator: Final = CohereV2ChatConfig().get_model_response_iterator( + streaming_response=iter(json.dumps(event) for event in events), sync_stream=True + ) + + calls: Final[dict[int, dict[str, str]]] = {} + for parsed in iterator: + tool_use = parsed["tool_use"] + if tool_use is None: + continue + call = calls.setdefault(tool_use["index"], {"id": "", "name": "", "arguments": ""}) + call["id"] += tool_use["id"] or "" + call["name"] += tool_use["function"]["name"] or "" + call["arguments"] += tool_use["function"]["arguments"] or "" + + assert calls == { + 0: {"id": "get_weather_a1", "name": "get_weather", "arguments": '{"city": "Paris"}'}, + 1: {"id": "get_time_b2", "name": "get_time", "arguments": "{}"}, + }