llm/llm_factory.py

from typing import Dict, List

from .huggingfacehub.hf_model import HF_Mistaril, HF_TinyLlama
from .llamacpp.lc_model import LC_TinyLlama, LC_Phi3


class LLM_Factory:

    # trigger = {"model_type": "execution_type"} -> {"hf": "small"}
    @staticmethod
    def create_llm(prompt_entity: str, prompt_id: int, trigger: Dict[str, str]):

        if trigger.keys() == "hf" and trigger.values() == "effective":
            model = HF_Mistaril(prompt_entity=prompt_entity, prompt_id=prompt_id)
        elif trigger.keys() == "hf" and trigger.values() == "small":
            model = HF_TinyLlama(prompt_entity=prompt_entity, prompt_id=prompt_id)
        elif trigger.keys() == "lc" and trigger.values() == "effective":
            model = LC_Phi3(prompt_entity=prompt_entity, prompt_id=prompt_id)
        elif trigger.keys() == "lc" and trigger.values() == "small":
            model = LC_TinyLlama(prompt_entity=prompt_entity, prompt_id=prompt_id)
        else:
            model = None

        return model