fix(anthropic): keep Mistral thinking text on streaming /v1/messages (#49059)

* fix(anthropic): keep Mistral thinking text on streaming /v1/messages

* test(anthropic): type the signed thinking stream test helpers

---------

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

litellm/llms/anthropic/messages/adapters/streaming_iterator.py

@@ -1,6 +1,7 @@
 # What is this?
 ## Translates OpenAI call to Anthropic `/v1/messages` format
 import copy
+import io
 import itertools
 import json
 import traceback
@@ -181,6 +182,44 @@ class _ToolCallPart:
         )
 
 
+@dataclass(frozen=True, slots=True)
+class _ThinkingPiece:
+    text: str
+    signature: str
+    index: int | None
+
+    @staticmethod
+    def of(chunk: "ModelResponseStream") -> "_ThinkingPiece | None":
+        choices: Final = _optional_attr_sequence(chunk, "choices")
+        if len(choices) != 1 or _optional_attr(choices[0], "finish_reason") is not None:
+            return None
+        delta: Final = _optional_attr(choices[0], "delta")
+        if _optional_attr(delta, "content") or _optional_attr(delta, "tool_calls"):
+            return None
+        blocks: Final = _optional_attr_sequence(delta, "thinking_blocks")
+        if not blocks:
+            reasoning: Final = _optional_str_attr(delta, "reasoning_content")
+            return _ThinkingPiece(text=reasoning, signature="", index=None) if reasoning else None
+        block: Final = as_str_mapping(blocks[0]) if len(blocks) == 1 else None
+        if block is None or block.get("type") != "thinking":
+            return None
+        text: Final = block.get("thinking")
+        signature: Final = block.get("signature")
+        index: Final = block.get("index")
+        return _ThinkingPiece(
+            text=text if isinstance(text, str) else "",
+            signature=signature if isinstance(signature, str) else "",
+            index=index if isinstance(index, int) else None,
+        )
+
+    def block(self, *, text: str, signature: str) -> ChatCompletionThinkingBlock:
+        if self.index is None:
+            return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature or None)
+        return ChatCompletionThinkingBlock(
+            type="thinking", thinking=text, signature=signature or None, index=self.index
+        )
+
+
 class _CombinedChunkSplitter:
     """
     Splits a streaming chunk that carries BOTH response content and a
@@ -209,6 +248,10 @@ class _CombinedChunkSplitter:
         self._buffer: deque[ModelResponseStream] = deque()
         self._held_calls: Mapping[int, tuple[ModelResponseStream, _ToolCallPart]] = MappingProxyType({})
         self._live_call_index: int | None = None
+        self._held_signature: ModelResponseStream | None = None
+        self._unsigned_thinking = io.StringIO()
+        self._unsigned_thinking_length = 0
+        self._thinking_index: int | None = None
         self._closed = False
 
     async def aclose(self) -> None:
@@ -458,7 +501,45 @@ class _CombinedChunkSplitter:
             for split_chunk in self._split_by_payload_kind(combined_chunk)  # comprehension-ok: flattens nested splits
         )
         block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks)))
-        return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, block_chunks)))
+        signed_chunks: Final = tuple(itertools.chain.from_iterable(map(self._defer_signature, block_chunks)))
+        return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, signed_chunks)))
+
+    def _end_thinking_run(self) -> "tuple[ModelResponseStream, ...]":
+        held: Final = self._held_signature
+        self._held_signature = None
+        self._unsigned_thinking = io.StringIO()
+        self._unsigned_thinking_length = 0
+        self._thinking_index = None
+        return (held,) if held is not None else ()
+
+    def _defer_signature(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]":
+        """Stream a signed thinking chunk's new text as thinking and send its signature once, when the block ends.
+
+        Anthropic-shaped producers stream unsigned text and then sign it, either with no text or
+        restating the unsigned text, which passes through untouched. Mistral signs every chunk while
+        each one still carries new text, which would otherwise collapse into ``signature_delta``
+        events and lose the text.
+        """
+        piece: Final = _ThinkingPiece.of(chunk)
+        released: Final = self._end_thinking_run() if piece is None or piece.index != self._thinking_index else ()
+        if piece is None:
+            return (*released, chunk)
+        self._thinking_index = piece.index
+        if not piece.signature:
+            self._unsigned_thinking_length += self._unsigned_thinking.write(piece.text)
+            return (*released, chunk)
+        self._held_signature = None
+        if not piece.text or self._restates_unsigned_thinking(piece.text):
+            return (*released, chunk)
+        signature_chunk: Final = self._with_thinking_block(
+            chunk, piece.block(text="", signature=piece.signature), keep_role=False
+        )
+        self._clear_usage(signature_chunk)
+        self._held_signature = signature_chunk
+        return (*released, self._with_thinking_block(chunk, piece.block(text=piece.text, signature=""), keep_role=True))
+
+    def _restates_unsigned_thinking(self, text: str) -> bool:
+        return len(text) == self._unsigned_thinking_length and text == self._unsigned_thinking.getvalue()
 
     def __iter__(self) -> "Iterator[ModelResponseStream]":
         return self
@@ -493,7 +574,7 @@ class _CombinedChunkSplitter:
             try:
                 chunk = next(self._sync_iter)
             except StopIteration:
-                self._buffer.extend(self._release_held_calls())
+                self._buffer.extend((*self._end_thinking_run(), *self._release_held_calls()))
                 if not self._buffer:
                     raise
                 break
@@ -512,7 +593,7 @@ class _CombinedChunkSplitter:
             try:
                 chunk = await self._async_iter.__anext__()
             except StopAsyncIteration:
-                self._buffer.extend(self._release_held_calls())
+                self._buffer.extend((*self._end_thinking_run(), *self._release_held_calls()))
                 if not self._buffer:
                     raise
                 break

tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_first_delta.py

@@ -21,7 +21,9 @@ into an open ``thinking`` block, crashing Anthropic SDK clients (Claude Code)
 with "Content block is not a text block".
 """
 
-from typing import List, Optional
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Final, List, Mapping, Optional, Sequence, Tuple
 from unittest.mock import MagicMock
 
 import pytest
@@ -30,6 +32,7 @@ import pytest
 from litellm.llms.anthropic.messages.adapters.streaming_iterator import (
     AnthropicStreamWrapper,
 )
+from litellm.types.llms.openai import ChatCompletionThinkingBlock
 from litellm.types.utils import (
     ChatCompletionDeltaToolCall,
     Delta,
@@ -558,11 +561,9 @@ async def test_empty_reasoning_delta_mid_thinking_block_is_suppressed_async():
 
 
 def _full_snapshot_signature_chunks() -> List[MagicMock]:
-    """Mirror litellm's real Anthropic streaming: incremental ``thinking_delta``
-    chunks (empty signature), then a terminal chunk whose ``thinking_blocks`` entry
-    re-states the *full accumulated thinking text* together with the signature
-    (anthropic/chat/handler.py builds the signature_delta event this way), then the
-    answer text.
+    """Incremental ``thinking_delta`` chunks (empty signature), then a terminal chunk
+    whose ``thinking_blocks`` entry re-states the *full accumulated thinking text*
+    together with the signature, then the answer text.
     """
     return [
         _thinking_chunk("Let me "),
@@ -1317,3 +1318,168 @@ async def test_signature_only_thinking_chunk_before_text_leaks_no_signature(is_a
     assert "signature" not in text_starts[0]
     assert _text_deltas(events) == ["Hello"]
     _assert_deltas_match_their_block_type(events)
+
+
+def _signed_thinking_block(
+    thinking: str, index: Optional[int] = None, signature: str = "mistral"
+) -> ChatCompletionThinkingBlock:
+    if index is None:
+        return ChatCompletionThinkingBlock(type="thinking", thinking=thinking, signature=signature)
+    return ChatCompletionThinkingBlock(type="thinking", thinking=thinking, signature=signature, index=index)
+
+
+def _signed_thinking_chunk(thinking: str, index: Optional[int] = None, signature: str = "mistral") -> MagicMock:
+    return _make_chunk(
+        Delta(
+            content=None,
+            reasoning_content=thinking,
+            thinking_blocks=[_signed_thinking_block(thinking, index, signature)],
+        )
+    )
+
+
+_SIGNED_THINKING_STREAMS: Final[Mapping[str, Tuple[MagicMock, ...]]] = MappingProxyType(
+    {
+        "unsigned-text-then-signed-new-text": (
+            _thinking_chunk("The user "),
+            _signed_thinking_chunk("wants 2+2. "),
+            _signed_thinking_chunk("That is 4."),
+            _make_chunk(Delta(content="4")),
+            _make_chunk(Delta(content=None), finish_reason="stop"),
+        ),
+        "then-text": (
+            _signed_thinking_chunk("The user "),
+            _signed_thinking_chunk("wants 2+2. "),
+            _signed_thinking_chunk("That is 4."),
+            _make_chunk(Delta(content="4")),
+            _make_chunk(Delta(content=None), finish_reason="stop"),
+        ),
+        "last-thought-shares-the-answer-chunk": (
+            _signed_thinking_chunk("The user "),
+            _signed_thinking_chunk("wants 2+2. "),
+            _make_chunk(
+                Delta(
+                    content="4",
+                    reasoning_content="That is 4.",
+                    thinking_blocks=[_signed_thinking_block("That is 4.")],
+                )
+            ),
+            _make_chunk(Delta(content=None), finish_reason="stop"),
+        ),
+        "blank-chunk-mid-thought": (
+            _signed_thinking_chunk("The user "),
+            _make_chunk(Delta(content="")),
+            _signed_thinking_chunk("wants 2+2. "),
+            _signed_thinking_chunk("That is 4."),
+            _make_chunk(Delta(content="4")),
+            _make_chunk(Delta(content=None), finish_reason="stop"),
+        ),
+        "then-tool-call": (
+            _signed_thinking_chunk("The user "),
+            _signed_thinking_chunk("wants 2+2. "),
+            _signed_thinking_chunk("That is 4."),
+            _tool_chunk("call_add", "add", '{"a": 2, "b": 2}'),
+            _make_chunk(Delta(content=None), finish_reason="tool_calls"),
+        ),
+        "thought-ends-with-finish": (
+            _signed_thinking_chunk("The user "),
+            _signed_thinking_chunk("wants 2+2. "),
+            _make_chunk(
+                Delta(
+                    content=None,
+                    reasoning_content="That is 4.",
+                    thinking_blocks=[_signed_thinking_block("That is 4.")],
+                ),
+                finish_reason="stop",
+            ),
+        ),
+        "stream-ends-mid-thought": (
+            _signed_thinking_chunk("The user "),
+            _signed_thinking_chunk("wants 2+2. "),
+            _signed_thinking_chunk("That is 4."),
+        ),
+    }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _AccumulatedThinking:
+    thinking: str
+    signature: str
+
+
+def _accumulate_thinking(start: dict, deltas: Sequence[dict]) -> _AccumulatedThinking:
+    signatures: Final = tuple(delta["signature"] for delta in deltas if delta["type"] == "signature_delta")
+    return _AccumulatedThinking(
+        thinking=(start.get("thinking") or "")
+        + "".join(delta["thinking"] for delta in deltas if delta["type"] == "thinking_delta"),
+        signature=signatures[-1] if signatures else (start.get("signature") or ""),
+    )
+
+
+def _thinking_starts(events: List[dict]) -> Mapping[int, dict]:
+    return MappingProxyType(
+        {
+            event["index"]: event["content_block"]
+            for event in events
+            if event.get("type") == "content_block_start" and event["content_block"]["type"] == "thinking"
+        }
+    )
+
+
+def _deltas_at(events: List[dict], index: int) -> Tuple[dict, ...]:
+    return tuple(
+        event["delta"] for event in events if event.get("type") == "content_block_delta" and event["index"] == index
+    )
+
+
+def _thinking_blocks_as_accumulated(events: List[dict]) -> Tuple[_AccumulatedThinking, ...]:
+    starts: Final = _thinking_starts(events)
+    return tuple(_accumulate_thinking(starts[index], _deltas_at(events, index)) for index in sorted(starts))
+
+
+def _last_delta_types_per_thinking_block(events: List[dict]) -> Tuple[str, ...]:
+    return tuple(_deltas_at(events, index)[-1]["type"] for index in sorted(_thinking_starts(events)))
+
+
+async def _events_for(chunks: Sequence[MagicMock], is_async: bool) -> List[dict]:
+    if is_async:
+        return await _drain_async(AnthropicStreamWrapper(completion_stream=_AsyncStream(list(chunks)), model="mistral-x"))
+    return _drain_sync(AnthropicStreamWrapper(completion_stream=iter(chunks), model="mistral-x"))
+
+
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.parametrize("stream", list(_SIGNED_THINKING_STREAMS), ids=list(_SIGNED_THINKING_STREAMS))
+@pytest.mark.asyncio
+async def test_signed_thinking_chunks_keep_all_text(is_async: bool, stream: str):
+    events: Final = await _events_for(_SIGNED_THINKING_STREAMS[stream], is_async)
+
+    blocks: Final = _thinking_blocks_as_accumulated(events)
+    assert "".join(block.thinking for block in blocks) == "The user wants 2+2. That is 4."
+    assert tuple(block.signature for block in blocks) == ("mistral",) * len(blocks)
+    assert len(_signature_deltas(events)) == len(blocks)
+    assert _last_delta_types_per_thinking_block(events) == ("signature_delta",) * len(blocks)
+    _assert_deltas_match_their_block_type(events)
+
+
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.asyncio
+async def test_signed_thinking_blocks_keep_their_own_text_and_signature(is_async: bool):
+    chunks: Final = (
+        _signed_thinking_chunk("First ", index=0, signature="sig-0"),
+        _signed_thinking_chunk("thought.", index=0, signature="sig-0"),
+        _signed_thinking_chunk("Second ", index=1, signature="sig-1"),
+        _signed_thinking_chunk("thought.", index=1, signature="sig-1"),
+        _make_chunk(Delta(content="Done")),
+        _make_chunk(Delta(content=None), finish_reason="stop"),
+    )
+
+    events: Final = await _events_for(chunks, is_async)
+
+    assert _thinking_blocks_as_accumulated(events) == (
+        _AccumulatedThinking(thinking="First thought.", signature="sig-0"),
+        _AccumulatedThinking(thinking="Second thought.", signature="sig-1"),
+    )
+    assert _last_delta_types_per_thinking_block(events) == ("signature_delta", "signature_delta")
+    assert _text_deltas(events) == ["Done"]
+    _assert_deltas_match_their_block_type(events)