Merge pull request #39965 from BerriAI/litellm_fix_oci_cohere_stream_tool_turn_dup

fix(oci): stream Cohere tool-calling answers once

litellm/llms/oci/chat/cohere.py

@@ -290,19 +290,14 @@ def handle_cohere_stream_chunk(
 ) -> ModelResponseStream:
     """Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream.
 
-    ``prior_tool_calls_emitted`` lets the caller signal whether tool calls
-    were already emitted in earlier chunks of the same stream. When set, the
-    terminal consolidation chunk's tool calls are suppressed (they would
-    duplicate prior deltas); otherwise they are passed through so a stream
-    that delivers tool calls only on the terminal chunk doesn't silently
-    drop them.
-
-    ``prior_text_emitted`` plays the analogous role for the ``text`` field:
-    when set, the terminal consolidation chunk's ``text`` is suppressed
-    (it would re-emit the full assembled response on top of prior deltas);
-    when unset (e.g. a degenerate stream that delivers the entire response
-    in a single SSE event carrying both ``chatHistory`` and ``finishReason``),
-    the text is passed through so the response content isn't silently lost.
+    OCI Cohere streams the answer as single-token ``text`` deltas, then restates
+    the whole assembled ``text`` on every chunk that carries ``toolCalls`` or
+    ``chatHistory`` (the tool-calls event and the terminal event). Once the
+    caller reports that earlier chunks already emitted text
+    (``prior_text_emitted``), those restatements are dropped so the client does
+    not see the answer twice; a stream whose only text lives on such a chunk
+    keeps it. ``prior_tool_calls_emitted`` plays the same role for the tool
+    calls the terminal ``chatHistory`` chunk repeats.
     """
     try:
         typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk)
@@ -315,33 +310,10 @@ def handle_cohere_stream_chunk(
     if typed_chunk.index is None:
         typed_chunk.index = 0
 
-    # OCI Cohere's terminal SSE event re-sends the full assembled response in
-    # `text` alongside a populated `chatHistory` and a non-null `finishReason`.
-    # Emitting that text would concatenate the whole response onto the
-    # already-streamed deltas. We require both signals to be present so that a
-    # future API change which adds `chatHistory` to intermediate chunks (or a
-    # rare early-populated case) doesn't silently drop legitimate token deltas.
-    is_terminal_consolidation: Final = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None
-    # On non-terminal text-free chunks (e.g. tool-call-only or keep-alive
-    # chunks) emit ``content=None`` rather than ``content=""`` so downstream
-    # stream-mergers that distinguish "no text in this delta" from "an
-    # explicitly empty text delta" behave correctly.
-    #
-    # We only suppress the terminal chunk's ``text`` when the caller has
-    # confirmed that text deltas were already emitted earlier — otherwise
-    # (e.g. a degenerate stream that delivers the whole response in a
-    # single SSE event), passing it through is the only chance to surface it.
-    text: Final[str | None] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text
-
-    # Tool calls on the terminal consolidation chunk (whether from
-    # `typed_chunk.toolCalls` or from `chatHistory`) typically restate what
-    # was already streamed in intermediate chunks. Re-emitting them would
-    # mint fresh `uuid4` IDs and cause downstream consumers to execute each
-    # tool call twice. We only suppress when the caller has confirmed that
-    # tool calls were already emitted earlier — otherwise (e.g. a short
-    # response that delivers tool calls exclusively on the terminal chunk),
-    # passing them through is the only chance to surface them.
-    cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls
+    restates_text: Final = typed_chunk.chatHistory is not None or typed_chunk.toolCalls is not None
+    restates_tool_calls: Final = typed_chunk.chatHistory is not None
+    text: Final[str | None] = None if (restates_text and prior_text_emitted) else typed_chunk.text
+    cohere_tool_calls: Final = None if (restates_tool_calls and prior_tool_calls_emitted) else typed_chunk.toolCalls
 
     tool_calls: list[dict[str, object]] | None = None
     if cohere_tool_calls:

tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py

@@ -445,6 +445,72 @@ class TestOCICohereToolCalls:
         assert result.choices[0].index == 0
         assert result.choices[0].finish_reason == "stop"  # COMPLETE is mapped to stop
 
+    _TOOL_TURN_TEXT = "I will use the tool to find out the weather in Paris."
+    _TOOL_TURN_DELTAS = [
+        "I", " will", " use", " the", " tool", " to", " find", " out", " the", " weather", " in", " Paris", ".",
+    ]
+    _TOOL_TURN_CALLS = [{"name": "get_weather", "parameters": {"city": "Paris"}}]
+    _TOOL_TURN_HISTORY = [
+        {"role": "USER", "message": "Briefly say what you will do, then find out the weather in Paris using the tool."},
+        {"role": "CHATBOT", "message": _TOOL_TURN_TEXT, "toolCalls": _TOOL_TURN_CALLS},
+    ]
+    _TOOL_TURN_TERMINAL_TOGETHER = [
+        {
+            "apiFormat": "COHERE",
+            "text": _TOOL_TURN_TEXT,
+            "chatHistory": _TOOL_TURN_HISTORY,
+            "finishReason": "COMPLETE",
+            "toolCalls": _TOOL_TURN_CALLS,
+        },
+    ]
+    _TOOL_TURN_TERMINAL_SPLIT = [
+        {
+            "apiFormat": "COHERE",
+            "text": _TOOL_TURN_TEXT,
+            "chatHistory": _TOOL_TURN_HISTORY,
+            "toolCalls": _TOOL_TURN_CALLS,
+        },
+        {"apiFormat": "COHERE", "finishReason": "COMPLETE"},
+    ]
+
+    @staticmethod
+    def _drain_cohere_stream(events):
+        wrapper = OCIStreamWrapper(
+            completion_stream=MagicMock(), model="cohere.command-a-03-2025", logging_obj=MagicMock()
+        )
+        chunks = [wrapper.chunk_creator(f"data: {json.dumps(event)}") for event in events]
+        content = "".join(chunk.choices[0].delta.content or "" for chunk in chunks)
+        tool_calls = [call for chunk in chunks for call in (chunk.choices[0].delta.tool_calls or [])]
+        finish_reasons = [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason]
+        return content, tool_calls, finish_reasons
+
+    @pytest.mark.parametrize("terminal_events", [_TOOL_TURN_TERMINAL_TOGETHER, _TOOL_TURN_TERMINAL_SPLIT])
+    def test_cohere_tool_turn_streams_the_answer_once(self, terminal_events):
+        """OCI restates the whole answer on the tool-calls chunk and again on the terminal chunk;
+        the client must read it exactly once, with one tool call and one finish reason."""
+        deltas = [{"apiFormat": "COHERE", "text": token} for token in self._TOOL_TURN_DELTAS]
+        tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS}
+
+        content, tool_calls, finish_reasons = self._drain_cohere_stream([*deltas, tool_calls_event, *terminal_events])
+
+        assert content == self._TOOL_TURN_TEXT
+        assert [(call["function"]["name"], call["function"]["arguments"]) for call in tool_calls] == [
+            ("get_weather", '{"city": "Paris"}')
+        ]
+        assert finish_reasons == ["stop"]
+
+    def test_cohere_tool_turn_without_preamble_deltas_keeps_the_only_text(self):
+        """When the tool-calls chunk carries the only copy of the text, dropping it would lose the answer."""
+        tool_calls_event = {"apiFormat": "COHERE", "text": self._TOOL_TURN_TEXT, "toolCalls": self._TOOL_TURN_CALLS}
+
+        content, tool_calls, finish_reasons = self._drain_cohere_stream(
+            [tool_calls_event, *self._TOOL_TURN_TERMINAL_TOGETHER]
+        )
+
+        assert content == self._TOOL_TURN_TEXT
+        assert len(tool_calls) == 1
+        assert finish_reasons == ["stop"]
+
     def test_cohere_parameter_mapping_excludes_tool_choice(self):
         """Test that tool_choice is excluded from Cohere parameter mapping"""
         config = OCIChatConfig()