"tests/vscode:/vscode.git/clone" did not exist on "db100c5cdebc7140b57cbb40b20b5a28d7bff386"
chat_utils.py 64.7 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
from vllm.tokenizers import TokenizerLike
53
from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
54
from vllm.transformers_utils.processor import cached_get_processor
55
from vllm.utils import random_uuid
56
from vllm.utils.collection_utils import is_list_of
57
from vllm.utils.func_utils import supports_kw
58
59
60
61
from vllm.utils.import_utils import LazyLoader

if TYPE_CHECKING:
    import torch
62
63

    from vllm.tokenizers.mistral import MistralTokenizer
64
65
else:
    torch = LazyLoader("torch", globals(), "torch")
66
67
68

logger = init_logger(__name__)

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

75

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


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


106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
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.
    """


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


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

141
142
143
144
145
146
147
148
149
150
151
152
    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
    }
    """
153

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


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

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

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


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

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

189
    audio_url: str | None
190
191


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

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


Julien Denize's avatar
Julien Denize committed
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
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."""


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


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

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

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

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

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

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

272

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


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

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

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

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

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

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

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

303

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

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

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

349
    return _is_attr_access(node, varname, key) if key else _is_var_access(node, varname)
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
377


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 = [
378
        varname for _, varname in _iter_nodes_assign_var_or_elems(root, "messages")
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
409
    ]

    # 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


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


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


440
def resolve_mistral_chat_template(
441
    chat_template: str | None,
442
    **kwargs: Any,
443
) -> str | None:
444
445
446
447
    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."
448
        )
449

450
451
    return None

452

453
_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()
454
455
456
457
458
459
460
461
462
"""
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(
463
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
464
    model_config: ModelConfig,
465
) -> str | None:
466
    cache_key = (tokenizer.name_or_path, model_config.trust_remote_code)
467
468
469
470
471
472
473
474
475
476
477
    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,
            ),
478
            trust_remote_code=model_config.trust_remote_code,
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
        )
        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


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

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

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

    return chat_template
543
544


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

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

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

574
    return detected_format
575
576
577


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

600
601

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

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

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

625
    return detected_format
626

627

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


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
669
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


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

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

680
        self._model_config = model_config
681

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

745
        mm_uuids = {}
746
747
748
749
        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
750
751
        if "audio_embeds" in uuids_by_modality:
            mm_uuids["audio"] = uuids_by_modality["audio_embeds"]
752
753
754
755
        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
756

757
758
        return mm_uuids

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


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

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

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

787
        return mm_inputs
788
789
790
791
792

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


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

798
799
800
801
802
803
804
805
        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()
        }
806
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
807
            raise ValueError("Mixing raw image and embedding inputs is not allowed")
808
809
        if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
            raise ValueError("Mixing raw audio and embedding inputs is not allowed")
810

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

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

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


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

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

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

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

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

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

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

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

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

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

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

889
890
891
892
893
894

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

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

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

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

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

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

915
    def parse_image_embeds(
916
        self,
917
918
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
919
    ) -> None:
920
921
922
923
924
925
        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`"
            )

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

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

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

940
        self._add_placeholder("image", placeholder)
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
966
    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)

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

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

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

979
    def parse_input_audio(
980
        self, input_audio: InputAudio | None, uuid: str | None = None
981
    ) -> None:
982
983
984
985
986
987
988
989
990
991
        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
992

993
        return self.parse_audio(audio_url, uuid)
994

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

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

1001
1002
1003
1004
1005
1006

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

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

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

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

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

1026
    def parse_image_embeds(
1027
        self,
1028
1029
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
1030
    ) -> None:
1031
1032
1033
1034
1035
1036
        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`"
            )

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

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

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

1053
        placeholder = self._tracker.add("image_embeds", future, uuid)
1054
        self._add_placeholder("image", placeholder)
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
1116
    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)

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

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

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

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

1135
    def parse_input_audio(
1136
        self, input_audio: InputAudio | None, uuid: str | None = None
1137
    ) -> None:
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
        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
1148

1149
        return self.parse_audio(audio_url, uuid)
1150

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

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

1161

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

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

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

1181
1182
1183
1184
1185
1186
1187
1188
            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}"
                )

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


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

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

1207
        return chat_template
1208

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

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

            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
1235

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


_cached_load_chat_template = lru_cache(_load_chat_template)


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


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

    return "\n".join(texts)


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

1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
    # 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

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

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

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


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

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

1331
# Define a mapping from part types to their corresponding parsing functions.
1332
MM_PARSER_MAP: dict[
1333
1334
1335
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
1336
1337
1338
    "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),
1339
    "output_text": lambda part: _TextParser(part).get("text", None),
1340
1341
1342
    "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),
1343
    "audio_embeds": lambda part: _AudioEmbedsParser(part).get("audio_embeds", None),
1344
    "image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
1345
1346
    "audio_url": lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
    "input_audio": lambda part: _InputAudioParser(part).get("input_audio", None),
1347
    "refusal": lambda part: _RefusalParser(part).get("refusal", None),
1348
    "video_url": lambda part: _VideoParser(part).get("video_url", {}).get("url", None),
1349
1350
1351
1352
}


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

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

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

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

        return part_type, content

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


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

1453

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

1464
    mm_parser = mm_tracker.create_parser()
1465
1466

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

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

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


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

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

1527
1528
1529
1530
1531
1532
    # 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)

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

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


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


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

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

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

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

1624

1625
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1626
1627
1628
1629
1630
1631
    # 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:
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
        if message["role"] == "assistant" and "tool_calls" in message:
            tool_calls = message.get("tool_calls")
            if not isinstance(tool_calls, list):
                continue

            if len(tool_calls) == 0:
                # Drop empty tool_calls to keep templates on the normal assistant path.
                message.pop("tool_calls", None)
                continue

            for item in tool_calls:
1643
1644
                # if arguments is None or empty string, set to {}
                if content := item["function"].get("arguments"):
1645
1646
                    if not isinstance(content, (dict, list)):
                        item["function"]["arguments"] = json.loads(content)
1647
1648
                else:
                    item["function"]["arguments"] = {}
1649
1650


1651
def parse_chat_messages(
1652
    messages: list[ChatCompletionMessageParam],
1653
    model_config: ModelConfig,
1654
    content_format: _ChatTemplateContentFormat,
1655
1656
) -> tuple[
    list[ConversationMessage],
1657
1658
    MultiModalDataDict | None,
    MultiModalUUIDDict | None,
1659
]:
1660
    conversation: list[ConversationMessage] = []
1661
    mm_tracker = MultiModalItemTracker(model_config)
1662
1663

    for msg in messages:
1664
1665
1666
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1667
            content_format,
1668
1669
1670
1671
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1672
            ),
1673
        )
1674

1675
        conversation.extend(sub_messages)
1676

1677
1678
    _postprocess_messages(conversation)

1679
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1680
1681


1682
def parse_chat_messages_futures(
1683
    messages: list[ChatCompletionMessageParam],
1684
    model_config: ModelConfig,
1685
    content_format: _ChatTemplateContentFormat,
1686
1687
) -> tuple[
    list[ConversationMessage],
1688
1689
    Awaitable[MultiModalDataDict | None],
    MultiModalUUIDDict | None,
1690
]:
1691
    conversation: list[ConversationMessage] = []
1692
    mm_tracker = AsyncMultiModalItemTracker(model_config)
1693
1694

    for msg in messages:
1695
1696
1697
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1698
            content_format,
1699
1700
1701
1702
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1703
            ),
1704
        )
1705
1706
1707

        conversation.extend(sub_messages)

1708
1709
    _postprocess_messages(conversation)

1710
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1711
1712


1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
# 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)


1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
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)


1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
@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)
    )


1760
def resolve_chat_template_kwargs(
1761
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
1762
1763
    chat_template: str,
    chat_template_kwargs: dict[str, Any],
1764
    raise_on_unexpected: bool = True,
1765
) -> dict[str, Any]:
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
    # 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}"
        )

1777
    fn_kw = {
1778
1779
        k
        for k in chat_template_kwargs
1780
1781
        if supports_kw(tokenizer.apply_chat_template, k, allow_var_kwargs=False)
    }
1782
    template_vars = _cached_resolve_chat_template_kwargs(chat_template)
1783
1784
1785
1786
1787

    # 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
1788
    return {k: v for k, v in chat_template_kwargs.items() if k in accept_vars}
1789
1790


1791
def apply_hf_chat_template(
1792
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
1793
    conversation: list[ConversationMessage],
1794
1795
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
1796
    *,
1797
    model_config: ModelConfig,
1798
    **kwargs: Any,
1799
) -> str:
1800
    hf_chat_template = resolve_hf_chat_template(
1801
1802
1803
        tokenizer,
        chat_template=chat_template,
        tools=tools,
1804
        model_config=model_config,
1805
    )
1806

1807
    if hf_chat_template is None:
1808
1809
1810
        raise ValueError(
            "As of transformers v4.44, default chat template is no longer "
            "allowed, so you must provide a chat template if the tokenizer "
1811
1812
            "does not define one."
        )
1813

1814
1815
1816
1817
1818
1819
    resolved_kwargs = resolve_chat_template_kwargs(
        tokenizer=tokenizer,
        chat_template=hf_chat_template,
        chat_template_kwargs=kwargs,
    )

1820
1821
1822
1823
1824
    try:
        return tokenizer.apply_chat_template(
            conversation=conversation,  # type: ignore[arg-type]
            tools=tools,  # type: ignore[arg-type]
            chat_template=hf_chat_template,
1825
            tokenize=False,
1826
            **resolved_kwargs,
1827
        )
1828

1829
1830
1831
1832
1833
1834
    # 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(
1835
1836
            "An error occurred in `transformers` while applying chat template"
        )
1837
        raise ValueError(str(e)) from e
1838

1839

1840
def apply_mistral_chat_template(
1841
    tokenizer: "MistralTokenizer",
1842
    messages: list[ChatCompletionMessageParam],
1843
1844
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
1845
    **kwargs: Any,
1846
) -> list[int]:
1847
1848
    from mistral_common.exceptions import MistralCommonException

1849
1850
1851
1852
1853
1854
    # 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,
    )
1855

1856
1857
1858
1859
1860
1861
1862
1863
1864
    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
1865
    # properly caught in the preprocessing_input step
1866
    except (AssertionError, MistralCommonException) as e:
1867
        raise ValueError(str(e)) from e
1868
1869
1870
1871
1872
1873
1874

    # 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(
1875
1876
            "An error occurred in `mistral_common` while applying chat template"
        )
1877
        raise ValueError(str(e)) from e
1878

1879

1880
1881
1882
def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
    idx = 0
    for msg in conversation:
1883
1884
1885
        if msg["role"] == "assistant":
            tool_calls = msg.get("tool_calls")
            idx += len(list(tool_calls)) if tool_calls is not None else 0  # noqa
1886
1887
1888
    return idx


1889
1890
1891
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}"
1892
1893
1894
    else:
        # by default return random
        return f"chatcmpl-tool-{random_uuid()}"