diff --git a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index 62a1e5f0a97..b19c0d8e2a9 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/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 diff --git a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_first_delta.py b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_first_delta.py index 32ec0ffae8e..a3cb6ed98f4 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_first_delta.py +++ b/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)