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

3
import asyncio
4
import codecs
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 cache, 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
32
# yapf: enable
# pydantic needs the TypedDict from typing_extensions
33
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
34
from typing_extensions import Required, TypeAlias, TypedDict
35

36
from vllm.config import ModelConfig
37
38
from vllm.logger import init_logger
from vllm.multimodal import 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
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
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
363
364
365
# 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"


def _resolve_chat_template_content_format(
    chat_template: Optional[str],
    given_format: ChatTemplateContentFormatOption,
    tokenizer: AnyTokenizer,
) -> _ChatTemplateContentFormat:
    if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
        tokenizer_chat_template = tokenizer.chat_template
    else:
        tokenizer_chat_template = None

    jinja_text: Optional[str]
    if isinstance(tokenizer_chat_template, str) and chat_template is None:
        jinja_text = tokenizer_chat_template
    elif (isinstance(tokenizer_chat_template, dict)
            and chat_template in tokenizer_chat_template):
        jinja_text = tokenizer_chat_template[chat_template]
    else:
        jinja_text = load_chat_template(chat_template, is_literal=True)

    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
def resolve_chat_template_content_format(
    chat_template: Optional[str],
    given_format: ChatTemplateContentFormatOption,
    tokenizer: AnyTokenizer,
) -> _ChatTemplateContentFormat:
    detected_format = _resolve_chat_template_content_format(
        chat_template,
        given_format,
        tokenizer,
    )

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

    return detected_format


366
ModalityStr = Literal["image", "audio", "video", "image_embeds"]
367
368
369
370
_T = TypeVar("_T")


class BaseMultiModalItemTracker(ABC, Generic[_T]):
371
372
373
374
375
376
377
    """
    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):
378
379
        super().__init__()

380
381
382
383
        self._model_config = model_config
        self._tokenizer = tokenizer
        self._allowed_items = (model_config.multimodal_config.limit_per_prompt
                               if model_config.multimodal_config else {})
384

385
        self._items_by_modality = defaultdict[str, list[_T]](list)
386

387
388
389
390
    @property
    def model_config(self) -> ModelConfig:
        return self._model_config

391
392
393
394
    @property
    def allowed_local_media_path(self):
        return self._model_config.allowed_local_media_path

395
    @staticmethod
396
    @cache
397
    def _cached_token_str(tokenizer: AnyTokenizer, token_index: int) -> str:
398
399
        return tokenizer.decode(token_index)

400
401
    def _placeholder_str(self, modality: ModalityStr,
                         current_count: int) -> Optional[str]:
402
403
        # TODO: Let user specify how to insert image tokens into prompt
        # (similar to chat template)
404
405
406
        hf_config = self._model_config.hf_config
        model_type = hf_config.model_type

407
408
409
        if modality in ("image", "image_embeds"):
            if model_type == "chatglm":
                return "<|begin_of_image|><|endoftext|><|end_of_image|>"
410
411
412
            if model_type == "phi3_v":
                # Workaround since this token is not defined in the tokenizer
                return f"<|image_{current_count}|>"
413
414
            if model_type == "phi4mm":
                return "<|endoftext10|>"  # 200010 (see vocab.json in hf model)
415
            if model_type in ("minicpmo", "minicpmv"):
416
                return "(<image>./</image>)"
417
            if model_type in ("blip-2", "fuyu", "paligemma", "pixtral"):
418
419
                # These models do not use image tokens in the prompt
                return None
420
421
            if model_type == "qwen":
                return f"Picture {current_count}: <img></img>"
422
            if model_type.startswith("llava"):
423
424
                return self._cached_token_str(self._tokenizer,
                                              hf_config.image_token_index)
425
426
            if model_type in ("chameleon", "deepseek_vl_v2", "internvl_chat",
                              "NVLM_D", "h2ovl_chat"):
427
                return "<image>"
428
429
            if model_type == "mllama":
                return "<|image|>"
Roger Wang's avatar
Roger Wang committed
430
            if model_type in ("qwen2_vl", "qwen2_5_vl"):
431
                return "<|vision_start|><|image_pad|><|vision_end|>"
432
433
            if model_type == "molmo":
                return ""
434
435
            if model_type == "idefics3":
                return "<image>"
436
437
            if model_type == "aria":
                return "<|fim_prefix|><|img|><|fim_suffix|>"
438
439
            if model_type == "gemma3":
                return "<start_of_image>"
440

441
            raise TypeError(f"Unknown {modality} model type: {model_type}")
442
443
        elif modality == "audio":
            if model_type == "ultravox":
444
                return "<|audio|>"
445
446
            if model_type == "phi4mm":
                return "<|endoftext11|>"  # 200011 (see vocab.json in hf model)
447
448
449
            if model_type == "qwen2_audio":
                return (f"Audio {current_count}: "
                        f"<|audio_bos|><|AUDIO|><|audio_eos|>")
450
451
            if model_type == "minicpmo":
                return "(<audio>./</audio>)"
452
            raise TypeError(f"Unknown model type: {model_type}")
453
        elif modality == "video":
Roger Wang's avatar
Roger Wang committed
454
            if model_type in ("qwen2_vl", "qwen2_5_vl"):
455
                return "<|vision_start|><|video_pad|><|vision_end|>"
456
457
            if model_type in ("minicpmo", "minicpmv"):
                return "(<video>./</video>)"
458
459
460
            if model_type.startswith("llava"):
                return self._cached_token_str(self._tokenizer,
                                              hf_config.video_token_index)
461
            raise TypeError(f"Unknown {modality} model type: {model_type}")
462
463
464
        else:
            raise TypeError(f"Unknown modality: {modality}")

465
466
467
468
469
470
    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.
        """
        allowed_count = self._allowed_items.get(modality, 1)
471
        current_count = len(self._items_by_modality[modality]) + 1
472
473
474
475
476
        if current_count > allowed_count:
            raise ValueError(
                f"At most {allowed_count} {modality}(s) may be provided in "
                "one request.")

477
        self._items_by_modality[modality].append(item)
478
479
480
481
482
483
484
485

        return self._placeholder_str(modality, current_count)

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


486
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
487
488

    def all_mm_data(self) -> Optional[MultiModalDataDict]:
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
        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]
        elif "image" in items_by_modality:
            mm_inputs["image"] = items_by_modality["image"] # A list of images
        elif "audio" in items_by_modality:
            mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
        elif "video" in items_by_modality:
            mm_inputs["video"] = items_by_modality["video"] # A list of videos
        return mm_inputs
510
511
512
513
514

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


515
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
516
517

    async def all_mm_data(self) -> Optional[MultiModalDataDict]:
518
519
520
521
        if not self._items_by_modality:
            return None
        mm_inputs = {}
        items_by_modality = {
522
523
524
                modality: await asyncio.gather(*items)
                for modality, items in self._items_by_modality.items()
            }
525

526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
        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]
        elif "image" in items_by_modality:
            mm_inputs["image"] = items_by_modality["image"] # A list of images
        elif "audio" in items_by_modality:
            mm_inputs["audio"] = items_by_modality["audio"] # A list of audios
        elif "video" in items_by_modality:
            mm_inputs["video"] = items_by_modality["video"] # A list of videos
        return mm_inputs
543
544
545
546
547
548
549
550
551
552
553

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


class BaseMultiModalContentParser(ABC):

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

        # multimodal placeholder_string : count
554
        self._placeholder_counts: dict[str, int] = defaultdict(lambda: 0)
555
556
557
558
559

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

560
    def mm_placeholder_counts(self) -> dict[str, int]:
561
562
563
564
565
566
        return dict(self._placeholder_counts)

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

567
568
569
570
571
    @abstractmethod
    def parse_image_embeds(self,
                           image_embeds: Union[str, dict[str, str]]) -> None:
        raise NotImplementedError

572
573
574
575
    @abstractmethod
    def parse_audio(self, audio_url: str) -> None:
        raise NotImplementedError

576
    @abstractmethod
577
    def parse_input_audio(self, input_audio: InputAudio) -> None:
578
579
        raise NotImplementedError

580
581
582
583
    @abstractmethod
    def parse_video(self, video_url: str) -> None:
        raise NotImplementedError

584
585
586
587
588
589
590
591

class MultiModalContentParser(BaseMultiModalContentParser):

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

        self._tracker = tracker

592
593
594
595
        self._connector = MediaConnector(
            allowed_local_media_path=tracker.allowed_local_media_path,
        )

596
    def parse_image(self, image_url: str) -> None:
597
        image = self._connector.fetch_image(image_url)
598
599
600
601

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

602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
    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)

617
    def parse_audio(self, audio_url: str) -> None:
618
        audio = self._connector.fetch_audio(audio_url)
619
620
621
622

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

623
624
625
626
    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}"
627

628
        return self.parse_audio(audio_url)
629

630
    def parse_video(self, video_url: str) -> None:
631
        video = self._connector.fetch_video(video_url)
632
633
634
635

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

636
637
638
639
640
641
642

class AsyncMultiModalContentParser(BaseMultiModalContentParser):

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

        self._tracker = tracker
643
644
645
        self._connector = MediaConnector(
            allowed_local_media_path=tracker.allowed_local_media_path,
        )
646
647

    def parse_image(self, image_url: str) -> None:
648
        image_coro = self._connector.fetch_image_async(image_url)
649
650
651
652

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

653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
    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)

672
    def parse_audio(self, audio_url: str) -> None:
673
        audio_coro = self._connector.fetch_audio_async(audio_url)
674
675
676

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

678
679
680
681
    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}"
682

683
        return self.parse_audio(audio_url)
684

685
    def parse_video(self, video_url: str) -> None:
686
        video = self._connector.fetch_video_async(video_url)
687
688
689
690

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

691

692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
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")


714
def load_chat_template(
715
716
717
718
    chat_template: Optional[Union[Path, str]],
    *,
    is_literal: bool = False,
) -> Optional[str]:
719
720
    if chat_template is None:
        return None
721
722
723
724
725
726
727
728

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

        return codecs.decode(chat_template, "unicode_escape")

729
    try:
730
        with open(chat_template) as f:
731
            return f.read()
732
    except OSError as e:
733
734
735
        if isinstance(chat_template, Path):
            raise

736
737
738
739
740
741
        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
742

743
744
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
745
        return load_chat_template(chat_template, is_literal=True)
746
747


748
# TODO: Let user specify how to insert multimodal tokens into prompt
749
# (similar to chat template)
750
def _get_full_multimodal_text_prompt(placeholder_counts: dict[str, int],
751
                                     text_prompt: str) -> str:
752
    """Combine multimodal prompts for a multimodal language model."""
753

754
    # Look through the text prompt to check for missing placeholders
755
    missing_placeholders: list[str] = []
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
    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])
772
773


774
775
776
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
_ImageParser = partial(cast, ChatCompletionContentPartImageParam)
777
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
778
_AudioParser = partial(cast, ChatCompletionContentPartAudioParam)
779
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
780
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
781
_VideoParser = partial(cast, ChatCompletionContentPartVideoParam)
782

783
_ContentPart: TypeAlias = Union[str, dict[str, str], InputAudio]
784

785
# Define a mapping from part types to their corresponding parsing functions.
786
MM_PARSER_MAP: dict[
787
788
789
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
790
791
792
793
    "text":
    lambda part: _TextParser(part).get("text", ""),
    "image_url":
    lambda part: _ImageParser(part).get("image_url", {}).get("url", ""),
794
795
    "image_embeds":
    lambda part: _ImageEmbedsParser(part).get("image_embeds", {}),
796
797
    "audio_url":
    lambda part: _AudioParser(part).get("audio_url", {}).get("url", ""),
798
799
    "input_audio":
    lambda part: _InputAudioParser(part).get("input_audio", {}),
800
801
    "refusal":
    lambda part: _RefusalParser(part).get("refusal", ""),
802
803
    "video_url":
    lambda part: _VideoParser(part).get("video_url", {}).get("url", ""),
804
805
806
807
}


def _parse_chat_message_content_mm_part(
808
        part: ChatCompletionContentPartParam) -> tuple[str, _ContentPart]:
809
    """
810
    Parses a given multi-modal content part based on its type.
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830

    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'
831
832
        # We only support 'auto', which is the default
        if part_type == "image_url" and part.get("detail", "auto") != "auto":
833
834
835
836
837
838
            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.
839
    # 'type' is required field by pydantic
840
841
842
843
844
845
846
847
848
    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", "")
849
        if part.get("input_audio") is not None:
850
            input_audio_params = cast(dict[str, str], part)
851
            return "input_audio", input_audio_params
852
853
854
855
        if part.get("video_url") is not None:
            video_params = cast(CustomChatCompletionContentSimpleVideoParam,
                                part)
            return "video_url", video_params.get("video_url", "")
856
857
858
859
860
861
862
863
864
        # 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",
865
                                       "image_embeds",
866
                                       "audio_url", "input_audio", "video_url")
867

868

869
870
871
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
872
    mm_tracker: BaseMultiModalItemTracker,
873
874
    *,
    wrap_dicts: bool,
875
) -> list[ConversationMessage]:
876
    content = list[_ContentPart]()
877

878
    mm_parser = mm_tracker.create_parser()
879
880

    for part in parts:
881
        parse_res = _parse_chat_message_content_part(
882
883
884
885
            part,
            mm_parser,
            wrap_dicts=wrap_dicts,
        )
886
887
        if parse_res:
            content.append(parse_res)
888

889
    if wrap_dicts:
890
        # Parsing wraps images and texts as interleaved dictionaries
891
        return [ConversationMessage(role=role,
892
                                    content=content)]  # type: ignore
893
    texts = cast(list[str], content)
894
895
896
897
898
899
900
901
902
    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(
903
904
905
906
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
907
) -> Optional[_ContentPart]:
908
909
910
911
912
913
914
915
    """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
916
        return part
917
918
919
920

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

921
    # if part_type is text/refusal/image_url/audio_url/video_url/input_audio but
922
923
924
    # content is empty, log a warning and skip
    if part_type in VALID_MESSAGE_CONTENT_MM_PART_TYPES and not content:
        logger.warning(
925
            "Skipping multimodal part (type: '%s') "
926
927
928
929
            "with empty / unparsable content.", part_type)
        return None

    if part_type in ("text", "refusal"):
930
931
932
933
934
        str_content = cast(str, content)
        if wrap_dicts:
            return {'type': 'text', 'text': str_content}
        else:
            return str_content
935
936

    if part_type == "image_url":
937
938
        str_content = cast(str, content)
        mm_parser.parse_image(str_content)
939
        return {'type': 'image'} if wrap_dicts else None
940
941
942
943
    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
944
    if part_type == "audio_url":
945
946
947
948
949
        str_content = cast(str, content)
        mm_parser.parse_audio(str_content)
        return {'type': 'audio'} if wrap_dicts else None

    if part_type == "input_audio":
950
        dict_content = cast(InputAudio, content)
951
        mm_parser.parse_input_audio(dict_content)
952
953
        return {'type': 'audio'} if wrap_dicts else None

954
    if part_type == "video_url":
955
956
        str_content = cast(str, content)
        mm_parser.parse_video(str_content)
957
958
        return {'type': 'video'} if wrap_dicts else None

959
    raise NotImplementedError(f"Unknown part type: {part_type}")
960
961


962
963
964
965
966
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


967
def _parse_chat_message_content(
968
969
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
970
    content_format: _ChatTemplateContentFormat,
971
) -> list[ConversationMessage]:
972
973
974
975
    role = message["role"]
    content = message.get("content")

    if content is None:
976
977
978
979
980
981
        content = []
    elif isinstance(content, str):
        content = [
            ChatCompletionContentPartTextParam(type="text", text=content)
        ]
    result = _parse_chat_message_content_parts(
982
983
        role,
        content,  # type: ignore
984
        mm_tracker,
985
        wrap_dicts=(content_format == "openai"),
986
    )
987

988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
    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

1004

1005
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
    # 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"])


1020
def parse_chat_messages(
1021
    messages: list[ChatCompletionMessageParam],
1022
    model_config: ModelConfig,
1023
    tokenizer: AnyTokenizer,
1024
    content_format: _ChatTemplateContentFormat,
1025
1026
) -> tuple[list[ConversationMessage], Optional[MultiModalDataDict]]:
    conversation: list[ConversationMessage] = []
1027
    mm_tracker = MultiModalItemTracker(model_config, tokenizer)
1028
1029

    for msg in messages:
1030
1031
1032
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1033
            content_format,
1034
        )
1035

1036
        conversation.extend(sub_messages)
1037

1038
1039
    _postprocess_messages(conversation)

1040
    return conversation, mm_tracker.all_mm_data()
1041
1042


1043
def parse_chat_messages_futures(
1044
    messages: list[ChatCompletionMessageParam],
1045
1046
    model_config: ModelConfig,
    tokenizer: AnyTokenizer,
1047
    content_format: _ChatTemplateContentFormat,
1048
1049
) -> tuple[list[ConversationMessage], Awaitable[Optional[MultiModalDataDict]]]:
    conversation: list[ConversationMessage] = []
1050
1051
1052
    mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)

    for msg in messages:
1053
1054
1055
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1056
            content_format,
1057
        )
1058
1059
1060

        conversation.extend(sub_messages)

1061
1062
    _postprocess_messages(conversation)

1063
1064
1065
    return conversation, mm_tracker.all_mm_data()


1066
1067
def apply_hf_chat_template(
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
1068
    conversation: list[ConversationMessage],
1069
1070
1071
1072
    chat_template: Optional[str],
    *,
    tokenize: bool = False,  # Different from HF's default
    **kwargs: Any,
1073
) -> str:
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
    if chat_template is None:
        chat_template = tokenizer.chat_template

    # FIXME: Temporary workaround for
    # https://huggingface.co/mistral-community/pixtral-12b/discussions/31
    if chat_template is None:
        try:
            processor = cached_get_processor(tokenizer.name_or_path)
            chat_template = processor.chat_template
        except Exception:
            pass

    if chat_template is None:
1087
1088
1089
1090
1091
        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.")

1092
1093
1094
1095
1096
1097
    return tokenizer.apply_chat_template(
        conversation=conversation,  # type: ignore[arg-type]
        chat_template=chat_template,
        tokenize=tokenize,
        **kwargs,
    )
1098
1099


1100
1101
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1102
    messages: list[ChatCompletionMessageParam],
1103
    chat_template: Optional[str] = None,
1104
    **kwargs: Any,
1105
) -> list[int]:
1106
    if chat_template is not None:
1107
        logger.warning_once(
1108
            "'chat_template' cannot be overridden for mistral tokenizer.")
1109
    if "add_generation_prompt" in kwargs:
1110
        logger.warning_once(
1111
1112
1113
            "'add_generation_prompt' is not supported for mistral tokenizer, "
            "so it will be ignored.")
    if "continue_final_message" in kwargs:
1114
        logger.warning_once(
1115
1116
            "'continue_final_message' is not supported for mistral tokenizer, "
            "so it will be ignored.")
1117

1118
1119
    return tokenizer.apply_chat_template(
        messages=messages,
1120
1121
        **kwargs,
    )