Refactor chat_utils and OpenAI modules for improved functionality and organization

- Updated import statements in chat_utils.py to replace the internal content format type with the public one for better clarity.
- Refactored api_server.py to streamline the chat template loading process by replacing the previous method with a direct function call.
- Enhanced serving_chat.py by reintroducing safe_apply_chat_template functions for better template handling.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/entrypoints/chat_utils.py

@@ -11,11 +11,11 @@ from vllm.entrypoints.chat_utils import (
     BaseMultiModalItemTracker,
     ChatCompletionContentPartParam,
     ChatCompletionMessageParam,
+    ChatTemplateContentFormat,
     ConversationMessage,
     MultiModalDataDict,
     MultiModalUUIDDict,
     _AssistantParser,
-    _ChatTemplateContentFormat,
     _ContentPart,
     _get_full_multimodal_text_prompt,
     _parse_chat_message_content_part,
@@ -128,7 +128,7 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
 def parse_chat_messages_futures(
     messages: list[ChatCompletionMessageParam],
     model_config: ModelConfig,
-    content_format: _ChatTemplateContentFormat,
+    content_format: ChatTemplateContentFormat,
     mm_processor_kwargs: dict[str, Any] | None = None,
 ) -> tuple[
     list[ConversationMessage],
@@ -161,7 +161,7 @@ def parse_chat_messages_futures(
 def _parse_chat_message_content(
     message: ChatCompletionMessageParam,
     mm_tracker: BaseMultiModalItemTracker,
-    content_format: _ChatTemplateContentFormat,
+    content_format: ChatTemplateContentFormat,
     interleave_strings: bool,
     mm_processor_kwargs: dict[str, Any] | None = None,
 ) -> list[ConversationMessage]:

vllm_omni/entrypoints/openai/api_server.py

@@ -19,6 +19,7 @@ from starlette.datastructures import State
 from starlette.routing import Route
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
+from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
@@ -33,13 +34,13 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
 # yapf: disable
 # yapf: enable
 from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse
-from vllm.entrypoints.openai.models.protocol import (
-    BaseModelPath,
+from vllm.entrypoints.openai.engine.protocol import (
+    ErrorResponse,
     ModelCard,
     ModelList,
     ModelPermission,
 )
+from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.orca_metrics import metrics_header
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
@@ -57,7 +58,6 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
 from vllm.entrypoints.utils import (
     load_aware_call,
-    process_chat_template,
     process_lora_modules,
     with_cancellation,
 )
@@ -398,11 +398,7 @@ async def omni_init_app_state(
         supported_tasks = set(await engine_client.get_supported_tasks())
     logger.info("Supported tasks: %s", supported_tasks)
 
-    resolved_chat_template = await process_chat_template(
-        args.chat_template,
-        engine_client,
-        vllm_config.model_config if vllm_config is not None else None,
-    )
+    resolved_chat_template = load_chat_template(args.chat_template)
 
     if args.tool_server == "demo":
         tool_server: ToolServer | None = DemoToolServer()

vllm_omni/entrypoints/openai/serving_chat.py

@@ -24,11 +24,8 @@ 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.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
@@ -64,6 +61,15 @@ from vllm.entrypoints.utils import should_include_usage
 from vllm.inputs.data import PromptType, TokensPrompt
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
+from vllm.renderers.hf import (
+    resolve_chat_template_content_format,
+)
+from vllm.renderers.hf import (
+    safe_apply_chat_template as apply_hf_chat_template,
+)
+from vllm.renderers.mistral import (
+    safe_apply_chat_template as apply_mistral_chat_template,
+)
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.tokenizers.mistral import (
@@ -390,6 +396,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tokenizer=tokenizer,
                 conversation=conversation,
                 model_config=model_config,
+                tokenize=False,
                 **_chat_template_kwargs,
             )