fix(cohere): emit streamed tool calls from v2 tool-call-start and tool-call-delta chunks (#47794)

* fix(cohere): emit streamed tool calls from v2 tool-call-start and tool-call-delta chunks

* fix(cohere): type _parse_tool_call_delta without growing basedpyright counts

The reassigned Final local and the untyped chunk lookups added basedpyright
errors past the base ceiling. Narrow them with casts; per-rule counts for the
file are now at or below the merge-base.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

* fix(cohere): parse streamed tool calls with a typed model and map Cohere TOOL_CALL to tool_calls

Read the streamed tool call through a frozen pydantic model instead of casts,
which the strict and type-discipline gates refuse. Map Cohere's TOOL_CALL and
STOP_SEQUENCE finish reasons, so a stream that reads Cohere's message-end still
ends tool_calls. The regression test now drives litellm.acompletion with two
parallel streamed calls and checks the rebuilt calls and the final finish reason

* fix(cohere): close a streamed tool call that sent no arguments with an empty JSON object

---------

Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/litellm_core_utils/core_helpers.py

@@ -315,6 +315,8 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = {
     "model_context_window_exceeded": "length",
     # Cohere
     "COMPLETE": "stop",
+    "STOP_SEQUENCE": "stop",
+    "TOOL_CALL": "tool_calls",
     "ERROR_TOXIC": "content_filter",
     "ERROR": "stop",
     # HuggingFace / Together AI

litellm/llms/cohere/common_utils.py

@@ -3,6 +3,7 @@ from typing import Final, Literal
 
 from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.types.llms.cohere import CohereV2StreamedToolCall
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import (
     ChatCompletionToolCallChunk,
@@ -222,6 +223,7 @@ class CohereV2ModelResponseIterator:
         self.content_blocks: list = []
         self.tool_index = -1
         self.json_mode = json_mode
+        self.tool_call_indexes_with_arguments: frozenset[int] = frozenset()
 
     def _parse_content_delta(self, chunk: dict) -> str:
         """Parse content-delta chunks to extract text."""
@@ -234,20 +236,23 @@ 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:
+        """Parse tool-call-start and tool-call-delta chunks to extract tool calls."""
+        tool_call: Final = CohereV2StreamedToolCall.model_validate(chunk)
+        if tool_call.arguments:
+            self.tool_call_indexes_with_arguments = self.tool_call_indexes_with_arguments | {index}
+        return {
+            "id": tool_call.id,
+            "type": "function",
+            "function": {"name": tool_call.name, "arguments": tool_call.arguments},
+            "index": index,
+        }
+
+    def _parse_tool_call_end(self, index: int) -> ChatCompletionToolCallChunk | None:
+        """Close a tool call that streamed no arguments with an empty JSON object, as OpenAI streams one."""
+        if index in self.tool_call_indexes_with_arguments:
+            return None
+        return {"id": None, "type": "function", "function": {"name": None, "arguments": "{}"}, "index": index}
 
     def _parse_tool_plan_delta(self, chunk: dict) -> dict | None:
         """Parse tool-plan-delta events to extract tool plan."""
@@ -301,7 +306,8 @@ 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 call end: chunk.type == "tool-call-end" -> "{}" arguments when the call streamed none
         - 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 +327,12 @@ 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-call-end"):
+                tool_use = (
+                    self._parse_tool_call_end(index)
+                    if chunk_type == "tool-call-end"
+                    else 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":

litellm/types/llms/cohere.py

@@ -1,5 +1,6 @@
 from typing import Literal
 
+from pydantic import AliasPath, BaseModel, ConfigDict, Field
 from typing_extensions import Required, TypedDict
 
 
@@ -98,3 +99,15 @@ class CohereV2ChatResponse(TypedDict):
     message: CohereV2ChatResponseMessage
     usage: CohereV2ChatResponseUsage
     logprobs: CohereV2ChatResponseLogProbs
+
+
+class CohereV2StreamedToolCall(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    id: str | None = Field(default=None, validation_alias=AliasPath("delta", "message", "tool_calls", "id"))
+    name: str | None = Field(
+        default=None, validation_alias=AliasPath("delta", "message", "tool_calls", "function", "name")
+    )
+    arguments: str = Field(
+        default="", validation_alias=AliasPath("delta", "message", "tool_calls", "function", "arguments")
+    )

tests/unit/litellm_core_utils/test_core_helpers.py

@@ -319,6 +319,12 @@ class TestMapFinishReasonCohere:
     def test_error(self):
         assert map_finish_reason("ERROR") == "stop"
 
+    def test_tool_call(self):
+        assert map_finish_reason("TOOL_CALL") == "tool_calls"
+
+    def test_stop_sequence(self):
+        assert map_finish_reason("STOP_SEQUENCE") == "stop"
+
 
 class TestMapFinishReasonHuggingFace:
     def test_eos_token(self):

tests/unit/llms/cohere/chat/test_cohere_transformation.py

@@ -1,3 +1,4 @@
+import itertools
 import json
 from collections.abc import Callable, Iterator
 from typing import Final
@@ -653,9 +654,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 +704,123 @@ 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 _cohere_v2_tool_call_events(index: int, call_id: str, city: str) -> tuple[dict[str, object], ...]:
+    return (
+        {
+            "type": "tool-call-start",
+            "index": index,
+            "delta": {
+                "message": {
+                    "tool_calls": {
+                        "id": call_id,
+                        "type": "function",
+                        "function": {"name": "get_weather", "arguments": ""},
+                    }
+                }
+            },
+        },
+        {
+            "type": "tool-call-delta",
+            "index": index,
+            "delta": {"message": {"tool_calls": {"function": {"arguments": '{"city": '}}}},
+        },
+        {
+            "type": "tool-call-delta",
+            "index": index,
+            "delta": {"message": {"tool_calls": {"function": {"arguments": json.dumps(city) + "}"}}}},
+        },
+        {"type": "tool-call-end", "index": index},
+    )
+
+
+@pytest.mark.usefixtures("_cohere_httpx_transport")
+@pytest.mark.respx(assert_all_called=True)
+async def test_cohere_v2_stream_returns_parallel_tool_calls(respx_mock: respx.MockRouter) -> None:
+    stream_events: Final = (
+        {"type": "message-start", "id": "c0a2b6e1", "delta": {"message": {"role": "assistant"}}},
+        *_cohere_v2_tool_call_events(0, "get_weather_paris", "Paris"),
+        *_cohere_v2_tool_call_events(1, "get_weather_tokyo", "Tokyo"),
+        {
+            "type": "message-end",
+            "delta": {"finish_reason": "TOOL_CALL", "usage": {"tokens": {"input_tokens": 40, "output_tokens": 30}}},
+        },
+    )
+    respx_mock.post(COHERE_V2_CHAT_URL).mock(
+        return_value=httpx.Response(200, content="".join(json.dumps(event) + "\n" for event in stream_events).encode())
+    )
+
+    response: Final = await litellm.acompletion(
+        model="cohere_chat/v2/command-a-plus-05-2026",
+        messages=[{"role": "user", "content": "What is the weather in Paris and in Tokyo?"}],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_weather",
+                    "parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
+                },
+            }
+        ],
+        stream=True,
+        api_key="cohere-test-key",
+    )
+    chunks: Final = [chunk async for chunk in response]
+    rebuilt: Final = litellm.stream_chunk_builder(chunks)
+
+    assert rebuilt is not None
+    assert [
+        (call.id, call.function.name, json.loads(call.function.arguments))
+        for call in rebuilt.choices[0].message.tool_calls
+    ] == [
+        ("get_weather_paris", "get_weather", {"city": "Paris"}),
+        ("get_weather_tokyo", "get_weather", {"city": "Tokyo"}),
+    ]
+    assert chunks[-1].choices[0].finish_reason == "tool_calls"
+
+
+@pytest.mark.usefixtures("_cohere_httpx_transport")
+@pytest.mark.respx(assert_all_called=True)
+async def test_cohere_v2_stream_closes_argument_less_tool_call_with_empty_object(respx_mock: respx.MockRouter) -> None:
+    stream_events: Final = (
+        {"type": "message-start", "id": "5d1e9f20", "delta": {"message": {"role": "assistant"}}},
+        {
+            "type": "tool-call-start",
+            "index": 0,
+            "delta": {
+                "message": {
+                    "tool_calls": {
+                        "id": "get_server_time_0",
+                        "type": "function",
+                        "function": {"name": "get_server_time", "arguments": ""},
+                    }
+                }
+            },
+        },
+        {"type": "tool-call-end", "index": 0},
+        {"type": "message-end", "delta": {"finish_reason": "TOOL_CALL"}},
+    )
+    respx_mock.post(COHERE_V2_CHAT_URL).mock(
+        return_value=httpx.Response(200, content="".join(json.dumps(event) + "\n" for event in stream_events).encode())
+    )
+
+    response: Final = await litellm.acompletion(
+        model="cohere_chat/v2/command-a-plus-05-2026",
+        messages=[{"role": "user", "content": "What time is it on the server?"}],
+        tools=[
+            {
+                "type": "function",
+                "function": {"name": "get_server_time", "parameters": {"type": "object", "properties": {}}},
+            }
+        ],
+        stream=True,
+        api_key="cohere-test-key",
+    )
+    deltas_per_chunk: Final = [chunk.choices[0].delta.tool_calls or () async for chunk in response]
+    streamed_calls: Final = tuple(itertools.chain.from_iterable(deltas_per_chunk))
+
+    assert [(call.id, call.function.name) for call in streamed_calls if call.id] == [
+        ("get_server_time_0", "get_server_time")
+    ]
+    assert json.loads("".join(call.function.arguments for call in streamed_calls)) == {}