# imports
from loguru import logger
# LLM modules
from langchain_community.llms.huggingface_hub import HuggingFaceHub
from langchain_community.llms.ollama import Ollama
from langchain_openai import ChatOpenAI, AzureChatOpenAI
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
# local imports
import settings_template as settings


class LLMCreator():
    """
    LLM class to import into other modules
    """
    def __init__(self, llm_type=None, llm_model_type=None, local_api_url=None, azureopenai_api_version=None) -> None:
        self.llm_type = settings.LLM_TYPE if llm_type is None else llm_type
        self.llm_model_type = settings.LLM_MODEL_TYPE if llm_model_type is None else llm_model_type
        self.local_api_url = settings.API_URL if local_api_url is None else local_api_url
        self.azureopenai_api_version = settings.AZUREOPENAI_API_VERSION \
            if azureopenai_api_version is None and settings.AZUREOPENAI_API_VERSION is not None \
            else azureopenai_api_version


    def get_llm(self):
        """
        returns, based on settings, the llm object
        """
        # if llm_type is "chatopenai"
        if self.llm_type == "chatopenai":
            # default llm_model_type value is "gpt-3.5-turbo"
            self.llm_model_type = "gpt-3.5-turbo"
            if self.llm_model_type == "gpt35_16":
                self.llm_model_type = "gpt-3.5-turbo-16k"
            elif self.llm_model_type == "gpt4":
                self.llm_model_type = "gpt-4"
            self.llm = ChatOpenAI(
                client=None,
                model=self.llm_model_type,
                temperature=0,
            )
        # else, if llm_type is "huggingface"
        elif self.llm_type == "huggingface":
            # default value is llama-2, with maximum output length 512
            self.llm_model_type = "meta-llama/Llama-2-7b-chat-hf"
            max_length = 512
            if self.llm_model_type == 'GoogleFlan':
                self.llm_model_type = 'google/flan-t5-base'
                max_length = 512
            self.llm = HuggingFaceHub(repo_id=self.llm_model_type,
                                      model_kwargs={"temperature": 0.1,
                                                    "max_length": max_length}
                                      )
        # else, if llm_type is "local_llm"
        elif self.llm_type == "local_llm":
            logger.info("Use Local LLM")
            logger.info("Retrieving " + self.llm_model_type)
            # If API URL is defined, use it
            if self.local_api_url is not None:
                logger.info("Using local api url " + self.local_api_url)
                self.llm = Ollama(
                    model=self.llm_model_type,
                    base_url=self.local_api_url,
                    callback_manager=CallbackManager([StreamingStdOutCallbackHandler()])
                )
            else:
                self.llm = Ollama(
                    model=self.llm_model_type,
                    callback_manager=CallbackManager([StreamingStdOutCallbackHandler()])
                )
            logger.info("Retrieved " + self.llm_model_type)
        # else, if llm_type is "azureopenai"
        elif self.llm_type == "azureopenai":
            logger.info("Use Azure OpenAI LLM")
            logger.info("Retrieving " + self.llm_model_type)
            self.llm = AzureChatOpenAI(
                azure_deployment=self.llm_model_type,
                azure_endpoint=self.local_api_url,
                api_version=self.azureopenai_api_version,
            )
            logger.info("Retrieved " + self.llm_model_type)
        
        return self.llm

langchain_openai.AzureChatOpenAI langchain_openai.ChatOpenAI langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler langchain_community.llms.huggingface_hub.HuggingFaceHub
[((1610, 1675), 'langchain_openai.ChatOpenAI', 'ChatOpenAI', ([], {'client': 'None', 'model': 'self.llm_model_type', 'temperature': '(0)'}), '(client=None, model=self.llm_model_type, temperature=0)\n', (1620, 1675), False, 'from langchain_openai import ChatOpenAI, AzureChatOpenAI\n'), ((2163, 2272), 'langchain_community.llms.huggingface_hub.HuggingFaceHub', 'HuggingFaceHub', ([], {'repo_id': 'self.llm_model_type', 'model_kwargs': "{'temperature': 0.1, 'max_length': max_length}"}), "(repo_id=self.llm_model_type, model_kwargs={'temperature': \n    0.1, 'max_length': max_length})\n", (2177, 2272), False, 'from langchain_community.llms.huggingface_hub import HuggingFaceHub\n'), ((2495, 2523), 'loguru.logger.info', 'logger.info', (['"""Use Local LLM"""'], {}), "('Use Local LLM')\n", (2506, 2523), False, 'from loguru import logger\n'), ((2536, 2584), 'loguru.logger.info', 'logger.info', (["('Retrieving ' + self.llm_model_type)"], {}), "('Retrieving ' + self.llm_model_type)\n", (2547, 2584), False, 'from loguru import logger\n'), ((3206, 3253), 'loguru.logger.info', 'logger.info', (["('Retrieved ' + self.llm_model_type)"], {}), "('Retrieved ' + self.llm_model_type)\n", (3217, 3253), False, 'from loguru import logger\n'), ((2692, 2748), 'loguru.logger.info', 'logger.info', (["('Using local api url ' + self.local_api_url)"], {}), "('Using local api url ' + self.local_api_url)\n", (2703, 2748), False, 'from loguru import logger\n'), ((3356, 3391), 'loguru.logger.info', 'logger.info', (['"""Use Azure OpenAI LLM"""'], {}), "('Use Azure OpenAI LLM')\n", (3367, 3391), False, 'from loguru import logger\n'), ((3404, 3452), 'loguru.logger.info', 'logger.info', (["('Retrieving ' + self.llm_model_type)"], {}), "('Retrieving ' + self.llm_model_type)\n", (3415, 3452), False, 'from loguru import logger\n'), ((3476, 3611), 'langchain_openai.AzureChatOpenAI', 'AzureChatOpenAI', ([], {'azure_deployment': 'self.llm_model_type', 'azure_endpoint': 'self.local_api_url', 'api_version': 'self.azureopenai_api_version'}), '(azure_deployment=self.llm_model_type, azure_endpoint=self.\n    local_api_url, api_version=self.azureopenai_api_version)\n', (3491, 3611), False, 'from langchain_openai import ChatOpenAI, AzureChatOpenAI\n'), ((3682, 3729), 'loguru.logger.info', 'logger.info', (["('Retrieved ' + self.llm_model_type)"], {}), "('Retrieved ' + self.llm_model_type)\n", (3693, 3729), False, 'from loguru import logger\n'), ((2934, 2966), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (2964, 2966), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n'), ((3141, 3173), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (3171, 3173), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n')]