"tests/kernels/attention/test_cache.py" did not exist on "c413c41cda0f9359e7a12bb674c0f87bf41798c5"
chat_utils.py 63.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 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.func_utils import supports_kw
57
58
59

logger = init_logger(__name__)

60
61
62
63
64
65
MODALITY_PLACEHOLDERS_MAP = {
    "image": "<##IMAGE##>",
    "audio": "<##AUDIO##>",
    "video": "<##VIDEO##>",
}

66

67
68
69
70
71
72
73
74
75
76
77
78
79
80
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."""


81
class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):
82
    image_embeds: str | dict[str, str] | None
83
84
85
86
87
88
89
    """
    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."""
90
    uuid: str | None
91
92
93
94
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
95
96


97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
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.
    """


113
114
115
116
117
118
119
120
121
122
123
124
125
126
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."""


127
128
129
130
class PILImage(BaseModel):
    """
    A PIL.Image.Image object.
    """
131

132
133
134
135
136
137
138
139
140
141
142
143
    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
    }
    """
144

145
146
    image_pil: PILImage | None
    uuid: str | None
147
148
149
150
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
151
152


153
154
155
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.
156

157
158
159
160
161
    Example:
    {
        "image_url": "https://example.com/image.jpg"
    }
    """
162

163
164
    image_url: str | None
    uuid: str | None
165
166
167
168
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
169
170
171
172


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

174
175
176
177
178
    Example:
    {
        "audio_url": "https://example.com/audio.mp3"
    }
    """
179

180
    audio_url: str | None
181
182


183
184
185
186
187
188
189
190
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"
    }
    """
191

192
193
    video_url: str | None
    uuid: str | None
194
195
196
197
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
198
199


Julien Denize's avatar
Julien Denize committed
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
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."""


221
222
223
224
225
226
227
228
229
ChatCompletionContentPartParam: TypeAlias = (
    OpenAIChatCompletionContentPartParam
    | ChatCompletionContentPartAudioParam
    | ChatCompletionContentPartInputAudioParam
    | ChatCompletionContentPartVideoParam
    | ChatCompletionContentPartRefusalParam
    | CustomChatCompletionContentPILImageParam
    | CustomChatCompletionContentSimpleImageParam
    | ChatCompletionContentPartImageEmbedsParam
230
    | ChatCompletionContentPartAudioEmbedsParam
231
232
233
234
235
    | CustomChatCompletionContentSimpleAudioParam
    | CustomChatCompletionContentSimpleVideoParam
    | str
    | CustomThinkCompletionContentParam
)
236
237
238
239


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

241
242
243
    role: Required[str]
    """The role of the message's author."""

244
    content: str | list[ChatCompletionContentPartParam]
245
246
247
248
249
250
251
252
253
    """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.
    """

254
    tool_call_id: str | None
255
256
    """Tool call that this message is responding to."""

257
    tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None
258
259
    """The tool calls generated by the model, such as function calls."""

260
261
262
    reasoning: str | None
    """The reasoning content for interleaved thinking."""

263

264
265
266
267
268
ChatCompletionMessageParam: TypeAlias = (
    OpenAIChatCompletionMessageParam
    | CustomChatCompletionMessageParam
    | OpenAIHarmonyMessage
)
269
270


271
# TODO: Make fields ReadOnly once mypy supports it
272
273
274
275
class ConversationMessage(TypedDict, total=False):
    role: Required[str]
    """The role of the message's author."""

276
    content: str | None | list[dict[str, str]]
277
278
    """The contents of the message"""

279
    tool_call_id: str | None
280
281
    """Tool call that this message is responding to."""

282
    name: str | None
283
284
    """The name of the function to call"""

285
    tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None
286
    """The tool calls generated by the model, such as function calls."""
287

288
289
290
291
292
293
    reasoning: str | None
    """The reasoning content for interleaved thinking."""

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

294

295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
# 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):
311
312
313
314
315
        return (
            _is_var_access(node.node, varname)
            and isinstance(node.arg, jinja2.nodes.Const)
            and node.arg.value == key
        )
316
317
318
319
320
321
322
323
324
325

    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,
326
    key: str | None = None,
327
328
) -> bool:
    if isinstance(node, jinja2.nodes.Filter):
329
        return node.node is not None and _is_var_or_elems_access(
330
331
            node.node, varname, key
        )
332
333
334
    if isinstance(node, jinja2.nodes.Test):
        return _is_var_or_elems_access(node.node, varname, key)

335
    if isinstance(node, jinja2.nodes.Getitem) and isinstance(
336
337
        node.arg, jinja2.nodes.Slice
    ):
338
339
        return _is_var_or_elems_access(node.node, varname, key)

340
    return _is_attr_access(node, varname, key) if key else _is_var_access(node, varname)
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368


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 = [
369
        varname for _, varname in _iter_nodes_assign_var_or_elems(root, "messages")
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
    ]

    # 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


401
def _try_extract_ast(chat_template: str) -> jinja2.nodes.Template | None:
402
403
404
405
406
407
408
409
    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


410
@lru_cache(maxsize=32)
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
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"


431
def resolve_mistral_chat_template(
432
    chat_template: str | None,
433
    **kwargs: Any,
434
) -> str | None:
435
436
437
438
    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."
439
        )
440

441
442
    return None

443

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


489
def resolve_hf_chat_template(
490
491
492
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
493
494
    *,
    model_config: ModelConfig,
495
) -> str | None:
496
497
498
499
500
501
    # 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:
502
        chat_template = _try_get_processor_chat_template(tokenizer, model_config)
503
504
        if chat_template is not None:
            return chat_template
505
506
507
508
509

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

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

    return chat_template
534
535


536
def _resolve_chat_template_content_format(
537
538
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
539
    tokenizer: TokenizerLike | None,
540
    *,
541
    model_config: ModelConfig,
542
543
) -> _ChatTemplateContentFormat:
    if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
544
        hf_chat_template = resolve_hf_chat_template(
545
546
547
            tokenizer,
            chat_template=chat_template,
            tools=tools,
548
            model_config=model_config,
549
        )
550
    else:
551
552
        hf_chat_template = None

553
554
555
556
557
    jinja_text = (
        hf_chat_template
        if isinstance(hf_chat_template, str)
        else load_chat_template(chat_template, is_literal=True)
    )
558

559
560
561
562
563
    detected_format = (
        "string"
        if jinja_text is None
        else _detect_content_format(jinja_text, default="string")
    )
564

565
    return detected_format
566
567
568


@lru_cache
569
def _log_chat_template_content_format(
570
    chat_template: str | None,
571
    given_format: ChatTemplateContentFormatOption,
572
573
    detected_format: ChatTemplateContentFormatOption,
):
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
    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,
        )

591
592

def resolve_chat_template_content_format(
593
594
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
595
    given_format: ChatTemplateContentFormatOption,
596
    tokenizer: TokenizerLike | None,
597
    *,
598
    model_config: ModelConfig,
599
) -> _ChatTemplateContentFormat:
600
601
602
    if given_format != "auto":
        return given_format

603
604
605
606
    detected_format = _resolve_chat_template_content_format(
        chat_template,
        tools,
        tokenizer,
607
        model_config=model_config,
608
609
610
611
612
613
614
615
    )

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

616
    return detected_format
617

618

619
ModalityStr = Literal["image", "audio", "video", "image_embeds", "audio_embeds"]
620
621
622
623
_T = TypeVar("_T")


class BaseMultiModalItemTracker(ABC, Generic[_T]):
624
625
626
627
628
629
    """
    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.
    """

630
    def __init__(self, model_config: ModelConfig):
631
632
        super().__init__()

633
        self._model_config = model_config
634

635
636
        self._items_by_modality = defaultdict[str, list[_T | None]](list)
        self._uuids_by_modality = defaultdict[str, list[str | None]](list)
637

638
    @property
639
640
    def model_config(self) -> ModelConfig:
        return self._model_config
641

642
    @cached_property
643
    def model_cls(self) -> type[SupportsMultiModal]:
644
        from vllm.model_executor.model_loader import get_model_cls
645

646
        model_cls = get_model_cls(self.model_config)
647
        return cast(type[SupportsMultiModal], model_cls)
648

649
650
    @property
    def allowed_local_media_path(self):
651
        return self._model_config.allowed_local_media_path
652

653
654
    @property
    def allowed_media_domains(self):
655
        return self._model_config.allowed_media_domains
656

657
658
659
660
    @property
    def mm_registry(self):
        return MULTIMODAL_REGISTRY

661
662
    @cached_property
    def mm_processor(self):
663
        return self.mm_registry.create_processor(self.model_config)
664

665
    def add(
666
667
        self,
        modality: ModalityStr,
668
669
670
        item: _T | None,
        uuid: str | None = None,
    ) -> str | None:
671
672
673
        """
        Add a multi-modal item to the current prompt and returns the
        placeholder string to use, if any.
674
675

        An optional uuid can be added which serves as a unique identifier of the
676
        media.
677
        """
678
        input_modality = modality.replace("_embeds", "")
679
        num_items = len(self._items_by_modality[modality]) + 1
680

681
        self.mm_processor.validate_num_items(input_modality, num_items)
682

683
        self._items_by_modality[modality].append(item)
684
        self._uuids_by_modality[modality].append(uuid)
685

686
        return self.model_cls.get_placeholder_str(modality, num_items)
687

688
    def all_mm_uuids(self) -> MultiModalUUIDDict | None:
689
690
691
692
693
        if not self._items_by_modality:
            return None
        mm_uuids = {}
        uuids_by_modality = dict(self._uuids_by_modality)
        if "image" in uuids_by_modality and "image_embeds" in uuids_by_modality:
694
            raise ValueError("Mixing raw image and embedding inputs is not allowed")
695
696
697
698
699

        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
700
701
        if "audio_embeds" in uuids_by_modality:
            mm_uuids["audio"] = uuids_by_modality["audio_embeds"]
702
703
704
705
706
707
        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
        return mm_uuids

708
709
710
711
712
    @abstractmethod
    def create_parser(self) -> "BaseMultiModalContentParser":
        raise NotImplementedError


713
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
714
    def all_mm_data(self) -> MultiModalDataDict | None:
715
716
717
718
719
        if not self._items_by_modality:
            return None
        mm_inputs = {}
        items_by_modality = dict(self._items_by_modality)
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
720
            raise ValueError("Mixing raw image and embedding inputs is not allowed")
721
722
        if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
            raise ValueError("Mixing raw audio and embedding inputs is not allowed")
723
724
725

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
726
727
728
            mm_inputs["image"] = (
                image_embeds_lst if len(image_embeds_lst) != 1 else image_embeds_lst[0]
            )
729
        if "image" in items_by_modality:
730
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
731
732
        if "audio_embeds" in items_by_modality:
            audio_embeds_lst = items_by_modality["audio_embeds"]
733
734
735
            mm_inputs["audio"] = (
                audio_embeds_lst if len(audio_embeds_lst) != 1 else audio_embeds_lst[0]
            )
736
        if "audio" in items_by_modality:
737
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
738
        if "video" in items_by_modality:
739
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
740
        return mm_inputs
741
742
743
744
745

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


746
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
747
    async def all_mm_data(self) -> MultiModalDataDict | None:
748
749
750
        if not self._items_by_modality:
            return None
        mm_inputs = {}
751
752
753
754
755
756
757
758
759
        items_by_modality = {}
        for modality, items in self._items_by_modality.items():
            coros = []
            for item in items:
                if item is not None:
                    coros.append(item)
                else:
                    coros.append(asyncio.sleep(0))
            items_by_modality[modality] = await asyncio.gather(*coros)
760

761
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
762
            raise ValueError("Mixing raw image and embedding inputs is not allowed")
763
764
        if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
            raise ValueError("Mixing raw audio and embedding inputs is not allowed")
765
766
767

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
768
769
770
            mm_inputs["image"] = (
                image_embeds_lst if len(image_embeds_lst) != 1 else image_embeds_lst[0]
            )
771
        if "image" in items_by_modality:
772
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
773
774
        if "audio_embeds" in items_by_modality:
            audio_embeds_lst = items_by_modality["audio_embeds"]
775
776
777
            mm_inputs["audio"] = (
                audio_embeds_lst if len(audio_embeds_lst) != 1 else audio_embeds_lst[0]
            )
778
        if "audio" in items_by_modality:
779
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
780
        if "video" in items_by_modality:
781
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
782
        return mm_inputs
783
784
785
786
787
788
789
790
791

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


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

792
        # stores model placeholders list with corresponding
793
794
795
796
797
798
799
        # general MM placeholder:
        # {
        #   "<##IMAGE##>": ["<image>", "<image>", "<image>"],
        #   "<##AUDIO##>": ["<audio>", "<audio>"]
        # }
        self._placeholder_storage: dict[str, list] = defaultdict(list)

800
    def _add_placeholder(self, modality: ModalityStr, placeholder: str | None):
801
        mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]
802
        if placeholder:
803
            self._placeholder_storage[mod_placeholder].append(placeholder)
804

805
806
    def mm_placeholder_storage(self) -> dict[str, list]:
        return dict(self._placeholder_storage)
807
808

    @abstractmethod
809
    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
810
811
        raise NotImplementedError

812
    @abstractmethod
813
    def parse_image_embeds(
814
        self,
815
816
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
817
    ) -> None:
818
819
        raise NotImplementedError

820
    @abstractmethod
821
    def parse_image_pil(
822
        self, image_pil: Image.Image | None, uuid: str | None = None
823
    ) -> None:
824
825
        raise NotImplementedError

826
    @abstractmethod
827
    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
828
829
        raise NotImplementedError

830
    @abstractmethod
831
    def parse_input_audio(
832
        self, input_audio: InputAudio | None, uuid: str | None = None
833
    ) -> None:
834
835
        raise NotImplementedError

836
837
838
839
840
841
842
843
    @abstractmethod
    def parse_audio_embeds(
        self,
        audio_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
    ) -> None:
        raise NotImplementedError

844
    @abstractmethod
845
    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
846
847
        raise NotImplementedError

848
849
850
851
852
853

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

        self._tracker = tracker
854
855
856
        multimodal_config = self._tracker.model_config.multimodal_config
        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)

857
858
        self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
859
            media_io_kwargs=media_io_kwargs,
860
            allowed_local_media_path=tracker.allowed_local_media_path,
861
            allowed_media_domains=tracker.allowed_media_domains,
862
863
        )

864
865
    @property
    def model_config(self) -> ModelConfig:
866
        return self._tracker.model_config
867

868
    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
869
        image = self._connector.fetch_image(image_url) if image_url else None
870

871
        placeholder = self._tracker.add("image", image, uuid)
872
        self._add_placeholder("image", placeholder)
873

874
    def parse_image_embeds(
875
        self,
876
877
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
878
    ) -> None:
879
880
881
882
883
884
        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`"
            )

885
886
887
888
889
        if isinstance(image_embeds, dict):
            embeds = {
                k: self._connector.fetch_image_embedding(v)
                for k, v in image_embeds.items()
            }
890
            placeholder = self._tracker.add("image_embeds", embeds, uuid)
891
892
893

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

896
897
898
        if image_embeds is None:
            placeholder = self._tracker.add("image_embeds", None, uuid)

899
        self._add_placeholder("image", placeholder)
900

901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
    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)

926
    def parse_image_pil(
927
        self, image_pil: Image.Image | None, uuid: str | None = None
928
929
    ) -> None:
        placeholder = self._tracker.add("image", image_pil, uuid)
930
        self._add_placeholder("image", placeholder)
931

932
    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
933
        audio = self._connector.fetch_audio(audio_url) if audio_url else None
934

935
        placeholder = self._tracker.add("audio", audio, uuid)
936
        self._add_placeholder("audio", placeholder)
937

938
    def parse_input_audio(
939
        self, input_audio: InputAudio | None, uuid: str | None = None
940
    ) -> None:
941
942
943
944
945
946
947
948
949
950
        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
951

952
        return self.parse_audio(audio_url, uuid)
953

954
    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
955
        video = self._connector.fetch_video(video_url=video_url) if video_url else None
956

957
        placeholder = self._tracker.add("video", video, uuid)
958
        self._add_placeholder("video", placeholder)
959

960
961
962
963
964
965

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

        self._tracker = tracker
966
967
        multimodal_config = self._tracker.model_config.multimodal_config
        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)
968
969
        self._connector: MediaConnector = MEDIA_CONNECTOR_REGISTRY.load(
            envs.VLLM_MEDIA_CONNECTOR,
970
            media_io_kwargs=media_io_kwargs,
971
            allowed_local_media_path=tracker.allowed_local_media_path,
972
            allowed_media_domains=tracker.allowed_media_domains,
973
        )
974

975
976
    @property
    def model_config(self) -> ModelConfig:
977
        return self._tracker.model_config
978

979
    def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
980
        image_coro = self._connector.fetch_image_async(image_url) if image_url else None
981

982
        placeholder = self._tracker.add("image", image_coro, uuid)
983
        self._add_placeholder("image", placeholder)
984

985
    def parse_image_embeds(
986
        self,
987
988
        image_embeds: str | dict[str, str] | None,
        uuid: str | None = None,
989
    ) -> None:
990
991
992
993
994
995
        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`"
            )

996
        future: asyncio.Future[str | dict[str, str] | None] = asyncio.Future()
997
998
999
1000
1001
1002
1003
1004
1005

        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):
1006
            embedding = self._connector.fetch_image_embedding(image_embeds)
1007
1008
            future.set_result(embedding)

1009
1010
1011
        if image_embeds is None:
            future.set_result(None)

1012
        placeholder = self._tracker.add("image_embeds", future, uuid)
1013
        self._add_placeholder("image", placeholder)
1014

1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
    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)

1076
    def parse_image_pil(
1077
        self, image_pil: Image.Image | None, uuid: str | None = None
1078
    ) -> None:
1079
        future: asyncio.Future[Image.Image | None] = asyncio.Future()
1080
1081
1082
1083
        if image_pil:
            future.set_result(image_pil)
        else:
            future.set_result(None)
1084

1085
        placeholder = self._tracker.add("image", future, uuid)
1086
        self._add_placeholder("image", placeholder)
1087

1088
    def parse_audio(self, audio_url: str | None, uuid: str | None = None) -> None:
1089
        audio_coro = self._connector.fetch_audio_async(audio_url) if audio_url else None
1090

1091
        placeholder = self._tracker.add("audio", audio_coro, uuid)
1092
        self._add_placeholder("audio", placeholder)
1093

1094
    def parse_input_audio(
1095
        self, input_audio: InputAudio | None, uuid: str | None = None
1096
    ) -> None:
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
        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
1107

1108
        return self.parse_audio(audio_url, uuid)
1109

1110
    def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
1111
1112
1113
1114
1115
        video = (
            self._connector.fetch_video_async(video_url=video_url)
            if video_url
            else None
        )
1116

1117
        placeholder = self._tracker.add("video", video, uuid)
1118
        self._add_placeholder("video", placeholder)
1119

1120

1121
def validate_chat_template(chat_template: Path | str | None):
1122
1123
1124
1125
1126
    """Raises if the provided chat template appears invalid."""
    if chat_template is None:
        return

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

    elif isinstance(chat_template, str):
        JINJA_CHARS = "{}\n"
1131
1132
1133
1134
        if (
            not any(c in chat_template for c in JINJA_CHARS)
            and not Path(chat_template).exists()
        ):
1135
1136
1137
            # Try to find the template in the built-in templates directory
            from vllm.transformers_utils.chat_templates.registry import (
                CHAT_TEMPLATES_DIR,
1138
            )
1139

1140
1141
1142
1143
1144
1145
1146
1147
            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}"
                )

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


1152
def _load_chat_template(
1153
    chat_template: Path | str | None,
1154
1155
    *,
    is_literal: bool = False,
1156
) -> str | None:
1157
1158
    if chat_template is None:
        return None
1159
1160
1161

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

1166
        return chat_template
1167

1168
    try:
1169
        with open(chat_template) as f:
1170
            return f.read()
1171
    except OSError as e:
1172
1173
1174
        if isinstance(chat_template, Path):
            raise

1175
1176
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template for c in JINJA_CHARS):
1177
1178
1179
            # Try to load from the built-in templates directory
            from vllm.transformers_utils.chat_templates.registry import (
                CHAT_TEMPLATES_DIR,
1180
            )
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193

            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
1194

1195
1196
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
1197
1198
1199
1200
1201
1202
1203
        return _load_chat_template(chat_template, is_literal=True)


_cached_load_chat_template = lru_cache(_load_chat_template)


def load_chat_template(
1204
    chat_template: Path | str | None,
1205
1206
    *,
    is_literal: bool = False,
1207
) -> str | None:
1208
    return _cached_load_chat_template(chat_template, is_literal=is_literal)
1209
1210


1211
1212
1213
def _get_interleaved_text_prompt(
    placeholder_storage: dict[str, list], texts: list[str]
) -> str:
1214
1215
1216
1217
1218
1219
1220
    for idx, elem in enumerate(texts):
        if elem in placeholder_storage:
            texts[idx] = placeholder_storage[elem].pop(0)

    return "\n".join(texts)


1221
# TODO: Let user specify how to insert multimodal tokens into prompt
1222
# (similar to chat template)
1223
1224
1225
1226
1227
def _get_full_multimodal_text_prompt(
    placeholder_storage: dict[str, list],
    texts: list[str],
    interleave_strings: bool,
) -> str:
1228
    """Combine multimodal prompts for a multimodal language model."""
1229

1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
    # 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

1247
    # Look through the text prompt to check for missing placeholders
1248
    missing_placeholders: list[str] = []
1249
1250
1251
1252
1253
    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:
1254
1255
1256
1257
            logger.error(
                "Placeholder count is negative! "
                "Ensure that the 'interleave_strings' flag is disabled "
                "(current value: %s) "
1258
1259
                "when manually placing image placeholders.",
                interleave_strings,
1260
1261
            )
            logger.debug("Input prompt: %s", text_prompt)
1262
1263
            raise ValueError(
                f"Found more '{placeholder}' placeholders in input prompt than "
1264
1265
                "actual multimodal data items."
            )
1266

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

1269
1270
    # NOTE: Default behaviour: we always add missing placeholders
    # at the front of the prompt, if interleave_strings=False
1271
    return "\n".join(missing_placeholders + [text_prompt])
1272
1273


1274
1275
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
1276
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
1277
_AudioEmbedsParser = partial(cast, ChatCompletionContentPartAudioEmbedsParam)
1278
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
1279
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
1280
_PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
Julien Denize's avatar
Julien Denize committed
1281
_ThinkParser = partial(cast, CustomThinkCompletionContentParam)
1282
1283
1284
1285
# Need to validate url objects
_ImageParser = TypeAdapter(ChatCompletionContentPartImageParam).validate_python
_AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python
_VideoParser = TypeAdapter(ChatCompletionContentPartVideoParam).validate_python
1286

1287
_ResponsesInputImageParser = TypeAdapter(ResponseInputImageParam).validate_python
1288
_ContentPart: TypeAlias = str | dict[str, str] | InputAudio | PILImage
1289

1290
# Define a mapping from part types to their corresponding parsing functions.
1291
MM_PARSER_MAP: dict[
1292
1293
1294
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
1295
1296
1297
    "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),
1298
    "output_text": lambda part: _TextParser(part).get("text", None),
1299
1300
1301
    "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),
1302
    "audio_embeds": lambda part: _AudioEmbedsParser(part).get("audio_embeds", None),
1303
    "image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
1304
1305
    "audio_url": lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
    "input_audio": lambda part: _InputAudioParser(part).get("input_audio", None),
1306
    "refusal": lambda part: _RefusalParser(part).get("refusal", None),
1307
    "video_url": lambda part: _VideoParser(part).get("video_url", {}).get("url", None),
1308
1309
1310
1311
}


def _parse_chat_message_content_mm_part(
1312
1313
    part: ChatCompletionContentPartParam,
) -> tuple[str, _ContentPart]:
1314
    """
1315
    Parses a given multi-modal content part based on its type.
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328

    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(
1329
1330
        part, dict
    )  # This is needed to avoid mypy errors: part.get() from str
1331
    part_type = part.get("type", None)
1332
    uuid = part.get("uuid", None)
1333

1334
    if isinstance(part_type, str) and part_type in MM_PARSER_MAP and uuid is None:  # noqa: E501
1335
1336
1337
        content = MM_PARSER_MAP[part_type](part)

        # Special case for 'image_url.detail'
1338
1339
        # We only support 'auto', which is the default
        if part_type == "image_url" and part.get("detail", "auto") != "auto":
1340
            logger.warning(
1341
                "'image_url.detail' is currently not supported and will be ignored."
1342
            )
1343
1344
1345
1346

        return part_type, content

    # Handle missing 'type' but provided direct URL fields.
1347
    # 'type' is required field by pydantic
1348
1349
    if part_type is None or uuid is not None:
        if "image_url" in part:
1350
            image_params = cast(CustomChatCompletionContentSimpleImageParam, part)
1351
1352
1353
1354
1355
1356
1357
1358
            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.
1359
            image_params = cast(  # type: ignore
1360
1361
1362
1363
1364
1365
                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.
1366
            image_params = cast(  # type: ignore
1367
1368
1369
1370
                ChatCompletionContentPartImageEmbedsParam, part
            )
            image_embeds = image_params.get("image_embeds", None)
            return "image_embeds", image_embeds
1371
1372
1373
1374
1375
1376
1377
        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
1378
        if "audio_url" in part:
1379
1380
1381
            audio_params = cast(  # type: ignore[assignment]
                CustomChatCompletionContentSimpleAudioParam, part
            )
1382
1383
1384
1385
1386
1387
            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
1388
        if part.get("input_audio") is not None:
1389
            input_audio_params = cast(dict[str, str], part)
1390
            return "input_audio", input_audio_params
1391
        if "video_url" in part:
1392
            video_params = cast(CustomChatCompletionContentSimpleVideoParam, part)
1393
1394
1395
1396
1397
1398
            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
1399
1400
1401
1402
1403
1404
1405
1406
        # 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"


1407
PART_TYPES_TO_SKIP_NONE_CONTENT = (
1408
1409
1410
    "text",
    "refusal",
)
1411

1412

1413
1414
1415
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
1416
    mm_tracker: BaseMultiModalItemTracker,
1417
1418
    *,
    wrap_dicts: bool,
1419
    interleave_strings: bool,
1420
) -> list[ConversationMessage]:
1421
    content = list[_ContentPart]()
1422

1423
    mm_parser = mm_tracker.create_parser()
1424
1425

    for part in parts:
1426
        parse_res = _parse_chat_message_content_part(
1427
1428
1429
            part,
            mm_parser,
            wrap_dicts=wrap_dicts,
1430
            interleave_strings=interleave_strings,
1431
        )
1432
1433
        if parse_res:
            content.append(parse_res)
1434

1435
    if wrap_dicts:
1436
        # Parsing wraps images and texts as interleaved dictionaries
1437
        return [ConversationMessage(role=role, content=content)]  # type: ignore
1438
    texts = cast(list[str], content)
1439
1440
    mm_placeholder_storage = mm_parser.mm_placeholder_storage()
    if mm_placeholder_storage:
1441
1442
1443
        text_prompt = _get_full_multimodal_text_prompt(
            mm_placeholder_storage, texts, interleave_strings
        )
1444
1445
1446
    else:
        text_prompt = "\n".join(texts)

1447
1448
1449
1450
    return [ConversationMessage(role=role, content=text_prompt)]


def _parse_chat_message_content_part(
1451
1452
1453
1454
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
1455
    interleave_strings: bool,
1456
) -> _ContentPart | None:
1457
1458
1459
1460
1461
1462
1463
1464
    """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
1465
        return part
1466
1467
    # Handle structured dictionary parts
    part_type, content = _parse_chat_message_content_mm_part(part)
1468
    # if part_type is text/refusal/image_url/audio_url/video_url/input_audio but
1469
    # content is None, log a warning and skip
1470
    if part_type in PART_TYPES_TO_SKIP_NONE_CONTENT and content is None:
1471
        logger.warning(
1472
            "Skipping multimodal part '%s' (type: '%s') "
1473
1474
1475
1476
            "with empty / unparsable content.",
            part,
            part_type,
        )
1477
1478
        return None

1479
    if part_type in ("text", "input_text", "output_text", "refusal", "thinking"):
1480
1481
        str_content = cast(str, content)
        if wrap_dicts:
1482
            return {"type": "text", "text": str_content}
1483
1484
        else:
            return str_content
1485

1486
1487
1488
1489
1490
1491
    # 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)

1492
    modality = None
1493
    if part_type == "image_pil":
1494
        image_content = cast(Image.Image, content) if content is not None else None
1495
        mm_parser.parse_image_pil(image_content, uuid)
1496
        modality = "image"
1497
    elif part_type in ("image_url", "input_image"):
1498
        str_content = cast(str, content)
1499
        mm_parser.parse_image(str_content, uuid)
1500
1501
        modality = "image"
    elif part_type == "image_embeds":
1502
        content = cast(str | dict[str, str], content) if content is not None else None
1503
        mm_parser.parse_image_embeds(content, uuid)
1504
        modality = "image"
1505
1506
1507
1508
    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"
1509
    elif part_type == "audio_url":
1510
        str_content = cast(str, content)
1511
        mm_parser.parse_audio(str_content, uuid)
1512
1513
        modality = "audio"
    elif part_type == "input_audio":
1514
        dict_content = cast(InputAudio, content)
1515
        mm_parser.parse_input_audio(dict_content, uuid)
1516
1517
        modality = "audio"
    elif part_type == "video_url":
1518
        str_content = cast(str, content)
1519
        mm_parser.parse_video(str_content, uuid)
1520
1521
1522
        modality = "video"
    else:
        raise NotImplementedError(f"Unknown part type: {part_type}")
1523

1524
1525
1526
    return (
        {"type": modality}
        if wrap_dicts
1527
        else (MODALITY_PLACEHOLDERS_MAP[modality] if interleave_strings else None)
1528
    )
1529
1530


1531
1532
1533
1534
1535
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


1536
def _parse_chat_message_content(
1537
1538
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
1539
    content_format: _ChatTemplateContentFormat,
1540
    interleave_strings: bool,
1541
) -> list[ConversationMessage]:
1542
1543
    role = message["role"]
    content = message.get("content")
1544
    reasoning = message.get("reasoning") or message.get("reasoning_content")
1545

1546
    if content is None:
1547
1548
        content = []
    elif isinstance(content, str):
1549
        content = [ChatCompletionContentPartTextParam(type="text", text=content)]
1550
    result = _parse_chat_message_content_parts(
1551
1552
        role,
        content,  # type: ignore
1553
        mm_tracker,
1554
        wrap_dicts=(content_format == "openai"),
1555
        interleave_strings=interleave_strings,
1556
    )
1557

1558
    for result_msg in result:
1559
        if role == "assistant":
1560
1561
            parsed_msg = _AssistantParser(message)

1562
1563
1564
            # The 'tool_calls' is not None check ensures compatibility.
            # It's needed only if downstream code doesn't strictly
            # follow the OpenAI spec.
1565
            if "tool_calls" in parsed_msg and parsed_msg["tool_calls"] is not None:
1566
                result_msg["tool_calls"] = list(parsed_msg["tool_calls"])
1567
1568
1569
1570
1571
1572
            # 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
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
        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

1583

1584
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1585
1586
1587
1588
1589
1590
    # 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:
1591
1592
1593
1594
1595
        if (
            message["role"] == "assistant"
            and "tool_calls" in message
            and isinstance(message["tool_calls"], list)
        ):
1596
            for item in message["tool_calls"]:
1597
1598
                # if arguments is None or empty string, set to {}
                if content := item["function"].get("arguments"):
1599
1600
                    if not isinstance(content, (dict, list)):
                        item["function"]["arguments"] = json.loads(content)
1601
1602
                else:
                    item["function"]["arguments"] = {}
1603
1604


1605
def parse_chat_messages(
1606
    messages: list[ChatCompletionMessageParam],
1607
    model_config: ModelConfig,
1608
    content_format: _ChatTemplateContentFormat,
1609
1610
) -> tuple[
    list[ConversationMessage],
1611
1612
    MultiModalDataDict | None,
    MultiModalUUIDDict | None,
1613
]:
1614
    conversation: list[ConversationMessage] = []
1615
    mm_tracker = MultiModalItemTracker(model_config)
1616
1617

    for msg in messages:
1618
1619
1620
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1621
            content_format,
1622
1623
1624
1625
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1626
            ),
1627
        )
1628

1629
        conversation.extend(sub_messages)
1630

1631
1632
    _postprocess_messages(conversation)

1633
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1634
1635


1636
def parse_chat_messages_futures(
1637
    messages: list[ChatCompletionMessageParam],
1638
    model_config: ModelConfig,
1639
    content_format: _ChatTemplateContentFormat,
1640
1641
) -> tuple[
    list[ConversationMessage],
1642
1643
    Awaitable[MultiModalDataDict | None],
    MultiModalUUIDDict | None,
1644
]:
1645
    conversation: list[ConversationMessage] = []
1646
    mm_tracker = AsyncMultiModalItemTracker(model_config)
1647
1648

    for msg in messages:
1649
1650
1651
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1652
            content_format,
1653
1654
1655
1656
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1657
            ),
1658
        )
1659
1660
1661

        conversation.extend(sub_messages)

1662
1663
    _postprocess_messages(conversation)

1664
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1665
1666


1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
# 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)


1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
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)


1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
@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)
    )


1714
def resolve_chat_template_kwargs(
1715
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
1716
1717
    chat_template: str,
    chat_template_kwargs: dict[str, Any],
1718
    raise_on_unexpected: bool = True,
1719
) -> dict[str, Any]:
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
    # 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}"
        )

1731
    fn_kw = {
1732
1733
        k
        for k in chat_template_kwargs
1734
1735
        if supports_kw(tokenizer.apply_chat_template, k, allow_var_kwargs=False)
    }
1736
    template_vars = _cached_resolve_chat_template_kwargs(chat_template)
1737
1738
1739
1740
1741

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


1745
def apply_hf_chat_template(
1746
    tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast,
1747
    conversation: list[ConversationMessage],
1748
1749
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
1750
    *,
1751
    model_config: ModelConfig,
1752
    **kwargs: Any,
1753
) -> str:
1754
    hf_chat_template = resolve_hf_chat_template(
1755
1756
1757
        tokenizer,
        chat_template=chat_template,
        tools=tools,
1758
        model_config=model_config,
1759
    )
1760

1761
    if hf_chat_template is None:
1762
1763
1764
        raise ValueError(
            "As of transformers v4.44, default chat template is no longer "
            "allowed, so you must provide a chat template if the tokenizer "
1765
1766
            "does not define one."
        )
1767

1768
1769
1770
1771
1772
1773
    resolved_kwargs = resolve_chat_template_kwargs(
        tokenizer=tokenizer,
        chat_template=hf_chat_template,
        chat_template_kwargs=kwargs,
    )

1774
1775
1776
1777
1778
    try:
        return tokenizer.apply_chat_template(
            conversation=conversation,  # type: ignore[arg-type]
            tools=tools,  # type: ignore[arg-type]
            chat_template=hf_chat_template,
1779
            tokenize=False,
1780
            **resolved_kwargs,
1781
        )
1782

1783
1784
1785
1786
1787
1788
    # 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(
1789
1790
            "An error occurred in `transformers` while applying chat template"
        )
1791
        raise ValueError(str(e)) from e
1792

1793

1794
1795
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1796
    messages: list[ChatCompletionMessageParam],
1797
1798
    chat_template: str | None,
    tools: list[dict[str, Any]] | None,
1799
    **kwargs: Any,
1800
) -> list[int]:
1801
1802
    from mistral_common.exceptions import MistralCommonException

1803
1804
1805
1806
1807
1808
    # 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,
    )
1809

1810
1811
1812
1813
1814
1815
1816
1817
1818
    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
1819
    # properly caught in the preprocessing_input step
1820
    except (AssertionError, MistralCommonException) as e:
1821
        raise ValueError(str(e)) from e
1822
1823
1824
1825
1826
1827
1828

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

1833

1834
1835
1836
def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
    idx = 0
    for msg in conversation:
1837
1838
1839
        if msg["role"] == "assistant":
            tool_calls = msg.get("tool_calls")
            idx += len(list(tool_calls)) if tool_calls is not None else 0  # noqa
1840
1841
1842
    return idx


1843
1844
1845
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}"
1846
1847
1848
    else:
        # by default return random
        return f"chatcmpl-tool-{random_uuid()}"