fix(bedrock): find the DeepSeek R1 end-of-thinking marker however it is chunked

The streaming iterator ended the thinking phase by comparing a whole chunk against
"</think>", so it only worked when the marker arrived as a chunk of its own. Bedrock
does not promise one token per chunk, and when the marker is split across chunks or
glued to the text on either side nothing flips has_finished_thinking. Every later
chunk stays filed as reasoning, the raw marker leaks into reasoning_content, and the
turn reaches the client with no content at all.

Carry a buffer across chunks and split on the marker wherever it lands. Text that
could still be the start of it is held back until the next chunk decides, and
released if the stream ends first, so nothing is lost either way. This is what
transform_response already does off one string via _parse_content_for_reasoning.

The iterator had no tests; the ones added here cover the marker at every
fragmentation, plus a chunk-size invariance check over the same generation.
This commit is contained in:
谭成 2026-08-22 15:09:10 +08:00
parent b9bff0998c
commit fe8b1aa17a
2 changed files with 149 additions and 6 deletions

View file

@ -76,10 +76,44 @@ class AmazonDeepSeekR1Config(AmazonLlamaConfig):
return response
_END_OF_THINKING: Final = "</think>"
class AmazonDeepseekR1ResponseIterator(BaseModelResponseIterator):
def __init__(self, streaming_response: Any, sync_stream: bool) -> None:
super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
self.has_finished_thinking = False
self.held_back = ""
def _split_on_end_of_thinking(self, generated_content: str, is_last_chunk: bool) -> tuple[str, str]:
"""Split a chunk of the thinking phase into (reasoning, content).
``</think>`` is not guaranteed to arrive as a chunk of its own: it can be glued to the
text on either side, or split across chunks. Matching the whole marker against one chunk
misses both, leaving every later chunk filed as reasoning and ``content`` empty for the
entire turn. Text that could still be the start of the marker is held back until the next
chunk decides it, and released if the stream ends first.
"""
buffered: Final = self.held_back + generated_content
reasoning, marker, content = buffered.partition(_END_OF_THINKING)
if marker:
verbose_logger.debug("Deepseek r1: </think> received, setting has_finished_thinking to True")
self.has_finished_thinking = True
self.held_back = ""
return reasoning, content
if is_last_chunk:
self.held_back = ""
return buffered, ""
partial: Final = next(
(
length
for length in range(min(len(buffered), len(_END_OF_THINKING) - 1), 0, -1)
if buffered.endswith(_END_OF_THINKING[:length])
),
0,
)
self.held_back = buffered[len(buffered) - partial :] if partial else ""
return buffered[: len(buffered) - partial] if partial else buffered, ""
def chunk_parser(self, chunk: dict) -> ModelResponseStream:
"""
@ -88,10 +122,11 @@ class AmazonDeepseekR1ResponseIterator(BaseModelResponseIterator):
try:
typed_chunk: Final = AmazonDeepSeekR1StreamingResponse(**chunk)
generated_content = typed_chunk["generation"]
if generated_content == "</think>" and not self.has_finished_thinking:
verbose_logger.debug("Deepseek r1: </think> received, setting has_finished_thinking to True")
generated_content = ""
self.has_finished_thinking = True
reasoning_delta: str = ""
if not self.has_finished_thinking:
reasoning_delta, generated_content = self._split_on_end_of_thinking(
generated_content, is_last_chunk=typed_chunk["stop_reason"] is not None
)
prompt_token_count: Final = typed_chunk.get("prompt_token_count") or 0
generation_token_count: Final = typed_chunk.get("generation_token_count") or 0
@ -106,8 +141,8 @@ class AmazonDeepseekR1ResponseIterator(BaseModelResponseIterator):
StreamingChoices(
finish_reason=typed_chunk["stop_reason"],
delta=Delta(
content=(generated_content if self.has_finished_thinking else None),
reasoning_content=(generated_content if not self.has_finished_thinking else None),
content=generated_content if self.has_finished_thinking else None,
reasoning_content=reasoning_delta or None,
),
)
],

View file

@ -0,0 +1,108 @@
import os
import sys
import pytest
# Ensure the project root is on the import path so `litellm` can be imported when
# tests are executed from any working directory.
sys.path.insert(0, os.path.abspath("../../../../../.."))
from litellm.llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import (
AmazonDeepseekR1ResponseIterator,
)
REASONING = "Let me think about this."
ANSWER = "The answer is 4."
def _drain(generations: list[str]) -> tuple[str, str]:
"""Feed one turn through the iterator and return its (reasoning, content) totals."""
iterator = AmazonDeepseekR1ResponseIterator(streaming_response=None, sync_stream=True)
reasoning_parts: list[str] = []
content_parts: list[str] = []
for position, generation in enumerate(generations):
chunk = iterator.chunk_parser(
{
"generation": generation,
"stop_reason": "stop" if position == len(generations) - 1 else None,
"prompt_token_count": 1,
"generation_token_count": 1,
}
)
delta = chunk.choices[0].delta
reasoning_parts.append(getattr(delta, "reasoning_content", None) or "")
content_parts.append(getattr(delta, "content", None) or "")
return "".join(reasoning_parts), "".join(content_parts)
@pytest.mark.parametrize(
"generations",
[
pytest.param([REASONING, "</think>", ANSWER], id="marker_alone"),
pytest.param([REASONING, "</", "think>", ANSWER], id="marker_split_in_two"),
pytest.param([REASONING, *"</think>", ANSWER], id="marker_split_per_character"),
pytest.param([REASONING, f"</think>{ANSWER}"], id="marker_glued_to_answer"),
pytest.param([f"{REASONING}</think>", ANSWER], id="marker_glued_to_reasoning"),
pytest.param([f"{REASONING}</think>{ANSWER}"], id="whole_turn_in_one_chunk"),
pytest.param(["Let me think ", "about this.", "</think>", "The answer ", "is 4."], id="both_sides_fragmented"),
],
)
def test_end_of_thinking_is_found_however_the_marker_is_chunked(generations):
"""`</think>` is only routed correctly when it lands as a chunk of its own.
Bedrock does not promise one token per chunk, so the marker can arrive split across chunks or
glued to the text on either side. Comparing a whole chunk against `"</think>"` misses both, and
since nothing else flips `has_finished_thinking`, every later chunk stays filed as reasoning and
the turn reaches the client with no content at all.
"""
reasoning, content = _drain(generations)
assert reasoning == REASONING
assert content == ANSWER
def test_reasoning_and_content_do_not_depend_on_where_the_stream_was_cut():
"""The same generation must assemble identically at every chunk size."""
whole = f"{REASONING}</think>{ANSWER}"
results = {
size: _drain([whole[i : i + size] for i in range(0, len(whole), size)])
for size in (1, 2, 3, 5, 8, 13, len(whole))
}
assert set(results.values()) == {(REASONING, ANSWER)}
def test_text_resembling_the_marker_is_not_swallowed():
"""Holding back a possible marker prefix must not eat text that never completes one."""
assert _drain(["a < b and c </ d", "</think>", ANSWER]) == ("a < b and c </ d", ANSWER)
def test_unterminated_thinking_is_released_at_the_end_of_the_stream():
"""A turn that never closes its thinking still owes the client every token it produced."""
assert _drain(["still thinking ", "</thi"]) == ("still thinking </thi", "")
def test_marker_in_the_answer_is_left_alone_once_thinking_ended():
"""Only the first `</think>` ends the thinking phase; a later one is ordinary content."""
assert _drain([REASONING, "</think>", "write </think> to close the block"]) == (
REASONING,
"write </think> to close the block",
)
def test_usage_and_finish_reason_still_come_from_the_chunk():
iterator = AmazonDeepseekR1ResponseIterator(streaming_response=None, sync_stream=True)
chunk = iterator.chunk_parser(
{
"generation": "done",
"stop_reason": "stop",
"prompt_token_count": 11,
"generation_token_count": 7,
}
)
assert chunk.choices[0].finish_reason == "stop"
assert chunk.usage["prompt_tokens"] == 11
assert chunk.usage["completion_tokens"] == 7
assert chunk.usage["total_tokens"] == 18