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

#set_verbose(True)
#set_debug(True)

repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
#repo_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
#repo_id = "HuggingFaceH4/zephyr-7b-beta"

template = """[INST]Tu eres Harry Potter, el mago más hábil de todo el mundo mágico.
Responde amablemente a la consulta del usuario basado en la información disponible.
Si no sabes la respuesta, pide al usuario que intente reformular su consulta.

{question}
[/INST]
"""

prompt = PromptTemplate.from_template(template)

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

llm_chain = (
    prompt
    | llm
)


def respond(message, history):
#    history_langchain_format = []
#
#    for human, ai in history:
#        history_langchain_format.append(HumanMessage(content=human))
#        history_langchain_format.append(AIMessage(content=ai))
#
#    history_langchain_format.append(HumanMessage(content=message))

    #print(message)
    response = llm_chain.invoke(message)
    #print(response)

    return response


demo = gr.ChatInterface(
    respond
)


if __name__ == "__main__":
    demo.launch()