from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI
from langsmith import traceable
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain.prompts.chat import SystemMessagePromptTemplate, HumanMessagePromptTemplate
import gradio as gr

embeddings = OpenAIEmbeddings()
vectordb = Chroma(persist_directory="./ai_experto_upc_v4", embedding_function=embeddings)
llm = ChatOpenAI(temperature=0, model_name='gpt-4o-2024-08-06')

general_system_template = r"""
Eres un experto que te diriges a estudiantes de 15 a 17 años de colegio que no saben sobre la UPC. Explica lo que te preguntan como si fueras un profesor, de manera simple, descriptiva e insightful, inspiralos a elegir la UPC como universidad. Si te preguntan por la competencia, responde elegamente, indicando que deben de haber muy buenas opciones, e indica que hace a la UPC especial. 
Proporciona mensajes breves y concisos con bullets y negritas. 
Toma los siguientes documentos de contexto {context} y responde únicamente basado en este contexto, no inventes. 
No menciones ningún código de curso al mencionar sus nombres (ejemplo: [PS1223])
Cuando el alumno pregunte por historias de alumnos o directores sobre su carrera, envía los links provistos en el contexto de los alumnos, directores de carrera y/o brochures.
Envía el link al brochure de la carrera cuando el alumno quiera más información sobre la carrera.
"""

general_user_template = "Pregunta:```{question}```"
messages = [
    SystemMessagePromptTemplate.from_template(general_system_template),
    HumanMessagePromptTemplate.from_template(general_user_template)
]
qa_prompt = ChatPromptTemplate.from_messages(messages)

def create_memory():
    return ConversationBufferMemory(memory_key='chat_history', return_messages=True)

@traceable
def pdf_qa(query, memory):
    function = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectordb.as_retriever(search_kwargs={'k': 8}),
        combine_docs_chain_kwargs={'prompt': qa_prompt},
        memory=memory
    )
    return function({"question": query})

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(label="Experto UPC", value=[[None, '''
    ¡Hola! Soy Experto UPC. Estoy aquí para ayudarte a qué conozcas cómo es la experiencia universitaria en la UPC y cómo puede potenciarte. 
    
    Puedes hacerme preguntas como:
    
    - ¿Qué hace que mi carrera sea especial en la UPC?
    - ¿Qué beneficios tengo como deportista?
    - ¿Puedo conocer historias de estudiantes y/o directores de carrera de cómo es la carrera que me interesa en la UPC?
    
    Estoy aquí para ayudarte a tomar una decisión informada. ¡Empecemos! 
    ''']])
    
    msg = gr.Textbox(placeholder="Escribe aquí", label='')
    submit = gr.Button("Enviar")
    memory_state = gr.State(create_memory)

    def user(query, chat_history, memory):
        print("User query:", query)
        print("Chat history:", chat_history)

        result = pdf_qa(query, memory)

        chat_history.append((query, result["answer"]))
        print("Updated chat history:", chat_history)

        return gr.update(value=""), chat_history, memory

    submit.click(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=False)
    msg.submit(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=False)

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