Merge a4c7d57c2c into dddbaa803d

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":

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": "{}"},
+    }