chat_utils.py 42.9 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import asyncio
4
import json
5
from abc import ABC, abstractmethod
6
from collections import defaultdict, deque
7
from collections.abc import Awaitable, Iterable
8
from functools import cache, lru_cache, partial
9
from pathlib import Path
10
11
from typing import (Any, Callable, Generic, Literal, Optional, TypeVar, Union,
                    cast)
12

13
14
import jinja2.nodes
import transformers.utils.chat_template_utils as hf_chat_utils
15
16
# yapf conflicts with isort for this block
# yapf: disable
17
from openai.types.chat import (ChatCompletionAssistantMessageParam,
18
19
                               ChatCompletionContentPartImageParam,
                               ChatCompletionContentPartInputAudioParam)
20
21
from openai.types.chat import (
    ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam)
22
23
from openai.types.chat import (ChatCompletionContentPartRefusalParam,
                               ChatCompletionContentPartTextParam)
24
25
from openai.types.chat import (
    ChatCompletionMessageParam as OpenAIChatCompletionMessageParam)
26
27
from openai.types.chat import (ChatCompletionMessageToolCallParam,
                               ChatCompletionToolMessageParam)
28
29
from openai.types.chat.chat_completion_content_part_input_audio_param import (
    InputAudio)
30
31
# yapf: enable
# pydantic needs the TypedDict from typing_extensions
32
33
from transformers import (PreTrainedTokenizer, PreTrainedTokenizerFast,
                          ProcessorMixin)
34
from typing_extensions import Required, TypeAlias, TypedDict
35

36
from vllm.config import ModelConfig
37
from vllm.logger import init_logger
38
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict
39
from vllm.multimodal.utils import MediaConnector
40
from vllm.transformers_utils.processor import cached_get_processor
41
from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
42
43
44
45

logger = init_logger(__name__)


46
47
48
49
50
51
52
53
54
55
56
57
58
59
class AudioURL(TypedDict, total=False):
    url: Required[str]
    """
    Either a URL of the audio or a data URL with base64 encoded audio data.
    """


class ChatCompletionContentPartAudioParam(TypedDict, total=False):
    audio_url: Required[AudioURL]

    type: Required[Literal["audio_url"]]
    """The type of the content part."""


60
61
62
63
64
65
66
67
68
69
70
class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):
    image_embeds: Required[Union[str, dict[str, str]]]
    """
    The image embeddings. It can be either:
    - A single base64 string.
    - A dictionary where each value is a base64 string.
    """
    type: Required[Literal["image_embeds"]]
    """The type of the content part."""


71
72
73
74
75
76
77
78
79
80
81
82
83
84
class VideoURL(TypedDict, total=False):
    url: Required[str]
    """
    Either a URL of the video or a data URL with base64 encoded video data.
    """


class ChatCompletionContentPartVideoParam(TypedDict, total=False):
    video_url: Required[VideoURL]

    type: Required[Literal["video_url"]]
    """The type of the content part."""


85
86
87
class CustomChatCompletionContentSimpleImageParam(TypedDict, total=False):
    """A simpler version of the param that only accepts a plain image_url.
    This is supported by OpenAI API, although it is not documented.
88

89
90
91
92
93
94
95
96
97
98
    Example:
    {
        "image_url": "https://example.com/image.jpg"
    }
    """
    image_url: Required[str]


class CustomChatCompletionContentSimpleAudioParam(TypedDict, total=False):
    """A simpler version of the param that only accepts a plain audio_url.
99

100
101
102
103
104
105
106
107
    Example:
    {
        "audio_url": "https://example.com/audio.mp3"
    }
    """
    audio_url: Required[str]


108
109
110
111
112
113
114
115
116
117
118
class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False):
    """A simpler version of the param that only accepts a plain audio_url.

    Example:
    {
        "video_url": "https://example.com/video.mp4"
    }
    """
    video_url: Required[str]


119
120
ChatCompletionContentPartParam: TypeAlias = Union[
    OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam,
121
    ChatCompletionContentPartInputAudioParam,
122
    ChatCompletionContentPartVideoParam, ChatCompletionContentPartRefusalParam,
123
    CustomChatCompletionContentSimpleImageParam,
124
    ChatCompletionContentPartImageEmbedsParam,
125
126
    CustomChatCompletionContentSimpleAudioParam,
    CustomChatCompletionContentSimpleVideoParam, str]
127
128
129
130
131
132
133


class CustomChatCompletionMessageParam(TypedDict, total=False):
    """Enables custom roles in the Chat Completion API."""
    role: Required[str]
    """The role of the message's author."""

134
    content: Union[str, list[ChatCompletionContentPartParam]]
135
136
137
138
139
140
141
142
143
    """The contents of the message."""

    name: str
    """An optional name for the participant.

    Provides the model information to differentiate between participants of the
    same role.
    """

144
145
146
147
148
149
    tool_call_id: Optional[str]
    """Tool call that this message is responding to."""

    tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
    """The tool calls generated by the model, such as function calls."""

150
151
152
153
154

ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam,
                                   CustomChatCompletionMessageParam]


155
# TODO: Make fields ReadOnly once mypy supports it
156
157
158
159
class ConversationMessage(TypedDict, total=False):
    role: Required[str]
    """The role of the message's author."""

160
    content: Union[Optional[str], list[dict[str, str]]]
161
162
163
164
165
166
167
168
169
170
    """The contents of the message"""

    tool_call_id: Optional[str]
    """Tool call that this message is responding to."""

    name: Optional[str]
    """The name of the function to call"""

    tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]]
    """The tool calls generated by the model, such as function calls."""
171
172


173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
# Passed in by user
ChatTemplateContentFormatOption = Literal["auto", "string", "openai"]

# Used internally
_ChatTemplateContentFormat = Literal["string", "openai"]


def _is_var_access(node: jinja2.nodes.Node, varname: str) -> bool:
    if isinstance(node, jinja2.nodes.Name):
        return node.ctx == "load" and node.name == varname

    return False


def _is_attr_access(node: jinja2.nodes.Node, varname: str, key: str) -> bool:
    if isinstance(node, jinja2.nodes.Getitem):
        return (_is_var_access(node.node, varname)
                and isinstance(node.arg, jinja2.nodes.Const)
                and node.arg.value == key)

    if isinstance(node, jinja2.nodes.Getattr):
        return _is_var_access(node.node, varname) and node.attr == key

    return False


def _is_var_or_elems_access(
    node: jinja2.nodes.Node,
    varname: str,
    key: Optional[str] = None,
) -> bool:
    if isinstance(node, jinja2.nodes.Filter):
        return (node.node is not None
                and _is_var_or_elems_access(node.node, varname, key))
    if isinstance(node, jinja2.nodes.Test):
        return _is_var_or_elems_access(node.node, varname, key)

    if (isinstance(node, jinja2.nodes.Getitem)
            and isinstance(node.arg, jinja2.nodes.Slice)):
        return _is_var_or_elems_access(node.node, varname, key)

    # yapf: disable
    return (
        _is_attr_access(node, varname, key) if key
        else _is_var_access(node, varname)
    ) # yapf: enable


def _iter_nodes_assign_var_or_elems(root: jinja2.nodes.Node, varname: str):
    # Global variable that is implicitly defined at the root
    yield root, varname

    # Iterative BFS
    related_varnames = deque([varname])
    while related_varnames:
        related_varname = related_varnames.popleft()

        for assign_ast in root.find_all(jinja2.nodes.Assign):
            lhs = assign_ast.target
            rhs = assign_ast.node

            if _is_var_or_elems_access(rhs, related_varname):
                assert isinstance(lhs, jinja2.nodes.Name)
                yield assign_ast, lhs.name

                # Avoid infinite looping for self-assignment
                if lhs.name != related_varname:
                    related_varnames.append(lhs.name)


# NOTE: The proper way to handle this is to build a CFG so that we can handle
# the scope in which each variable is defined, but that is too complicated
def _iter_nodes_assign_messages_item(root: jinja2.nodes.Node):
    messages_varnames = [
        varname
        for _, varname in _iter_nodes_assign_var_or_elems(root, "messages")
    ]

    # Search for {%- for message in messages -%} loops
    for loop_ast in root.find_all(jinja2.nodes.For):
        loop_iter = loop_ast.iter
        loop_target = loop_ast.target

        for varname in messages_varnames:
            if _is_var_or_elems_access(loop_iter, varname):
                assert isinstance(loop_target, jinja2.nodes.Name)
                yield loop_ast, loop_target.name
                break


def _iter_nodes_assign_content_item(root: jinja2.nodes.Node):
    message_varnames = [
        varname for _, varname in _iter_nodes_assign_messages_item(root)
    ]

    # Search for {%- for content in message['content'] -%} loops
    for loop_ast in root.find_all(jinja2.nodes.For):
        loop_iter = loop_ast.iter
        loop_target = loop_ast.target

        for varname in message_varnames:
            if _is_var_or_elems_access(loop_iter, varname, "content"):
                assert isinstance(loop_target, jinja2.nodes.Name)
                yield loop_ast, loop_target.name
                break


def _try_extract_ast(chat_template: str) -> Optional[jinja2.nodes.Template]:
    try:
        jinja_compiled = hf_chat_utils._compile_jinja_template(chat_template)
        return jinja_compiled.environment.parse(chat_template)
    except Exception:
        logger.exception("Error when compiling Jinja template")
        return None


def _detect_content_format(
    chat_template: str,
    *,
    default: _ChatTemplateContentFormat,
) -> _ChatTemplateContentFormat:
    jinja_ast = _try_extract_ast(chat_template)
    if jinja_ast is None:
        return default

    try:
        next(_iter_nodes_assign_content_item(jinja_ast))
    except StopIteration:
        return "string"
    except Exception:
        logger.exception("Error when parsing AST of Jinja template")
        return default
    else:
        return "openai"


309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
def resolve_mistral_chat_template(
    chat_template: Optional[str],
    **kwargs: Any,
) -> Optional[str]:
    if chat_template is not None:
        logger.warning_once(
            "'chat_template' cannot be overridden for mistral tokenizer.")
    if "add_generation_prompt" in kwargs:
        logger.warning_once(
            "'add_generation_prompt' is not supported for mistral tokenizer, "
            "so it will be ignored.")
    if "continue_final_message" in kwargs:
        logger.warning_once(
            "'continue_final_message' is not supported for mistral tokenizer, "
            "so it will be ignored.")
    return None

def resolve_hf_chat_template(
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
    *,
    trust_remote_code: bool,
) -> Optional[str]:
    # 1st priority: The given chat template
    if chat_template is not None:
        return chat_template

    # 2nd priority: AutoProcessor chat template, unless tool calling is enabled
    if tools is None:
        try:
            processor = cached_get_processor(
                tokenizer.name_or_path,
                processor_cls=(PreTrainedTokenizer, PreTrainedTokenizerFast,
                               ProcessorMixin),
                trust_remote_code=trust_remote_code,
            )
            if isinstance(processor, ProcessorMixin) and \
                processor.chat_template is not None:
                return processor.chat_template
        except Exception:
            logger.debug("Failed to load AutoProcessor chat template for %s",
                        tokenizer.name_or_path, exc_info=True)

    # 3rd priority: AutoTokenizer chat template
    try:
        return tokenizer.get_chat_template(chat_template, tools=tools)
    except Exception:
        logger.debug("Failed to load AutoTokenizer chat template for %s",
                     tokenizer.name_or_path, exc_info=True)

    return None


363
364
def _resolve_chat_template_content_format(
    chat_template: Optional[str],
365
    tools: Optional[list[dict[str, Any]]],
366
367
    given_format: ChatTemplateContentFormatOption,
    tokenizer: AnyTokenizer,
368
369
    *,
    trust_remote_code: bool,
370
371
) -> _ChatTemplateContentFormat:
    if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
372
        hf_chat_template = resolve_hf_chat_template(
373
374
375
376
377
            tokenizer,
            chat_template=chat_template,
            trust_remote_code=trust_remote_code,
            tools=tools,
        )
378
    else:
379
380
381
382
        hf_chat_template = None

    jinja_text = (hf_chat_template if isinstance(hf_chat_template, str)
                  else load_chat_template(chat_template, is_literal=True))
383
384
385
386
387
388
389
390

    detected_format = ("string" if jinja_text is None else
                       _detect_content_format(jinja_text, default="string"))

    return detected_format if given_format == "auto" else given_format


@lru_cache
391
def _log_chat_template_content_format(
392
393
    chat_template: Optional[str],
    given_format: ChatTemplateContentFormatOption,
394
395
    detected_format: ChatTemplateContentFormatOption,
):
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
    logger.info(
        "Detected the chat template content format to be '%s'. "
        "You can set `--chat-template-content-format` to override this.",
        detected_format,
    )

    if given_format != "auto" and given_format != detected_format:
        logger.warning(
            "You specified `--chat-template-content-format %s` "
            "which is different from the detected format '%s'. "
            "If our automatic detection is incorrect, please consider "
            "opening a GitHub issue so that we can improve it: "
            "https://github.com/vllm-project/vllm/issues/new/choose",
            given_format,
            detected_format,
        )

413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435

def resolve_chat_template_content_format(
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
    given_format: ChatTemplateContentFormatOption,
    tokenizer: AnyTokenizer,
    *,
    trust_remote_code: bool = False,
) -> _ChatTemplateContentFormat:
    detected_format = _resolve_chat_template_content_format(
        chat_template,
        tools,
        given_format,
        tokenizer,
        trust_remote_code=trust_remote_code,
    )

    _log_chat_template_content_format(
        chat_template,
        given_format=given_format,
        detected_format=detected_format,
    )

436
437
438
    return detected_format


439
ModalityStr = Literal["image", "audio", "video", "image_embeds"]
440
441
442
443
_T = TypeVar("_T")


class BaseMultiModalItemTracker(ABC, Generic[_T]):
444
445
446
447
448
449
450
    """
    Tracks multi-modal items in a given request and ensures that the number
    of multi-modal items in a given request does not exceed the configured
    maximum per prompt.
    """

    def __init__(self, model_config: ModelConfig, tokenizer: AnyTokenizer):
451
452
        super().__init__()

453
454
        self._model_config = model_config
        self._tokenizer = tokenizer
455

456
        self._items_by_modality = defaultdict[str, list[_T]](list)
457

458
459
460
461
    @property
    def model_config(self) -> ModelConfig:
        return self._model_config

462
463
464
465
    @property
    def allowed_local_media_path(self):
        return self._model_config.allowed_local_media_path

466
467
468
469
    @property
    def mm_registry(self):
        return MULTIMODAL_REGISTRY

470
    @staticmethod
471
    @cache
472
    def _cached_token_str(tokenizer: AnyTokenizer, token_index: int) -> str:
473
474
        return tokenizer.decode(token_index)

475
476
    def _placeholder_str(self, modality: ModalityStr,
                         current_count: int) -> Optional[str]:
477
478
        # TODO: Let user specify how to insert image tokens into prompt
        # (similar to chat template)
479
480
481
        hf_config = self._model_config.hf_config
        model_type = hf_config.model_type

482
483
484
        if modality in ("image", "image_embeds"):
            if model_type == "chatglm":
                return "<|begin_of_image|><|endoftext|><|end_of_image|>"
485
486
487
            if model_type == "phi3_v":
                # Workaround since this token is not defined in the tokenizer
                return f"<|image_{current_count}|>"
488
489
            if model_type == "phi4mm":
                return "<|endoftext10|>"  # 200010 (see vocab.json in hf model)
490
            if model_type in ("minicpmo", "minicpmv"):
491
                return "(<image>./</image>)"
492
493
            if model_type in ("blip-2", "florence2", "fuyu", "paligemma",
                              "pixtral", "mistral3"):
494
495
                # These models do not use image tokens in the prompt
                return None
496
497
            if model_type == "qwen":
                return f"Picture {current_count}: <img></img>"
498
            if model_type.startswith("llava"):
499
500
                return self._cached_token_str(self._tokenizer,
                                              hf_config.image_token_index)
Jennifer Zhao's avatar
Jennifer Zhao committed
501
502
            if model_type in ("aya_vision", "chameleon", "deepseek_vl_v2",
                              "internvl_chat", "skywork_chat", "NVLM_D",
503
                              "h2ovl_chat", "idefics3", "smolvlm"):
504
                return "<image>"
505
            if model_type in ("mllama", "llama4"):
506
                return "<|image|>"
Roger Wang's avatar
Roger Wang committed
507
            if model_type in ("qwen2_vl", "qwen2_5_vl"):
508
                return "<|vision_start|><|image_pad|><|vision_end|>"
509
510
            if model_type == "molmo":
                return ""
511
512
            if model_type == "aria":
                return "<|fim_prefix|><|img|><|fim_suffix|>"
513
514
            if model_type == "gemma3":
                return "<start_of_image>"
515

516
            raise TypeError(f"Unknown {modality} model type: {model_type}")
517
518
        elif modality == "audio":
            if model_type == "ultravox":
519
                return "<|audio|>"
520
521
            if model_type == "phi4mm":
                return "<|endoftext11|>"  # 200011 (see vocab.json in hf model)
522
523
524
            if model_type == "qwen2_audio":
                return (f"Audio {current_count}: "
                        f"<|audio_bos|><|AUDIO|><|audio_eos|>")
525
526
            if model_type == "minicpmo":
                return "(<audio>./</audio>)"
527
            raise TypeError(f"Unknown model type: {model_type}")
528
        elif modality == "video":
Roger Wang's avatar
Roger Wang committed
529
            if model_type in ("qwen2_vl", "qwen2_5_vl"):
530
                return "<|vision_start|><|video_pad|><|vision_end|>"
531
532
            if model_type in ("minicpmo", "minicpmv"):
                return "(<video>./</video>)"
533
534
535
            if model_type.startswith("llava"):
                return self._cached_token_str(self._tokenizer,
                                              hf_config.video_token_index)
536
            raise TypeError(f"Unknown {modality} model type: {model_type}")
537
538
539
        else:
            raise TypeError(f"Unknown modality: {modality}")

540
541
542
543
544
    def add(self, modality: ModalityStr, item: _T) -> Optional[str]:
        """
        Add a multi-modal item to the current prompt and returns the
        placeholder string to use, if any.
        """
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
        mm_registry = self.mm_registry
        model_config = self.model_config

        input_modality = modality.replace("_embeds", "")

        if mm_registry.has_processor(model_config):
            mm_processor = mm_registry.create_processor(model_config)
            allowed_counts = mm_processor.info.get_allowed_mm_limits()
            allowed_count = allowed_counts.get(input_modality, 0)
        else:
            mm_config = model_config.multimodal_config
            if mm_config is None:
                msg = "This model does not support multi-modal inputs"
                raise ValueError(msg)

            allowed_count = mm_config.get_limit_per_prompt(input_modality)

562
        current_count = len(self._items_by_modality[modality]) + 1
563
564
565
        if current_count > allowed_count:
            raise ValueError(
                f"At most {allowed_count} {modality}(s) may be provided in "
566
567
                "one request. You can set `--limit-mm-per-prompt` to "
                "increase this limit if the model supports it.")
568

569
        self._items_by_modality[modality].append(item)
570
571
572
573
574
575
576
577

        return self._placeholder_str(modality, current_count)

    @abstractmethod
    def create_parser(self) -> "BaseMultiModalContentParser":
        raise NotImplementedError


578
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
579
580

    def all_mm_data(self) -> Optional[MultiModalDataDict]:
581
582
583
584
585
586
587
588
589
590
591
592
593
594
        if not self._items_by_modality:
            return None
        mm_inputs = {}
        items_by_modality = dict(self._items_by_modality)
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
            raise ValueError(\
                "Mixing raw image and embedding inputs is not allowed")

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
            if len(image_embeds_lst) > 1:
                raise ValueError(\
                    "Only one message can have {'type': 'image_embeds'}")
            mm_inputs["image"] = image_embeds_lst[0]
595
        if "image" in items_by_modality:
596
            mm_inputs["image"] = items_by_modality["image"] # A list of images
597
        if "audio" in items_by_modality:
598
            mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
599
        if "video" in items_by_modality:
600
601
            mm_inputs["video"] = items_by_modality["video"] # A list of videos
        return mm_inputs
602
603
604
605
606

    def create_parser(self) -> "BaseMultiModalContentParser":
        return MultiModalContentParser(self)


607
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
608
609

    async def all_mm_data(self) -> Optional[MultiModalDataDict]:
610
611
612
613
        if not self._items_by_modality:
            return None
        mm_inputs = {}
        items_by_modality = {
614
615
616
                modality: await asyncio.gather(*items)
                for modality, items in self._items_by_modality.items()
            }
617

618
619
620
621
622
623
624
625
626
627
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
            raise ValueError(
                "Mixing raw image and embedding inputs is not allowed")

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
            if len(image_embeds_lst) > 1:
                raise ValueError(
                    "Only one message can have {'type': 'image_embeds'}")
            mm_inputs["image"] = image_embeds_lst[0]
628
        if "image" in items_by_modality:
629
            mm_inputs["image"] = items_by_modality["image"] # A list of images
630
        if "audio" in items_by_modality:
631
            mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
632
        if "video" in items_by_modality:
633
634
            mm_inputs["video"] = items_by_modality["video"] # A list of videos
        return mm_inputs
635
636
637
638
639
640
641
642
643
644
645

    def create_parser(self) -> "BaseMultiModalContentParser":
        return AsyncMultiModalContentParser(self)


class BaseMultiModalContentParser(ABC):

    def __init__(self) -> None:
        super().__init__()

        # multimodal placeholder_string : count
646
        self._placeholder_counts: dict[str, int] = defaultdict(lambda: 0)
647
648
649
650
651

    def _add_placeholder(self, placeholder: Optional[str]):
        if placeholder:
            self._placeholder_counts[placeholder] += 1

652
    def mm_placeholder_counts(self) -> dict[str, int]:
653
654
655
656
657
658
        return dict(self._placeholder_counts)

    @abstractmethod
    def parse_image(self, image_url: str) -> None:
        raise NotImplementedError

659
660
661
662
663
    @abstractmethod
    def parse_image_embeds(self,
                           image_embeds: Union[str, dict[str, str]]) -> None:
        raise NotImplementedError

664
665
666
667
    @abstractmethod
    def parse_audio(self, audio_url: str) -> None:
        raise NotImplementedError

668
    @abstractmethod
669
    def parse_input_audio(self, input_audio: InputAudio) -> None:
670
671
        raise NotImplementedError

672
673
674
675
    @abstractmethod
    def parse_video(self, video_url: str) -> None:
        raise NotImplementedError

676
677
678
679
680
681
682
683

class MultiModalContentParser(BaseMultiModalContentParser):

    def __init__(self, tracker: MultiModalItemTracker) -> None:
        super().__init__()

        self._tracker = tracker

684
685
686
687
        self._connector = MediaConnector(
            allowed_local_media_path=tracker.allowed_local_media_path,
        )

688
    def parse_image(self, image_url: str) -> None:
689
        image = self._connector.fetch_image(image_url)
690
691
692
693

        placeholder = self._tracker.add("image", image)
        self._add_placeholder(placeholder)

694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
    def parse_image_embeds(self,
                           image_embeds: Union[str, dict[str, str]]) -> None:
        if isinstance(image_embeds, dict):
            embeds = {
                k: self._connector.fetch_image_embedding(v)
                for k, v in image_embeds.items()
            }
            placeholder = self._tracker.add("image_embeds", embeds)

        if isinstance(image_embeds, str):
            embedding = self._connector.fetch_image_embedding(image_embeds)
            placeholder = self._tracker.add("image_embeds", embedding)

        self._add_placeholder(placeholder)

709
    def parse_audio(self, audio_url: str) -> None:
710
        audio = self._connector.fetch_audio(audio_url)
711
712
713
714

        placeholder = self._tracker.add("audio", audio)
        self._add_placeholder(placeholder)

715
716
717
718
    def parse_input_audio(self, input_audio: InputAudio) -> None:
        audio_data = input_audio.get("data", "")
        audio_format = input_audio.get("format", "")
        audio_url = f"data:audio/{audio_format};base64,{audio_data}"
719

720
        return self.parse_audio(audio_url)
721

722
    def parse_video(self, video_url: str) -> None:
723
        video = self._connector.fetch_video(video_url)
724
725
726
727

        placeholder = self._tracker.add("video", video)
        self._add_placeholder(placeholder)

728
729
730
731
732
733
734

class AsyncMultiModalContentParser(BaseMultiModalContentParser):

    def __init__(self, tracker: AsyncMultiModalItemTracker) -> None:
        super().__init__()

        self._tracker = tracker
735
736
737
        self._connector = MediaConnector(
            allowed_local_media_path=tracker.allowed_local_media_path,
        )
738
739

    def parse_image(self, image_url: str) -> None:
740
        image_coro = self._connector.fetch_image_async(image_url)
741
742
743
744

        placeholder = self._tracker.add("image", image_coro)
        self._add_placeholder(placeholder)

745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
    def parse_image_embeds(self,
                           image_embeds: Union[str, dict[str, str]]) -> None:
        future: asyncio.Future[Union[str, dict[str, str]]] = asyncio.Future()

        if isinstance(image_embeds, dict):
            embeds = {
                k: self._connector.fetch_image_embedding(v)
                for k, v in image_embeds.items()
            }
            future.set_result(embeds)

        if isinstance(image_embeds, str):
            embedding = self._connector.\
                fetch_image_embedding(image_embeds)
            future.set_result(embedding)

        placeholder = self._tracker.add("image_embeds", future)
        self._add_placeholder(placeholder)

764
    def parse_audio(self, audio_url: str) -> None:
765
        audio_coro = self._connector.fetch_audio_async(audio_url)
766
767
768

        placeholder = self._tracker.add("audio", audio_coro)
        self._add_placeholder(placeholder)
769

770
771
772
773
    def parse_input_audio(self, input_audio: InputAudio) -> None:
        audio_data = input_audio.get("data", "")
        audio_format = input_audio.get("format", "")
        audio_url = f"data:audio/{audio_format};base64,{audio_data}"
774

775
        return self.parse_audio(audio_url)
776

777
    def parse_video(self, video_url: str) -> None:
778
        video = self._connector.fetch_video_async(video_url)
779
780
781
782

        placeholder = self._tracker.add("video", video)
        self._add_placeholder(placeholder)

783

784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
def validate_chat_template(chat_template: Optional[Union[Path, str]]):
    """Raises if the provided chat template appears invalid."""
    if chat_template is None:
        return

    elif isinstance(chat_template, Path) and not chat_template.exists():
        raise FileNotFoundError(
            "the supplied chat template path doesn't exist")

    elif isinstance(chat_template, str):
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template
                   for c in JINJA_CHARS) and not Path(chat_template).exists():
            raise ValueError(
                f"The supplied chat template string ({chat_template}) "
                f"appears path-like, but doesn't exist!")

    else:
        raise TypeError(
            f"{type(chat_template)} is not a valid chat template type")


806
def _load_chat_template(
807
808
809
810
    chat_template: Optional[Union[Path, str]],
    *,
    is_literal: bool = False,
) -> Optional[str]:
811
812
    if chat_template is None:
        return None
813
814
815
816
817
818

    if is_literal:
        if isinstance(chat_template, Path):
            raise TypeError("chat_template is expected to be read directly "
                            "from its value")

819
        return chat_template
820

821
    try:
822
        with open(chat_template) as f:
823
            return f.read()
824
    except OSError as e:
825
826
827
        if isinstance(chat_template, Path):
            raise

828
829
830
831
832
833
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template for c in JINJA_CHARS):
            msg = (f"The supplied chat template ({chat_template}) "
                   f"looks like a file path, but it failed to be "
                   f"opened. Reason: {e}")
            raise ValueError(msg) from e
834

835
836
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
837
838
839
840
841
842
843
844
845
846
847
848
        return _load_chat_template(chat_template, is_literal=True)


_cached_load_chat_template = lru_cache(_load_chat_template)


def load_chat_template(
    chat_template: Optional[Union[Path, str]],
    *,
    is_literal: bool = False,
) -> Optional[str]:
    return _cached_load_chat_template(chat_template, is_literal=is_literal)
849
850


851
# TODO: Let user specify how to insert multimodal tokens into prompt
852
# (similar to chat template)
853
def _get_full_multimodal_text_prompt(placeholder_counts: dict[str, int],
854
                                     text_prompt: str) -> str:
855
    """Combine multimodal prompts for a multimodal language model."""
856

857
    # Look through the text prompt to check for missing placeholders
858
    missing_placeholders: list[str] = []
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
    for placeholder in placeholder_counts:

        # For any existing placeholder in the text prompt, we leave it as is
        placeholder_counts[placeholder] -= text_prompt.count(placeholder)

        if placeholder_counts[placeholder] < 0:
            raise ValueError(
                f"Found more '{placeholder}' placeholders in input prompt than "
                "actual multimodal data items.")

        missing_placeholders.extend([placeholder] *
                                    placeholder_counts[placeholder])

    # NOTE: For now we always add missing placeholders at the front of
    # the prompt. This may change to be customizable in the future.
    return "\n".join(missing_placeholders + [text_prompt])
875
876


877
878
879
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
_ImageParser = partial(cast, ChatCompletionContentPartImageParam)
880
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
881
_AudioParser = partial(cast, ChatCompletionContentPartAudioParam)
882
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
883
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
884
_VideoParser = partial(cast, ChatCompletionContentPartVideoParam)
885

886
_ContentPart: TypeAlias = Union[str, dict[str, str], InputAudio]
887

888
# Define a mapping from part types to their corresponding parsing functions.
889
MM_PARSER_MAP: dict[
890
891
892
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
893
    "text":
894
    lambda part: _TextParser(part).get("text", None),
895
    "image_url":
896
    lambda part: _ImageParser(part).get("image_url", {}).get("url", None),
897
    "image_embeds":
898
    lambda part: _ImageEmbedsParser(part).get("image_embeds", None),
899
    "audio_url":
900
    lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
901
    "input_audio":
902
    lambda part: _InputAudioParser(part).get("input_audio", None),
903
    "refusal":
904
    lambda part: _RefusalParser(part).get("refusal", None),
905
    "video_url":
906
    lambda part: _VideoParser(part).get("video_url", {}).get("url", None),
907
908
909
910
}


def _parse_chat_message_content_mm_part(
911
        part: ChatCompletionContentPartParam) -> tuple[str, _ContentPart]:
912
    """
913
    Parses a given multi-modal content part based on its type.
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933

    Args:
        part: A dict containing the content part, with a potential 'type' field.

    Returns:
        A tuple (part_type, content) where:
        - part_type: Type of the part (e.g., 'text', 'image_url').
        - content: Parsed content (e.g., text, image URL).

    Raises:
        ValueError: If the 'type' field is missing and no direct URL is found.
    """
    assert isinstance(
        part, dict)  # This is needed to avoid mypy errors: part.get() from str
    part_type = part.get("type", None)

    if isinstance(part_type, str) and part_type in MM_PARSER_MAP:
        content = MM_PARSER_MAP[part_type](part)

        # Special case for 'image_url.detail'
934
935
        # We only support 'auto', which is the default
        if part_type == "image_url" and part.get("detail", "auto") != "auto":
936
937
938
939
940
941
            logger.warning("'image_url.detail' is currently not supported "
                           "and will be ignored.")

        return part_type, content

    # Handle missing 'type' but provided direct URL fields.
942
    # 'type' is required field by pydantic
943
944
945
946
947
948
949
950
951
    if part_type is None:
        if part.get("image_url") is not None:
            image_params = cast(CustomChatCompletionContentSimpleImageParam,
                                part)
            return "image_url", image_params.get("image_url", "")
        if part.get("audio_url") is not None:
            audio_params = cast(CustomChatCompletionContentSimpleAudioParam,
                                part)
            return "audio_url", audio_params.get("audio_url", "")
952
        if part.get("input_audio") is not None:
953
            input_audio_params = cast(dict[str, str], part)
954
            return "input_audio", input_audio_params
955
956
957
958
        if part.get("video_url") is not None:
            video_params = cast(CustomChatCompletionContentSimpleVideoParam,
                                part)
            return "video_url", video_params.get("video_url", "")
959
960
961
962
963
964
965
966
967
        # Raise an error if no 'type' or direct URL is found.
        raise ValueError("Missing 'type' field in multimodal part.")

    if not isinstance(part_type, str):
        raise ValueError("Invalid 'type' field in multimodal part.")
    return part_type, "unknown part_type content"


VALID_MESSAGE_CONTENT_MM_PART_TYPES = ("text", "refusal", "image_url",
968
                                       "image_embeds",
969
                                       "audio_url", "input_audio", "video_url")
970

971

972
973
974
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
975
    mm_tracker: BaseMultiModalItemTracker,
976
977
    *,
    wrap_dicts: bool,
978
) -> list[ConversationMessage]:
979
    content = list[_ContentPart]()
980

981
    mm_parser = mm_tracker.create_parser()
982
983

    for part in parts:
984
        parse_res = _parse_chat_message_content_part(
985
986
987
988
            part,
            mm_parser,
            wrap_dicts=wrap_dicts,
        )
989
990
        if parse_res:
            content.append(parse_res)
991

992
    if wrap_dicts:
993
        # Parsing wraps images and texts as interleaved dictionaries
994
        return [ConversationMessage(role=role,
995
                                    content=content)]  # type: ignore
996
    texts = cast(list[str], content)
997
998
999
1000
1001
1002
1003
1004
1005
    text_prompt = "\n".join(texts)
    mm_placeholder_counts = mm_parser.mm_placeholder_counts()
    if mm_placeholder_counts:
        text_prompt = _get_full_multimodal_text_prompt(mm_placeholder_counts,
                                                       text_prompt)
    return [ConversationMessage(role=role, content=text_prompt)]


def _parse_chat_message_content_part(
1006
1007
1008
1009
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
1010
) -> Optional[_ContentPart]:
1011
1012
1013
1014
1015
1016
1017
1018
    """Parses a single part of a conversation. If wrap_dicts is True,
    structured dictionary pieces for texts and images will be
    wrapped in dictionaries, i.e., {"type": "text", "text", ...} and
    {"type": "image"}, respectively. Otherwise multimodal data will be
    handled by mm_parser, and texts will be returned as strings to be joined
    with multimodal placeholders.
    """
    if isinstance(part, str):  # Handle plain text parts
1019
        return part
1020
1021
1022
1023

    # Handle structured dictionary parts
    part_type, content = _parse_chat_message_content_mm_part(part)

1024
    # if part_type is text/refusal/image_url/audio_url/video_url/input_audio but
1025
1026
    # content is None, log a warning and skip
    if part_type in VALID_MESSAGE_CONTENT_MM_PART_TYPES and content is None:
1027
        logger.warning(
1028
1029
            "Skipping multimodal part '%s' (type: '%s') "
            "with empty / unparsable content.", part, part_type)
1030
1031
1032
        return None

    if part_type in ("text", "refusal"):
1033
1034
1035
1036
1037
        str_content = cast(str, content)
        if wrap_dicts:
            return {'type': 'text', 'text': str_content}
        else:
            return str_content
1038
1039

    if part_type == "image_url":
1040
1041
        str_content = cast(str, content)
        mm_parser.parse_image(str_content)
1042
        return {'type': 'image'} if wrap_dicts else None
1043
1044
1045
1046
    if part_type == "image_embeds":
        content = cast(Union[str, dict[str, str]], content)
        mm_parser.parse_image_embeds(content)
        return {'type': 'image'} if wrap_dicts else None
1047
    if part_type == "audio_url":
1048
1049
1050
1051
1052
        str_content = cast(str, content)
        mm_parser.parse_audio(str_content)
        return {'type': 'audio'} if wrap_dicts else None

    if part_type == "input_audio":
1053
        dict_content = cast(InputAudio, content)
1054
        mm_parser.parse_input_audio(dict_content)
1055
1056
        return {'type': 'audio'} if wrap_dicts else None

1057
    if part_type == "video_url":
1058
1059
        str_content = cast(str, content)
        mm_parser.parse_video(str_content)
1060
1061
        return {'type': 'video'} if wrap_dicts else None

1062
    raise NotImplementedError(f"Unknown part type: {part_type}")
1063
1064


1065
1066
1067
1068
1069
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


1070
def _parse_chat_message_content(
1071
1072
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
1073
    content_format: _ChatTemplateContentFormat,
1074
) -> list[ConversationMessage]:
1075
1076
1077
1078
    role = message["role"]
    content = message.get("content")

    if content is None:
1079
1080
1081
1082
1083
1084
        content = []
    elif isinstance(content, str):
        content = [
            ChatCompletionContentPartTextParam(type="text", text=content)
        ]
    result = _parse_chat_message_content_parts(
1085
1086
        role,
        content,  # type: ignore
1087
        mm_tracker,
1088
        wrap_dicts=(content_format == "openai"),
1089
    )
1090

1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
    for result_msg in result:
        if role == 'assistant':
            parsed_msg = _AssistantParser(message)

            if "tool_calls" in parsed_msg:
                result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
        elif role == "tool":
            parsed_msg = _ToolParser(message)
            if "tool_call_id" in parsed_msg:
                result_msg["tool_call_id"] = parsed_msg["tool_call_id"]

        if "name" in message and isinstance(message["name"], str):
            result_msg["name"] = message["name"]

    return result

1107

1108
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
    # per the Transformers docs & maintainers, tool call arguments in
    # assistant-role messages with tool_calls need to be dicts not JSON str -
    # this is how tool-use chat templates will expect them moving forwards
    # so, for messages that have tool_calls, parse the string (which we get
    # from openAI format) to dict
    for message in messages:
        if (message["role"] == "assistant" and "tool_calls" in message
                and isinstance(message["tool_calls"], list)):

            for item in message["tool_calls"]:
                item["function"]["arguments"] = json.loads(
                    item["function"]["arguments"])


1123
def parse_chat_messages(
1124
    messages: list[ChatCompletionMessageParam],
1125
    model_config: ModelConfig,
1126
    tokenizer: AnyTokenizer,
1127
    content_format: _ChatTemplateContentFormat,
1128
1129
) -> tuple[list[ConversationMessage], Optional[MultiModalDataDict]]:
    conversation: list[ConversationMessage] = []
1130
    mm_tracker = MultiModalItemTracker(model_config, tokenizer)
1131
1132

    for msg in messages:
1133
1134
1135
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1136
            content_format,
1137
        )
1138

1139
        conversation.extend(sub_messages)
1140

1141
1142
    _postprocess_messages(conversation)

1143
    return conversation, mm_tracker.all_mm_data()
1144
1145


1146
def parse_chat_messages_futures(
1147
    messages: list[ChatCompletionMessageParam],
1148
1149
    model_config: ModelConfig,
    tokenizer: AnyTokenizer,
1150
    content_format: _ChatTemplateContentFormat,
1151
1152
) -> tuple[list[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]:
    conversation: list[ConversationMessage] = []
1153
1154
1155
    mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)

    for msg in messages:
1156
1157
1158
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1159
            content_format,
1160
        )
1161
1162
1163

        conversation.extend(sub_messages)

1164
1165
    _postprocess_messages(conversation)

1166
1167
1168
    return conversation, mm_tracker.all_mm_data()


1169
1170
def apply_hf_chat_template(
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
1171
    conversation: list[ConversationMessage],
1172
    chat_template: Optional[str],
1173
    tools: Optional[list[dict[str, Any]]],
1174
    *,
1175
    trust_remote_code: bool = False,
1176
1177
    tokenize: bool = False,  # Different from HF's default
    **kwargs: Any,
1178
) -> str:
1179
    hf_chat_template = resolve_hf_chat_template(
1180
1181
1182
1183
1184
        tokenizer,
        chat_template=chat_template,
        tools=tools,
        trust_remote_code=trust_remote_code,
    )
1185

1186
    if hf_chat_template is None:
1187
1188
1189
1190
1191
        raise ValueError(
            "As of transformers v4.44, default chat template is no longer "
            "allowed, so you must provide a chat template if the tokenizer "
            "does not define one.")

1192
1193
    return tokenizer.apply_chat_template(
        conversation=conversation,  # type: ignore[arg-type]
1194
1195
        tools=tools,  # type: ignore[arg-type]
        chat_template=hf_chat_template,
1196
1197
1198
        tokenize=tokenize,
        **kwargs,
    )
1199
1200


1201
1202
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1203
    messages: list[ChatCompletionMessageParam],
1204
1205
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
1206
    **kwargs: Any,
1207
) -> list[int]:
1208
1209
1210
1211
1212
1213
    # The return value of resolve_mistral_chat_template is always None,
    # and we won't use it.
    resolve_mistral_chat_template(
        chat_template=chat_template,
        **kwargs,
    )
1214

1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
    try:
        return tokenizer.apply_chat_template(
            messages=messages,
            tools=tools,
            **kwargs,
        )
    # mistral-common uses assert statements to stop processing of input
    # if input does not comply with the expected format.
    # We convert those assertion errors to ValueErrors so they can be
    # are properly caught in the preprocessing_input step
    except AssertionError as e:
        raise ValueError from e