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

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

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

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

logger = init_logger(__name__)


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


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

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


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


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


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

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


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

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


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

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


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

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


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


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

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

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

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

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

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

150
151
152
153
154

ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam,
                                   CustomChatCompletionMessageParam]


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

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

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

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

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


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

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


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

    return False


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

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

    return False


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

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

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


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

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

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

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

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


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

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

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


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

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

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


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


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

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


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

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

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

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

    return None


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

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

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

    return detected_format if given_format == "auto" else given_format


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

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

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

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

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

436
437
438
    return detected_format


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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

            allowed_count = mm_config.get_limit_per_prompt(input_modality)

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

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

        return self._placeholder_str(modality, current_count)

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


580
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
581
582

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

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

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


609
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
610
611

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

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

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

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


class BaseMultiModalContentParser(ABC):

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

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

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

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

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

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

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

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

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

678
679
680
681
682
683
684
685

class MultiModalContentParser(BaseMultiModalContentParser):

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

        self._tracker = tracker

686
687
688
689
        self._connector = MediaConnector(
            allowed_local_media_path=tracker.allowed_local_media_path,
        )

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

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

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

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

        self._add_placeholder(placeholder)

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

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

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

722
        return self.parse_audio(audio_url)
723

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

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

730
731
732
733
734
735
736

class AsyncMultiModalContentParser(BaseMultiModalContentParser):

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

        self._tracker = tracker
737
738
739
        self._connector = MediaConnector(
            allowed_local_media_path=tracker.allowed_local_media_path,
        )
740
741

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

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

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

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

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

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

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

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

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

777
        return self.parse_audio(audio_url)
778

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

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

785

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

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

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

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


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

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

821
        return chat_template
822

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

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

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


_cached_load_chat_template = lru_cache(_load_chat_template)


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


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

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

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

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

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

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


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

888
_ContentPart: TypeAlias = Union[str, dict[str, str], InputAudio]
889

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


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

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

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

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

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

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

        return part_type, content

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

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


VALID_MESSAGE_CONTENT_MM_PART_TYPES = ("text", "refusal", "image_url",
970
                                       "image_embeds",
971
                                       "audio_url", "input_audio", "video_url")
972

973

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

983
    mm_parser = mm_tracker.create_parser()
984
985

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

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


def _parse_chat_message_content_part(
1008
1009
1010
1011
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
1012
) -> Optional[_ContentPart]:
1013
1014
1015
1016
1017
1018
1019
1020
    """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
1021
        return part
1022
1023
1024
1025

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

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

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

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

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

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

1064
    raise NotImplementedError(f"Unknown part type: {part_type}")
1065
1066


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


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

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

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

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

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

    return result

1109

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

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


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

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

1141
        conversation.extend(sub_messages)
1142

1143
1144
    _postprocess_messages(conversation)

1145
    return conversation, mm_tracker.all_mm_data()
1146
1147


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

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

        conversation.extend(sub_messages)

1166
1167
    _postprocess_messages(conversation)

1168
1169
1170
    return conversation, mm_tracker.all_mm_data()


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

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

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


1203
1204
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1205
    messages: list[ChatCompletionMessageParam],
1206
1207
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
1208
    **kwargs: Any,
1209
) -> list[int]:
1210
1211
1212
1213
1214
1215
    # 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,
    )
1216

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