flows/chat_interface.py

import gradio as gr
import os
from knowledge_base.build_kb import crear_knowledge_base
from knowledge_base.chat_initializer import inicializar_chat
import time

def interfaz_chat(session_id_state):

    with gr.Column():
        gr.Markdown("## 🧠 Chat con IA sobre la Knowledge Base")
        mensaje_estado_chat = gr.Markdown("Consulta los datos a la base de datos y crea la base de conocimiento antes de inicializar el chat.")

        with gr.Row():
            boton_crear_kb = gr.Button("📦 Crear y Descargar Knowledge Base")
            boton_inicializar = gr.Button("💬 Inicializar Chat")

        archivo_kb = gr.File(label="📁 Descarga la Knowledge Base aquí", visible=False)

        # Contenedor del chat (inicialmente oculto)
        with gr.Column(visible=False) as contenedor_chat:
            interfaz_chat = gr.ChatInterface(fn=lambda msg, hist: "⚠️ Inicializa el chat primero para comenzar.", type="messages")    

        # -------------------------
        # LÓGICA: CREAR KB
        # -------------------------

        def crear_kb_wrapper(session_id):
            for archivo, mensaje in crear_knowledge_base(session_id):
                yield archivo, mensaje

        boton_crear_kb.click(
            fn=crear_kb_wrapper,
            inputs=[session_id_state],
            outputs=[archivo_kb, mensaje_estado_chat]
        )

        # -------------------------
        # LÓGICA: INICIALIZAR CHAT
        # -------------------------

        def mostrar_mensaje_cargando():
            time.sleep(0.3)
            return "⏳ Cargando documentos y generando embeddings... Esto puede tardar unos minutos."

        def cargar_chat(session_id):
            chain, mensaje = inicializar_chat(session_id)
            if chain:
                def responder(msg, hist):
                    result = chain.invoke({"question": msg})
                    return {"role": "assistant", "content": result["answer"]}
                interfaz_chat.fn = responder
                return mensaje, gr.update(visible=True)
            else:
                return mensaje, gr.update(visible=False)

        boton_inicializar.click(
            fn=mostrar_mensaje_cargando,
            inputs=[],
            outputs=[mensaje_estado_chat]
        ).then(
            fn=cargar_chat,
            inputs=[session_id_state],
            outputs=[mensaje_estado_chat, contenedor_chat]
        )

    return mensaje_estado_chat, contenedor_chat