chat_utils.py 64.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import asyncio
5
import inspect
6
import json
7
from abc import ABC, abstractmethod
8
from collections import Counter, defaultdict, deque
9
from collections.abc import Awaitable, Callable, Iterable
10
from functools import cached_property, lru_cache, partial
11
from pathlib import Path
12
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar, cast
13

14
15
16
import jinja2
import jinja2.ext
import jinja2.meta
17
import jinja2.nodes
18
19
import jinja2.parser
import jinja2.sandbox
20
import transformers.utils.chat_template_utils as hf_chat_utils
21
from openai.types.chat import (
22
23
24
25
26
27
28
29
30
31
32
    ChatCompletionAssistantMessageParam,
    ChatCompletionContentPartImageParam,
    ChatCompletionContentPartInputAudioParam,
    ChatCompletionContentPartRefusalParam,
    ChatCompletionContentPartTextParam,
    ChatCompletionMessageToolCallParam,
    ChatCompletionToolMessageParam,
)
from openai.types.chat import (
    ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam,
)
33
from openai.types.chat import (
34
35
36
    ChatCompletionMessageParam as OpenAIChatCompletionMessageParam,
)
from openai.types.chat.chat_completion_content_part_input_audio_param import InputAudio
37
from openai.types.responses import ResponseInputImageParam
38
from openai_harmony import Message as OpenAIHarmonyMessage
39
40
from PIL import Image
from pydantic import BaseModel, ConfigDict, TypeAdapter
41
42
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast, ProcessorMixin

43
# pydantic needs the TypedDict from typing_extensions
44
from typing_extensions import Required, TypedDict
45

46
from vllm import envs
47
from vllm.config import ModelConfig
48
from vllm.logger import init_logger
49
from vllm.model_executor.models import SupportsMultiModal
50
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict, MultiModalUUIDDict
51
from vllm.multimodal.utils import MEDIA_CONNECTOR_REGISTRY, MediaConnector
52
53
from vllm.tokenizers import TokenizerLike
from vllm.tokenizers.mistral import MistralTokenizer
54
from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
55
from vllm.transformers_utils.processor import cached_get_processor
56
from vllm.utils import random_uuid
57
from vllm.utils.collection_utils import is_list_of
58
from vllm.utils.func_utils import supports_kw
59
60
61
62
63
64
from vllm.utils.import_utils import LazyLoader

if TYPE_CHECKING:
    import torch
else:
    torch = LazyLoader("torch", globals(), "torch")
65
66
67

logger = init_logger(__name__)

68
69
70
71
72
73
MODALITY_PLACEHOLDERS_MAP = {
    "image": "<##IMAGE##>",
    "audio": "<##AUDIO##>",
    "video": "<##VIDEO##>",
}

74

75
76
77
78
79
80
81
82
83
84
85
86
87
88
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."""


89
class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):
90
    image_embeds: str | dict[str, str] | None
91
92
93
94
95
96
97
    """
    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."""
98
    uuid: str | None
99
100
101
102
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
103
104


105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False):
    audio_embeds: str | dict[str, str] | None
    """
    The audio embeddings. It can be either:
    - A single base64 string representing a serialized torch tensor.
    - A dictionary where each value is a base64 string.
    """
    type: Required[Literal["audio_embeds"]]
    """The type of the content part."""
    uuid: str | None
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """


121
122
123
124
125
126
127
128
129
130
131
132
133
134
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."""


135
136
137
138
class PILImage(BaseModel):
    """
    A PIL.Image.Image object.
    """
139

140
141
142
143
144
145
146
147
148
149
150
151
    image_pil: Image.Image
    model_config = ConfigDict(arbitrary_types_allowed=True)


class CustomChatCompletionContentPILImageParam(TypedDict, total=False):
    """A simpler version of the param that only accepts a PIL image.

    Example:
    {
        "image_pil": ImageAsset('cherry_blossom').pil_image
    }
    """
152

153
154
    image_pil: PILImage | None
    uuid: str | None
155
156
157
158
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
159
160


161
162
163
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.
164

165
166
167
168
169
    Example:
    {
        "image_url": "https://example.com/image.jpg"
    }
    """
170

171
172
    image_url: str | None
    uuid: str | None
173
174
175
176
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
177
178
179
180


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

182
183
184
185
186
    Example:
    {
        "audio_url": "https://example.com/audio.mp3"
    }
    """
187

188
    audio_url: str | None
189
190


191
192
193
194
195
196
197
198
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"
    }
    """
199

200
201
    video_url: str | None
    uuid: str | None
202
203
204
205
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
206
207


Julien Denize's avatar
Julien Denize committed
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
class CustomThinkCompletionContentParam(TypedDict, total=False):
    """A Think Completion Content Param that accepts a plain text and a boolean.

    Example:
    {
        "thinking": "I am thinking about the answer",
        "closed": True,
        "type": "thinking"
    }
    """

    thinking: Required[str]
    """The thinking content."""

    closed: bool
    """Whether the thinking is closed."""

    type: Required[Literal["thinking"]]
    """The thinking type."""


229
230
231
232
233
234
235
236
237
ChatCompletionContentPartParam: TypeAlias = (
    OpenAIChatCompletionContentPartParam
    | ChatCompletionContentPartAudioParam
    | ChatCompletionContentPartInputAudioParam
    | ChatCompletionContentPartVideoParam
    | ChatCompletionContentPartRefusalParam
    | CustomChatCompletionContentPILImageParam
    | CustomChatCompletionContentSimpleImageParam
    | ChatCompletionContentPartImageEmbedsParam
238
    | ChatCompletionContentPartAudioEmbedsParam
239
240
241
242
243
    | CustomChatCompletionContentSimpleAudioParam
    | CustomChatCompletionContentSimpleVideoParam
    | str
    | CustomThinkCompletionContentParam
)
244
245
246
247


class CustomChatCompletionMessageParam(TypedDict, total=False):
    """Enables custom roles in the Chat Completion API."""
248

249
250
251
    role: Required[str]
    """The role of the message's author."""

252
    content: str | list[ChatCompletionContentPartParam]
253
254
255
256
257
258
259
260
261
    """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.
    """

262
    tool_call_id: str | None
263
264
    """Tool call that this message is responding to."""

265
    tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None
266
267
    """The tool calls generated by the model, such as function calls."""

268
269
270
    reasoning: str | None
    """The reasoning content for interleaved thinking."""

271

272
273
274
275
276
ChatCompletionMessageParam: TypeAlias = (
    OpenAIChatCompletionMessageParam
    | CustomChatCompletionMessageParam
    | OpenAIHarmonyMessage
)
277
278


279
# TODO: Make fields ReadOnly once mypy supports it
280
281
282
283
class ConversationMessage(TypedDict, total=False):
    role: Required[str]
    """The role of the message's author."""

284
    content: str | None | list[dict[str, str]]
285
286
    """The contents of the message"""

287
    tool_call_id: str | None
288
289
    """Tool call that this message is responding to."""

290
    name: str | None
291
292
    """The name of the function to call"""

293
    tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None
294
    """The tool calls generated by the model, such as function calls."""
295

296
297
298
299
300
301
    reasoning: str | None
    """The reasoning content for interleaved thinking."""

    reasoning_content: str | None
    """Deprecated: The reasoning content for interleaved thinking."""

302

303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
# 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):
319
320
321
322
323
        return (
            _is_var_access(node.node, varname)
            and isinstance(node.arg, jinja2.nodes.Const)
            and node.arg.value == key
        )
324
325
326
327
328
329
330
331
332
333

    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,
334
    key: str | None = None,
335
336
) -> bool:
    if isinstance(node, jinja2.nodes.Filter):
337
        return node.node is not None and _is_var_or_elems_access(
338
339
            node.node, varname, key
        )
340
341
342
    if isinstance(node, jinja2.nodes.Test):
        return _is_var_or_elems_access(node.node, varname, key)

343
    if isinstance(node, jinja2.nodes.Getitem) and isinstance(
344
345
        node.arg, jinja2.nodes.Slice
    ):
346
347
        return _is_var_or_elems_access(node.node, varname, key)

348
    return _is_attr_access(node, varname, key) if key else _is_var_access(node, varname)
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376


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 = [
377
        varname for _, varname in _iter_nodes_assign_var_or_elems(root, "messages")
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
    ]

    # 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


409
def _try_extract_ast(chat_template: str) -> jinja2.nodes.Template | None:
410
411
412
413
414
415
416
417
    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


418
@lru_cache(maxsize=32)
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
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"


439
def resolve_mistral_chat_template(
440
    chat_template: str | None,
441
    **kwargs: Any,
442
) -> str | None:
443
444
445
446
    if chat_template is not None or kwargs.get("chat_template_kwargs") is not None:
        raise ValueError(
            "'chat_template' or 'chat_template_kwargs' cannot be overridden "
            "for mistral tokenizer."
447
        )
448

449
450
    return None

451

452
_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()
453
454
455
456
457
458
459
460
461
"""
Used in `_try_get_processor_chat_template` to avoid calling
`cached_get_processor` again if the processor fails to be loaded.

This is needed because `lru_cache` does not cache when an exception happens.
"""


def _try_get_processor_chat_template(
462
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
463
    model_config: ModelConfig,
464
) -> str | None:
465
    cache_key = (tokenizer.name_or_path, model_config.trust_remote_code)
466
467
468
469
470
471
472
473
474
475
476
    if cache_key in _PROCESSOR_CHAT_TEMPLATES:
        return _PROCESSOR_CHAT_TEMPLATES[cache_key]

    try:
        processor = cached_get_processor(
            tokenizer.name_or_path,
            processor_cls=(
                PreTrainedTokenizer,
                PreTrainedTokenizerFast,
                ProcessorMixin,
            ),
477
            trust_remote_code=model_config.trust_remote_code,
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
        )
        if (
            isinstance(processor, ProcessorMixin)
            and hasattr(processor, "chat_template")
            and (chat_template := processor.chat_template) is not None
        ):
            _PROCESSOR_CHAT_TEMPLATES[cache_key] = chat_template
            return chat_template
    except Exception:
        logger.debug(
            "Failed to load AutoProcessor chat template for %s",
            tokenizer.name_or_path,
            exc_info=True,
        )

    _PROCESSOR_CHAT_TEMPLATES[cache_key] = None
    return None


497
def resolve_hf_chat_template(
498
499
500
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
501
502
    *,
    model_config: ModelConfig,
503
) -> str | None:
504
505
506
507
508
509
    # 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:
510
        chat_template = _try_get_processor_chat_template(tokenizer, model_config)
511
512
        if chat_template is not None:
            return chat_template
513
514
515
516
517

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

524
    # 4th priority: Predefined fallbacks
525
526
    path = get_chat_template_fallback_path(
        model_type=model_config.hf_config.model_type,
527
        tokenizer_name_or_path=model_config.tokenizer,
528
529
    )
    if path is not None:
530
        logger.info_once(
531
532
533
534
            "Loading chat template fallback for %s as there isn't one "
            "defined on HF Hub.",
            tokenizer.name_or_path,
        )
535
536
        chat_template = load_chat_template(path)
    else:
537
        logger.debug_once(
538
539
            "There is no chat template fallback for %s", tokenizer.name_or_path
        )
540
541

    return chat_template
542
543


544
def _resolve_chat_template_content_format(
545
546
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
547
    tokenizer: TokenizerLike | None,
548
    *,
549
    model_config: ModelConfig,
550
551
) -> _ChatTemplateContentFormat:
    if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
552
        hf_chat_template = resolve_hf_chat_template(
553
554
555
            tokenizer,
            chat_template=chat_template,
            tools=tools,
556
            model_config=model_config,
557
        )
558
    else:
559
560
        hf_chat_template = None

561
562
563
564
565
    jinja_text = (
        hf_chat_template
        if isinstance(hf_chat_template, str)
        else load_chat_template(chat_template, is_literal=True)
    )
566

567
568
569
570
571
    detected_format = (
        "string"
        if jinja_text is None
        else _detect_content_format(jinja_text, default="string")
    )
572

573
    return detected_format
574
575
576


@lru_cache
577
def _log_chat_template_content_format(
578
    chat_template: str | None,
579
    given_format: ChatTemplateContentFormatOption,
580
581
    detected_format: ChatTemplateContentFormatOption,
):
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
    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,
        )

599
600

def resolve_chat_template_content_format(
601
602
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
603
    given_format: ChatTemplateContentFormatOption,
604
    tokenizer: TokenizerLike | None,
605
    *,
606
    model_config: ModelConfig,
607
) -> _ChatTemplateContentFormat:
608
609
610
    if given_format != "auto":
        return given_format

611
612
613
614
    detected_format = _resolve_chat_template_content_format(
        chat_template,
        tools,
        tokenizer,
615
        model_config=model_config,
616
617
618
619
620
621
622
623
    )

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

624
    return detected_format
625

626

627
ModalityStr = Literal["image", "audio", "video", "image_embeds", "audio_embeds"]
628
629
630
_T = TypeVar("_T")


631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
def _extract_embeds(tensors: list[torch.Tensor]):
    if len(tensors) == 0:
        return tensors

    if len(tensors) == 1:
        tensors[0]._is_single_item = True  # type: ignore
        return tensors[0]  # To keep backwards compatibility for single item input

    first_shape = tensors[0].shape
    if all(t.shape == first_shape for t in tensors):
        return torch.stack(tensors)

    return tensors


def _get_embeds_data(items_by_modality: dict[str, list[Any]], modality: str):
    embeds_key = f"{modality}_embeds"
    embeds = items_by_modality[embeds_key]

    if len(embeds) == 0:
        return embeds
    if is_list_of(embeds, torch.Tensor):
        return _extract_embeds(embeds)
    if is_list_of(embeds, dict):
        if not embeds:
            return {}

        first_keys = set(embeds[0].keys())
        if any(set(item.keys()) != first_keys for item in embeds[1:]):
            raise ValueError(
                "All dictionaries in the list of embeddings must have the same keys."
            )

        return {k: _extract_embeds([item[k] for item in embeds]) for k in first_keys}

    return embeds


669
class BaseMultiModalItemTracker(ABC, Generic[_T]):
670
671
672
673
674
675
    """
    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.
    """

676
    def __init__(self, model_config: ModelConfig):
677
678
        super().__init__()

679
        self._model_config = model_config
680

681
682
        self._items_by_modality = defaultdict[str, list[_T | None]](list)
        self._uuids_by_modality = defaultdict[str, list[str | None]](list)
683

684
    @property
685
686
    def model_config(self) -> ModelConfig:
        return self._model_config
687

688
    @cached_property
689
    def model_cls(self) -> type[SupportsMultiModal]:
690
        from vllm.model_executor.model_loader import get_model_cls
691

692
        model_cls = get_model_cls(self.model_config)
693
        return cast(type[SupportsMultiModal], model_cls)
694

695
696
    @property
    def allowed_local_media_path(self):
697
        return self._model_config.allowed_local_media_path
698

699
700
    @property
    def allowed_media_domains(self):
701
        return self._model_config.allowed_media_domains
702

703
704
705
706
    @property
    def mm_registry(self):
        return MULTIMODAL_REGISTRY

707
708
    @cached_property
    def mm_processor(self):
709
        return self.mm_registry.create_processor(self.model_config)
710

711
    def add(
712
713
        self,
        modality: ModalityStr,
714
715
716
        item: _T | None,
        uuid: str | None = None,
    ) -> str | None:
717
718
719
        """
        Add a multi-modal item to the current prompt and returns the
        placeholder string to use, if any.
720
721

        An optional uuid can be added which serves as a unique identifier of the
722
        media.
723
        """
724
        input_modality = modality.replace("_embeds", "")
725
        num_items = len(self._items_by_modality[modality]) + 1
726

727
        self.mm_processor.validate_num_items(input_modality, num_items)
728

729
        self._items_by_modality[modality].append(item)
730
        self._uuids_by_modality[modality].append(uuid)
731

732
        return self.model_cls.get_placeholder_str(modality, num_items)
733

734
    def all_mm_uuids(self) -> MultiModalUUIDDict | None:
735
736
        if not self._items_by_modality:
            return None
737

738
739
        uuids_by_modality = dict(self._uuids_by_modality)
        if "image" in uuids_by_modality and "image_embeds" in uuids_by_modality:
740
            raise ValueError("Mixing raw image and embedding inputs is not allowed")
741
742
        if "audio" in uuids_by_modality and "audio_embeds" in uuids_by_modality:
            raise ValueError("Mixing raw audio and embedding inputs is not allowed")
743

744
        mm_uuids = {}
745
746
747
748
        if "image_embeds" in uuids_by_modality:
            mm_uuids["image"] = uuids_by_modality["image_embeds"]
        if "image" in uuids_by_modality:
            mm_uuids["image"] = uuids_by_modality["image"]  # UUIDs of images
749
750
        if "audio_embeds" in uuids_by_modality:
            mm_uuids["audio"] = uuids_by_modality["audio_embeds"]
751
752
753
754
        if "audio" in uuids_by_modality:
            mm_uuids["audio"] = uuids_by_modality["audio"]  # UUIDs of audios
        if "video" in uuids_by_modality:
            mm_uuids["video"] = uuids_by_modality["video"]  # UUIDs of videos
755

756
757
        return mm_uuids

758
759
760
761
762
    @abstractmethod
    def create_parser(self) -> "BaseMultiModalContentParser":
        raise NotImplementedError


763
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
764
    def all_mm_data(self) -> MultiModalDataDict | None:
765
766
        if not self._items_by_modality:
            return None
767

768
769
        items_by_modality = dict(self._items_by_modality)
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
770
            raise ValueError("Mixing raw image and embedding inputs is not allowed")
771
772
        if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
            raise ValueError("Mixing raw audio and embedding inputs is not allowed")
773

774
        mm_inputs = {}
775
        if "image_embeds" in items_by_modality:
776
            mm_inputs["image"] = _get_embeds_data(items_by_modality, "image")
777
        if "image" in items_by_modality:
778
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
779
        if "audio_embeds" in items_by_modality:
780
            mm_inputs["audio"] = _get_embeds_data(items_by_modality, "audio")
781
        if "audio" in items_by_modality:
782
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
783
        if "video" in items_by_modality:
784
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
785

786
        return mm_inputs
787
788
789
790
791

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


792
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
793
    async def all_mm_data(self) -> MultiModalDataDict | None:
794
795
        if not self._items_by_modality:
            return None
796

797
798
799
800
801
802
803
804
        coros_by_modality = {
            modality: [item or asyncio.sleep(0) for item in items]
            for modality, items in self._items_by_modality.items()
        }
        items_by_modality: dict[str, list[object | None]] = {
            modality: await asyncio.gather(*coros)
            for modality, coros in coros_by_modality.items()
        }
805
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
806
            raise ValueError("Mixing raw image and embedding inputs is not allowed")
807
808
        if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
            raise ValueError("Mixing raw audio and embedding inputs is not allowed")
809

810
        mm_inputs = {}
811
        if "image_embeds" in items_by_modality:
812
            mm_inputs["image"] = _get_embeds_data(items_by_modality, "image")
813
        if "image" in items_by_modality:
814
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
815
        if "audio_embeds" in items_by_modality:
816
            mm_inputs["audio"] = _get_embeds_data(items_by_modality, "audio")
817
        if "audio" in items_by_modality:
818
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
819
        if "video" in items_by_modality:
820
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
821

822
        return mm_inputs
823
824
825
826
827
828
829
830
831

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


class BaseMultiModalContentParser(ABC):
    def __init__(self) -> None:
        super().__init__()

832
        # stores model placeholders list with corresponding
833
834
835
836
837
838
839
        # general MM placeholder:
        # {
        #   "<##IMAGE##>": ["<image>", "<image>", "<image>"],
        #   "<##AUDIO##>": ["<audio>", "<audio>"]
        # }
        self._placeholder_storage: dict[str, list] = defaultdict(list)

840
    def _add_placeholder(self, modality: ModalityStr, placeholder: str | None):
841
        mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]
842
        if placeholder:
843
            self._placeholder_storage[mod_placeholder].append(placeholder)
844

845
846
    def mm_placeholder_storage(self) -> dict[str, list]:
        return dict(self._placeholder_storage)
847
848

    @abstractmethod
849
    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
850
851
        raise NotImplementedError

852
    @abstractmethod
853
    def parse_image_embeds(
854
        self,
855
856
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
857
    ) -> None:
858
859
        raise NotImplementedError

860
    @abstractmethod
861
    def parse_image_pil(
862
        self, image_pil: Image.Image | None, uuid: str | None = None
863
    ) -> None:
864
865
        raise NotImplementedError

866
    @abstractmethod
867
    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
868
869
        raise NotImplementedError

870
    @abstractmethod
871
    def parse_input_audio(
872
        self, input_audio: InputAudio | None, uuid: str | None = None
873
    ) -> None:
874
875
        raise NotImplementedError

876
877
878
879
880
881
882
883
    @abstractmethod
    def parse_audio_embeds(
        self,
        audio_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
    ) -> None:
        raise NotImplementedError

884
    @abstractmethod
885
    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
886
887
        raise NotImplementedError

888
889
890
891
892
893

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

        self._tracker = tracker
894
895
896
        multimodal_config = self._tracker.model_config.multimodal_config
        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)

897
898
        self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
899
            media_io_kwargs=media_io_kwargs,
900
            allowed_local_media_path=tracker.allowed_local_media_path,
901
            allowed_media_domains=tracker.allowed_media_domains,
902
903
        )

904
905
    @property
    def model_config(self) -> ModelConfig:
906
        return self._tracker.model_config
907

908
    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
909
        image = self._connector.fetch_image(image_url) if image_url else None
910

911
        placeholder = self._tracker.add("image", image, uuid)
912
        self._add_placeholder("image", placeholder)
913

914
    def parse_image_embeds(
915
        self,
916
917
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
918
    ) -> None:
919
920
921
922
923
924
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise ValueError(
                "You must set `--enable-mm-embeds` to input `image_embeds`"
            )

925
926
927
928
929
        if isinstance(image_embeds, dict):
            embeds = {
                k: self._connector.fetch_image_embedding(v)
                for k, v in image_embeds.items()
            }
930
            placeholder = self._tracker.add("image_embeds", embeds, uuid)
931
932
933

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

936
937
938
        if image_embeds is None:
            placeholder = self._tracker.add("image_embeds", None, uuid)

939
        self._add_placeholder("image", placeholder)
940

941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
    def parse_audio_embeds(
        self,
        audio_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise ValueError(
                "You must set `--enable-mm-embeds` to input `audio_embeds`"
            )

        if isinstance(audio_embeds, dict):
            embeds = {
                k: self._connector.fetch_audio_embedding(v)
                for k, v in audio_embeds.items()
            }
            placeholder = self._tracker.add("audio_embeds", embeds, uuid)
        elif isinstance(audio_embeds, str):
            embedding = self._connector.fetch_audio_embedding(audio_embeds)
            placeholder = self._tracker.add("audio_embeds", embedding, uuid)
        else:
            placeholder = self._tracker.add("audio_embeds", None, uuid)

        self._add_placeholder("audio", placeholder)

966
    def parse_image_pil(
967
        self, image_pil: Image.Image | None, uuid: str | None = None
968
969
    ) -> None:
        placeholder = self._tracker.add("image", image_pil, uuid)
970
        self._add_placeholder("image", placeholder)
971

972
    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
973
        audio = self._connector.fetch_audio(audio_url) if audio_url else None
974

975
        placeholder = self._tracker.add("audio", audio, uuid)
976
        self._add_placeholder("audio", placeholder)
977

978
    def parse_input_audio(
979
        self, input_audio: InputAudio | None, uuid: str | None = None
980
    ) -> None:
981
982
983
984
985
986
987
988
989
990
        if input_audio:
            audio_data = input_audio.get("data", "")
            audio_format = input_audio.get("format", "")
            if audio_data:
                audio_url = f"data:audio/{audio_format};base64,{audio_data}"
            else:
                # If a UUID is provided, audio data may be empty.
                audio_url = None
        else:
            audio_url = None
991

992
        return self.parse_audio(audio_url, uuid)
993

994
    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
995
        video = self._connector.fetch_video(video_url=video_url) if video_url else None
996

997
        placeholder = self._tracker.add("video", video, uuid)
998
        self._add_placeholder("video", placeholder)
999

1000
1001
1002
1003
1004
1005

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

        self._tracker = tracker
1006
1007
        multimodal_config = self._tracker.model_config.multimodal_config
        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)
1008
1009
        self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
1010
            media_io_kwargs=media_io_kwargs,
1011
            allowed_local_media_path=tracker.allowed_local_media_path,
1012
            allowed_media_domains=tracker.allowed_media_domains,
1013
        )
1014

1015
1016
    @property
    def model_config(self) -> ModelConfig:
1017
        return self._tracker.model_config
1018

1019
    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
1020
        image_coro = self._connector.fetch_image_async(image_url) if image_url else None
1021

1022
        placeholder = self._tracker.add("image", image_coro, uuid)
1023
        self._add_placeholder("image", placeholder)
1024

1025
    def parse_image_embeds(
1026
        self,
1027
1028
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
1029
    ) -> None:
1030
1031
1032
1033
1034
1035
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise ValueError(
                "You must set `--enable-mm-embeds` to input `image_embeds`"
            )

1036
        future: asyncio.Future[str | dict[str, str] | None] = asyncio.Future()
1037
1038
1039
1040
1041
1042
1043
1044
1045

        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):
1046
            embedding = self._connector.fetch_image_embedding(image_embeds)
1047
1048
            future.set_result(embedding)

1049
1050
1051
        if image_embeds is None:
            future.set_result(None)

1052
        placeholder = self._tracker.add("image_embeds", future, uuid)
1053
        self._add_placeholder("image", placeholder)
1054

1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
    def parse_audio_embeds(
        self,
        audio_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
    ) -> None:
        mm_config = self.model_config.get_multimodal_config()
        if not mm_config.enable_mm_embeds:
            raise ValueError(
                "You must set `--enable-mm-embeds` to input `audio_embeds`"
            )

        logger.info(
            "🎵 Parsing audio_embeds: type=%s, uuid=%s, is_dict=%s, "
            "is_str=%s, is_none=%s",
            type(audio_embeds).__name__,
            uuid,
            isinstance(audio_embeds, dict),
            isinstance(audio_embeds, str),
            audio_embeds is None,
        )

        future: asyncio.Future[str | dict[str, str] | None] = asyncio.Future()

        if isinstance(audio_embeds, dict):
            logger.info(
                "🎵 Processing dict audio_embeds with %d entries",
                len(audio_embeds),
            )
            embeds = {
                k: self._connector.fetch_audio_embedding(v)
                for k, v in audio_embeds.items()
            }
            future.set_result(embeds)
            logger.info(
                "🎵 Successfully loaded %d audio embeddings from dict",
                len(embeds),
            )

        if isinstance(audio_embeds, str):
            base64_size = len(audio_embeds)
            logger.info(
                "🎵 Processing base64 audio_embeds: %d chars (%.2f KB)",
                base64_size,
                base64_size / 1024,
            )
            embedding = self._connector.fetch_audio_embedding(audio_embeds)
            future.set_result(embedding)
            logger.info(
                "🎵 Successfully loaded audio embedding tensor: shape=%s, dtype=%s",
                embedding.shape,
                embedding.dtype,
            )

        if audio_embeds is None:
            logger.info("🎵 Audio embeds is None (UUID-only reference)")
            future.set_result(None)

        placeholder = self._tracker.add("audio_embeds", future, uuid)
        self._add_placeholder("audio", placeholder)
        logger.info("🎵 Added audio_embeds placeholder with uuid=%s", uuid)

1116
    def parse_image_pil(
1117
        self, image_pil: Image.Image | None, uuid: str | None = None
1118
    ) -> None:
1119
        future: asyncio.Future[Image.Image | None] = asyncio.Future()
1120
1121
1122
1123
        if image_pil:
            future.set_result(image_pil)
        else:
            future.set_result(None)
1124

1125
        placeholder = self._tracker.add("image", future, uuid)
1126
        self._add_placeholder("image", placeholder)
1127

1128
    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
1129
        audio_coro = self._connector.fetch_audio_async(audio_url) if audio_url else None
1130

1131
        placeholder = self._tracker.add("audio", audio_coro, uuid)
1132
        self._add_placeholder("audio", placeholder)
1133

1134
    def parse_input_audio(
1135
        self, input_audio: InputAudio | None, uuid: str | None = None
1136
    ) -> None:
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
        if input_audio:
            audio_data = input_audio.get("data", "")
            audio_format = input_audio.get("format", "")
            if audio_data:
                audio_url = f"data:audio/{audio_format};base64,{audio_data}"
            else:
                # If a UUID is provided, audio data may be empty.
                audio_url = None
        else:
            audio_url = None
1147

1148
        return self.parse_audio(audio_url, uuid)
1149

1150
    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
1151
1152
1153
1154
1155
        video = (
            self._connector.fetch_video_async(video_url=video_url)
            if video_url
            else None
        )
1156

1157
        placeholder = self._tracker.add("video", video, uuid)
1158
        self._add_placeholder("video", placeholder)
1159

1160

1161
def validate_chat_template(chat_template: Path | str | None):
1162
1163
1164
1165
1166
    """Raises if the provided chat template appears invalid."""
    if chat_template is None:
        return

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

    elif isinstance(chat_template, str):
        JINJA_CHARS = "{}\n"
1171
1172
1173
1174
        if (
            not any(c in chat_template for c in JINJA_CHARS)
            and not Path(chat_template).exists()
        ):
1175
1176
1177
            # Try to find the template in the built-in templates directory
            from vllm.transformers_utils.chat_templates.registry import (
                CHAT_TEMPLATES_DIR,
1178
            )
1179

1180
1181
1182
1183
1184
1185
1186
1187
            builtin_template_path = CHAT_TEMPLATES_DIR / chat_template
            if not builtin_template_path.exists():
                raise ValueError(
                    f"The supplied chat template string ({chat_template}) "
                    f"appears path-like, but doesn't exist! "
                    f"Tried: {chat_template} and {builtin_template_path}"
                )

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


1192
def _load_chat_template(
1193
    chat_template: Path | str | None,
1194
1195
    *,
    is_literal: bool = False,
1196
) -> str | None:
1197
1198
    if chat_template is None:
        return None
1199
1200
1201

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

1206
        return chat_template
1207

1208
    try:
1209
        with open(chat_template) as f:
1210
            return f.read()
1211
    except OSError as e:
1212
1213
1214
        if isinstance(chat_template, Path):
            raise

1215
1216
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template for c in JINJA_CHARS):
1217
1218
1219
            # Try to load from the built-in templates directory
            from vllm.transformers_utils.chat_templates.registry import (
                CHAT_TEMPLATES_DIR,
1220
            )
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233

            builtin_template_path = CHAT_TEMPLATES_DIR / chat_template
            try:
                with open(builtin_template_path) as f:
                    return f.read()
            except OSError:
                msg = (
                    f"The supplied chat template ({chat_template}) "
                    f"looks like a file path, but it failed to be opened. "
                    f"Tried: {chat_template} and {builtin_template_path}. "
                    f"Reason: {e}"
                )
                raise ValueError(msg) from e
1234

1235
1236
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
1237
1238
1239
1240
1241
1242
1243
        return _load_chat_template(chat_template, is_literal=True)


_cached_load_chat_template = lru_cache(_load_chat_template)


def load_chat_template(
1244
    chat_template: Path | str | None,
1245
1246
    *,
    is_literal: bool = False,
1247
) -> str | None:
1248
    return _cached_load_chat_template(chat_template, is_literal=is_literal)
1249
1250


1251
1252
1253
def _get_interleaved_text_prompt(
    placeholder_storage: dict[str, list], texts: list[str]
) -> str:
1254
1255
1256
1257
1258
1259
1260
    for idx, elem in enumerate(texts):
        if elem in placeholder_storage:
            texts[idx] = placeholder_storage[elem].pop(0)

    return "\n".join(texts)


1261
# TODO: Let user specify how to insert multimodal tokens into prompt
1262
# (similar to chat template)
1263
1264
1265
1266
1267
def _get_full_multimodal_text_prompt(
    placeholder_storage: dict[str, list],
    texts: list[str],
    interleave_strings: bool,
) -> str:
1268
    """Combine multimodal prompts for a multimodal language model."""
1269

1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
    # flatten storage to make it looks like
    # {
    #   "<|image|>": 2,
    #   "<|audio|>": 1
    # }
    placeholder_counts = Counter(
        [v for elem in placeholder_storage.values() for v in elem]
    )

    if interleave_strings:
        text_prompt = _get_interleaved_text_prompt(placeholder_storage, texts)
    else:
        text_prompt = "\n".join(texts)

    # Pass interleaved text further in case the user used image placeholders
    # himself, but forgot to disable the 'interleave_strings' flag

1287
    # Look through the text prompt to check for missing placeholders
1288
    missing_placeholders: list[str] = []
1289
1290
1291
1292
1293
    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:
1294
1295
1296
1297
            logger.error(
                "Placeholder count is negative! "
                "Ensure that the 'interleave_strings' flag is disabled "
                "(current value: %s) "
1298
1299
                "when manually placing image placeholders.",
                interleave_strings,
1300
1301
            )
            logger.debug("Input prompt: %s", text_prompt)
1302
1303
            raise ValueError(
                f"Found more '{placeholder}' placeholders in input prompt than "
1304
1305
                "actual multimodal data items."
            )
1306

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

1309
1310
    # NOTE: Default behaviour: we always add missing placeholders
    # at the front of the prompt, if interleave_strings=False
1311
    return "\n".join(missing_placeholders + [text_prompt])
1312
1313


1314
1315
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
1316
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
1317
_AudioEmbedsParser = partial(cast, ChatCompletionContentPartAudioEmbedsParam)
1318
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
1319
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
1320
_PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
Julien Denize's avatar
Julien Denize committed
1321
_ThinkParser = partial(cast, CustomThinkCompletionContentParam)
1322
1323
1324
1325
# Need to validate url objects
_ImageParser = TypeAdapter(ChatCompletionContentPartImageParam).validate_python
_AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python
_VideoParser = TypeAdapter(ChatCompletionContentPartVideoParam).validate_python
1326

1327
_ResponsesInputImageParser = TypeAdapter(ResponseInputImageParam).validate_python
1328
_ContentPart: TypeAlias = str | dict[str, str] | InputAudio | PILImage
1329

1330
# Define a mapping from part types to their corresponding parsing functions.
1331
MM_PARSER_MAP: dict[
1332
1333
1334
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
1335
1336
1337
    "text": lambda part: _TextParser(part).get("text", None),
    "thinking": lambda part: _ThinkParser(part).get("thinking", None),
    "input_text": lambda part: _TextParser(part).get("text", None),
1338
    "output_text": lambda part: _TextParser(part).get("text", None),
1339
1340
1341
    "input_image": lambda part: _ResponsesInputImageParser(part).get("image_url", None),
    "image_url": lambda part: _ImageParser(part).get("image_url", {}).get("url", None),
    "image_embeds": lambda part: _ImageEmbedsParser(part).get("image_embeds", None),
1342
    "audio_embeds": lambda part: _AudioEmbedsParser(part).get("audio_embeds", None),
1343
    "image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
1344
1345
    "audio_url": lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
    "input_audio": lambda part: _InputAudioParser(part).get("input_audio", None),
1346
    "refusal": lambda part: _RefusalParser(part).get("refusal", None),
1347
    "video_url": lambda part: _VideoParser(part).get("video_url", {}).get("url", None),
1348
1349
1350
1351
}


def _parse_chat_message_content_mm_part(
1352
1353
    part: ChatCompletionContentPartParam,
) -> tuple[str, _ContentPart]:
1354
    """
1355
    Parses a given multi-modal content part based on its type.
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368

    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(
1369
1370
        part, dict
    )  # This is needed to avoid mypy errors: part.get() from str
1371
    part_type = part.get("type", None)
1372
    uuid = part.get("uuid", None)
1373

1374
    if isinstance(part_type, str) and part_type in MM_PARSER_MAP and uuid is None:  # noqa: E501
1375
1376
1377
        content = MM_PARSER_MAP[part_type](part)

        # Special case for 'image_url.detail'
1378
1379
        # We only support 'auto', which is the default
        if part_type == "image_url" and part.get("detail", "auto") != "auto":
1380
            logger.warning(
1381
                "'image_url.detail' is currently not supported and will be ignored."
1382
            )
1383
1384
1385
1386

        return part_type, content

    # Handle missing 'type' but provided direct URL fields.
1387
    # 'type' is required field by pydantic
1388
1389
    if part_type is None or uuid is not None:
        if "image_url" in part:
1390
            image_params = cast(CustomChatCompletionContentSimpleImageParam, part)
1391
1392
1393
1394
1395
1396
1397
1398
            image_url = image_params.get("image_url", None)
            if isinstance(image_url, dict):
                # Can potentially happen if user provides a uuid
                # with url as a dict of {"url": url}
                image_url = image_url.get("url", None)
            return "image_url", image_url
        if "image_pil" in part:
            # "image_pil" could be None if UUID is provided.
1399
            image_params = cast(  # type: ignore
1400
1401
1402
1403
1404
1405
                CustomChatCompletionContentPILImageParam, part
            )
            image_pil = image_params.get("image_pil", None)
            return "image_pil", image_pil
        if "image_embeds" in part:
            # "image_embeds" could be None if UUID is provided.
1406
            image_params = cast(  # type: ignore
1407
1408
1409
1410
                ChatCompletionContentPartImageEmbedsParam, part
            )
            image_embeds = image_params.get("image_embeds", None)
            return "image_embeds", image_embeds
1411
1412
1413
1414
1415
1416
1417
        if "audio_embeds" in part:
            # "audio_embeds" could be None if UUID is provided.
            audio_params = cast(  # type: ignore[assignment]
                ChatCompletionContentPartAudioEmbedsParam, part
            )
            audio_embeds = audio_params.get("audio_embeds", None)
            return "audio_embeds", audio_embeds
1418
        if "audio_url" in part:
1419
1420
1421
            audio_params = cast(  # type: ignore[assignment]
                CustomChatCompletionContentSimpleAudioParam, part
            )
1422
1423
1424
1425
1426
1427
            audio_url = audio_params.get("audio_url", None)
            if isinstance(audio_url, dict):
                # Can potentially happen if user provides a uuid
                # with url as a dict of {"url": url}
                audio_url = audio_url.get("url", None)
            return "audio_url", audio_url
1428
        if part.get("input_audio") is not None:
1429
            input_audio_params = cast(dict[str, str], part)
1430
            return "input_audio", input_audio_params
1431
        if "video_url" in part:
1432
            video_params = cast(CustomChatCompletionContentSimpleVideoParam, part)
1433
1434
1435
1436
1437
1438
            video_url = video_params.get("video_url", None)
            if isinstance(video_url, dict):
                # Can potentially happen if user provides a uuid
                # with url as a dict of {"url": url}
                video_url = video_url.get("url", None)
            return "video_url", video_url
1439
1440
1441
1442
1443
1444
1445
1446
        # 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"


1447
PART_TYPES_TO_SKIP_NONE_CONTENT = (
1448
1449
1450
    "text",
    "refusal",
)
1451

1452

1453
1454
1455
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
1456
    mm_tracker: BaseMultiModalItemTracker,
1457
1458
    *,
    wrap_dicts: bool,
1459
    interleave_strings: bool,
1460
) -> list[ConversationMessage]:
1461
    content = list[_ContentPart]()
1462

1463
    mm_parser = mm_tracker.create_parser()
1464
1465

    for part in parts:
1466
        parse_res = _parse_chat_message_content_part(
1467
1468
1469
            part,
            mm_parser,
            wrap_dicts=wrap_dicts,
1470
            interleave_strings=interleave_strings,
1471
        )
1472
1473
        if parse_res:
            content.append(parse_res)
1474

1475
    if wrap_dicts:
1476
        # Parsing wraps images and texts as interleaved dictionaries
1477
        return [ConversationMessage(role=role, content=content)]  # type: ignore
1478
    texts = cast(list[str], content)
1479
1480
    mm_placeholder_storage = mm_parser.mm_placeholder_storage()
    if mm_placeholder_storage:
1481
1482
1483
        text_prompt = _get_full_multimodal_text_prompt(
            mm_placeholder_storage, texts, interleave_strings
        )
1484
1485
1486
    else:
        text_prompt = "\n".join(texts)

1487
1488
1489
1490
    return [ConversationMessage(role=role, content=text_prompt)]


def _parse_chat_message_content_part(
1491
1492
1493
1494
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
1495
    interleave_strings: bool,
1496
) -> _ContentPart | None:
1497
1498
1499
1500
1501
1502
1503
1504
    """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
1505
        return part
1506
1507
    # Handle structured dictionary parts
    part_type, content = _parse_chat_message_content_mm_part(part)
1508
    # if part_type is text/refusal/image_url/audio_url/video_url/input_audio but
1509
    # content is None, log a warning and skip
1510
    if part_type in PART_TYPES_TO_SKIP_NONE_CONTENT and content is None:
1511
        logger.warning(
1512
            "Skipping multimodal part '%s' (type: '%s') "
1513
1514
1515
1516
            "with empty / unparsable content.",
            part,
            part_type,
        )
1517
1518
        return None

1519
    if part_type in ("text", "input_text", "output_text", "refusal", "thinking"):
1520
1521
        str_content = cast(str, content)
        if wrap_dicts:
1522
            return {"type": "text", "text": str_content}
1523
1524
        else:
            return str_content
1525

1526
1527
1528
1529
1530
1531
    # For media items, if a user has provided one, use it. Otherwise, insert
    # a placeholder empty uuid.
    uuid = part.get("uuid", None)
    if uuid is not None:
        uuid = str(uuid)

1532
    modality = None
1533
    if part_type == "image_pil":
1534
        image_content = cast(Image.Image, content) if content is not None else None
1535
        mm_parser.parse_image_pil(image_content, uuid)
1536
        modality = "image"
1537
    elif part_type in ("image_url", "input_image"):
1538
        str_content = cast(str, content)
1539
        mm_parser.parse_image(str_content, uuid)
1540
1541
        modality = "image"
    elif part_type == "image_embeds":
1542
        content = cast(str | dict[str, str], content) if content is not None else None
1543
        mm_parser.parse_image_embeds(content, uuid)
1544
        modality = "image"
1545
1546
1547
1548
    elif part_type == "audio_embeds":
        content = cast(str | dict[str, str], content) if content is not None else None
        mm_parser.parse_audio_embeds(content, uuid)
        modality = "audio"
1549
    elif part_type == "audio_url":
1550
        str_content = cast(str, content)
1551
        mm_parser.parse_audio(str_content, uuid)
1552
1553
        modality = "audio"
    elif part_type == "input_audio":
1554
        dict_content = cast(InputAudio, content)
1555
        mm_parser.parse_input_audio(dict_content, uuid)
1556
1557
        modality = "audio"
    elif part_type == "video_url":
1558
        str_content = cast(str, content)
1559
        mm_parser.parse_video(str_content, uuid)
1560
1561
1562
        modality = "video"
    else:
        raise NotImplementedError(f"Unknown part type: {part_type}")
1563

1564
1565
1566
    return (
        {"type": modality}
        if wrap_dicts
1567
        else (MODALITY_PLACEHOLDERS_MAP[modality] if interleave_strings else None)
1568
    )
1569
1570


1571
1572
1573
1574
1575
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


1576
def _parse_chat_message_content(
1577
1578
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
1579
    content_format: _ChatTemplateContentFormat,
1580
    interleave_strings: bool,
1581
) -> list[ConversationMessage]:
1582
1583
    role = message["role"]
    content = message.get("content")
1584
    reasoning = message.get("reasoning") or message.get("reasoning_content")
1585

1586
    if content is None:
1587
1588
        content = []
    elif isinstance(content, str):
1589
        content = [ChatCompletionContentPartTextParam(type="text", text=content)]
1590
    result = _parse_chat_message_content_parts(
1591
1592
        role,
        content,  # type: ignore
1593
        mm_tracker,
1594
        wrap_dicts=(content_format == "openai"),
1595
        interleave_strings=interleave_strings,
1596
    )
1597

1598
    for result_msg in result:
1599
        if role == "assistant":
1600
1601
            parsed_msg = _AssistantParser(message)

1602
1603
1604
            # The 'tool_calls' is not None check ensures compatibility.
            # It's needed only if downstream code doesn't strictly
            # follow the OpenAI spec.
1605
            if "tool_calls" in parsed_msg and parsed_msg["tool_calls"] is not None:
1606
                result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
1607
1608
1609
1610
1611
1612
            # Include reasoning if present for interleaved thinking.
            if reasoning is not None:
                result_msg["reasoning"] = cast(str, reasoning)
                result_msg["reasoning_content"] = cast(
                    str, reasoning
                )  # keep compatibility
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
        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

1623

1624
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1625
1626
1627
1628
1629
1630
    # 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:
1631
1632
1633
1634
1635
        if (
            message["role"] == "assistant"
            and "tool_calls" in message
            and isinstance(message["tool_calls"], list)
        ):
1636
            for item in message["tool_calls"]:
1637
1638
                # if arguments is None or empty string, set to {}
                if content := item["function"].get("arguments"):
1639
1640
                    if not isinstance(content, (dict, list)):
                        item["function"]["arguments"] = json.loads(content)
1641
1642
                else:
                    item["function"]["arguments"] = {}
1643
1644


1645
def parse_chat_messages(
1646
    messages: list[ChatCompletionMessageParam],
1647
    model_config: ModelConfig,
1648
    content_format: _ChatTemplateContentFormat,
1649
1650
) -> tuple[
    list[ConversationMessage],
1651
1652
    MultiModalDataDict | None,
    MultiModalUUIDDict | None,
1653
]:
1654
    conversation: list[ConversationMessage] = []
1655
    mm_tracker = MultiModalItemTracker(model_config)
1656
1657

    for msg in messages:
1658
1659
1660
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1661
            content_format,
1662
1663
1664
1665
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1666
            ),
1667
        )
1668

1669
        conversation.extend(sub_messages)
1670

1671
1672
    _postprocess_messages(conversation)

1673
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1674
1675


1676
def parse_chat_messages_futures(
1677
    messages: list[ChatCompletionMessageParam],
1678
    model_config: ModelConfig,
1679
    content_format: _ChatTemplateContentFormat,
1680
1681
) -> tuple[
    list[ConversationMessage],
1682
1683
    Awaitable[MultiModalDataDict | None],
    MultiModalUUIDDict | None,
1684
]:
1685
    conversation: list[ConversationMessage] = []
1686
    mm_tracker = AsyncMultiModalItemTracker(model_config)
1687
1688

    for msg in messages:
1689
1690
1691
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1692
            content_format,
1693
1694
1695
1696
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1697
            ),
1698
        )
1699
1700
1701

        conversation.extend(sub_messages)

1702
1703
    _postprocess_messages(conversation)

1704
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1705
1706


1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
# adapted from https://github.com/huggingface/transformers/blob/v4.56.2/src/transformers/utils/chat_template_utils.py#L398-L412
# only preserve the parse function used to resolve chat template kwargs
class AssistantTracker(jinja2.ext.Extension):
    tags = {"generation"}

    def parse(self, parser: jinja2.parser.Parser) -> jinja2.nodes.CallBlock:
        lineno = next(parser.stream).lineno
        body = parser.parse_statements(["name:endgeneration"], drop_needle=True)
        call = self.call_method("_generation_support")
        call_block = jinja2.nodes.CallBlock(call, [], [], body)
        return call_block.set_lineno(lineno)


1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
def _resolve_chat_template_kwargs(
    chat_template: str,
):
    env = jinja2.sandbox.ImmutableSandboxedEnvironment(
        trim_blocks=True,
        lstrip_blocks=True,
        extensions=[AssistantTracker, jinja2.ext.loopcontrols],
    )
    parsed_content = env.parse(chat_template)
    template_vars = jinja2.meta.find_undeclared_variables(parsed_content)
    return template_vars


_cached_resolve_chat_template_kwargs = lru_cache(_resolve_chat_template_kwargs)


1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
@lru_cache
def _get_hf_base_chat_template_params() -> frozenset[str]:
    # Get standard parameters from HuggingFace's base tokenizer class.
    # This dynamically extracts parameters from PreTrainedTokenizer's
    # apply_chat_template method, ensuring compatibility with tokenizers
    # that use **kwargs to receive standard parameters.

    # Read signature from HF's base class - the single source of truth
    base_sig = inspect.signature(PreTrainedTokenizer.apply_chat_template)
    # Exclude VAR_KEYWORD (**kwargs) and VAR_POSITIONAL (*args) placeholders
    return frozenset(
        p.name
        for p in base_sig.parameters.values()
        if p.kind
        not in (inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL)
    )


1754
def resolve_chat_template_kwargs(
1755
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
1756
1757
    chat_template: str,
    chat_template_kwargs: dict[str, Any],
1758
    raise_on_unexpected: bool = True,
1759
) -> dict[str, Any]:
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
    # We exclude chat_template from kwargs here, because
    # chat template has been already resolved at this stage
    unexpected_vars = {"chat_template", "tokenize"}
    if raise_on_unexpected and (
        unexpected_in_kwargs := unexpected_vars & chat_template_kwargs.keys()
    ):
        raise ValueError(
            "Found unexpected chat template kwargs from request: "
            f"{unexpected_in_kwargs}"
        )

1771
    fn_kw = {
1772
1773
        k
        for k in chat_template_kwargs
1774
1775
        if supports_kw(tokenizer.apply_chat_template, k, allow_var_kwargs=False)
    }
1776
    template_vars = _cached_resolve_chat_template_kwargs(chat_template)
1777
1778
1779
1780
1781

    # Allow standard HF parameters even if tokenizer uses **kwargs to receive them
    hf_base_params = _get_hf_base_chat_template_params()

    accept_vars = (fn_kw | template_vars | hf_base_params) - unexpected_vars
1782
    return {k: v for k, v in chat_template_kwargs.items() if k in accept_vars}
1783
1784


1785
def apply_hf_chat_template(
1786
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
1787
    conversation: list[ConversationMessage],
1788
1789
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
1790
    *,
1791
    model_config: ModelConfig,
1792
    **kwargs: Any,
1793
) -> str:
1794
    hf_chat_template = resolve_hf_chat_template(
1795
1796
1797
        tokenizer,
        chat_template=chat_template,
        tools=tools,
1798
        model_config=model_config,
1799
    )
1800

1801
    if hf_chat_template is None:
1802
1803
1804
        raise ValueError(
            "As of transformers v4.44, default chat template is no longer "
            "allowed, so you must provide a chat template if the tokenizer "
1805
1806
            "does not define one."
        )
1807

1808
1809
1810
1811
1812
1813
    resolved_kwargs = resolve_chat_template_kwargs(
        tokenizer=tokenizer,
        chat_template=hf_chat_template,
        chat_template_kwargs=kwargs,
    )

1814
1815
1816
1817
1818
    try:
        return tokenizer.apply_chat_template(
            conversation=conversation,  # type: ignore[arg-type]
            tools=tools,  # type: ignore[arg-type]
            chat_template=hf_chat_template,
1819
            tokenize=False,
1820
            **resolved_kwargs,
1821
        )
1822

1823
1824
1825
1826
1827
1828
    # External library exceptions can sometimes occur despite the framework's
    # internal exception management capabilities.
    except Exception as e:
        # Log and report any library-related exceptions for further
        # investigation.
        logger.exception(
1829
1830
            "An error occurred in `transformers` while applying chat template"
        )
1831
        raise ValueError(str(e)) from e
1832

1833

1834
1835
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1836
    messages: list[ChatCompletionMessageParam],
1837
1838
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
1839
    **kwargs: Any,
1840
) -> list[int]:
1841
1842
    from mistral_common.exceptions import MistralCommonException

1843
1844
1845
1846
1847
1848
    # 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,
    )
1849

1850
1851
1852
1853
1854
1855
1856
1857
1858
    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
1859
    # properly caught in the preprocessing_input step
1860
    except (AssertionError, MistralCommonException) as e:
1861
        raise ValueError(str(e)) from e
1862
1863
1864
1865
1866
1867
1868

    # External library exceptions can sometimes occur despite the framework's
    # internal exception management capabilities.
    except Exception as e:
        # Log and report any library-related exceptions for further
        # investigation.
        logger.exception(
1869
1870
            "An error occurred in `mistral_common` while applying chat template"
        )
1871
        raise ValueError(str(e)) from e
1872

1873

1874
1875
1876
def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
    idx = 0
    for msg in conversation:
1877
1878
1879
        if msg["role"] == "assistant":
            tool_calls = msg.get("tool_calls")
            idx += len(list(tool_calls)) if tool_calls is not None else 0  # noqa
1880
1881
1882
    return idx


1883
1884
1885
def make_tool_call_id(id_type: str = "random", func_name=None, idx=None):
    if id_type == "kimi_k2":
        return f"functions.{func_name}:{idx}"
1886
1887
1888
    else:
        # by default return random
        return f"chatcmpl-tool-{random_uuid()}"