"vllm/vscode:/vscode.git/clone" did not exist on "e2b31243c092e9f4ade5ffe4bf9a5d5ddae06ca7"
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
from vllm.tokenizers import MistralTokenizer, 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
62
63
from vllm.utils.import_utils import LazyLoader

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

logger = init_logger(__name__)

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

73

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


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


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


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


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

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

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


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

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

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


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

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

187
    audio_url: str | None
188
189


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

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


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


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


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

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

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

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

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

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

270

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


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

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

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

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

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

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

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

301

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

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

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

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


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

    # 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


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


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


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

448
449
    return None

450

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


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

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

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

    return chat_template
541
542


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

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

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

572
    return detected_format
573
574
575


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

598
599

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

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

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

623
    return detected_format
624

625

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


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


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

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

678
        self._model_config = model_config
679

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

755
756
        return mm_uuids

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


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

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

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

785
        return mm_inputs
786
787
788
789
790

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


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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

887
888
889
890
891
892

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

991
        return self.parse_audio(audio_url, uuid)
992

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

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

999
1000
1001
1002
1003
1004

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1147
        return self.parse_audio(audio_url, uuid)
1148

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

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

1159

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

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

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

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

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


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

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

1205
        return chat_template
1206

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

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

            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
1233

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


_cached_load_chat_template = lru_cache(_load_chat_template)


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


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

    return "\n".join(texts)


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

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

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

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

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


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

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

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


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

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

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

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

        return part_type, content

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


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

1451

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

1462
    mm_parser = mm_tracker.create_parser()
1463
1464

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

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

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


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

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

1525
1526
1527
1528
1529
1530
    # 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)

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

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


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


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

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

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

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

1622

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


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

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

1668
        conversation.extend(sub_messages)
1669

1670
1671
    _postprocess_messages(conversation)

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


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

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

        conversation.extend(sub_messages)

1701
1702
    _postprocess_messages(conversation)

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


1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
# 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)


1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
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)


1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
@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)
    )


1753
def resolve_chat_template_kwargs(
1754
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
1755
1756
    chat_template: str,
    chat_template_kwargs: dict[str, Any],
1757
    raise_on_unexpected: bool = True,
1758
) -> dict[str, Any]:
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
    # 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}"
        )

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

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


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

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

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

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

1822
1823
1824
1825
1826
1827
    # 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(
1828
1829
            "An error occurred in `transformers` while applying chat template"
        )
1830
        raise ValueError(str(e)) from e
1831

1832

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

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

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

    # 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(
1868
1869
            "An error occurred in `mistral_common` while applying chat template"
        )
1870
        raise ValueError(str(e)) from e
1871

1872

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


1882
1883
1884
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}"
1885
1886
1887
    else:
        # by default return random
        return f"chatcmpl-tool-{random_uuid()}"