import gradio as gr
from langchain_huggingface import HuggingFaceEndpoint, HuggingFacePipeline
from langchain_core.prompts import PromptTemplate
from langchain.globals import set_verbose, set_debug
import os

def isDevelopmentEnv():
    return "DEVELOPMENT" in os.environ



def initPrompt():
    template = """[INST]Tu eres Harry Potter, el estudiante de magia más hábil de todo el mundo mágico.
    Responde amablemente a la consulta del usuario basado en la información disponible y a las siguientes reglas:
    1. Si no sabes la respuesta, pide al usuario que intente reformular su consulta.
    2. Responde siempre en idioma Español.
    3. Da respuestas únicamente relacionadas al mundo mágico.

    Consulta: {question}
    [/INST]
    """

    prompt = PromptTemplate.from_template(template)
    return prompt



def initLLM():
    """
    Inicializamos el modelo LLM.

    Otros modelos que podríamos usar:
        - meta-llama/Meta-Llama-3.1-8B-Instruct
        - HuggingFaceH4/zephyr-7b-beta
    """

    model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"

    llm = HuggingFaceEndpoint(
        repo_id = model_id,
        task = "text-generation",
        temperature = 0.5,
        model_kwargs = {
            "min_length": 200,
            "max_length": 2000,
            "num_return_sequences": 1
        }
    )

    return llm



def respond(message, history):
    response = ""

    try:
        response = llm_chain.invoke(message)
    except:
        raise gradio.Error("Se ha producido un error al interactuar con el modelo LLM.", duratio=5)

    return response



if __name__ == "__main__":
    """
    Entrypoint de la app.
    """

    if isDevelopmentEnv():
        set_verbose(True)
        set_debug(True)

    prompt = initPrompt()

    llm = initLLM()

    llm_chain = (
        prompt
        | llm
    )

    demo = gr.ChatInterface(
        fn = respond,
        title = "Hola 👋! Soy Harry Potter ⚡",
        description = "Intentaré responder cualquier consulta relacionada a Hogwarts, animales fantásticos y al mundo mágico en general. Adelante!"
    )

    demo.launch()