fix(router): keep content blocks on a trailing assistant turn when extending it with the partial output

A trailing assistant message whose content is a list of blocks was stringified before the partial output was appended, so the fallback received the Python repr of the blocks as its prefill. The partial now lands as one more text block, and a content shape that cannot be extended surfaces the original stream error instead of a garbled continuation.
This commit is contained in:
mateo-berri 2026-10-01 18:11:49 -07:00
parent 197a58b6d2
commit d2ab01a88b
2 changed files with 35 additions and 12 deletions

View file

@ -475,6 +475,7 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"})
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
@ -2997,12 +2998,13 @@ class Router:
getattr(complete_response_object_usage, "completion_tokens", 0) or 0
)
reduced_ceilings: Final = self._continuation_output_ceilings(initial_kwargs, emitted_tokens)
if reduced_ceilings is None:
self._raise_original_mid_stream_error(e)
initial_kwargs.update(reduced_ceilings)
initial_kwargs["messages"] = self._build_completion_continuation_input(
continuation_messages: Final = self._build_completion_continuation_input(
messages, e.generated_content
)
if reduced_ceilings is None or continuation_messages is None:
self._raise_original_mid_stream_error(e)
initial_kwargs.update(reduced_ceilings)
initial_kwargs["messages"] = continuation_messages
initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER
else:
initial_kwargs["messages"] = messages
@ -3243,17 +3245,23 @@ class Router:
@staticmethod
def _build_completion_continuation_input(
messages: Sequence[Mapping[str, str]],
messages: Sequence[Mapping[str, object]],
generated_content: str,
) -> Sequence[Mapping[str, object]]:
) -> Sequence[Mapping[str, object]] | None:
"""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."""
request never ends in two assistant messages. None when that turn's content has a shape this cannot extend."""
last: Final = messages[-1] if messages else None
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]
if last is None or last.get("role") != "assistant":
return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}]
content: Final = last.get("content")
if content is None or isinstance(content, str):
return [*messages[:-1], {**last, "content": f"{content or ''}{generated_content}", "prefix": True}]
try:
blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content)
except ValidationError:
return None
extended: Final = [*blocks, {"type": "text", "text": generated_content}]
return [*messages[:-1], {**last, "content": extended, "prefix": True}]
@staticmethod
def _build_responses_continuation_input(

View file

@ -2901,6 +2901,21 @@ def test_build_completion_continuation_input_folds_into_trailing_plain_assistant
assert built[-1] == {"role": "assistant", "content": "Sure, here it is", "prefix": True}
def test_build_completion_continuation_input_keeps_structured_assistant_content():
"""Content blocks on a trailing assistant turn stay blocks: the partial lands as one more text
block instead of a stringified list, and a content shape that cannot be extended declines."""
messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}]}]
built = litellm.Router._build_completion_continuation_input(messages, "here it is")
assert built is not None
assert [m["role"] for m in built] == ["user", "assistant"]
assert built[-1] == {
"role": "assistant",
"content": [{"type": "text", "text": "Sure, "}, {"type": "text", "text": "here it is"}],
"prefix": True,
}
assert litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42}], "x") is None
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."""