From fe8b1aa17a67a22b8334c3be68b7b06071ae4c76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B0=AD=E6=88=90?= Date: Sat, 22 Aug 2026 15:09:10 +0800 Subject: [PATCH 1/2] 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 "", 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. --- .../amazon_deepseek_transformation.py | 47 +++++++- .../test_amazon_deepseek_transformation.py | 108 ++++++++++++++++++ 2 files changed, 149 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py index d86c756ca99..2a4daa94a20 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py @@ -76,10 +76,44 @@ class AmazonDeepSeekR1Config(AmazonLlamaConfig): return response +_END_OF_THINKING: Final = "" + + 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). + + ```` 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: 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 == "" and not self.has_finished_thinking: - verbose_logger.debug("Deepseek r1: 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, ), ) ], diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py new file mode 100644 index 00000000000..0f95022136c --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py @@ -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, "", ANSWER], id="marker_alone"), + pytest.param([REASONING, "", ANSWER], id="marker_split_in_two"), + pytest.param([REASONING, *"", ANSWER], id="marker_split_per_character"), + pytest.param([REASONING, f"{ANSWER}"], id="marker_glued_to_answer"), + pytest.param([f"{REASONING}", ANSWER], id="marker_glued_to_reasoning"), + pytest.param([f"{REASONING}{ANSWER}"], id="whole_turn_in_one_chunk"), + pytest.param(["Let me think ", "about this.", "", "The answer ", "is 4."], id="both_sides_fragmented"), + ], +) +def test_end_of_thinking_is_found_however_the_marker_is_chunked(generations): + """`` 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 `""` 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}{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 ", ANSWER]) == ("a < b and c ` ends the thinking phase; a later one is ordinary content.""" + assert _drain([REASONING, "", "write to close the block"]) == ( + REASONING, + "write 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 From d2df7c65d6a850cde49e6432d88f4e6d45bfa0a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B0=AD=E6=88=90?= Date: Sat, 22 Aug 2026 15:57:56 +0800 Subject: [PATCH 2/2] test(bedrock): import litellm directly instead of patching sys.path pytest's rootdir already makes litellm importable, and TQ003 counts the insert against the test-quality budget. --- .../test_amazon_deepseek_transformation.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py index 0f95022136c..3cd31fa4437 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py @@ -1,12 +1,5 @@ -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, )