From 292f926fcd5b4882e3cda1e28c6fa9211a7860d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:56:17 -0700 Subject: [PATCH 1/2] fix(oci): stream Cohere tool-calling answers once OCI Cohere restates the whole assistant text on the chunk that carries the tool calls and again on the terminal chunk that carries chatHistory. Only the terminal restatement was dropped, so a tool-calling turn streamed the text twice. Treat a toolCalls-bearing chunk as a restatement too and drop its text once deltas were already emitted. Resolves LIT-6819 --- litellm/llms/oci/chat/cohere.py | 52 ++++------------- .../oci/chat/test_oci_cohere_tool_calls.py | 57 +++++++++++++++++++ 2 files changed, 69 insertions(+), 40 deletions(-) diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 6e9bb83b0a0..1b494ebad47 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/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: diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py index 5dd44d72d68..97b2fd88cb9 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -445,6 +445,63 @@ 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() From 43bb55d84934762a775bef3c32f6ed5ad8bc5ce6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:38:23 -0700 Subject: [PATCH 2/2] style(oci): wrap the Cohere tool-turn test fixtures to 120 columns --- .../llms/oci/chat/test_oci_cohere_tool_calls.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py index 97b2fd88cb9..729a2d25f41 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_cohere_tool_calls.py @@ -446,7 +446,9 @@ class TestOCICohereToolCalls: 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_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."}, @@ -462,13 +464,20 @@ class TestOCICohereToolCalls: }, ] _TOOL_TURN_TERMINAL_SPLIT = [ - {"apiFormat": "COHERE", "text": _TOOL_TURN_TEXT, "chatHistory": _TOOL_TURN_HISTORY, "toolCalls": _TOOL_TURN_CALLS}, + { + "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()) + 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 [])]