chat_utils.py 53.1 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
def resolve_hf_chat_template(
425
426
427
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
428
429
    *,
    model_config: ModelConfig,
430
431
432
433
434
435
436
437
438
439
) -> Optional[str]:
    # 1st priority: The given chat template
    if chat_template is not None:
        return chat_template

    # 2nd priority: AutoProcessor chat template, unless tool calling is enabled
    if tools is None:
        try:
            processor = cached_get_processor(
                tokenizer.name_or_path,
440
441
442
443
444
                processor_cls=(
                    PreTrainedTokenizer,
                    PreTrainedTokenizerFast,
                    ProcessorMixin,
                ),
445
                trust_remote_code=model_config.trust_remote_code,
446
            )
447
448
449
450
451
            if (
                isinstance(processor, ProcessorMixin)
                and hasattr(processor, "chat_template")
                and processor.chat_template is not None
            ):
452
453
                return processor.chat_template
        except Exception:
454
455
456
457
458
            logger.debug(
                "Failed to load AutoProcessor chat template for %s",
                tokenizer.name_or_path,
                exc_info=True,
            )  # noqa: E501
459
460
461
462
463

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

470
471
472
473
474
475
    # 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:
476
477
478
479
480
        logger.info(
            "Loading chat template fallback for %s as there isn't one "
            "defined on HF Hub.",
            tokenizer.name_or_path,
        )
481
482
        chat_template = load_chat_template(path)
    else:
483
484
485
        logger.debug(
            "There is no chat template fallback for %s", tokenizer.name_or_path
        )
486
487

    return chat_template
488
489


490
491
def _resolve_chat_template_content_format(
    chat_template: Optional[str],
492
    tools: Optional[list[dict[str, Any]]],
493
    tokenizer: AnyTokenizer,
494
495
    *,
    model_config: ModelConfig,
496
497
) -> _ChatTemplateContentFormat:
    if isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)):
498
        hf_chat_template = resolve_hf_chat_template(
499
500
501
            tokenizer,
            chat_template=chat_template,
            tools=tools,
502
            model_config=model_config,
503
        )
504
    else:
505
506
        hf_chat_template = None

507
508
509
510
511
    jinja_text = (
        hf_chat_template
        if isinstance(hf_chat_template, str)
        else load_chat_template(chat_template, is_literal=True)
    )
512

513
514
515
516
517
    detected_format = (
        "string"
        if jinja_text is None
        else _detect_content_format(jinja_text, default="string")
    )
518

519
    return detected_format
520
521
522


@lru_cache
523
def _log_chat_template_content_format(
524
525
    chat_template: Optional[str],
    given_format: ChatTemplateContentFormatOption,
526
527
    detected_format: ChatTemplateContentFormatOption,
):
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
    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,
        )

545
546
547
548
549
550

def resolve_chat_template_content_format(
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
    given_format: ChatTemplateContentFormatOption,
    tokenizer: AnyTokenizer,
551
552
    *,
    model_config: ModelConfig,
553
) -> _ChatTemplateContentFormat:
554
555
556
    if given_format != "auto":
        return given_format

557
558
559
560
    detected_format = _resolve_chat_template_content_format(
        chat_template,
        tools,
        tokenizer,
561
        model_config=model_config,
562
563
564
565
566
567
568
569
    )

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

570
    return detected_format
571

572

573
ModalityStr = Literal["image", "audio", "video", "image_embeds"]
574
575
576
577
_T = TypeVar("_T")


class BaseMultiModalItemTracker(ABC, Generic[_T]):
578
579
580
581
582
583
584
    """
    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):
585
586
        super().__init__()

587
588
        self._model_config = model_config
        self._tokenizer = tokenizer
589

590
        self._items_by_modality = defaultdict[str, list[Optional[_T]]](list)
591
        self._uuids_by_modality = defaultdict[str, list[Optional[str]]](list)
592

593
594
595
596
    @property
    def model_config(self) -> ModelConfig:
        return self._model_config

597
    @cached_property
598
    def model_cls(self) -> type[SupportsMultiModal]:
599
        from vllm.model_executor.model_loader import get_model_cls
600

601
602
        model_cls = get_model_cls(self.model_config)
        return cast(type[SupportsMultiModal], model_cls)
603

604
605
606
607
    @property
    def allowed_local_media_path(self):
        return self._model_config.allowed_local_media_path

608
609
610
611
    @property
    def mm_registry(self):
        return MULTIMODAL_REGISTRY

612
613
614
615
    @cached_property
    def mm_processor(self):
        return self.mm_registry.create_processor(self.model_config)

616
    def add(
617
618
619
620
        self,
        modality: ModalityStr,
        item: Optional[_T],
        uuid: Optional[str] = None,
621
    ) -> Optional[str]:
622
623
624
        """
        Add a multi-modal item to the current prompt and returns the
        placeholder string to use, if any.
625
626

        An optional uuid can be added which serves as a unique identifier of the
627
        media.
628
        """
629
        input_modality = modality.replace("_embeds", "")
630
        num_items = len(self._items_by_modality[modality]) + 1
631

632
        self.mm_processor.validate_num_items(input_modality, num_items)
633

634
        self._items_by_modality[modality].append(item)
635
        self._uuids_by_modality[modality].append(uuid)
636

637
        return self.model_cls.get_placeholder_str(modality, num_items)
638

639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
    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

664
665
666
667
668
    @abstractmethod
    def create_parser(self) -> "BaseMultiModalContentParser":
        raise NotImplementedError


669
class MultiModalItemTracker(BaseMultiModalItemTracker[object]):
670
    def all_mm_data(self) -> Optional[MultiModalDataDict]:
671
672
673
674
675
        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:
676
677
678
            raise ValueError(
                "Mixing raw image and embedding inputs is not allowed"
            )
679
680
681
682

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
            if len(image_embeds_lst) > 1:
683
684
685
                raise ValueError(
                    "Only one message can have {'type': 'image_embeds'}"
                )
686
            mm_inputs["image"] = image_embeds_lst[0]
687
        if "image" in items_by_modality:
688
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
689
        if "audio" in items_by_modality:
690
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
691
        if "video" in items_by_modality:
692
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
693
        return mm_inputs
694
695
696
697
698

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


699
class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):
700
    async def all_mm_data(self) -> Optional[MultiModalDataDict]:
701
702
703
        if not self._items_by_modality:
            return None
        mm_inputs = {}
704
705
706
707
708
709
710
711
712
        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)
713

714
715
        if "image" in items_by_modality and "image_embeds" in items_by_modality:
            raise ValueError(
716
717
                "Mixing raw image and embedding inputs is not allowed"
            )
718
719
720
721
722

        if "image_embeds" in items_by_modality:
            image_embeds_lst = items_by_modality["image_embeds"]
            if len(image_embeds_lst) > 1:
                raise ValueError(
723
724
                    "Only one message can have {'type': 'image_embeds'}"
                )
725
            mm_inputs["image"] = image_embeds_lst[0]
726
        if "image" in items_by_modality:
727
            mm_inputs["image"] = items_by_modality["image"]  # A list of images
728
        if "audio" in items_by_modality:
729
            mm_inputs["audio"] = items_by_modality["audio"]  # A list of audios
730
        if "video" in items_by_modality:
731
            mm_inputs["video"] = items_by_modality["video"]  # A list of videos
732
        return mm_inputs
733
734
735
736
737
738
739
740
741

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


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

742
        # stores model placeholders list with corresponding
743
744
745
746
747
748
749
        # general MM placeholder:
        # {
        #   "<##IMAGE##>": ["<image>", "<image>", "<image>"],
        #   "<##AUDIO##>": ["<audio>", "<audio>"]
        # }
        self._placeholder_storage: dict[str, list] = defaultdict(list)

750
751
752
    def _add_placeholder(
        self, modality: ModalityStr, placeholder: Optional[str]
    ):
753
        mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]
754
        if placeholder:
755
            self._placeholder_storage[mod_placeholder].append(placeholder)
756

757
758
    def mm_placeholder_storage(self) -> dict[str, list]:
        return dict(self._placeholder_storage)
759
760

    @abstractmethod
761
762
    def parse_image(
        self, image_url: Optional[str], uuid: Optional[str] = None) -> None:
763
764
        raise NotImplementedError

765
    @abstractmethod
766
    def parse_image_embeds(
767
        self,
768
        image_embeds: Union[str, dict[str, str], None],
769
        uuid: Optional[str] = None,
770
    ) -> None:
771
772
        raise NotImplementedError

773
    @abstractmethod
774
    def parse_image_pil(
775
        self, image_pil: Optional[Image.Image], uuid: Optional[str] = None
776
    ) -> None:
777
778
        raise NotImplementedError

779
    @abstractmethod
780
781
782
    def parse_audio(
        self, audio_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
783
784
        raise NotImplementedError

785
    @abstractmethod
786
    def parse_input_audio(
787
        self, input_audio: Optional[InputAudio], uuid: Optional[str] = None
788
    ) -> None:
789
790
        raise NotImplementedError

791
    @abstractmethod
792
793
794
    def parse_video(
        self, video_url: Optional[str], uuid: Optional[str] = None
    ) -> None:
795
796
        raise NotImplementedError

797
798
799
800
801
802
803

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

        self._tracker = tracker

804
        self._connector = MediaConnector(
805
            media_io_kwargs=self._tracker._model_config.media_io_kwargs,
806
807
808
            allowed_local_media_path=tracker.allowed_local_media_path,
        )

809
810
811
812
    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
813

814
        placeholder = self._tracker.add("image", image, uuid)
815
        self._add_placeholder("image", placeholder)
816

817
    def parse_image_embeds(
818
        self,
819
        image_embeds: Union[str, dict[str, str], None],
820
        uuid: Optional[str] = None,
821
    ) -> None:
822
823
824
825
826
        if isinstance(image_embeds, dict):
            embeds = {
                k: self._connector.fetch_image_embedding(v)
                for k, v in image_embeds.items()
            }
827
            placeholder = self._tracker.add("image_embeds", embeds, uuid)
828
829
830

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

833
834
835
        if image_embeds is None:
            placeholder = self._tracker.add("image_embeds", None, uuid)

836
        self._add_placeholder("image", placeholder)
837

838
    def parse_image_pil(
839
        self, image_pil: Optional[Image.Image], uuid: Optional[str] = None
840
841
    ) -> None:
        placeholder = self._tracker.add("image", image_pil, uuid)
842
        self._add_placeholder("image", placeholder)
843

844
845
846
847
    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
848

849
        placeholder = self._tracker.add("audio", audio, uuid)
850
        self._add_placeholder("audio", placeholder)
851

852
    def parse_input_audio(
853
        self, input_audio: Optional[InputAudio], uuid: Optional[str] = None
854
    ) -> None:
855
856
857
858
859
860
861
862
863
864
        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
865

866
        return self.parse_audio(audio_url, uuid)
867

868
869
870
871
872
873
874
875
    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
        )
876

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

880
881
882
883
884
885

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

        self._tracker = tracker
886
        self._connector = MediaConnector(
887
            media_io_kwargs=self._tracker._model_config.media_io_kwargs,
888
            allowed_local_media_path=tracker.allowed_local_media_path,
889
        )
890

891
892
893
894
895
896
    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
        )
897

898
        placeholder = self._tracker.add("image", image_coro, uuid)
899
        self._add_placeholder("image", placeholder)
900

901
    def parse_image_embeds(
902
        self,
903
        image_embeds: Union[str, dict[str, str], None],
904
        uuid: Optional[str] = None,
905
    ) -> None:
906
907
908
        future: asyncio.Future[Union[str, dict[str, str], None]] = (
            asyncio.Future()
        )
909
910
911
912
913
914
915
916
917

        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):
918
            embedding = self._connector.fetch_image_embedding(image_embeds)
919
920
            future.set_result(embedding)

921
922
923
        if image_embeds is None:
            future.set_result(None)

924
        placeholder = self._tracker.add("image_embeds", future, uuid)
925
        self._add_placeholder("image", placeholder)
926

927
    def parse_image_pil(
928
        self, image_pil: Optional[Image.Image], uuid: Optional[str] = None
929
    ) -> None:
930
931
932
933
934
        future: asyncio.Future[Optional[Image.Image]] = asyncio.Future()
        if image_pil:
            future.set_result(image_pil)
        else:
            future.set_result(None)
935

936
        placeholder = self._tracker.add("image", future, uuid)
937
        self._add_placeholder("image", placeholder)
938

939
940
941
942
943
944
    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
        )
945

946
        placeholder = self._tracker.add("audio", audio_coro, uuid)
947
        self._add_placeholder("audio", placeholder)
948

949
    def parse_input_audio(
950
        self, input_audio: Optional[InputAudio], uuid: Optional[str] = None
951
    ) -> None:
952
953
954
955
956
957
958
959
960
961
        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
962

963
        return self.parse_audio(audio_url, uuid)
964

965
966
967
968
969
970
971
972
    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
        )
973

974
        placeholder = self._tracker.add("video", video, uuid)
975
        self._add_placeholder("video", placeholder)
976

977

978
979
980
981
982
983
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():
984
        raise FileNotFoundError("the supplied chat template path doesn't exist")
985
986
987

    elif isinstance(chat_template, str):
        JINJA_CHARS = "{}\n"
988
989
990
991
        if (
            not any(c in chat_template for c in JINJA_CHARS)
            and not Path(chat_template).exists()
        ):
992
993
            raise ValueError(
                f"The supplied chat template string ({chat_template}) "
994
995
                f"appears path-like, but doesn't exist!"
            )
996
997
998

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


1003
def _load_chat_template(
1004
1005
1006
1007
    chat_template: Optional[Union[Path, str]],
    *,
    is_literal: bool = False,
) -> Optional[str]:
1008
1009
    if chat_template is None:
        return None
1010
1011
1012

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

1017
        return chat_template
1018

1019
    try:
1020
        with open(chat_template) as f:
1021
            return f.read()
1022
    except OSError as e:
1023
1024
1025
        if isinstance(chat_template, Path):
            raise

1026
1027
        JINJA_CHARS = "{}\n"
        if not any(c in chat_template for c in JINJA_CHARS):
1028
1029
1030
1031
1032
            msg = (
                f"The supplied chat template ({chat_template}) "
                f"looks like a file path, but it failed to be "
                f"opened. Reason: {e}"
            )
1033
            raise ValueError(msg) from e
1034

1035
1036
        # If opening a file fails, set chat template to be args to
        # ensure we decode so our escape are interpreted correctly
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
        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)
1049
1050


1051
1052
1053
def _get_interleaved_text_prompt(
    placeholder_storage: dict[str, list], texts: list[str]
) -> str:
1054
1055
1056
1057
1058
1059
1060
    for idx, elem in enumerate(texts):
        if elem in placeholder_storage:
            texts[idx] = placeholder_storage[elem].pop(0)

    return "\n".join(texts)


1061
# TODO: Let user specify how to insert multimodal tokens into prompt
1062
# (similar to chat template)
1063
1064
1065
1066
1067
def _get_full_multimodal_text_prompt(
    placeholder_storage: dict[str, list],
    texts: list[str],
    interleave_strings: bool,
) -> str:
1068
    """Combine multimodal prompts for a multimodal language model."""
1069

1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
    # 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

1087
    # Look through the text prompt to check for missing placeholders
1088
    missing_placeholders: list[str] = []
1089
1090
1091
1092
1093
    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:
1094
1095
1096
1097
            logger.error(
                "Placeholder count is negative! "
                "Ensure that the 'interleave_strings' flag is disabled "
                "(current value: %s) "
1098
1099
                "when manually placing image placeholders.",
                interleave_strings,
1100
1101
            )
            logger.debug("Input prompt: %s", text_prompt)
1102
1103
            raise ValueError(
                f"Found more '{placeholder}' placeholders in input prompt than "
1104
1105
                "actual multimodal data items."
            )
1106

1107
1108
1109
        missing_placeholders.extend(
            [placeholder] * placeholder_counts[placeholder]
        )
1110

1111
1112
    # NOTE: Default behaviour: we always add missing placeholders
    # at the front of the prompt, if interleave_strings=False
1113
    return "\n".join(missing_placeholders + [text_prompt])
1114
1115


1116
1117
# No need to validate using Pydantic again
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
1118
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
1119
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
1120
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
1121
_PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
Julien Denize's avatar
Julien Denize committed
1122
_ThinkParser = partial(cast, CustomThinkCompletionContentParam)
1123
1124
1125
1126
# Need to validate url objects
_ImageParser = TypeAdapter(ChatCompletionContentPartImageParam).validate_python
_AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python
_VideoParser = TypeAdapter(ChatCompletionContentPartVideoParam).validate_python
1127

1128
_ResponsesInputImageParser = TypeAdapter(
1129
1130
    ResponseInputImageParam
).validate_python
1131
_ContentPart: TypeAlias = Union[str, dict[str, str], InputAudio, PILImage]
1132

1133
# Define a mapping from part types to their corresponding parsing functions.
1134
MM_PARSER_MAP: dict[
1135
1136
1137
    str,
    Callable[[ChatCompletionContentPartParam], _ContentPart],
] = {
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
    "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
    ),
1150
    "image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
    "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),
1161
1162
1163
1164
}


def _parse_chat_message_content_mm_part(
1165
1166
    part: ChatCompletionContentPartParam,
) -> tuple[str, _ContentPart]:
1167
    """
1168
    Parses a given multi-modal content part based on its type.
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181

    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(
1182
1183
        part, dict
    )  # This is needed to avoid mypy errors: part.get() from str
1184
    part_type = part.get("type", None)
1185
    uuid = part.get("uuid", None)
1186

1187
    if isinstance(part_type, str) and part_type in MM_PARSER_MAP and uuid is None: # noqa: E501
1188
1189
1190
        content = MM_PARSER_MAP[part_type](part)

        # Special case for 'image_url.detail'
1191
1192
        # We only support 'auto', which is the default
        if part_type == "image_url" and part.get("detail", "auto") != "auto":
1193
1194
1195
1196
            logger.warning(
                "'image_url.detail' is currently not supported "
                "and will be ignored."
            )
1197
1198
1199
1200

        return part_type, content

    # Handle missing 'type' but provided direct URL fields.
1201
    # 'type' is required field by pydantic
1202
1203
    if part_type is None or uuid is not None:
        if "image_url" in part:
1204
1205
1206
            image_params = cast(
                CustomChatCompletionContentSimpleImageParam, part
            )
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
            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:
1228
1229
1230
            audio_params = cast(
                CustomChatCompletionContentSimpleAudioParam, part
            )
1231
1232
1233
1234
1235
1236
            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
1237
        if part.get("input_audio") is not None:
1238
            input_audio_params = cast(dict[str, str], part)
1239
            return "input_audio", input_audio_params
1240
        if "video_url" in part:
1241
1242
1243
            video_params = cast(
                CustomChatCompletionContentSimpleVideoParam, part
            )
1244
1245
1246
1247
1248
1249
            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
1250
1251
1252
1253
1254
1255
1256
1257
        # 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"


1258
PART_TYPES_TO_SKIP_NONE_CONTENT = (
1259
1260
1261
    "text",
    "refusal",
)
1262

1263

1264
1265
1266
def _parse_chat_message_content_parts(
    role: str,
    parts: Iterable[ChatCompletionContentPartParam],
1267
    mm_tracker: BaseMultiModalItemTracker,
1268
1269
    *,
    wrap_dicts: bool,
1270
    interleave_strings: bool,
1271
) -> list[ConversationMessage]:
1272
    content = list[_ContentPart]()
1273

1274
    mm_parser = mm_tracker.create_parser()
1275
1276

    for part in parts:
1277
        parse_res = _parse_chat_message_content_part(
1278
1279
1280
            part,
            mm_parser,
            wrap_dicts=wrap_dicts,
1281
            interleave_strings=interleave_strings,
1282
        )
1283
1284
        if parse_res:
            content.append(parse_res)
1285

1286
    if wrap_dicts:
1287
        # Parsing wraps images and texts as interleaved dictionaries
1288
        return [ConversationMessage(role=role, content=content)]  # type: ignore
1289
    texts = cast(list[str], content)
1290
1291
    mm_placeholder_storage = mm_parser.mm_placeholder_storage()
    if mm_placeholder_storage:
1292
1293
1294
        text_prompt = _get_full_multimodal_text_prompt(
            mm_placeholder_storage, texts, interleave_strings
        )
1295
1296
1297
    else:
        text_prompt = "\n".join(texts)

1298
1299
1300
1301
    return [ConversationMessage(role=role, content=text_prompt)]


def _parse_chat_message_content_part(
1302
1303
1304
1305
    part: ChatCompletionContentPartParam,
    mm_parser: BaseMultiModalContentParser,
    *,
    wrap_dicts: bool,
1306
    interleave_strings: bool,
1307
) -> Optional[_ContentPart]:
1308
1309
1310
1311
1312
1313
1314
1315
    """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
1316
        return part
1317
1318
    # Handle structured dictionary parts
    part_type, content = _parse_chat_message_content_mm_part(part)
1319
    # if part_type is text/refusal/image_url/audio_url/video_url/input_audio but
1320
    # content is None, log a warning and skip
1321
    if part_type in PART_TYPES_TO_SKIP_NONE_CONTENT and content is None:
1322
        logger.warning(
1323
            "Skipping multimodal part '%s' (type: '%s') "
1324
1325
1326
1327
            "with empty / unparsable content.",
            part,
            part_type,
        )
1328
1329
        return None

Julien Denize's avatar
Julien Denize committed
1330
    if part_type in ("text", "input_text", "refusal", "thinking"):
1331
1332
        str_content = cast(str, content)
        if wrap_dicts:
1333
            return {"type": "text", "text": str_content}
1334
1335
        else:
            return str_content
1336

1337
1338
1339
1340
1341
1342
    # 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)

1343
    modality = None
1344
    if part_type == "image_pil":
1345
1346
1347
1348
        if content is not None:
            image_content = cast(Image.Image, content)
        else:
            image_content = None
1349
        mm_parser.parse_image_pil(image_content, uuid)
1350
        modality = "image"
1351
    elif part_type in ("image_url", "input_image"):
1352
        str_content = cast(str, content)
1353
        mm_parser.parse_image(str_content, uuid)
1354
1355
        modality = "image"
    elif part_type == "image_embeds":
1356
1357
1358
1359
        if content is not None:
            content = cast(Union[str, dict[str, str]], content)
        else:
            content = None
1360
        mm_parser.parse_image_embeds(content, uuid)
1361
1362
        modality = "image"
    elif part_type == "audio_url":
1363
        str_content = cast(str, content)
1364
        mm_parser.parse_audio(str_content, uuid)
1365
1366
        modality = "audio"
    elif part_type == "input_audio":
1367
        dict_content = cast(InputAudio, content)
1368
        mm_parser.parse_input_audio(dict_content, uuid)
1369
1370
        modality = "audio"
    elif part_type == "video_url":
1371
        str_content = cast(str, content)
1372
        mm_parser.parse_video(str_content, uuid)
1373
1374
1375
        modality = "video"
    else:
        raise NotImplementedError(f"Unknown part type: {part_type}")
1376

1377
1378
1379
1380
1381
1382
    return (
        {"type": modality}
        if wrap_dicts
        else (
            MODALITY_PLACEHOLDERS_MAP[modality] if interleave_strings else None
        )
1383
    )
1384
1385


1386
1387
1388
1389
1390
# No need to validate using Pydantic again
_AssistantParser = partial(cast, ChatCompletionAssistantMessageParam)
_ToolParser = partial(cast, ChatCompletionToolMessageParam)


1391
def _parse_chat_message_content(
1392
1393
    message: ChatCompletionMessageParam,
    mm_tracker: BaseMultiModalItemTracker,
1394
    content_format: _ChatTemplateContentFormat,
1395
    interleave_strings: bool,
1396
) -> list[ConversationMessage]:
1397
1398
1399
1400
    role = message["role"]
    content = message.get("content")

    if content is None:
1401
1402
1403
1404
1405
1406
        content = []
    elif isinstance(content, str):
        content = [
            ChatCompletionContentPartTextParam(type="text", text=content)
        ]
    result = _parse_chat_message_content_parts(
1407
1408
        role,
        content,  # type: ignore
1409
        mm_tracker,
1410
        wrap_dicts=(content_format == "openai"),
1411
        interleave_strings=interleave_strings,
1412
    )
1413

1414
    for result_msg in result:
1415
        if role == "assistant":
1416
1417
            parsed_msg = _AssistantParser(message)

1418
1419
1420
            # The 'tool_calls' is not None check ensures compatibility.
            # It's needed only if downstream code doesn't strictly
            # follow the OpenAI spec.
1421
1422
1423
1424
            if (
                "tool_calls" in parsed_msg
                and parsed_msg["tool_calls"] is not None
            ):
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
                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

1436

1437
def _postprocess_messages(messages: list[ConversationMessage]) -> None:
1438
1439
1440
1441
1442
1443
    # 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:
1444
1445
1446
1447
1448
        if (
            message["role"] == "assistant"
            and "tool_calls" in message
            and isinstance(message["tool_calls"], list)
        ):
1449
1450
            for item in message["tool_calls"]:
                item["function"]["arguments"] = json.loads(
1451
1452
                    item["function"]["arguments"]
                )
1453
1454


1455
def parse_chat_messages(
1456
    messages: list[ChatCompletionMessageParam],
1457
    model_config: ModelConfig,
1458
    tokenizer: AnyTokenizer,
1459
    content_format: _ChatTemplateContentFormat,
1460
1461
1462
1463
1464
) -> tuple[
    list[ConversationMessage],
    Optional[MultiModalDataDict],
    Optional[MultiModalUUIDDict],
]:
1465
    conversation: list[ConversationMessage] = []
1466
    mm_tracker = MultiModalItemTracker(model_config, tokenizer)
1467
1468

    for msg in messages:
1469
1470
1471
        sub_messages = _parse_chat_message_content(
            msg,
            mm_tracker,
1472
            content_format,
1473
1474
1475
1476
            interleave_strings=(
                content_format == "string"
                and model_config.multimodal_config is not None
                and model_config.multimodal_config.interleave_mm_strings
1477
            ),
1478
        )
1479

1480
        conversation.extend(sub_messages)
1481

1482
1483
    _postprocess_messages(conversation)

1484
    return conversation, mm_tracker.all_mm_data(), mm_tracker.all_mm_uuids()
1485
1486


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

    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
1513

        conversation.extend(sub_messages)

1514
1515
    _postprocess_messages(conversation)

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


1519
1520
def apply_hf_chat_template(
    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
1521
    conversation: list[ConversationMessage],
1522
    chat_template: Optional[str],
1523
    tools: Optional[list[dict[str, Any]]],
1524
    *,
1525
    model_config: ModelConfig,
1526
1527
    tokenize: bool = False,  # Different from HF's default
    **kwargs: Any,
1528
) -> str:
1529
    hf_chat_template = resolve_hf_chat_template(
1530
1531
1532
        tokenizer,
        chat_template=chat_template,
        tools=tools,
1533
        model_config=model_config,
1534
    )
1535

1536
    if hf_chat_template is None:
1537
1538
1539
        raise ValueError(
            "As of transformers v4.44, default chat template is no longer "
            "allowed, so you must provide a chat template if the tokenizer "
1540
1541
            "does not define one."
        )
1542

1543
1544
1545
1546
1547
1548
1549
1550
    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,
        )
1551

1552
1553
1554
1555
1556
1557
    # 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(
1558
1559
            "An error occurred in `transformers` while applying chat template"
        )
1560
        raise ValueError(str(e)) from e
1561

1562

1563
1564
def apply_mistral_chat_template(
    tokenizer: MistralTokenizer,
1565
    messages: list[ChatCompletionMessageParam],
1566
1567
    chat_template: Optional[str],
    tools: Optional[list[dict[str, Any]]],
1568
    **kwargs: Any,
1569
) -> list[int]:
1570
1571
    from mistral_common.exceptions import MistralCommonException

1572
1573
1574
1575
1576
1577
    # 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,
    )
1578

1579
1580
1581
1582
1583
1584
1585
1586
1587
    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
1588
    # properly caught in the preprocessing_input step
1589
    except (AssertionError, MistralCommonException) as e:
1590
        raise ValueError(str(e)) from e
1591
1592
1593
1594
1595
1596
1597

    # 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(
1598
1599
            "An error occurred in `mistral_common` while applying chat template"
        )
1600
        raise ValueError(str(e)) from e
1601

1602

1603
1604
1605
def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
    idx = 0
    for msg in conversation:
1606
1607
1608
        if msg["role"] == "assistant":
            tool_calls = msg.get("tool_calls")
            idx += len(list(tool_calls)) if tool_calls is not None else 0  # noqa
1609
1610
1611
    return idx


1612
1613
1614
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}"
1615
1616
1617
    else:
        # by default return random
        return f"chatcmpl-tool-{random_uuid()}"