diff --git a/litellm/router.py b/litellm/router.py index 2d3d6971e3b..bb8485e0ef5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d59674f1c57..3d592737e1e 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -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."""