From a4c7d57c2c6b2dbc043ad7abcbad36b1241b475f Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Wed, 21 Oct 2026 00:53:56 -0400 Subject: [PATCH] fix(cohere): emit streamed tool calls from v2 tool-call-start and tool-call-delta chunks --- litellm/llms/cohere/common_utils.py | 35 +++++----- .../cohere/chat/test_cohere_transformation.py | 69 ++++++++++++++++++- 2 files changed, 84 insertions(+), 20 deletions(-) diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index c964e60fac7..e5f741ccd3d 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -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.""" + tool_call: Final = chunk.get("delta", {}).get("message", {}).get("tool_calls") or {} + if isinstance(tool_call, list): + tool_call = tool_call[0] if tool_call else {} + function: Final = tool_call.get("function") or {} + return { + "id": 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 087a5254ae0..5893f62f1a8 100644 --- a/tests/unit/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/unit/llms/cohere/chat/test_cohere_transformation.py @@ -613,9 +613,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( @@ -665,3 +663,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": "{}"}, + }