llama-index

1
import base64
2
from typing import Any, Dict, Union
3

4
from llama_index.legacy.llms import ChatMessage, MessageRole
5

6

7
def is_gemini_model(model: str) -> bool:
8
    return model.startswith("gemini")
9

10

11
def create_gemini_client(model: str) -> Any:
12
    from vertexai.preview.generative_models import GenerativeModel
13

14
    return GenerativeModel(model_name=model)
15

16

17
def convert_chat_message_to_gemini_content(
18
    message: ChatMessage, is_history: bool = True
19
) -> Any:
20
    from vertexai.preview.generative_models import Content, Part
21

22
    def _convert_gemini_part_to_prompt(part: Union[str, Dict]) -> Part:
23
        from vertexai.preview.generative_models import Image, Part
24

25
        if isinstance(part, str):
26
            return Part.from_text(part)
27

28
        if not isinstance(part, Dict):
29
            raise ValueError(
30
                f"Message's content is expected to be a dict, got {type(part)}!"
31
            )
32
        if part["type"] == "text":
33
            return Part.from_text(part["text"])
34
        elif part["type"] == "image_url":
35
            path = part["image_url"]
36
            if path.startswith("gs://"):
37
                raise ValueError("Only local image path is supported!")
38
            elif path.startswith("data:image/jpeg;base64,"):
39
                image = Image.from_bytes(base64.b64decode(path[23:]))
40
            else:
41
                image = Image.load_from_file(path)
42
        else:
43
            raise ValueError("Only text and image_url types are supported!")
44
        return Part.from_image(image)
45

46
    raw_content = message.content
47
    if raw_content is None:
48
        raw_content = ""
49
    if isinstance(raw_content, str):
50
        raw_content = [raw_content]
51
    parts = [_convert_gemini_part_to_prompt(part) for part in raw_content]
52
    if is_history:
53
        return Content(
54
            role="user" if message.role == MessageRole.USER else "model",
55
            parts=parts,
56
        )
57
    else:
58
        return parts
59