mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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:
parent
b9bff0998c
commit
fe8b1aa17a
2 changed files with 149 additions and 6 deletions
|
|
@ -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,
|
||||
),
|
||||
)
|
||||
],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue