[Frontend] Implement Batched Chat Completions (#5317)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/e2e/online_serving/test_flux_2_dev_expansion.py

@@ -136,3 +136,36 @@ def test_flux_2_dev(
     }
 
     openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_flux_2_dev_feature_cases(MODEL),
+    indirect=True,
+)
+def test_flux_2_dev_batched_chat_completions(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """Validate batched chat completions for diffusion models."""
+    messages = [
+        [{"role": "user", "content": PROMPT}],
+        [{"role": "user", "content": "A sunset over the ocean."}],
+    ]
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                "num_inference_steps": 2,
+                "height": 512,
+                "width": 512,
+                "seed": 42,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert resp.success
+    choices = resp.json_body["choices"]
+    assert len(choices) == len(messages)

tests/e2e/online_serving/test_qwen3_omni.py

@@ -250,3 +250,101 @@ def test_completions_rejected_for_thinker_talker(omni_server, openai_client) ->
         err_code=400,
     )
     assert not responses[0].success
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_batched_completions_text(omni_server, openai_client) -> None:
+    """Ensure that we can make a batch chat completions request (text only)."""
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": [
+                    [{"role": "user", "content": "What color is the sky? Answer in one word."}],
+                    [{"role": "user", "content": "What is 2+2? Answer in one word."}],
+                ],
+                "modalities": ["text"],
+                "max_tokens": 20,
+            },
+        },
+    )
+    assert responses
+    resp = responses[0]
+    assert resp.success
+    body = resp.json_body
+    choices = body["choices"]
+    assert len(choices) == 2
+    for choice in choices:
+        assert choice["message"]["content"]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_batched_completions_audio_out(omni_server, openai_client) -> None:
+    """Ensure that we can make a batch chat completions request (audio + text)."""
+    messages = [
+        [{"role": "user", "content": "Say hello."}],
+        [{"role": "user", "content": "Say goodbye."}],
+    ]
+    num_messages = len(messages)
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                "modalities": ["text", "audio"],
+                "max_tokens": 200,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert resp.success
+    choices = resp.json_body["choices"]
+
+    assert len(choices) == num_messages
+
+    # Every message in the batch should contain both the audio and text content
+    for choice in choices:
+        assert choice["message"]["audio"]["data"]
+        assert choice["message"]["content"]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize(
+    "sad_opts",
+    [
+        {"stream": "not-a-bool"},  # Fails validation
+        {"n": 2},  # vLLM Batched endpoints doesn't support n > 1
+    ],
+)
+def test_batched_completions_with_bad_values(omni_server, openai_client, sad_opts) -> None:
+    """Ensure that a bad values are correctly handled as 400s."""
+    messages = [
+        [{"role": "user", "content": "Say hello."}],
+        [{"role": "user", "content": "Say goodbye."}],
+    ]
+    responses = openai_client.send_batched_chat_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "messages": messages,
+                **sad_opts,
+            },
+        },
+    )
+    assert responses and len(responses) == 1
+    resp = responses[0]
+    assert not responses[0].success
+    assert resp.status_code == 400

tests/entrypoints/openai/test_batch_chat_completions.py

@@ -0,0 +1,150 @@
+import asyncio
+from unittest.mock import AsyncMock
+
+import pytest
+from openai.types.chat import ChatCompletionUserMessageParam
+from openai.types.chat.chat_completion_audio import ChatCompletionAudio
+from starlette.requests import Request
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    BatchChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+    ChatMessage,
+)
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
+
+from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+MSG_BATCH: list[list[ChatCompletionMessageParam]] = [
+    [ChatCompletionUserMessageParam(role="user", content="What color is the sky?")],
+    [ChatCompletionUserMessageParam(role="user", content="What is 2+2?")],
+]
+collapse = OmniOpenAIServingChatBatch._maybe_collapse_choices
+
+
+# Helpers for testing creating a __new__ serving instance and ensuring IDs are correct
+def _make_raw_request(headers: list[tuple[bytes, bytes]]):
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/v1/chat/completions/batch",
+        "headers": headers,
+    }
+    return Request(scope)
+
+
+def _make_completion(text):
+    return ChatCompletionResponse(
+        model="test-model",
+        choices=[
+            ChatCompletionResponseChoice(
+                index=0,
+                message=ChatMessage(role="assistant", content=text),
+            )
+        ],
+        usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0),
+    )
+
+
+def _make_handler(responses):
+    handler = OmniOpenAIServingChatBatch.__new__(OmniOpenAIServingChatBatch)
+    handler.create_chat_completion = AsyncMock(side_effect=responses)
+    return handler
+
+
+# Helpers for creating choices of different modality types
+def _text_choice(text="hello"):
+    return ChatCompletionResponseChoice(
+        index=0,
+        message=ChatMessage(role="assistant", content=text),
+    )
+
+
+def _audio_choice():
+    audio = ChatCompletionAudio(id="a1", data="base64audio", expires_at=0, transcript="")
+    return ChatCompletionResponseChoice(
+        index=0,
+        message=ChatMessage(role="assistant", content=None, audio=audio),
+    )
+
+
+def _image_choice():
+    content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]
+    msg = ChatMessage.model_construct(role="assistant")
+    object.__setattr__(msg, "content", content)
+    return ChatCompletionResponseChoice.model_construct(
+        index=0,
+        message=msg,
+    )
+
+
+def _get_subrequest_ids(handler):
+    return [call.args[0].request_id for call in handler.create_chat_completion.call_args_list]
+
+
+@pytest.mark.parametrize("has_header", [True, False])
+def test_subrequest_ids_are_unique(has_header: bool):
+    """Ensure that a submitted batch request creates unique subrequest IDs"""
+    header = None if not has_header else b"user-123"
+    headers = [] if not header else [(b"x-request-id", header)]
+
+    handler = _make_handler([_make_completion("hello"), _make_completion("world")])
+    request = BatchChatCompletionRequest(
+        model="test-model",
+        messages=MSG_BATCH,
+        seed=42,
+    )
+    result = asyncio.run(handler.create_batch_chat_completion(request, _make_raw_request(headers)))
+    assert not isinstance(result, ErrorResponse)
+    assert result.id.startswith("chatcmpl-batch")
+
+    # Ensure request IDs are unique and that if we had a header, it's in all reqs
+    req_ids = _get_subrequest_ids(handler)
+    assert len(set(req_ids)) == len(req_ids) == 2
+
+    if header is not None:
+        assert all(header.decode("utf-8") in req_id for req_id in req_ids)
+
+
+### Tests for ensuring that the (currently batch specific) choice collapsing is correct
+def test_single_text_passthrough():
+    result = collapse([_text_choice()])
+    assert result.message.content == "hello"
+
+
+def test_single_image_passthrough():
+    result = collapse([_image_choice()])
+    assert isinstance(result.message.content, list)
+    assert result.message.content[0]["type"] == "image_url"
+
+
+def test_text_plus_audio():
+    result = collapse([_text_choice(), _audio_choice()])
+    assert result.message.content == "hello"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_audio_plus_text_order_independent():
+    result = collapse([_audio_choice(), _text_choice()])
+    assert result.message.content == "hello"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_image_plus_audio():
+    result = collapse([_image_choice(), _audio_choice()])
+    assert isinstance(result.message.content, list)
+    assert result.message.content[0]["type"] == "image_url"
+    assert result.message.audio.data == "base64audio"
+
+
+def test_two_content_choices_raises():
+    with pytest.raises(ValueError, match="Multiple content choices cannot be set"):
+        collapse([_text_choice(), _text_choice()])
+
+
+def test_three_choices_raises():
+    with pytest.raises(ValueError, match="got 3"):
+        collapse([_text_choice(), _audio_choice(), _text_choice()])

tests/helpers/runtime.py

@@ -1252,6 +1252,28 @@ class OpenAIClientHandler:
         )
         return [resp]
 
+    def send_batched_chat_completions_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """Post to batched chat completions."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
     def send_completions_http_request(
         self,
         request_config: dict[str, Any],

vllm_omni/config/endpoint_policy.py

@@ -21,7 +21,6 @@ class RouteTarget(NamedTuple):
 class OmniServingCapability(Enum):
     """Serving capabilities that pipelines can shut down."""
 
-    CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"}))
     COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
 
     @property
@@ -39,16 +38,6 @@ class EndpointRestriction:
     reason: str
 
 
-# Routes that are not supported for any model, but are supported in vLLM.
-# This is only temporary to avoid 500s for batched chat completions.
-UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = (
-    EndpointRestriction(
-        OmniServingCapability.CHAT_COMPLETIONS_BATCH,
-        "Batched chat completions are not yet supported by vLLM Omni.",
-    ),
-)
-
-
 def build_rejection_handler(reason: str):
     """Build a rejection handler for a given endpoint for the provided reason."""
 
@@ -71,11 +60,7 @@ def shutdown_unsupported_routes(
     """
     from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
 
-    # Generally these should not overlap since there is no point. If they do,
-    # we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages.
-    restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES)
-
-    for end_restrict in restricted_endpoints:
+    for end_restrict in endpoint_restrictions:
         capability = end_restrict.capability
         # Remove the route from the app
         _remove_route_from_app(app, capability.path, capability.methods)

vllm_omni/entrypoints/openai/api_server.py

@@ -39,6 +39,7 @@ from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, Tool
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
 from vllm.entrypoints.openai.chat_completion.protocol import (
+    BatchChatCompletionRequest,
     ChatCompletionRequest,
     ChatCompletionResponse,
 )
@@ -95,6 +96,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes
 from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
 from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
 from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_endpoint
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
@@ -513,6 +515,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         # OMNI: Remove upstream routes that we override with omni-specific handlers
         _remove_route_from_app(app, "/v1/chat/completions", {"POST"})
+        _remove_route_from_app(app, "/v1/chat/completions/batch", {"POST"})
         _remove_route_from_app(app, "/v1/models", {"GET"})  # Remove upstream /v1/models to use omni's handler
         app.include_router(router)
 
@@ -777,6 +780,10 @@ async def omni_init_app_state(
             diffusion_engine=engine_client,  # type: ignore
             model_name=model_name,
         )
+        state.openai_serving_chat_batch = OmniOpenAIServingChatBatch.for_diffusion(
+            diffusion_engine=engine_client,  # type: ignore
+            model_name=model_name,
+        )
 
         # audio related
         state.openai_serving_speech = None
@@ -950,30 +957,33 @@ async def omni_init_app_state(
         if "generate" in supported_tasks
         else None
     )
-    state.openai_serving_chat = (
-        OmniOpenAIServingChat(
-            engine_client,
-            state.openai_serving_models,
-            args.response_role,
-            online_renderer=state.online_renderer,
-            request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            default_chat_template_kwargs=args.default_chat_template_kwargs,
-            trust_request_chat_template=args.trust_request_chat_template,
-            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
-            enable_auto_tools=args.enable_auto_tool_choice,
-            exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
-            tool_parser=args.tool_call_parser,
-            reasoning_parser=args.structured_outputs_config.reasoning_parser,
-            enable_prompt_tokens_details=args.enable_prompt_tokens_details,
-            enable_force_include_usage=args.enable_force_include_usage,
-            enable_log_outputs=args.enable_log_outputs,
-            enable_log_deltas=args.enable_log_deltas,
-        )
-        if "generate" in supported_tasks
-        else None
+
+    _chat_kwargs = dict(
+        engine_client=engine_client,
+        models=state.openai_serving_models,
+        response_role=args.response_role,
+        online_renderer=state.online_renderer,
+        request_logger=request_logger,
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        default_chat_template_kwargs=args.default_chat_template_kwargs,
+        trust_request_chat_template=args.trust_request_chat_template,
+        return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+        enable_auto_tools=args.enable_auto_tool_choice,
+        exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
+        tool_parser=args.tool_call_parser,
+        reasoning_parser=args.structured_outputs_config.reasoning_parser,
+        enable_prompt_tokens_details=args.enable_prompt_tokens_details,
+        enable_force_include_usage=args.enable_force_include_usage,
+        enable_log_outputs=args.enable_log_outputs,
+        enable_log_deltas=args.enable_log_deltas,
     )
+
+    state.openai_serving_chat = OmniOpenAIServingChat(**_chat_kwargs) if "generate" in supported_tasks else None
+    state.openai_serving_chat_batch = (
+        OmniOpenAIServingChatBatch(**_chat_kwargs) if "generate" in supported_tasks else None
+    )
+
     # Warm up chat template processing to avoid first-request latency
     if state.openai_serving_chat is not None:
         state.openai_serving_chat.warmup()
@@ -1162,6 +1172,10 @@ def Omnichat(request: Request) -> OmniOpenAIServingChat | None:
     return request.app.state.openai_serving_chat
 
 
+def OmniBatchChat(request: Request) -> OmniOpenAIServingChatBatch | None:
+    return request.app.state.openai_serving_chat_batch
+
+
 def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
     return request.app.state.openai_serving_speech
 
@@ -1245,6 +1259,45 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
     return StreamingResponse(content=generator, media_type="text/event-stream")
 
 
+@router.post(
+    "/v1/chat/completions/batch",
+    dependencies=[Depends(validate_json_request)],
+    responses={
+        HTTPStatus.OK.value: {},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_IMPLEMENTED.value: {"model": ErrorResponse},
+    },
+)
+@with_cancellation
+@load_aware_call
+async def create_batch_chat_completion(request: BatchChatCompletionRequest, raw_request: Request):
+    handler = OmniBatchChat(raw_request)
+    if handler is None:
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
+        if base_server is None:
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND.value,
+                detail="The model does not support Chat Completions API",
+            )
+        return base_server.create_error_response(message="The model does not support Chat Completions API")
+    try:
+        result = await handler.create_batch_chat_completion(request, raw_request)
+    except (EngineGenerateError, EngineDeadError) as exc:
+        return _create_engine_error_json_response(raw_request, exc)
+    except Exception as e:
+        logger.exception("Batched chat completion failed: %s", e)
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
+
+    if isinstance(result, ErrorResponse):
+        return JSONResponse(
+            content=result.model_dump(),
+            status_code=result.error.code if result.error else 400,
+        )
+    return JSONResponse(content=result.model_dump(mode="json"))
+
+
 _remove_route_from_router(router, "/v1/audio/speech", {"POST"})
 
 

vllm_omni/entrypoints/openai/batch_serving.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import time
+
+from fastapi import Request
+from pydantic import ValidationError
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    BatchChatCompletionRequest,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+)
+from vllm.entrypoints.openai.engine.protocol import (
+    ErrorResponse,
+    UsageInfo,
+)
+from vllm.logger import init_logger
+
+from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+logger = init_logger(__name__)
+
+
+class OmniOpenAIServingChatBatch(OmniOpenAIServingChat):
+    @staticmethod
+    def _maybe_collapse_choices(choices: list[ChatCompletionResponseChoice]):
+        """Collapse content + audio choices into one choice.
+
+        NOTE: This is a hack for models that produce separate text and
+        audio choices to ensure we have 1:1 mapping between number of
+        inputs in the batch and responses.
+        """
+        content_choice, audio_choice = None, None
+        if len(choices) == 1:
+            return choices[0]
+        if len(choices) != 2:
+            raise ValueError(f"Expected 1 or 2 choices to collapse, got {len(choices)}. ")
+
+        for choice in choices:
+            if choice.message.audio:
+                if audio_choice is not None:
+                    raise ValueError("Multiple audio choices cannot be set in one completion message")
+                audio_choice = choice
+            else:
+                if content_choice is not None:
+                    raise ValueError("Multiple content choices cannot be set in one completion message")
+                content_choice = choice
+
+        if content_choice is None or audio_choice is None:
+            raise ValueError(
+                "Expected one content and one audio choice, but got: "
+                f"content={content_choice is not None}, audio={audio_choice is not None}"
+            )
+        content_choice.message.audio = audio_choice.message.audio
+        return content_choice
+
+    @staticmethod
+    def _get_subrequest_ids(base_id, request: BatchChatCompletionRequest) -> list[str]:
+        """Get the request ID for each entry in the batch request to be resubmitted to chat completions."""
+        return [f"{base_id}-idx-{idx}" for idx in range(len(request.messages))]
+
+    async def create_batch_chat_completion(
+        self,
+        request: BatchChatCompletionRequest,
+        raw_request: Request,
+    ) -> ChatCompletionResponse | ErrorResponse:
+        """Given a request, submit each request to chat completions & collect the results."""
+        model = ""
+        enabled_streaming = False
+        chat_requests: list[ChatCompletionRequest] = []
+        choices: list[ChatCompletionResponseChoice] = []
+        total_prompt_tokens = 0
+        total_completion_tokens = 0
+
+        # Get the base request from the raw_request before maybe mutating the header
+        base_id = OmniOpenAIServingChatBatch._base_request_id(
+            raw_request,
+            request.request_id,
+        )
+        batch_req_id = f"chatcmpl-batch-{base_id}"
+        sub_req_ids = self._get_subrequest_ids(base_id, request)
+
+        # Remove the raw request header if it exists now that we have subreq IDs,
+        # since this takes priority over the object request_ids
+        if raw_request.headers.get("X-Request-Id"):
+            copy_headers = raw_request.headers.mutablecopy()
+            del copy_headers["X-Request-Id"]
+            raw_request._headers = copy_headers
+
+        for idx, msg in enumerate(request.messages):
+            # Update the current request ID to avoid subrequest collisions
+            request.request_id = sub_req_ids[idx]
+            try:
+                chat_cmp_request = request.to_chat_completion_request(msg)
+            except ValidationError as e:
+                return self._create_error_response(
+                    f"Message {idx} could not be converted to a chat completion request: {e}",
+                )
+            # Streaming isn't supported for batch chat completions,
+            # so always set to False & warn if it was requested.
+            enabled_streaming |= chat_cmp_request.stream
+            chat_cmp_request.stream = False
+            chat_requests.append(chat_cmp_request)
+
+        if enabled_streaming:
+            logger.warning("Streaming is not supported for batched chat completions; ignoring stream=True.")
+
+        # Submit each chat completion request as a task, then gather results.
+        # TODO (Alex): optimize this
+        tasks = [asyncio.create_task(self.create_chat_completion(c, raw_request)) for c in chat_requests]
+        try:
+            results = await asyncio.gather(*tasks)
+        finally:
+            # Ensure we cancel remaining tasks if needed, e.g., early exit due to bad behavior
+            for t in tasks:
+                if not t.done():
+                    t.cancel()
+
+        for i, resp in enumerate(results):
+            if isinstance(resp, ErrorResponse):
+                return resp
+            completion: ChatCompletionResponse = resp
+            model = completion.model
+            # FIXME (Alex): We should probably handle this in chat completions,
+            # not here, but we need to ensure streaming is properly handled
+            try:
+                collapsed_choice = self._maybe_collapse_choices(completion.choices)
+            except ValueError as e:
+                return self._create_error_response(
+                    f"Failed to collapse choices with error: {e}",
+                )
+            collapsed_choice.index = i
+            choices.append(collapsed_choice)
+            if completion.usage:
+                total_prompt_tokens += completion.usage.prompt_tokens
+                total_completion_tokens += completion.usage.completion_tokens
+
+        usage = UsageInfo(
+            prompt_tokens=total_prompt_tokens,
+            completion_tokens=total_completion_tokens,
+            total_tokens=total_prompt_tokens + total_completion_tokens,
+        )
+
+        return ChatCompletionResponse(
+            id=batch_req_id,
+            created=int(time.time()),
+            model=model,
+            choices=choices,
+            usage=usage,
+        )