diff --git a/litellm/router.py b/litellm/router.py index 849edfc66a7..2d3d6971e3b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3246,10 +3246,11 @@ class Router: messages: Sequence[Mapping[str, str]], generated_content: str, ) -> Sequence[Mapping[str, object]]: - """Append the partial output as an assistant prefill, extending an existing trailing prefill in place.""" + """Append the partial output as an assistant prefill, or extend a trailing assistant turn in place so the + request never ends in two assistant messages.""" last: Final = messages[-1] if messages else None - if last is not None and last.get("role") == "assistant" and last.get("prefix"): - merged: Final = {**last, "content": str(last.get("content") or "") + generated_content} + if last is not None and last.get("role") == "assistant": + merged: Final = {**last, "content": str(last.get("content") or "") + generated_content, "prefix": True} return [*messages[:-1], merged] prefill: Final = {"role": "assistant", "content": generated_content, "prefix": True} return [*messages, prefill] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 8676e2079f4..d59674f1c57 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2892,6 +2892,15 @@ def test_build_completion_continuation_input_folds_into_existing_prefill(): assert twice[-1] == {"role": "assistant", "content": "part one part two", "prefix": True} +def test_build_completion_continuation_input_folds_into_trailing_plain_assistant_turn(): + """A request that already ends in an assistant message takes the partial as its prefill + instead of gaining a second assistant turn, which prefill providers reject.""" + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Sure, "}] + built = litellm.Router._build_completion_continuation_input(messages, "here it is") + assert [m["role"] for m in built] == ["user", "assistant"] + assert built[-1] == {"role": "assistant", "content": "Sure, here it is", "prefix": True} + + def test_continuation_output_ceilings_reduces_by_emitted_tokens(): """A continuation must complete within the caller's original allowance, so each output ceiling is reduced by the tokens already emitted."""