From b15d4e6a64f80725a9b33e3fa4bba2fbf100d44c Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 7 Aug 2026 16:12:09 -0600 Subject: [PATCH] [Frontend] Implement Batched Chat Completions (#5317) Signed-off-by: Alex Brooks --- .../test_flux_2_dev_expansion.py | 33 ++++ tests/e2e/online_serving/test_qwen3_omni.py | 98 ++++++++++++ .../openai/test_batch_chat_completions.py | 150 +++++++++++++++++ tests/helpers/runtime.py | 22 +++ vllm_omni/config/endpoint_policy.py | 17 +- vllm_omni/entrypoints/openai/api_server.py | 99 +++++++++--- vllm_omni/entrypoints/openai/batch_serving.py | 151 ++++++++++++++++++ 7 files changed, 531 insertions(+), 39 deletions(-) create mode 100644 tests/entrypoints/openai/test_batch_chat_completions.py create mode 100644 vllm_omni/entrypoints/openai/batch_serving.py diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index 9bc3308b8..19a14ce39 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/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) diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index ffddafbc0..a54d2fe8a 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/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 diff --git a/tests/entrypoints/openai/test_batch_chat_completions.py b/tests/entrypoints/openai/test_batch_chat_completions.py new file mode 100644 index 000000000..9fc72b2c1 --- /dev/null +++ b/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()]) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 5c1601cd0..b01c152f5 100644 --- a/tests/helpers/runtime.py +++ b/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], diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py index 5a1449614..d22c0acc1 100644 --- a/vllm_omni/config/endpoint_policy.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 048fea3c5..791f9063a 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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"}) diff --git a/vllm_omni/entrypoints/openai/batch_serving.py b/vllm_omni/entrypoints/openai/batch_serving.py new file mode 100644 index 000000000..bc5ef36ca --- /dev/null +++ b/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, + )