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..3cd31fa4437 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_deepseek_transformation.py @@ -0,0 +1,101 @@ +import pytest + +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