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

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

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

38
from vllm.config import ModelConfig
39
from vllm.logger import init_logger
40
41
from vllm.model_executor.model_loader import get_model_cls
from vllm.model_executor.models import SupportsMultiModal
42
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict
43
from vllm.multimodal.utils import MediaConnector
44
45
46
47
# yapf: disable
from vllm.transformers_utils.chat_templates import (
    get_chat_template_fallback_path)
# yapf: enable
48
from vllm.transformers_utils.processor import cached_get_processor
49
from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
50
from vllm.utils import deprecate_kwargs, random_uuid
51
52
53
54

logger = init_logger(__name__)


55
56
57
58
59
60
61
62
63
64
65
66
67
68
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."""


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


80
81
82
83
84
85
86
87
88
89
90
91
92
93
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."""


94
95
96
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.
97

98
99
100
101
102
103
104
105
106
107
    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.
108

109
110
111
112
113
114
115
116
    Example:
    {
        "audio_url": "https://example.com/audio.mp3"
    }
    """
    audio_url: Required[str]


117
118
119
120
121
122
123
124
125
126
127
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]


128
129
ChatCompletionContentPartParam: TypeAlias = Union[
    OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam,
130
    ChatCompletionContentPartInputAudioParam,
131
    ChatCompletionContentPartVideoParam, ChatCompletionContentPartRefusalParam,
132
    CustomChatCompletionContentSimpleImageParam,
133
    ChatCompletionContentPartImageEmbedsParam,
134
135
    CustomChatCompletionContentSimpleAudioParam,
    CustomChatCompletionContentSimpleVideoParam, str]
136
137
138
139
140
141
142


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

143
    content: Union[str, list[ChatCompletionContentPartParam]]
144
145
146
147
148
149
150
151
152
    """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.
    """

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

159
160
161
162
163

ChatCompletionMessageParam = Union[OpenAIChatCompletionMessageParam,
                                   CustomChatCompletionMessageParam]


164
# TODO: Make fields ReadOnly once mypy supports it
165
166
167
168
class ConversationMessage(TypedDict, total=False):
    role: Required[str]
    """The role of the message's author."""

169
    content: Union[Optional[str], list[dict[str, str]]]
170
171
172
173
174
175
176
177
178
179
    """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."""
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
# 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


298
@lru_cache(maxsize=32)
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
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"


319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
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

336
337
338
339
@deprecate_kwargs(
    "trust_remote_code",
    additional_message="Please use `model_config.trust_remote_code` instead.",
)
340
def resolve_hf_chat_template(
341
342
343
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
344
345
    *,
    model_config: ModelConfig,
346
    trust_remote_code: Optional[bool] = None,
347
348
349
350
351
352
353
354
355
356
357
358
) -> 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),
359
                trust_remote_code=model_config.trust_remote_code,
360
361
            )
            if isinstance(processor, ProcessorMixin) and \
362
                hasattr(processor, 'chat_template') and \
363
364
365
                processor.chat_template is not None:
                return processor.chat_template
        except Exception:
366
            logger.debug("Failed to load AutoProcessor chat template for %s", tokenizer.name_or_path, exc_info=True)  # noqa: E501
367
368
369
370
371
372
373
374

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

375
376
377
378
379
380
381
382
383
384
385
386
387
388
    # 4th priority: Predefined fallbacks
    path = get_chat_template_fallback_path(
        model_type=model_config.hf_config.model_type,
        tokenizer_name_or_path=model_config.tokenizer,
    )
    if path is not None:
        logger.info("Loading chat template fallback for %s as there isn't one "
                    "defined on HF Hub.", tokenizer.name_or_path)
        chat_template = load_chat_template(path)
    else:
        logger.debug("There is no chat template fallback for %s",
                     tokenizer.name_or_path)

    return chat_template
389
390


391
392
def _resolve_chat_template_content_format(
    chat_template: Optional[str],
393
    tools: Optional[list[dict[str, Any]]],
394
    tokenizer: AnyTokenizer,
395
396
    *,
    model_config: ModelConfig,
397
398
) -> _ChatTemplateContentFormat:
    if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
399
        hf_chat_template = resolve_hf_chat_template(
400
401
402
            tokenizer,
            chat_template=chat_template,
            tools=tools,
403
            model_config=model_config,
404
        )
405
    else:
406
407
408
409
        hf_chat_template = None

    jinja_text = (hf_chat_template if isinstance(hf_chat_template, str)
                  else load_chat_template(chat_template, is_literal=True))
410
411
412
413

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

414
    return detected_format
415
416
417


@lru_cache
418
def _log_chat_template_content_format(
419
420
    chat_template: Optional[str],
    given_format: ChatTemplateContentFormatOption,
421
422
    detected_format: ChatTemplateContentFormatOption,
):
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
    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,
        )

440

441
442
443
444
@deprecate_kwargs(
    "trust_remote_code",
    additional_message="Please use `model_config.trust_remote_code` instead.",
)
445
446
447
448
449
def resolve_chat_template_content_format(
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
    given_format: ChatTemplateContentFormatOption,
    tokenizer: AnyTokenizer,
450
451
452
    *,
    model_config: ModelConfig,
    trust_remote_code: Optional[bool] = None,
453
) -> _ChatTemplateContentFormat:
454
455
456
    if given_format != "auto":
        return given_format

457
458
459
460
    detected_format = _resolve_chat_template_content_format(
        chat_template,
        tools,
        tokenizer,
461
        model_config=model_config,
462
463
464
465
466
467
468
469
    )

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

470
    return detected_format
471

472
473


474
ModalityStr = Literal["image", "audio", "video", "image_embeds"]
475
476
477
478
_T = TypeVar("_T")


class BaseMultiModalItemTracker(ABC, Generic[_T]):
479
480
481
482
483
484
485
    """
    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):
486
487
        super().__init__()

488
489
        self._model_config = model_config
        self._tokenizer = tokenizer
490

491
        self._items_by_modality = defaultdict[str, list[_T]](list)
492

493
494
495
496
    @property
    def model_config(self) -> ModelConfig:
        return self._model_config

497
498
499
500
    @cached_property
    def model_cls(self):
        return get_model_cls(self.model_config)

501
502
503
504
    @property
    def allowed_local_media_path(self):
        return self._model_config.allowed_local_media_path

505
506
507
508
    @property
    def mm_registry(self):
        return MULTIMODAL_REGISTRY

509
510
511
512
513
    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.
        """
514
515
        mm_registry = self.mm_registry
        model_config = self.model_config
516
        model_cls = cast(SupportsMultiModal, self.model_cls)
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531

        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)

532
        current_count = len(self._items_by_modality[modality]) + 1
533
534
535
        if current_count > allowed_count:
            raise ValueError(
                f"At most {allowed_count} {modality}(s) may be provided in "
536
537
                "one request. You can set `--limit-mm-per-prompt` to "
                "increase this limit if the model supports it.")
538

539
        self._items_by_modality[modality].append(item)
540

541
        return model_cls.get_placeholder_str(modality, current_count)
542
543
544
545
546
547

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


548
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
549
550

    def all_mm_data(self) -> Optional[MultiModalDataDict]:
551
552
553
554
555
556
557
558
559
560
561
562
563
564
        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]
565
        if "image" in items_by_modality:
566
            mm_inputs["image"] = items_by_modality["image"] # A list of images
567
        if "audio" in items_by_modality:
568
            mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
569
        if "video" in items_by_modality:
570
571
            mm_inputs["video"] = items_by_modality["video"] # A list of videos
        return mm_inputs
572
573
574
575
576

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


577
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
578
579

    async def all_mm_data(self) -> Optional[MultiModalDataDict]:
580
581
582
583
        if not self._items_by_modality:
            return None
        mm_inputs = {}
        items_by_modality = {
584
585
586
                modality: await asyncio.gather(*items)
                for modality, items in self._items_by_modality.items()
            }
587

588
589
590
591
592
593
594
595
596
597
        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]
598
        if "image" in items_by_modality:
599
            mm_inputs["image"] = items_by_modality["image"] # A list of images
600
        if "audio" in items_by_modality:
601
            mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
602
        if "video" in items_by_modality:
603
604
            mm_inputs["video"] = items_by_modality["video"] # A list of videos
        return mm_inputs
605
606
607
608
609
610
611
612
613
614
615

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


class BaseMultiModalContentParser(ABC):

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

        # multimodal placeholder_string : count
616
        self._placeholder_counts: dict[str, int] = defaultdict(lambda: 0)
617
618
619
620
621

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

622
    def mm_placeholder_counts(self) -> dict[str, int]:
623
624
625
626
627
628
        return dict(self._placeholder_counts)

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

629
630
631
632
633
    @abstractmethod
    def parse_image_embeds(self,
                           image_embeds: Union[str, dict[str, str]]) -> None:
        raise NotImplementedError

634
635
636
637
    @abstractmethod
    def parse_audio(self, audio_url: str) -> None:
        raise NotImplementedError

638
    @abstractmethod
639
    def parse_input_audio(self, input_audio: InputAudio) -> None:
640
641
        raise NotImplementedError

642
643
644
645
    @abstractmethod
    def parse_video(self, video_url: str) -> None:
        raise NotImplementedError

646
647
648
649
650
651
652
653

class MultiModalContentParser(BaseMultiModalContentParser):

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

        self._tracker = tracker

654
        self._connector = MediaConnector(
655
            media_io_kwargs=self._tracker._model_config.media_io_kwargs,
656
657
658
            allowed_local_media_path=tracker.allowed_local_media_path,
        )

659
    def parse_image(self, image_url: str) -> None:
660
        image = self._connector.fetch_image(image_url)
661
662
663
664

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

665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
    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)

680
    def parse_audio(self, audio_url: str) -> None:
681
        audio = self._connector.fetch_audio(audio_url)
682
683
684
685

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

686
687
688
689
    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}"
690

691
        return self.parse_audio(audio_url)
692

693
    def parse_video(self, video_url: str) -> None:
694
        video = self._connector.fetch_video(video_url=video_url)
695
696
697
698

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

699
700
701
702
703
704
705

class AsyncMultiModalContentParser(BaseMultiModalContentParser):

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

        self._tracker = tracker
706
        self._connector = MediaConnector(
707
708
            media_io_kwargs=self._tracker._model_config.media_io_kwargs,
            allowed_local_media_path=tracker.allowed_local_media_path
709
        )
710
711

    def parse_image(self, image_url: str) -> None:
712
        image_coro = self._connector.fetch_image_async(image_url)
713
714
715
716

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

717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
    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)

736
    def parse_audio(self, audio_url: str) -> None:
737
        audio_coro = self._connector.fetch_audio_async(audio_url)
738
739
740

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

742
743
744
745
    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}"
746

747
        return self.parse_audio(audio_url)
748

749
    def parse_video(self, video_url: str) -> None:
750
        video = self._connector.fetch_video_async(video_url=video_url)
751
752
753
754

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

755

756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
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")


778
def _load_chat_template(
779
780
781
782
    chat_template: Optional[Union[Path, str]],
    *,
    is_literal: bool = False,
) -> Optional[str]:
783
784
    if chat_template is None:
        return None
785
786
787
788
789
790

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

791
        return chat_template
792

793
    try:
794
        with open(chat_template) as f:
795
            return f.read()
796
    except OSError as e:
797
798
799
        if isinstance(chat_template, Path):
            raise

800
801
802
803
804
805
        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
806

807
808
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
809
810
811
812
813
814
815
816
817
818
819
820
        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)
821
822


823
# TODO: Let user specify how to insert multimodal tokens into prompt
824
# (similar to chat template)
825
def _get_full_multimodal_text_prompt(placeholder_counts: dict[str, int],
826
                                     text_prompt: str) -> str:
827
    """Combine multimodal prompts for a multimodal language model."""
828

829
    # Look through the text prompt to check for missing placeholders
830
    missing_placeholders: list[str] = []
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
    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])
847
848


849
850
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
851
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
852
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
853
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
854
855
856
857
# Need to validate url objects
_ImageParser = TypeAdapter(ChatCompletionContentPartImageParam).validate_python
_AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python
_VideoParser = TypeAdapter(ChatCompletionContentPartVideoParam).validate_python
858

859
_ContentPart: TypeAlias = Union[str, dict[str, str], InputAudio]
860

861
# Define a mapping from part types to their corresponding parsing functions.
862
MM_PARSER_MAP: dict[
863
864
865
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
866
    "text":
867
    lambda part: _TextParser(part).get("text", None),
868
    "image_url":
869
    lambda part: _ImageParser(part).get("image_url", {}).get("url", None),
870
    "image_embeds":
871
    lambda part: _ImageEmbedsParser(part).get("image_embeds", None),
872
    "audio_url":
873
    lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
874
    "input_audio":
875
    lambda part: _InputAudioParser(part).get("input_audio", None),
876
    "refusal":
877
    lambda part: _RefusalParser(part).get("refusal", None),
878
    "video_url":
879
    lambda part: _VideoParser(part).get("video_url", {}).get("url", None),
880
881
882
883
}


def _parse_chat_message_content_mm_part(
884
        part: ChatCompletionContentPartParam) -> tuple[str, _ContentPart]:
885
    """
886
    Parses a given multi-modal content part based on its type.
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906

    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'
907
908
        # We only support 'auto', which is the default
        if part_type == "image_url" and part.get("detail", "auto") != "auto":
909
910
911
912
913
914
            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.
915
    # 'type' is required field by pydantic
916
917
918
919
920
921
922
923
924
    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", "")
925
        if part.get("input_audio") is not None:
926
            input_audio_params = cast(dict[str, str], part)
927
            return "input_audio", input_audio_params
928
929
930
931
        if part.get("video_url") is not None:
            video_params = cast(CustomChatCompletionContentSimpleVideoParam,
                                part)
            return "video_url", video_params.get("video_url", "")
932
933
934
935
936
937
938
939
940
        # 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",
941
                                       "image_embeds",
942
                                       "audio_url", "input_audio", "video_url")
943

944

945
946
947
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
948
    mm_tracker: BaseMultiModalItemTracker,
949
950
    *,
    wrap_dicts: bool,
951
) -> list[ConversationMessage]:
952
    content = list[_ContentPart]()
953

954
    mm_parser = mm_tracker.create_parser()
955
956

    for part in parts:
957
        parse_res = _parse_chat_message_content_part(
958
959
960
961
            part,
            mm_parser,
            wrap_dicts=wrap_dicts,
        )
962
963
        if parse_res:
            content.append(parse_res)
964

965
    if wrap_dicts:
966
        # Parsing wraps images and texts as interleaved dictionaries
967
        return [ConversationMessage(role=role,
968
                                    content=content)]  # type: ignore
969
    texts = cast(list[str], content)
970
971
972
973
974
975
976
977
978
    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(
979
980
981
982
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
983
) -> Optional[_ContentPart]:
984
985
986
987
988
989
990
991
    """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
992
        return part
993
994
995
996

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

997
    # if part_type is text/refusal/image_url/audio_url/video_url/input_audio but
998
999
    # content is None, log a warning and skip
    if part_type in VALID_MESSAGE_CONTENT_MM_PART_TYPES and content is None:
1000
        logger.warning(
1001
1002
            "Skipping multimodal part '%s' (type: '%s') "
            "with empty / unparsable content.", part, part_type)
1003
1004
1005
        return None

    if part_type in ("text", "refusal"):
1006
1007
1008
1009
1010
        str_content = cast(str, content)
        if wrap_dicts:
            return {'type': 'text', 'text': str_content}
        else:
            return str_content
1011
1012

    if part_type == "image_url":
1013
1014
        str_content = cast(str, content)
        mm_parser.parse_image(str_content)
1015
        return {'type': 'image'} if wrap_dicts else None
1016
1017
1018
1019
    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
1020
    if part_type == "audio_url":
1021
1022
1023
1024
1025
        str_content = cast(str, content)
        mm_parser.parse_audio(str_content)
        return {'type': 'audio'} if wrap_dicts else None

    if part_type == "input_audio":
1026
        dict_content = cast(InputAudio, content)
1027
        mm_parser.parse_input_audio(dict_content)
1028
1029
        return {'type': 'audio'} if wrap_dicts else None

1030
    if part_type == "video_url":
1031
1032
        str_content = cast(str, content)
        mm_parser.parse_video(str_content)
1033
1034
        return {'type': 'video'} if wrap_dicts else None

1035
    raise NotImplementedError(f"Unknown part type: {part_type}")
1036
1037


1038
1039
1040
1041
1042
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


1043
def _parse_chat_message_content(
1044
1045
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
1046
    content_format: _ChatTemplateContentFormat,
1047
) -> list[ConversationMessage]:
1048
1049
1050
1051
    role = message["role"]
    content = message.get("content")

    if content is None:
1052
1053
1054
1055
1056
1057
        content = []
    elif isinstance(content, str):
        content = [
            ChatCompletionContentPartTextParam(type="text", text=content)
        ]
    result = _parse_chat_message_content_parts(
1058
1059
        role,
        content,  # type: ignore
1060
        mm_tracker,
1061
        wrap_dicts=(content_format == "openai"),
1062
    )
1063

1064
1065
1066
1067
    for result_msg in result:
        if role == 'assistant':
            parsed_msg = _AssistantParser(message)

1068
1069
1070
1071
1072
            # The 'tool_calls' is not None check ensures compatibility.
            # It's needed only if downstream code doesn't strictly
            # follow the OpenAI spec.
            if ("tool_calls" in parsed_msg
                and parsed_msg["tool_calls"] is not None):
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
                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

1084

1085
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
    # 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"])


1100
def parse_chat_messages(
1101
    messages: list[ChatCompletionMessageParam],
1102
    model_config: ModelConfig,
1103
    tokenizer: AnyTokenizer,
1104
    content_format: _ChatTemplateContentFormat,
1105
1106
) -> tuple[list[ConversationMessage], Optional[MultiModalDataDict]]:
    conversation: list[ConversationMessage] = []
1107
    mm_tracker = MultiModalItemTracker(model_config, tokenizer)
1108
1109

    for msg in messages:
1110
1111
1112
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1113
            content_format,
1114
        )
1115

1116
        conversation.extend(sub_messages)
1117

1118
1119
    _postprocess_messages(conversation)

1120
    return conversation, mm_tracker.all_mm_data()
1121
1122


1123
def parse_chat_messages_futures(
1124
    messages: list[ChatCompletionMessageParam],
1125
1126
    model_config: ModelConfig,
    tokenizer: AnyTokenizer,
1127
    content_format: _ChatTemplateContentFormat,
1128
1129
) -> tuple[list[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]:
    conversation: list[ConversationMessage] = []
1130
1131
1132
    mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)

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

        conversation.extend(sub_messages)

1141
1142
    _postprocess_messages(conversation)

1143
1144
1145
    return conversation, mm_tracker.all_mm_data()


1146
1147
1148
1149
@deprecate_kwargs(
    "trust_remote_code",
    additional_message="Please use `model_config.trust_remote_code` instead.",
)
1150
1151
def apply_hf_chat_template(
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
1152
    conversation: list[ConversationMessage],
1153
    chat_template: Optional[str],
1154
    tools: Optional[list[dict[str, Any]]],
1155
    *,
1156
    model_config: ModelConfig,
1157
    tokenize: bool = False,  # Different from HF's default
1158
1159
    # Deprecated, explicitly capture here so it doesn't slit into kwargs.
    trust_remote_code: Optional[bool] = None,
1160
    **kwargs: Any,
1161
) -> str:
1162
    hf_chat_template = resolve_hf_chat_template(
1163
1164
1165
        tokenizer,
        chat_template=chat_template,
        tools=tools,
1166
        model_config=model_config,
1167
    )
1168

1169
    if hf_chat_template is None:
1170
1171
1172
1173
1174
        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.")

1175
1176
1177
1178
1179
1180
1181
1182
1183
    try:

        return tokenizer.apply_chat_template(
            conversation=conversation,  # type: ignore[arg-type]
            tools=tools,  # type: ignore[arg-type]
            chat_template=hf_chat_template,
            tokenize=tokenize,
            **kwargs,
        )
1184

1185
1186
1187
1188
1189
1190
1191
1192
    # 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(
            "An error occurred in `transformers` while applying chat template")
1193
        raise ValueError(str(e)) from e
1194

1195
1196
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1197
    messages: list[ChatCompletionMessageParam],
1198
1199
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
1200
    **kwargs: Any,
1201
) -> list[int]:
1202
1203
    from mistral_common.exceptions import MistralCommonException

1204
1205
1206
1207
1208
1209
    # 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,
    )
1210

1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
    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
1221
    except (AssertionError, MistralCommonException) as e:
1222
        raise ValueError(str(e)) from e
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232

    # 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(
            "An error occurred in `mistral_common` while applying chat "
            "template")
1233
        raise ValueError(str(e)) from e
1234
1235

def random_tool_call_id() -> str:
1236
    return f"chatcmpl-tool-{random_uuid()}"