import gradio as gr
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain.globals import set_verbose, set_debug
import os

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



def initPrompt():
    system_prompt = """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.
    """

    prompt = ChatPromptTemplate.from_messages([
      ("system", system_prompt),
      ("human", "{input}"),
    ])

    return prompt



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

    Modelos:
        * llama3-8b-8192
        * mixtral-8x7b-32768
    """
    llm = ChatGroq(
        model="mixtral-8x7b-32768",
        temperature=0,
        max_tokens=None,
        timeout=None,
        max_retries=2,
        # other params...
    )

    return llm



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

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

    print(response)
    return response.content



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()