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..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 @@ -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()