llm_factory.py

from llama_index.llms.openai_like import OpenAILike
from llama_index.llms.openai import OpenAI
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from args import InterfaceChoice, Args


class LLMFactory():

    @classmethod
    def create(cls, system_prompt, temperature = None, max_tokens = None):
        if Args.INTERFACE == InterfaceChoice.OPENAILIKE:
            return cls._openailike_create(system_prompt, temperature, max_tokens)
        elif Args.INTERFACE == InterfaceChoice.OPENAI:
            return cls._openai_create(system_prompt, temperature, max_tokens)
        elif Args.INTERFACE == InterfaceChoice.HUGGINGFACE:
            return cls._hf_create(system_prompt, temperature, max_tokens)
        else:
            raise ValueError(f"Invalid interface choice: {Args.INTERFACE}")

    @staticmethod
    def _openailike_create(system_prompt, temperature=None, max_tokens=None):
        kwargs = {
            "model": Args.model_name,
            "api_base": Args.api_base,
            "api_key": Args.api_key,
            "system_prompt": system_prompt,
        }

        if temperature is not None:
            kwargs["temperature"] = temperature

        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens

        llm = OpenAILike(**kwargs)
        return llm

    @staticmethod
    def _openai_create(system_prompt, temperature = None, max_tokens = None):
        kwargs = {
            "model": Args.model_name,
            "api_key": Args.api_key,
            "system_prompt": system_prompt,
        }

        if temperature is not None:
            kwargs["temperature"] = temperature

        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens

        llm = OpenAI(**kwargs)
        return llm

    @staticmethod
    def _hf_create(system_prompt, temperature = None, max_tokens = None):
        kwargs = {
            "model_name": Args.model_name,
            "system_prompt": system_prompt,
            "token": Args.token,
        }

        if temperature is not None:
            kwargs["temperature"] = temperature

        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens

        llm = HuggingFaceInferenceAPI(**kwargs)
        return llm