import streamlit as st
from modelo import get_chain

#Menu Visual
st.markdown("<h1 style='text-align: center; color: yellow;'>Chatbot Códigos de Leyes</h1>", unsafe_allow_html=True) #mandar un texto en html
st.header("🤖🦾ChatBot entrenado usando los codigos que componen la ley. Actualizado en 30/01/2024.")

with st.chat_message(name="ai"): #assistant or ai
    st.write('🤖 Hola soy tu asistente del dia de hoy, en que te puedo ayudar')


if "mensajes" not in st.session_state:
    st.session_state.mensajes = []

for message in st.session_state.mensajes:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

pregunta = st.chat_input("Ingresa tu pregunta")
#Manejador del prompt, es un input y button a la vez
with st.spinner("Cargando modelo, por favor espere, mientras puede ingresar su pregunta"):
    chain = get_chain()  #windows

if pregunta:
    #Muestra el mensaje del usuario en el chat
    with st.chat_message(name="human"): #assistant or ai
        st.markdown(pregunta)

    st.session_state.mensajes.append({"role" : "human", "content": pregunta})
    with st.spinner("Procesando su respuesta, por favor espere"):
        respuesta = chain.invoke(pregunta)

    with st.chat_message(name="ai"): #assistant or ai
        st.markdown(respuesta)
    st.session_state.mensajes.append({"role" : "ai", "content": respuesta})