chat_utils.py 54.3 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 Counter, 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 openai.types.responses import ResponseInputImageParam
32
from openai_harmony import Message as OpenAIHarmonyMessage
33
34
from PIL import Image
from pydantic import BaseModel, ConfigDict, TypeAdapter
35
# yapf: enable
36
37
from transformers import (PreTrainedTokenizer, PreTrainedTokenizerFast,
                          ProcessorMixin)
38
# pydantic needs the TypedDict from typing_extensions
39
from typing_extensions import Required, TypeAlias, TypedDict
40

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

logger = init_logger(__name__)

57
58
59
60
61
62
MODALITY_PLACEHOLDERS_MAP = {
    "image": "<##IMAGE##>",
    "audio": "<##AUDIO##>",
    "video": "<##VIDEO##>",
}

63

64
65
66
67
68
69
70
71
72
73
74
75
76
77
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."""


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


94
95
96
97
98
99
100
101
102
103
104
105
106
107
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."""


108
109
110
111
class PILImage(BaseModel):
    """
    A PIL.Image.Image object.
    """
112

113
114
115
116
117
118
119
120
121
122
123
124
    image_pil: Image.Image
    model_config = ConfigDict(arbitrary_types_allowed=True)


class CustomChatCompletionContentPILImageParam(TypedDict, total=False):
    """A simpler version of the param that only accepts a PIL image.

    Example:
    {
        "image_pil": ImageAsset('cherry_blossom').pil_image
    }
    """
125

126
    image_pil: Optional[PILImage]
127
128
129
130
131
    uuid: Optional[str]
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
132
133


134
135
136
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.
137

138
139
140
141
142
    Example:
    {
        "image_url": "https://example.com/image.jpg"
    }
    """
143

144
    image_url: Optional[str]
145
146
147
148
149
    uuid: Optional[str]
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
150
151
152
153


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

155
156
157
158
159
    Example:
    {
        "audio_url": "https://example.com/audio.mp3"
    }
    """
160

161
    audio_url: Optional[str]
162
163


164
165
166
167
168
169
170
171
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"
    }
    """
172

173
    video_url: Optional[str]
174
175
176
177
178
    uuid: Optional[str]
    """
    User-provided UUID of a media. User must guarantee that it is properly
    generated and unique for different medias.
    """
179
180


Julien Denize's avatar
Julien Denize committed
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
class CustomThinkCompletionContentParam(TypedDict, total=False):
    """A Think Completion Content Param that accepts a plain text and a boolean.

    Example:
    {
        "thinking": "I am thinking about the answer",
        "closed": True,
        "type": "thinking"
    }
    """

    thinking: Required[str]
    """The thinking content."""

    closed: bool
    """Whether the thinking is closed."""

    type: Required[Literal["thinking"]]
    """The thinking type."""


202
ChatCompletionContentPartParam: TypeAlias = Union[
203
204
    OpenAIChatCompletionContentPartParam,
    ChatCompletionContentPartAudioParam,
205
    ChatCompletionContentPartInputAudioParam,
206
207
    ChatCompletionContentPartVideoParam,
    ChatCompletionContentPartRefusalParam,
208
    CustomChatCompletionContentPILImageParam,
209
    CustomChatCompletionContentSimpleImageParam,
210
    ChatCompletionContentPartImageEmbedsParam,
211
    CustomChatCompletionContentSimpleAudioParam,
212
213
214
215
    CustomChatCompletionContentSimpleVideoParam,
    str,
    CustomThinkCompletionContentParam,
]
216
217
218
219


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

221
222
223
    role: Required[str]
    """The role of the message's author."""

224
    content: Union[str, list[ChatCompletionContentPartParam]]
225
226
227
228
229
230
231
232
233
    """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.
    """

234
235
236
237
238
239
    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."""

240

241
242
243
244
245
ChatCompletionMessageParam = Union[
    OpenAIChatCompletionMessageParam,
    CustomChatCompletionMessageParam,
    OpenAIHarmonyMessage,
]
246
247


248
# TODO: Make fields ReadOnly once mypy supports it
249
250
251
252
class ConversationMessage(TypedDict, total=False):
    role: Required[str]
    """The role of the message's author."""

253
    content: Union[Optional[str], list[dict[str, str]]]
254
255
256
257
258
259
260
261
262
263
    """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."""
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):
298
299
        return node.node is not None and _is_var_or_elems_access(
            node.node, varname, key)
300
301
302
    if isinstance(node, jinja2.nodes.Test):
        return _is_var_or_elems_access(node.node, varname, key)

303
304
    if isinstance(node, jinja2.nodes.Getitem) and isinstance(
            node.arg, jinja2.nodes.Slice):
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
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
        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


382
@lru_cache(maxsize=32)
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
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"


403
404
405
406
407
408
def resolve_mistral_chat_template(
    chat_template: Optional[str],
    **kwargs: Any,
) -> Optional[str]:
    if chat_template is not None:
        logger.warning_once(
409
410
            "'chat_template' cannot be overridden for mistral tokenizer."
        )
411
412
413
    if "add_generation_prompt" in kwargs:
        logger.warning_once(
            "'add_generation_prompt' is not supported for mistral tokenizer, "
414
415
            "so it will be ignored."
        )
416
417
418
    if "continue_final_message" in kwargs:
        logger.warning_once(
            "'continue_final_message' is not supported for mistral tokenizer, "
419
420
            "so it will be ignored."
        )
421
422
    return None

423

424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], Optional[str]]()
"""
Used in `_try_get_processor_chat_template` to avoid calling
`cached_get_processor` again if the processor fails to be loaded.

This is needed because `lru_cache` does not cache when an exception happens.
"""


def _try_get_processor_chat_template(
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
    model_config: ModelConfig,
) -> Optional[str]:
    cache_key = (tokenizer.name_or_path, model_config.trust_remote_code)
    if cache_key in _PROCESSOR_CHAT_TEMPLATES:
        return _PROCESSOR_CHAT_TEMPLATES[cache_key]

    try:
        processor = cached_get_processor(
            tokenizer.name_or_path,
            processor_cls=(
                PreTrainedTokenizer,
                PreTrainedTokenizerFast,
                ProcessorMixin,
            ),
            trust_remote_code=model_config.trust_remote_code,
        )
        if (
            isinstance(processor, ProcessorMixin)
            and hasattr(processor, "chat_template")
            and (chat_template := processor.chat_template) is not None
        ):
            _PROCESSOR_CHAT_TEMPLATES[cache_key] = chat_template
            return chat_template
    except Exception:
        logger.debug(
            "Failed to load AutoProcessor chat template for %s",
            tokenizer.name_or_path,
            exc_info=True,
        )

    _PROCESSOR_CHAT_TEMPLATES[cache_key] = None
    return None


469
def resolve_hf_chat_template(
470
471
472
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
473
474
    *,
    model_config: ModelConfig,
475
476
477
478
479
480
481
) -> 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:
482
483
484
485
        chat_template = _try_get_processor_chat_template(tokenizer,
                                                         model_config)
        if chat_template is not None:
            return chat_template
486
487
488
489
490

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

497
498
499
500
501
502
    # 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:
503
504
505
506
507
        logger.info(
            "Loading chat template fallback for %s as there isn't one "
            "defined on HF Hub.",
            tokenizer.name_or_path,
        )
508
509
        chat_template = load_chat_template(path)
    else:
510
511
512
        logger.debug(
            "There is no chat template fallback for %s", tokenizer.name_or_path
        )
513
514

    return chat_template
515
516


517
518
def _resolve_chat_template_content_format(
    chat_template: Optional[str],
519
    tools: Optional[list[dict[str, Any]]],
520
    tokenizer: AnyTokenizer,
521
522
    *,
    model_config: ModelConfig,
523
524
) -> _ChatTemplateContentFormat:
    if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
525
        hf_chat_template = resolve_hf_chat_template(
526
527
528
            tokenizer,
            chat_template=chat_template,
            tools=tools,
529
            model_config=model_config,
530
        )
531
    else:
532
533
        hf_chat_template = None

534
535
536
537
538
    jinja_text = (
        hf_chat_template
        if isinstance(hf_chat_template, str)
        else load_chat_template(chat_template, is_literal=True)
    )
539

540
541
542
543
544
    detected_format = (
        "string"
        if jinja_text is None
        else _detect_content_format(jinja_text, default="string")
    )
545

546
    return detected_format
547
548
549


@lru_cache
550
def _log_chat_template_content_format(
551
552
    chat_template: Optional[str],
    given_format: ChatTemplateContentFormatOption,
553
554
    detected_format: ChatTemplateContentFormatOption,
):
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
    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,
        )

572
573
574
575
576
577

def resolve_chat_template_content_format(
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
    given_format: ChatTemplateContentFormatOption,
    tokenizer: AnyTokenizer,
578
579
    *,
    model_config: ModelConfig,
580
) -> _ChatTemplateContentFormat:
581
582
583
    if given_format != "auto":
        return given_format

584
585
586
587
    detected_format = _resolve_chat_template_content_format(
        chat_template,
        tools,
        tokenizer,
588
        model_config=model_config,
589
590
591
592
593
594
595
596
    )

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

597
    return detected_format
598

599

600
ModalityStr = Literal["image", "audio", "video", "image_embeds"]
601
602
603
604
_T = TypeVar("_T")


class BaseMultiModalItemTracker(ABC, Generic[_T]):
605
606
607
608
609
610
611
    """
    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):
612
613
        super().__init__()

614
615
        self._model_config = model_config
        self._tokenizer = tokenizer
616

617
        self._items_by_modality = defaultdict[str, list[Optional[_T]]](list)
618
        self._uuids_by_modality = defaultdict[str, list[Optional[str]]](list)
619

620
621
622
623
    @property
    def model_config(self) -> ModelConfig:
        return self._model_config

624
    @cached_property
625
    def model_cls(self) -> type[SupportsMultiModal]:
626
        from vllm.model_executor.model_loader import get_model_cls
627

628
629
        model_cls = get_model_cls(self.model_config)
        return cast(type[SupportsMultiModal], model_cls)
630

631
632
633
634
    @property
    def allowed_local_media_path(self):
        return self._model_config.allowed_local_media_path

635
636
637
638
    @property
    def mm_registry(self):
        return MULTIMODAL_REGISTRY

639
640
641
642
    @cached_property
    def mm_processor(self):
        return self.mm_registry.create_processor(self.model_config)

643
    def add(
644
645
646
647
        self,
        modality: ModalityStr,
        item: Optional[_T],
        uuid: Optional[str] = None,
648
    ) -> Optional[str]:
649
650
651
        """
        Add a multi-modal item to the current prompt and returns the
        placeholder string to use, if any.
652
653

        An optional uuid can be added which serves as a unique identifier of the
654
        media.
655
        """
656
        input_modality = modality.replace("_embeds", "")
657
        num_items = len(self._items_by_modality[modality]) + 1
658

659
        self.mm_processor.validate_num_items(input_modality, num_items)
660

661
        self._items_by_modality[modality].append(item)
662
        self._uuids_by_modality[modality].append(uuid)
663

664
        return self.model_cls.get_placeholder_str(modality, num_items)
665

666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
    def all_mm_uuids(self) -> Optional[MultiModalUUIDDict]:
        if not self._items_by_modality:
            return None
        mm_uuids = {}
        uuids_by_modality = dict(self._uuids_by_modality)
        if "image" in uuids_by_modality and "image_embeds" in uuids_by_modality:
            raise ValueError(
                "Mixing raw image and embedding inputs is not allowed"
            )

        if "image_embeds" in uuids_by_modality:
            image_embeds_uuids = uuids_by_modality["image_embeds"]
            if len(image_embeds_uuids) > 1:
                raise ValueError(
                    "Only one message can have {'type': 'image_embeds'}"
                )
            mm_uuids["image"] = uuids_by_modality["image_embeds"]
        if "image" in uuids_by_modality:
            mm_uuids["image"] = uuids_by_modality["image"]  # UUIDs of images
        if "audio" in uuids_by_modality:
            mm_uuids["audio"] = uuids_by_modality["audio"]  # UUIDs of audios
        if "video" in uuids_by_modality:
            mm_uuids["video"] = uuids_by_modality["video"]  # UUIDs of videos
        return mm_uuids

691
692
693
694
695
    @abstractmethod
    def create_parser(self) -> "BaseMultiModalContentParser":
        raise NotImplementedError


696
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
697
    def all_mm_data(self) -> Optional[MultiModalDataDict]:
698
699
700
701
702
        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:
703
704
705
            raise ValueError(
                "Mixing raw image and embedding inputs is not allowed"
            )
706
707
708
709

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
            if len(image_embeds_lst) > 1:
710
711
712
                raise ValueError(
                    "Only one message can have {'type': 'image_embeds'}"
                )
713
            mm_inputs["image"] = image_embeds_lst[0]
714
        if "image" in items_by_modality:
715
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
716
        if "audio" in items_by_modality:
717
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
718
        if "video" in items_by_modality:
719
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
720
        return mm_inputs
721
722
723
724
725

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


726
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
727
    async def all_mm_data(self) -> Optional[MultiModalDataDict]:
728
729
730
        if not self._items_by_modality:
            return None
        mm_inputs = {}
731
732
733
734
735
736
737
738
739
        items_by_modality = {}
        for modality, items in self._items_by_modality.items():
            coros = []
            for item in items:
                if item is not None:
                    coros.append(item)
                else:
                    coros.append(asyncio.sleep(0))
            items_by_modality[modality] = await asyncio.gather(*coros)
740

741
742
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
            raise ValueError(
743
744
                "Mixing raw image and embedding inputs is not allowed"
            )
745
746
747
748
749

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
            if len(image_embeds_lst) > 1:
                raise ValueError(
750
751
                    "Only one message can have {'type': 'image_embeds'}"
                )
752
            mm_inputs["image"] = image_embeds_lst[0]
753
        if "image" in items_by_modality:
754
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
755
        if "audio" in items_by_modality:
756
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
757
        if "video" in items_by_modality:
758
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
759
        return mm_inputs
760
761
762
763
764
765
766
767
768

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


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

769
        # stores model placeholders list with corresponding
770
771
772
773
774
775
776
        # general MM placeholder:
        # {
        #   "<##IMAGE##>": ["<image>", "<image>", "<image>"],
        #   "<##AUDIO##>": ["<audio>", "<audio>"]
        # }
        self._placeholder_storage: dict[str, list] = defaultdict(list)

777
778
779
    def _add_placeholder(
        self, modality: ModalityStr, placeholder: Optional[str]
    ):
780
        mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]
781
        if placeholder:
782
            self._placeholder_storage[mod_placeholder].append(placeholder)
783

784
785
    def mm_placeholder_storage(self) -> dict[str, list]:
        return dict(self._placeholder_storage)
786
787

    @abstractmethod
788
789
    def parse_image(
        self, image_url: Optional[str], uuid: Optional[str] = None) -> None:
790
791
        raise NotImplementedError

792
    @abstractmethod
793
    def parse_image_embeds(
794
        self,
795
        image_embeds: Union[str, dict[str, str], None],
796
        uuid: Optional[str] = None,
797
    ) -> None:
798
799
        raise NotImplementedError

800
    @abstractmethod
801
    def parse_image_pil(
802
        self, image_pil: Optional[Image.Image], uuid: Optional[str] = None
803
    ) -> None:
804
805
        raise NotImplementedError

806
    @abstractmethod
807
808
809
    def parse_audio(
        self, audio_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
810
811
        raise NotImplementedError

812
    @abstractmethod
813
    def parse_input_audio(
814
        self, input_audio: Optional[InputAudio], uuid: Optional[str] = None
815
    ) -> None:
816
817
        raise NotImplementedError

818
    @abstractmethod
819
820
821
    def parse_video(
        self, video_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
822
823
        raise NotImplementedError

824
825
826
827
828
829

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

        self._tracker = tracker
830
831
        multimodal_config = self._tracker.model_config.multimodal_config
        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)
832
        self._connector = MediaConnector(
833
            media_io_kwargs=media_io_kwargs,
834
835
836
            allowed_local_media_path=tracker.allowed_local_media_path,
        )

837
838
839
840
    def parse_image(
        self, image_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
        image = self._connector.fetch_image(image_url) if image_url else None
841

842
        placeholder = self._tracker.add("image", image, uuid)
843
        self._add_placeholder("image", placeholder)
844

845
    def parse_image_embeds(
846
        self,
847
        image_embeds: Union[str, dict[str, str], None],
848
        uuid: Optional[str] = None,
849
    ) -> None:
850
851
852
853
854
        if isinstance(image_embeds, dict):
            embeds = {
                k: self._connector.fetch_image_embedding(v)
                for k, v in image_embeds.items()
            }
855
            placeholder = self._tracker.add("image_embeds", embeds, uuid)
856
857
858

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

861
862
863
        if image_embeds is None:
            placeholder = self._tracker.add("image_embeds", None, uuid)

864
        self._add_placeholder("image", placeholder)
865

866
    def parse_image_pil(
867
        self, image_pil: Optional[Image.Image], uuid: Optional[str] = None
868
869
    ) -> None:
        placeholder = self._tracker.add("image", image_pil, uuid)
870
        self._add_placeholder("image", placeholder)
871

872
873
874
875
    def parse_audio(
        self, audio_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
        audio = self._connector.fetch_audio(audio_url) if audio_url else None
876

877
        placeholder = self._tracker.add("audio", audio, uuid)
878
        self._add_placeholder("audio", placeholder)
879

880
    def parse_input_audio(
881
        self, input_audio: Optional[InputAudio], uuid: Optional[str] = None
882
    ) -> None:
883
884
885
886
887
888
889
890
891
892
        if input_audio:
            audio_data = input_audio.get("data", "")
            audio_format = input_audio.get("format", "")
            if audio_data:
                audio_url = f"data:audio/{audio_format};base64,{audio_data}"
            else:
                # If a UUID is provided, audio data may be empty.
                audio_url = None
        else:
            audio_url = None
893

894
        return self.parse_audio(audio_url, uuid)
895

896
897
898
899
900
901
902
903
    def parse_video(
        self, video_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
        video = (
            self._connector.fetch_video(video_url=video_url)
            if video_url
            else None
        )
904

905
        placeholder = self._tracker.add("video", video, uuid)
906
        self._add_placeholder("video", placeholder)
907

908
909
910
911
912
913

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

        self._tracker = tracker
914
915
        multimodal_config = self._tracker.model_config.multimodal_config
        media_io_kwargs = getattr(multimodal_config, "media_io_kwargs", None)
916
        self._connector = MediaConnector(
917
            media_io_kwargs=media_io_kwargs,
918
            allowed_local_media_path=tracker.allowed_local_media_path,
919
        )
920

921
922
923
924
925
926
    def parse_image(
        self, image_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
        image_coro = (
            self._connector.fetch_image_async(image_url) if image_url else None
        )
927

928
        placeholder = self._tracker.add("image", image_coro, uuid)
929
        self._add_placeholder("image", placeholder)
930

931
    def parse_image_embeds(
932
        self,
933
        image_embeds: Union[str, dict[str, str], None],
934
        uuid: Optional[str] = None,
935
    ) -> None:
936
937
938
        future: asyncio.Future[Union[str, dict[str, str], None]] = (
            asyncio.Future()
        )
939
940
941
942
943
944
945
946
947

        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):
948
            embedding = self._connector.fetch_image_embedding(image_embeds)
949
950
            future.set_result(embedding)

951
952
953
        if image_embeds is None:
            future.set_result(None)

954
        placeholder = self._tracker.add("image_embeds", future, uuid)
955
        self._add_placeholder("image", placeholder)
956

957
    def parse_image_pil(
958
        self, image_pil: Optional[Image.Image], uuid: Optional[str] = None
959
    ) -> None:
960
961
962
963
964
        future: asyncio.Future[Optional[Image.Image]] = asyncio.Future()
        if image_pil:
            future.set_result(image_pil)
        else:
            future.set_result(None)
965

966
        placeholder = self._tracker.add("image", future, uuid)
967
        self._add_placeholder("image", placeholder)
968

969
970
971
972
973
974
    def parse_audio(
        self, audio_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
        audio_coro = (
            self._connector.fetch_audio_async(audio_url) if audio_url else None
        )
975

976
        placeholder = self._tracker.add("audio", audio_coro, uuid)
977
        self._add_placeholder("audio", placeholder)
978

979
    def parse_input_audio(
980
        self, input_audio: Optional[InputAudio], uuid: Optional[str] = None
981
    ) -> None:
982
983
984
985
986
987
988
989
990
991
        if input_audio:
            audio_data = input_audio.get("data", "")
            audio_format = input_audio.get("format", "")
            if audio_data:
                audio_url = f"data:audio/{audio_format};base64,{audio_data}"
            else:
                # If a UUID is provided, audio data may be empty.
                audio_url = None
        else:
            audio_url = None
992

993
        return self.parse_audio(audio_url, uuid)
994

995
996
997
998
999
1000
1001
1002
    def parse_video(
        self, video_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
        video = (
            self._connector.fetch_video_async(video_url=video_url)
            if video_url
            else None
        )
1003

1004
        placeholder = self._tracker.add("video", video, uuid)
1005
        self._add_placeholder("video", placeholder)
1006

1007

1008
1009
1010
1011
1012
1013
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():
1014
        raise FileNotFoundError("the supplied chat template path doesn't exist")
1015
1016
1017

    elif isinstance(chat_template, str):
        JINJA_CHARS = "{}\n"
1018
1019
1020
1021
        if (
            not any(c in chat_template for c in JINJA_CHARS)
            and not Path(chat_template).exists()
        ):
1022
1023
            raise ValueError(
                f"The supplied chat template string ({chat_template}) "
1024
1025
                f"appears path-like, but doesn't exist!"
            )
1026
1027
1028

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


1033
def _load_chat_template(
1034
1035
1036
1037
    chat_template: Optional[Union[Path, str]],
    *,
    is_literal: bool = False,
) -> Optional[str]:
1038
1039
    if chat_template is None:
        return None
1040
1041
1042

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

1047
        return chat_template
1048

1049
    try:
1050
        with open(chat_template) as f:
1051
            return f.read()
1052
    except OSError as e:
1053
1054
1055
        if isinstance(chat_template, Path):
            raise

1056
1057
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template for c in JINJA_CHARS):
1058
1059
1060
1061
1062
            msg = (
                f"The supplied chat template ({chat_template}) "
                f"looks like a file path, but it failed to be "
                f"opened. Reason: {e}"
            )
1063
            raise ValueError(msg) from e
1064

1065
1066
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
        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)
1079
1080


1081
1082
1083
def _get_interleaved_text_prompt(
    placeholder_storage: dict[str, list], texts: list[str]
) -> str:
1084
1085
1086
1087
1088
1089
1090
    for idx, elem in enumerate(texts):
        if elem in placeholder_storage:
            texts[idx] = placeholder_storage[elem].pop(0)

    return "\n".join(texts)


1091
# TODO: Let user specify how to insert multimodal tokens into prompt
1092
# (similar to chat template)
1093
1094
1095
1096
1097
def _get_full_multimodal_text_prompt(
    placeholder_storage: dict[str, list],
    texts: list[str],
    interleave_strings: bool,
) -> str:
1098
    """Combine multimodal prompts for a multimodal language model."""
1099

1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
    # flatten storage to make it looks like
    # {
    #   "<|image|>": 2,
    #   "<|audio|>": 1
    # }
    placeholder_counts = Counter(
        [v for elem in placeholder_storage.values() for v in elem]
    )

    if interleave_strings:
        text_prompt = _get_interleaved_text_prompt(placeholder_storage, texts)
    else:
        text_prompt = "\n".join(texts)

    # Pass interleaved text further in case the user used image placeholders
    # himself, but forgot to disable the 'interleave_strings' flag

1117
    # Look through the text prompt to check for missing placeholders
1118
    missing_placeholders: list[str] = []
1119
1120
1121
1122
1123
    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:
1124
1125
1126
1127
            logger.error(
                "Placeholder count is negative! "
                "Ensure that the 'interleave_strings' flag is disabled "
                "(current value: %s) "
1128
1129
                "when manually placing image placeholders.",
                interleave_strings,
1130
1131
            )
            logger.debug("Input prompt: %s", text_prompt)
1132
1133
            raise ValueError(
                f"Found more '{placeholder}' placeholders in input prompt than "
1134
1135
                "actual multimodal data items."
            )
1136

1137
1138
1139
        missing_placeholders.extend(
            [placeholder] * placeholder_counts[placeholder]
        )
1140

1141
1142
    # NOTE: Default behaviour: we always add missing placeholders
    # at the front of the prompt, if interleave_strings=False
1143
    return "\n".join(missing_placeholders + [text_prompt])
1144
1145


1146
1147
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
1148
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
1149
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
1150
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
1151
_PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
Julien Denize's avatar
Julien Denize committed
1152
_ThinkParser = partial(cast, CustomThinkCompletionContentParam)
1153
1154
1155
1156
# Need to validate url objects
_ImageParser = TypeAdapter(ChatCompletionContentPartImageParam).validate_python
_AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python
_VideoParser = TypeAdapter(ChatCompletionContentPartVideoParam).validate_python
1157

1158
_ResponsesInputImageParser = TypeAdapter(
1159
1160
    ResponseInputImageParam
).validate_python
1161
_ContentPart: TypeAlias = Union[str, dict[str, str], InputAudio, PILImage]
1162

1163
# Define a mapping from part types to their corresponding parsing functions.
1164
MM_PARSER_MAP: dict[
1165
1166
1167
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
    "text": lambda part: _TextParser(part).get("text", None),
    "thinking": lambda part: _ThinkParser(part).get("thinking", None),
    "input_text": lambda part: _TextParser(part).get("text", None),
    "input_image": lambda part: _ResponsesInputImageParser(part).get(
        "image_url", None
    ),
    "image_url": lambda part: _ImageParser(part)
    .get("image_url", {})
    .get("url", None),
    "image_embeds": lambda part: _ImageEmbedsParser(part).get(
        "image_embeds", None
    ),
1180
    "image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
    "audio_url": lambda part: _AudioParser(part)
    .get("audio_url", {})
    .get("url", None),
    "input_audio": lambda part: _InputAudioParser(part).get(
        "input_audio", None
    ),
    "refusal": lambda part: _RefusalParser(part).get("refusal", None),
    "video_url": lambda part: _VideoParser(part)
    .get("video_url", {})
    .get("url", None),
1191
1192
1193
1194
}


def _parse_chat_message_content_mm_part(
1195
1196
    part: ChatCompletionContentPartParam,
) -> tuple[str, _ContentPart]:
1197
    """
1198
    Parses a given multi-modal content part based on its type.
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211

    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(
1212
1213
        part, dict
    )  # This is needed to avoid mypy errors: part.get() from str
1214
    part_type = part.get("type", None)
1215
    uuid = part.get("uuid", None)
1216

1217
    if isinstance(part_type, str) and part_type in MM_PARSER_MAP and uuid is None: # noqa: E501
1218
1219
1220
        content = MM_PARSER_MAP[part_type](part)

        # Special case for 'image_url.detail'
1221
1222
        # We only support 'auto', which is the default
        if part_type == "image_url" and part.get("detail", "auto") != "auto":
1223
1224
1225
1226
            logger.warning(
                "'image_url.detail' is currently not supported "
                "and will be ignored."
            )
1227
1228
1229
1230

        return part_type, content

    # Handle missing 'type' but provided direct URL fields.
1231
    # 'type' is required field by pydantic
1232
1233
    if part_type is None or uuid is not None:
        if "image_url" in part:
1234
1235
1236
            image_params = cast(
                CustomChatCompletionContentSimpleImageParam, part
            )
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
            image_url = image_params.get("image_url", None)
            if isinstance(image_url, dict):
                # Can potentially happen if user provides a uuid
                # with url as a dict of {"url": url}
                image_url = image_url.get("url", None)
            return "image_url", image_url
        if "image_pil" in part:
            # "image_pil" could be None if UUID is provided.
            image_params = cast( # type: ignore 
                CustomChatCompletionContentPILImageParam, part
            )
            image_pil = image_params.get("image_pil", None)
            return "image_pil", image_pil
        if "image_embeds" in part:
            # "image_embeds" could be None if UUID is provided.
            image_params = cast( # type: ignore 
                ChatCompletionContentPartImageEmbedsParam, part
            )
            image_embeds = image_params.get("image_embeds", None)
            return "image_embeds", image_embeds
        if "audio_url" in part:
1258
1259
1260
            audio_params = cast(
                CustomChatCompletionContentSimpleAudioParam, part
            )
1261
1262
1263
1264
1265
1266
            audio_url = audio_params.get("audio_url", None)
            if isinstance(audio_url, dict):
                # Can potentially happen if user provides a uuid
                # with url as a dict of {"url": url}
                audio_url = audio_url.get("url", None)
            return "audio_url", audio_url
1267
        if part.get("input_audio") is not None:
1268
            input_audio_params = cast(dict[str, str], part)
1269
            return "input_audio", input_audio_params
1270
        if "video_url" in part:
1271
1272
1273
            video_params = cast(
                CustomChatCompletionContentSimpleVideoParam, part
            )
1274
1275
1276
1277
1278
1279
            video_url = video_params.get("video_url", None)
            if isinstance(video_url, dict):
                # Can potentially happen if user provides a uuid
                # with url as a dict of {"url": url}
                video_url = video_url.get("url", None)
            return "video_url", video_url
1280
1281
1282
1283
1284
1285
1286
1287
        # 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"


1288
PART_TYPES_TO_SKIP_NONE_CONTENT = (
1289
1290
1291
    "text",
    "refusal",
)
1292

1293

1294
1295
1296
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
1297
    mm_tracker: BaseMultiModalItemTracker,
1298
1299
    *,
    wrap_dicts: bool,
1300
    interleave_strings: bool,
1301
) -> list[ConversationMessage]:
1302
    content = list[_ContentPart]()
1303

1304
    mm_parser = mm_tracker.create_parser()
1305
1306

    for part in parts:
1307
        parse_res = _parse_chat_message_content_part(
1308
1309
1310
            part,
            mm_parser,
            wrap_dicts=wrap_dicts,
1311
            interleave_strings=interleave_strings,
1312
        )
1313
1314
        if parse_res:
            content.append(parse_res)
1315

1316
    if wrap_dicts:
1317
        # Parsing wraps images and texts as interleaved dictionaries
1318
        return [ConversationMessage(role=role, content=content)]  # type: ignore
1319
    texts = cast(list[str], content)
1320
1321
    mm_placeholder_storage = mm_parser.mm_placeholder_storage()
    if mm_placeholder_storage:
1322
1323
1324
        text_prompt = _get_full_multimodal_text_prompt(
            mm_placeholder_storage, texts, interleave_strings
        )
1325
1326
1327
    else:
        text_prompt = "\n".join(texts)

1328
1329
1330
1331
    return [ConversationMessage(role=role, content=text_prompt)]


def _parse_chat_message_content_part(
1332
1333
1334
1335
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
1336
    interleave_strings: bool,
1337
) -> Optional[_ContentPart]:
1338
1339
1340
1341
1342
1343
1344
1345
    """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
1346
        return part
1347
1348
    # Handle structured dictionary parts
    part_type, content = _parse_chat_message_content_mm_part(part)
1349
    # if part_type is text/refusal/image_url/audio_url/video_url/input_audio but
1350
    # content is None, log a warning and skip
1351
    if part_type in PART_TYPES_TO_SKIP_NONE_CONTENT and content is None:
1352
        logger.warning(
1353
            "Skipping multimodal part '%s' (type: '%s') "
1354
1355
1356
1357
            "with empty / unparsable content.",
            part,
            part_type,
        )
1358
1359
        return None

Julien Denize's avatar
Julien Denize committed
1360
    if part_type in ("text", "input_text", "refusal", "thinking"):
1361
1362
        str_content = cast(str, content)
        if wrap_dicts:
1363
            return {"type": "text", "text": str_content}
1364
1365
        else:
            return str_content
1366

1367
1368
1369
1370
1371
1372
    # For media items, if a user has provided one, use it. Otherwise, insert
    # a placeholder empty uuid.
    uuid = part.get("uuid", None)
    if uuid is not None:
        uuid = str(uuid)

1373
    modality = None
1374
    if part_type == "image_pil":
1375
1376
1377
1378
        if content is not None:
            image_content = cast(Image.Image, content)
        else:
            image_content = None
1379
        mm_parser.parse_image_pil(image_content, uuid)
1380
        modality = "image"
1381
    elif part_type in ("image_url", "input_image"):
1382
        str_content = cast(str, content)
1383
        mm_parser.parse_image(str_content, uuid)
1384
1385
        modality = "image"
    elif part_type == "image_embeds":
1386
1387
1388
1389
        if content is not None:
            content = cast(Union[str, dict[str, str]], content)
        else:
            content = None
1390
        mm_parser.parse_image_embeds(content, uuid)
1391
1392
        modality = "image"
    elif part_type == "audio_url":
1393
        str_content = cast(str, content)
1394
        mm_parser.parse_audio(str_content, uuid)
1395
1396
        modality = "audio"
    elif part_type == "input_audio":
1397
        dict_content = cast(InputAudio, content)
1398
        mm_parser.parse_input_audio(dict_content, uuid)
1399
1400
        modality = "audio"
    elif part_type == "video_url":
1401
        str_content = cast(str, content)
1402
        mm_parser.parse_video(str_content, uuid)
1403
1404
1405
        modality = "video"
    else:
        raise NotImplementedError(f"Unknown part type: {part_type}")
1406

1407
1408
1409
1410
1411
1412
    return (
        {"type": modality}
        if wrap_dicts
        else (
            MODALITY_PLACEHOLDERS_MAP[modality] if interleave_strings else None
        )
1413
    )
1414
1415


1416
1417
1418
1419
1420
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


1421
def _parse_chat_message_content(
1422
1423
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
1424
    content_format: _ChatTemplateContentFormat,
1425
    interleave_strings: bool,
1426
) -> list[ConversationMessage]:
1427
1428
1429
1430
    role = message["role"]
    content = message.get("content")

    if content is None:
1431
1432
1433
1434
1435
1436
        content = []
    elif isinstance(content, str):
        content = [
            ChatCompletionContentPartTextParam(type="text", text=content)
        ]
    result = _parse_chat_message_content_parts(
1437
1438
        role,
        content,  # type: ignore
1439
        mm_tracker,
1440
        wrap_dicts=(content_format == "openai"),
1441
        interleave_strings=interleave_strings,
1442
    )
1443

1444
    for result_msg in result:
1445
        if role == "assistant":
1446
1447
            parsed_msg = _AssistantParser(message)

1448
1449
1450
            # The 'tool_calls' is not None check ensures compatibility.
            # It's needed only if downstream code doesn't strictly
            # follow the OpenAI spec.
1451
1452
1453
1454
            if (
                "tool_calls" in parsed_msg
                and parsed_msg["tool_calls"] is not None
            ):
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
                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

1466

1467
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1468
1469
1470
1471
1472
1473
    # 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:
1474
1475
1476
1477
1478
        if (
            message["role"] == "assistant"
            and "tool_calls" in message
            and isinstance(message["tool_calls"], list)
        ):
1479
            for item in message["tool_calls"]:
1480
1481
1482
1483
1484
                # if arguments is None or empty string, set to {}
                if content := item["function"].get("arguments"):
                    item["function"]["arguments"] = json.loads(content)
                else:
                    item["function"]["arguments"] = {}
1485
1486


1487
def parse_chat_messages(
1488
    messages: list[ChatCompletionMessageParam],
1489
    model_config: ModelConfig,
1490
    tokenizer: AnyTokenizer,
1491
    content_format: _ChatTemplateContentFormat,
1492
1493
1494
1495
1496
) -> tuple[
    list[ConversationMessage],
    Optional[MultiModalDataDict],
    Optional[MultiModalUUIDDict],
]:
1497
    conversation: list[ConversationMessage] = []
1498
    mm_tracker = MultiModalItemTracker(model_config, tokenizer)
1499
1500

    for msg in messages:
1501
1502
1503
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1504
            content_format,
1505
1506
1507
1508
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1509
            ),
1510
        )
1511

1512
        conversation.extend(sub_messages)
1513

1514
1515
    _postprocess_messages(conversation)

1516
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1517
1518


1519
def parse_chat_messages_futures(
1520
    messages: list[ChatCompletionMessageParam],
1521
1522
    model_config: ModelConfig,
    tokenizer: AnyTokenizer,
1523
    content_format: _ChatTemplateContentFormat,
1524
1525
1526
1527
1528
) -> tuple[
    list[ConversationMessage],
    Awaitable[Optional[MultiModalDataDict]],
    Optional[MultiModalUUIDDict],
]:
1529
    conversation: list[ConversationMessage] = []
1530
1531
1532
    mm_tracker = AsyncMultiModalItemTracker(model_config, tokenizer)

    for msg in messages:
1533
1534
1535
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1536
            content_format,
1537
1538
1539
1540
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1541
            ),
1542
        )
1543
1544
1545

        conversation.extend(sub_messages)

1546
1547
    _postprocess_messages(conversation)

1548
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1549
1550


1551
1552
def apply_hf_chat_template(
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
1553
    conversation: list[ConversationMessage],
1554
    chat_template: Optional[str],
1555
    tools: Optional[list[dict[str, Any]]],
1556
    *,
1557
    model_config: ModelConfig,
1558
1559
    tokenize: bool = False,  # Different from HF's default
    **kwargs: Any,
1560
) -> str:
1561
    hf_chat_template = resolve_hf_chat_template(
1562
1563
1564
        tokenizer,
        chat_template=chat_template,
        tools=tools,
1565
        model_config=model_config,
1566
    )
1567

1568
    if hf_chat_template is None:
1569
1570
1571
        raise ValueError(
            "As of transformers v4.44, default chat template is no longer "
            "allowed, so you must provide a chat template if the tokenizer "
1572
1573
            "does not define one."
        )
1574

1575
1576
1577
1578
1579
1580
1581
1582
    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,
        )
1583

1584
1585
1586
1587
1588
1589
    # 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(
1590
1591
            "An error occurred in `transformers` while applying chat template"
        )
1592
        raise ValueError(str(e)) from e
1593

1594

1595
1596
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1597
    messages: list[ChatCompletionMessageParam],
1598
1599
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
1600
    **kwargs: Any,
1601
) -> list[int]:
1602
1603
    from mistral_common.exceptions import MistralCommonException

1604
1605
1606
1607
1608
1609
    # 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,
    )
1610

1611
1612
1613
1614
1615
1616
1617
1618
1619
    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
1620
    # properly caught in the preprocessing_input step
1621
    except (AssertionError, MistralCommonException) as e:
1622
        raise ValueError(str(e)) from e
1623
1624
1625
1626
1627
1628
1629

    # 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(
1630
1631
            "An error occurred in `mistral_common` while applying chat template"
        )
1632
        raise ValueError(str(e)) from e
1633

1634

1635
1636
1637
def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
    idx = 0
    for msg in conversation:
1638
1639
1640
        if msg["role"] == "assistant":
            tool_calls = msg.get("tool_calls")
            idx += len(list(tool_calls)) if tool_calls is not None else 0  # noqa
1641
1642
1643
    return idx


1644
1645
1646
def make_tool_call_id(id_type: str = "random", func_name=None, idx=None):
    if id_type == "kimi_k2":
        return f"functions.{func_name}:{idx}"
1647
1648
1649
    else:
        # by default return random
        return f"chatcmpl-tool-{random_uuid()}"