medrag_multi_modal/assistant/llm_client.py

import json
import os
from enum import Enum
from typing import Any, Optional, Union

import instructor
import weave
from PIL import Image

from ..utils import base64_encode_image


class ClientType(str, Enum):
    GEMINI = "gemini"
    MISTRAL = "mistral"
    OPENAI = "openai"


GOOGLE_MODELS = [
    "gemini-1.0-pro-latest",
    "gemini-1.0-pro",
    "gemini-pro",
    "gemini-1.0-pro-001",
    "gemini-1.0-pro-vision-latest",
    "gemini-pro-vision",
    "gemini-1.5-pro-latest",
    "gemini-1.5-pro-001",
    "gemini-1.5-pro-002",
    "gemini-1.5-pro",
    "gemini-1.5-pro-exp-0801",
    "gemini-1.5-pro-exp-0827",
    "gemini-1.5-flash-latest",
    "gemini-1.5-flash-001",
    "gemini-1.5-flash-001-tuning",
    "gemini-1.5-flash",
    "gemini-1.5-flash-exp-0827",
    "gemini-1.5-flash-002",
    "gemini-1.5-flash-8b",
    "gemini-1.5-flash-8b-001",
    "gemini-1.5-flash-8b-latest",
    "gemini-1.5-flash-8b-exp-0827",
    "gemini-1.5-flash-8b-exp-0924",
]

MISTRAL_MODELS = [
    "ministral-3b-latest",
    "ministral-8b-latest",
    "mistral-large-latest",
    "mistral-small-latest",
    "codestral-latest",
    "pixtral-12b-2409",
    "open-mistral-nemo",
    "open-codestral-mamba",
    "open-mistral-7b",
    "open-mixtral-8x7b",
    "open-mixtral-8x22b",
]

OPENAI_MODELS = ["gpt-4o", "gpt-4o-2024-08-06", "gpt-4o-mini", "gpt-4o-mini-2024-07-18"]


class LLMClient(weave.Model):
    """
    LLMClient is a class that interfaces with different large language model (LLM) providers
    such as Google Gemini, Mistral, and OpenAI. It abstracts the complexity of interacting with
    these different APIs and provides a unified interface for making predictions.

    Args:
        model_name (str): The name of the model to be used for predictions.
        client_type (Optional[ClientType]): The type of client (e.g., GEMINI, MISTRAL, OPENAI).
            If not provided, it is inferred from the model_name.
    """

    model_name: str
    client_type: Optional[ClientType]

    def __init__(self, model_name: str, client_type: Optional[ClientType] = None):
        if client_type is None:
            if model_name in GOOGLE_MODELS:
                client_type = ClientType.GEMINI
            elif model_name in MISTRAL_MODELS:
                client_type = ClientType.MISTRAL
            elif model_name in OPENAI_MODELS:
                client_type = ClientType.OPENAI
            else:
                raise ValueError(f"Invalid model name: {model_name}")
        super().__init__(model_name=model_name, client_type=client_type)

    @weave.op()
    def execute_gemini_sdk(
        self,
        user_prompt: Union[str, list[str]],
        system_prompt: Optional[Union[str, list[str]]] = None,
        schema: Optional[Any] = None,
    ) -> Union[str, Any]:
        import google.generativeai as genai
        from google.generativeai.types import HarmBlockThreshold, HarmCategory

        system_prompt = (
            [system_prompt] if isinstance(system_prompt, str) else system_prompt
        )
        user_prompt = [user_prompt] if isinstance(user_prompt, str) else user_prompt

        genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
        model = genai.GenerativeModel(self.model_name, system_instruction=system_prompt)
        generation_config = (
            None
            if schema is None
            else genai.GenerationConfig(
                response_mime_type="application/json", response_schema=schema
            )
        )
        response = model.generate_content(
            user_prompt,
            generation_config=generation_config,
            # This is necessary in order to answer questions about anatomy, sexual diseases,
            # medical devices, medicines, etc.
            safety_settings={
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
            },
        )
        return response.text if schema is None else json.loads(response.text)

    @weave.op()
    def execute_mistral_sdk(
        self,
        user_prompt: Union[str, list[str]],
        system_prompt: Optional[Union[str, list[str]]] = None,
        schema: Optional[Any] = None,
    ) -> Union[str, Any]:
        from mistralai import Mistral

        system_prompt = (
            [system_prompt] if isinstance(system_prompt, str) else system_prompt
        )
        user_prompt = [user_prompt] if isinstance(user_prompt, str) else user_prompt
        system_messages = [{"type": "text", "text": prompt} for prompt in system_prompt]
        user_messages = []
        for prompt in user_prompt:
            if isinstance(prompt, Image.Image):
                user_messages.append(
                    {
                        "type": "image_url",
                        "image_url": base64_encode_image(prompt, "image/png"),
                    }
                )
            else:
                user_messages.append({"type": "text", "text": prompt})
        messages = [
            {"role": "system", "content": system_messages},
            {"role": "user", "content": user_messages},
        ]

        client = Mistral(api_key=os.environ.get("MISTRAL_API_KEY"))
        client = instructor.from_mistral(client) if schema is not None else client

        if schema is None:
            raise NotImplementedError(
                "Mistral does not support structured output using a schema"
            )
        else:
            response = client.chat.complete(model=self.model_name, messages=messages)
            return response.choices[0].message.content

    @weave.op()
    def execute_openai_sdk(
        self,
        user_prompt: Union[str, list[str]],
        system_prompt: Optional[Union[str, list[str]]] = None,
        schema: Optional[Any] = None,
    ) -> Union[str, Any]:
        from openai import OpenAI

        system_prompt = (
            [system_prompt] if isinstance(system_prompt, str) else system_prompt
        )
        user_prompt = [user_prompt] if isinstance(user_prompt, str) else user_prompt

        system_messages = [
            {"role": "system", "content": prompt} for prompt in system_prompt
        ]
        user_messages = []
        for prompt in user_prompt:
            if isinstance(prompt, Image.Image):
                user_messages.append(
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": base64_encode_image(prompt, "image/png"),
                        },
                    },
                )
            else:
                user_messages.append({"type": "text", "text": prompt})
        messages = system_messages + [{"role": "user", "content": user_messages}]

        client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

        if schema is None:
            completion = client.chat.completions.create(
                model=self.model_name, messages=messages
            )
            return completion.choices[0].message.content

        completion = weave.op()(client.beta.chat.completions.parse)(
            model=self.model_name, messages=messages, response_format=schema
        )
        return completion.choices[0].message.parsed

    @weave.op()
    def predict(
        self,
        user_prompt: Union[str, list[str]],
        system_prompt: Optional[Union[str, list[str]]] = None,
        schema: Optional[Any] = None,
    ) -> Union[str, Any]:
        """
        Predicts the response from a language model based on the provided prompts and schema.

        This function determines the client type and calls the appropriate SDK execution function
        to get the response from the language model. It supports multiple client types including
        GEMINI, MISTRAL, and OPENAI. Depending on the client type, it calls the corresponding
        execution function with the provided user and system prompts, and an optional schema.

        Args:
            user_prompt (Union[str, list[str]]): The user prompt(s) to be sent to the language model.
            system_prompt (Optional[Union[str, list[str]]]): The system prompt(s) to be sent to the language model.
            schema (Optional[Any]): The schema to be used for parsing the response, if applicable.

        Returns:
            Union[str, Any]: The response from the language model, which could be a string or any other type
            depending on the schema provided.

        Raises:
            ValueError: If the client type is invalid.
        """
        if self.client_type == ClientType.GEMINI:
            return self.execute_gemini_sdk(user_prompt, system_prompt, schema)
        elif self.client_type == ClientType.MISTRAL:
            return self.execute_mistral_sdk(user_prompt, system_prompt, schema)
        elif self.client_type == ClientType.OPENAI:
            return self.execute_openai_sdk(user_prompt, system_prompt, schema)
        else:
            raise ValueError(f"Invalid client type: {self.client_type}")