import asyncio import base64 import json import time import uuid from collections.abc import AsyncGenerator, AsyncIterator, Sequence from datetime import datetime, timedelta, timezone from io import BytesIO from typing import Any, Callable, Optional, Union import jinja2 from fastapi import Request from pydantic import TypeAdapter try: import soundfile except ImportError: soundfile = None from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateContentFormatOption, ConversationMessage, apply_hf_chat_template, apply_mistral_chat_template, get_history_tool_calls_cnt, make_tool_call_id, resolve_chat_template_content_format, ) from vllm.entrypoints.harmony_utils import parse_chat_output from vllm.entrypoints.openai.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, ChatMessage, ErrorResponse, FunctionCall, FunctionDefinition, PromptTokenUsageInfo, RequestResponseMetadata, ToolCall, UsageInfo, ) from vllm.entrypoints.openai.serving_chat import OpenAIServingChat from vllm.entrypoints.openai.serving_engine import ( ChatLikeRequest, EngineTokensPrompt, RequestPrompt, ResponsesRequest, TextTokensPrompt, clamp_prompt_logprobs, is_list_of, ) from vllm.entrypoints.openai.tool_parsers import ToolParser from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall from vllm.inputs.data import PromptType from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer from vllm.transformers_utils.tokenizers import ( maybe_serialize_tool_calls, truncate_tool_call_ids, validate_request_params, ) from vllm.utils import as_list from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) class OmniOpenAIServingChat(OpenAIServingChat): async def create_chat_completion( self, request: ChatCompletionRequest, raw_request: Optional[Request] = None, ) -> Union[AsyncGenerator[str, None], ChatCompletionResponse, ErrorResponse]: """ Chat Completion API similar to OpenAI's API. See https://platform.openai.com/docs/api-reference/chat/create for the API specification. This API mimics the OpenAI Chat Completion API. """ error_check_ret = await self._check_model(request) if error_check_ret is not None: logger.error("Error with model %s", error_check_ret) return error_check_ret # If the engine is dead, raise the engine's DEAD_ERROR. # This is required for the streaming case, where we return a # success status before we actually start generating text :). if self.engine_client.errored: raise self.engine_client.dead_error try: lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True) model_name = self.models.model_name(lora_request) tokenizer = await self.engine_client.get_tokenizer() tool_parser = self.tool_parser if isinstance(tokenizer, MistralTokenizer): # because of issues with pydantic we need to potentially # re-serialize the tool_calls field of the request # for more info: see comment in `maybe_serialize_tool_calls` maybe_serialize_tool_calls(request) truncate_tool_call_ids(request) validate_request_params(request) if ( request.tool_choice == "auto" and not (self.enable_auto_tools and tool_parser is not None) and not isinstance(tokenizer, MistralTokenizer) and not self.use_harmony ): # for hf tokenizers, "auto" tools requires # --enable-auto-tool-choice and --tool-call-parser return self.create_error_response( '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set' ) if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none): tool_dicts = None else: tool_dicts = [tool.model_dump() for tool in request.tools] # Common case. request_chat_template = request.chat_template chat_template_kwargs = request.chat_template_kwargs if not self.trust_request_chat_template and ( request_chat_template is not None or (chat_template_kwargs and chat_template_kwargs.get("chat_template") is not None) ): return self.create_error_response( "Chat template is passed with request, but --trust-request-chat-template is not set. " "Refused request with untrusted chat template." ) ( conversation, request_prompts, engine_prompts, ) = await self._preprocess_chat( request, tokenizer, request.messages, chat_template=request_chat_template or self.chat_template, chat_template_content_format=self.chat_template_content_format, add_generation_prompt=request.add_generation_prompt, continue_final_message=request.continue_final_message, tool_dicts=tool_dicts, documents=request.documents, chat_template_kwargs=request.chat_template_kwargs, tool_parser=tool_parser, add_special_tokens=request.add_special_tokens, ) except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e: logger.exception("Error in preprocessing prompt inputs") return self.create_error_response(f"{e} {e.__cause__}") request_id = f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}" request_metadata = RequestResponseMetadata(request_id=request_id) if raw_request: raw_request.state.request_metadata = request_metadata # Schedule the request and get the result generator. generators: list[AsyncGenerator[RequestOutput, None]] = [] try: for i, engine_prompt in enumerate(engine_prompts): if hasattr(request, "sampling_params_list"): sampling_params_list = self._to_sampling_params_list(request.sampling_params_list) else: sampling_params_list = None self._log_inputs( request_id, request_prompts[i], params_list=sampling_params_list, lora_request=lora_request, ) trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers) generator = self.engine_client.generate( prompt=engine_prompt, request_id=request_id, sampling_params_list=sampling_params_list, lora_request=lora_request, trace_headers=trace_headers, priority=request.priority, ) generators.append(generator) except ValueError as e: # TODO: Use a vllm-specific Validation Error return self.create_error_response(str(e)) assert len(generators) == 1 (result_generator,) = generators # Streaming response if request.stream: return self.chat_completion_stream_generator( request, result_generator, request_id, model_name, conversation, tokenizer, request_metadata, enable_force_include_usage=self.enable_force_include_usage, ) try: return await self.chat_completion_full_generator( request, result_generator, request_id, model_name, conversation, tokenizer, request_metadata, ) except ValueError as e: # TODO: Use a vllm-specific Validation Error return self.create_error_response(str(e)) async def _preprocess_chat( self, request: Union[ChatLikeRequest, ResponsesRequest], tokenizer: AnyTokenizer, messages: list[ChatCompletionMessageParam], chat_template: Optional[str], chat_template_content_format: ChatTemplateContentFormatOption, add_generation_prompt: bool = True, continue_final_message: bool = False, tool_dicts: Optional[list[dict[str, Any]]] = None, documents: Optional[list[dict[str, str]]] = None, chat_template_kwargs: Optional[dict[str, Any]] = None, tool_parser: Optional[Callable[[AnyTokenizer], ToolParser]] = None, add_special_tokens: bool = False, ) -> tuple[ list[ConversationMessage], Sequence[RequestPrompt], list[EngineTokensPrompt], ]: model_config = self.model_config resolved_content_format = resolve_chat_template_content_format( chat_template, tool_dicts, chat_template_content_format, tokenizer, model_config=model_config, ) conversation, mm_data_future, mm_uuids = parse_chat_messages_futures( messages, model_config, tokenizer, content_format=resolved_content_format, mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), ) _chat_template_kwargs: dict[str, Any] = dict( chat_template=chat_template, add_generation_prompt=add_generation_prompt, continue_final_message=continue_final_message, tools=tool_dicts, documents=documents, ) _chat_template_kwargs.update(chat_template_kwargs or {}) request_prompt: Union[str, list[int]] if tokenizer is None: request_prompt = "placeholder" elif isinstance(tokenizer, MistralTokenizer): request_prompt = apply_mistral_chat_template( tokenizer, messages=messages, **_chat_template_kwargs, ) else: request_prompt = apply_hf_chat_template( tokenizer=tokenizer, conversation=conversation, model_config=model_config, **_chat_template_kwargs, ) mm_data = await mm_data_future # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser # is set, we want to prevent parsing a tool_call hallucinated by the LLM should_parse_tools = tool_parser is not None and ( hasattr(request, "tool_choice") and request.tool_choice != "none" ) if should_parse_tools: if not isinstance(request, ChatCompletionRequest): msg = "Tool usage is only supported for Chat Completions API" raise NotImplementedError(msg) request = tool_parser(tokenizer).adjust_request( # type: ignore request=request ) if tokenizer is None: assert isinstance(request_prompt, str), ( "Prompt has to be a string", "when the tokenizer is not initialised", ) prompt_inputs = TextTokensPrompt(prompt=request_prompt, prompt_token_ids=[1]) elif isinstance(request_prompt, str): prompt_inputs = await self._tokenize_prompt_input_async( request, tokenizer, request_prompt, add_special_tokens=add_special_tokens, ) else: # For MistralTokenizer assert is_list_of(request_prompt, int), "Prompt has to be either a string or a list of token ids" prompt_inputs = TextTokensPrompt( prompt=tokenizer.decode(request_prompt), prompt_token_ids=request_prompt, ) engine_prompt = EngineTokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"]) if mm_data is not None: engine_prompt["multi_modal_data"] = mm_data if mm_uuids is not None: engine_prompt["multi_modal_uuids"] = mm_uuids if request.mm_processor_kwargs is not None: engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs if hasattr(request, "cache_salt") and request.cache_salt is not None: engine_prompt["cache_salt"] = request.cache_salt return conversation, [request_prompt], [engine_prompt] def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]: final_sampling_params_list = [] for sampling_params in sampling_params_list: if isinstance(sampling_params, dict): final_sampling_params_list.append(SamplingParams(**sampling_params)) elif isinstance(sampling_params, SamplingParams): final_sampling_params_list.append(sampling_params) else: raise ValueError(f"Invalid sampling params: {sampling_params}") return final_sampling_params_list def _log_inputs( self, request_id: str, inputs: Union[RequestPrompt, PromptType], params_list: Optional[list[SamplingParams]], lora_request: Optional[LoRARequest], ) -> None: if self.request_logger is None: return prompt, prompt_token_ids, prompt_embeds = None, None, None if isinstance(inputs, str): prompt = inputs elif isinstance(inputs, list): prompt_token_ids = inputs else: prompt = getattr(inputs, "prompt", None) prompt_token_ids = getattr(inputs, "prompt_token_ids", None) logger.info( "Received request %s: prompt: %r, params_list: %s, prompt_token_ids: %s, prompt_embeds shape: %s, lora_request: %s.", # noqa: E501 request_id, prompt, params_list, prompt_token_ids, prompt_embeds.shape if prompt_embeds is not None else None, lora_request, ) async def chat_completion_full_generator( self, request: ChatCompletionRequest, result_generator: AsyncIterator[RequestOutput], request_id: str, model_name: str, conversation: list[ConversationMessage], tokenizer: AnyTokenizer, request_metadata: RequestResponseMetadata, ) -> Union[ErrorResponse, ChatCompletionResponse]: created_time = int(time.time()) final_res: Optional[RequestOutput] = None final_outputs: list[OmniRequestOutput] = [] try: async for res in result_generator: final_outputs.append(res) except asyncio.CancelledError: return self.create_error_response("Client disconnected") except ValueError as e: # TODO: Use a vllm-specific Validation Error return self.create_error_response(str(e)) assert final_outputs is not None choices: list[ChatCompletionResponseChoice] = [] usage = UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0) role = self.get_chat_request_role(request) prompt_logprobs = None prompt_token_ids = None kv_transfer_params = None for omni_outputs in final_outputs: choices_data = [] if omni_outputs.final_output_type == "text": ( choices_data, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params, ) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role) elif omni_outputs.final_output_type == "audio": choices_data = self._create_audio_choice(omni_outputs, role) else: logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}") continue choices.extend(choices_data) response = ChatCompletionResponse( id=request_id, created=created_time, model=model_name, choices=choices, usage=usage, prompt_logprobs=prompt_logprobs, prompt_token_ids=prompt_token_ids, kv_transfer_params=kv_transfer_params, ) # Log complete response if output logging is enabled if self.enable_log_outputs and self.request_logger: for choice in choices: output_text = "" if choice.message.content: output_text = choice.message.content elif choice.message.tool_calls: # For tool calls, log the function name and arguments tool_call_descriptions = [] for tc in choice.message.tool_calls: if hasattr(tc.function, "name") and hasattr(tc.function, "arguments"): tool_call_descriptions.append(f"{tc.function.name}({tc.function.arguments})") tool_calls_str = ", ".join(tool_call_descriptions) output_text = f"[tool_calls: {tool_calls_str}]" if output_text: # Get the corresponding output token IDs output_token_ids = None if choice.index < len(final_res.outputs): output_token_ids = final_res.outputs[choice.index].token_ids self.request_logger.log_outputs( request_id=request_id, outputs=output_text, output_token_ids=output_token_ids, finish_reason=choice.finish_reason, is_streaming=False, delta=False, ) return response def _create_text_choice( self, request: ChatCompletionRequest, omni_outputs: OmniRequestOutput, tokenizer: AnyTokenizer, conversation: list[ConversationMessage], role: str, ): final_res = omni_outputs.request_output if self.tool_call_id_type == "kimi_k2": history_tool_call_cnt = get_history_tool_calls_cnt(conversation) else: history_tool_call_cnt = 0 choices: list[ChatCompletionResponseChoice] = [] for output in final_res.outputs: token_ids = output.token_ids out_logprobs = output.logprobs tool_call_info = None if request.logprobs and request.top_logprobs is not None: assert out_logprobs is not None, "Did not output logprobs" logprobs = self._create_chat_logprobs( token_ids=token_ids, top_logprobs=out_logprobs, num_output_top_logprobs=request.top_logprobs, tokenizer=tokenizer, return_as_token_id=request.return_tokens_as_token_ids, ) else: logprobs = None if self.use_harmony: reasoning_content, content, _ = parse_chat_output(token_ids) if not request.include_reasoning: reasoning_content = None if self.tool_parser is not None: tool_parser = self.tool_parser(tokenizer) # NOTE: We use token_ids for openai tool parser tool_call_info = tool_parser.extract_tool_calls( "", request=request, token_ids=token_ids, # type: ignore ) content = tool_call_info.content message = ChatMessage( role=role, reasoning_content=reasoning_content, content=content, tool_calls=tool_call_info.tool_calls, ) else: message = ChatMessage( role=role, reasoning_content=reasoning_content, content=content, ) choice_data = ChatCompletionResponseChoice( index=output.index, message=message, logprobs=logprobs, finish_reason=( "tool_calls" if (tool_call_info is not None and tool_call_info.tools_called) else (output.finish_reason if output.finish_reason else "stop") ), stop_reason=output.stop_reason, ) choices.append(choice_data) continue if self.reasoning_parser: try: reasoning_parser = self.reasoning_parser(tokenizer) except RuntimeError as e: logger.exception("Error in reasoning parser creation.") return self.create_error_response(str(e)) # If the reasoning parser is enabled, # tool calls are extracted exclusively from the content. reasoning_content, content = reasoning_parser.extract_reasoning_content(output.text, request=request) if not request.include_reasoning: reasoning_content = None else: reasoning_content = None content = output.text auto_tools_called = False # if auto tools are not enabled, and a named tool choice using # outlines is not being used if (not self.enable_auto_tools or not self.tool_parser) and ( not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) and request.tool_choice != "required" ): message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content) # if the request uses tools and specified a tool choice elif request.tool_choice and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam: tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall message = ChatMessage( role=role, reasoning_content=reasoning_content, content="", tool_calls=[ tool_call_class( function=FunctionCall( name=request.tool_choice.function.name, arguments=content, ) ) ], ) elif request.tool_choice and request.tool_choice == "required": tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall # the fields of FunctionDefinition are a superset of the # tool call outputs and can be used for parsing assert content is not None tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(content) tool_call_ids = [] for tool_call in tool_calls: tool_call_ids.append( make_tool_call_id( id_type=self.tool_call_id_type, func_name=tool_call.name, idx=history_tool_call_cnt, ) ) history_tool_call_cnt += 1 message = ChatMessage( role=role, content="", tool_calls=[ tool_call_class( id=tool_call_ids[i], function=FunctionCall( name=tool_call.name, arguments=json.dumps(tool_call.parameters, ensure_ascii=False), ), ) for i, tool_call in enumerate(tool_calls) ], reasoning_content=reasoning_content, ) # if the request doesn't use tool choice # OR specifies to not use a tool elif not request.tool_choice or request.tool_choice == "none": message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content) # handle when there are tools and tool choice is auto elif ( request.tools and (request.tool_choice == "auto" or request.tool_choice is None) and self.enable_auto_tools and self.tool_parser ): try: tool_parser = self.tool_parser(tokenizer) except RuntimeError as e: logger.exception("Error in tool parser creation.") return self.create_error_response(str(e)) tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request) # In the OpenAI API the finish_reason is "tools_called" # if the tool choice is auto and the model produced a tool # call. The same is not true for named function calls auto_tools_called = tool_call_info.tools_called if tool_call_info.tools_called: message = ChatMessage( role=role, reasoning_content=reasoning_content, content=tool_call_info.content, tool_calls=tool_call_info.tool_calls, ) else: # FOR NOW make it a chat message; we will have to detect # the type to make it later. ret_content = content # try to use content return from tool parser first, # tool parser may do some modify for the content. if tool_call_info.content and len(tool_call_info.content) > 0: ret_content = tool_call_info.content message = ChatMessage( role=role, reasoning_content=reasoning_content, content=ret_content, ) # undetermined case that is still important to handle else: logger.error( "Error in chat_completion_full_generator - cannot determine if tools should be extracted. " "Returning a standard chat completion." ) message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content) choice_data = ChatCompletionResponseChoice( index=output.index, message=message, logprobs=logprobs, finish_reason=( "tool_calls" if auto_tools_called else output.finish_reason if output.finish_reason else "stop" ), stop_reason=output.stop_reason, token_ids=(as_list(output.token_ids) if request.return_token_ids else None), ) choices.append(choice_data) if request.echo: last_msg_content: Union[str, list[dict[str, str]]] = "" if conversation and "content" in conversation[-1] and conversation[-1].get("role") == role: last_msg_content = conversation[-1]["content"] or "" if isinstance(last_msg_content, list): last_msg_content = "\n".join(msg["text"] for msg in last_msg_content) for choice in choices: full_message = last_msg_content + (choice.message.content or "") choice.message.content = full_message assert final_res.prompt_token_ids is not None num_prompt_tokens = len(final_res.prompt_token_ids) if final_res.encoder_prompt_token_ids is not None: num_prompt_tokens += len(final_res.encoder_prompt_token_ids) num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs) usage = UsageInfo( prompt_tokens=num_prompt_tokens, completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) if self.enable_prompt_tokens_details and final_res.num_cached_tokens: usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens) prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs) prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None kv_transfer_params = final_res.kv_transfer_params return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params def _create_audio_choice(self, omni_outputs: OmniRequestOutput, role: str): choices: list[ChatCompletionResponseChoice] = [] final_res = omni_outputs.request_output audio_tensor = final_res.multimodal_output["audio"].detach().cpu().numpy() # Convert numpy array to WAV bytes and encode as base64 if soundfile is None: raise ImportError( "soundfile is required for audio generation. Please install it with: pip install soundfile" ) # Default sample rate for TTS models (typically 24000 Hz) # You may need to adjust this based on your model's configuration sample_rate = 24000 # Ensure audio is 1D (flatten if needed) if audio_tensor.ndim > 1: audio_tensor = audio_tensor.flatten() # Convert to WAV format and encode as base64 with BytesIO() as buffer: soundfile.write(buffer, audio_tensor, sample_rate, format="WAV") wav_bytes = buffer.getvalue() audio_base64 = base64.b64encode(wav_bytes).decode("utf-8") # Generate unique ID for the audio audio_id = f"audio-{uuid.uuid4().hex[:16]}" # Set expiration time (e.g., 24 hours from now) as Unix timestamp expires_at = int((datetime.now(timezone.utc) + timedelta(hours=24)).timestamp()) # Create OpenAIChatCompletionAudio object with all required fields audio_obj = OpenAIChatCompletionAudio( id=audio_id, data=audio_base64, expires_at=expires_at, transcript="", # Empty transcript if not available ) for output in final_res.outputs: choice_data = ChatCompletionResponseChoice( index=output.index, message=ChatMessage(role=role, audio=audio_obj), logprobs=None, finish_reason="stop", stop_reason=None, ) choices.append(choice_data) return choices