main gradio.py

import os
import gradio as gr
import shutil
from typing import List

from src.file_processor import chunk_pdfs, chunk_all_documents
from src.chroma_db import save_to_chroma_db, get_chroma_client
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaEmbeddings
from langchain_ollama import ChatOllama

# Initialize components - procesamiento condicional
def initialize_system(process_documents=True):
    """
    Inicializa el sistema RAG con opción de procesar documentos
    """
    if process_documents:
        print("Procesando documentos...")
        processed_documents = chunk_pdfs()

        print("Inicializando modelo de embeddings...")
        embedding_model = OllamaEmbeddings(
            model="nomic-embed-text"
        )

        print("Guardando documentos en la base de datos...")
        db = save_to_chroma_db(processed_documents, embedding_model)
        return db, embedding_model
    else:
        print("Saltando procesamiento de documentos...")
        print("Inicializando modelo de embeddings...")
        embedding_model = OllamaEmbeddings(
            model="nomic-embed-text"
        )

        # Intentar conectar con base de datos existente
        try:
            db = get_chroma_client()
            print("Conectado a base de datos existente")
            return db, embedding_model
        except Exception as e:
            print(f"Error conectando a base de datos existente: {e}")
            return None, embedding_model

# Estado global para controlar si los documentos están procesados
documents_processed = False
db = None
embedding_model = None

# Define the prompt template
PROMPT_TEMPLATE = """
Tienes que responder la siguiente pregunta basada en el contexto proporcionado:
{context}

Responde la siguiente pregunta: {question}

Proporciona una respuesta con un enfoque de análisis histórico, considerando las causas, consecuencias y evolución de los hechos descritos.
Sitúa los eventos en su marco temporal y geopolítico, y explica los factores sociales, económicos y políticos relevantes.
Evita opiniones o juicios de valor y no incluyas información que no esté sustentada en el contexto.
"""


prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)

# Initialize Ollama LLM model
model = ChatOllama(model="hf.co/unsloth/granite-4.0-h-small-GGUF:Q2_K_L")



def answer_question(question):
    """
    Función que responde preguntas basadas en el contexto de los documentos usando ChromaDB Docker
    """
    global documents_processed, db

    if not question.strip():
        return "Por favor ingresa una pregunta válida."

    if not documents_processed or db is None:
        return "❌ No hay documentos procesados disponibles. Por favor, procesa algunos documentos primero usando la opción de arriba."

    try:
        # Perform similarity search with the query
        docs = db.similarity_search_with_score(question, k=3)

        if not docs:
            return "No se encontraron documentos relevantes para responder tu pregunta."

        context = "\n\n---\n\n".join([doc.page_content for doc, _score in docs])

        # Generate the prompt
        prompt = prompt_template.format(context=context, question=question)

        # Get response from model
        response = model.invoke(prompt)

        return response.content if hasattr(response, 'content') else str(response)

    except Exception as e:
        return f"Error al procesar la pregunta: {str(e)}. Verifica que ChromaDB Docker esté funcionando en el puerto 8000."

# Definir constante para la carpeta de aportaciones
APORTACIONES_PATH = 'aportaciones'

def handle_file_upload(files) -> str:
    """
    Función que maneja la subida de archivos de los usuarios
    """
    if not files:
        return "😅 ¡Ups! No has seleccionado ningún archivo. ¡Inténtalo de nuevo!"

    success_count = 0
    error_count = 0
    error_messages = []

    # Crear carpeta aportaciones si no existe
    os.makedirs(APORTACIONES_PATH, exist_ok=True)

    for file_obj in files:
        try:
            # Obtener el nombre del archivo
            filename = os.path.basename(file_obj.name)

            # Crear ruta de destino
            destination_path = os.path.join(APORTACIONES_PATH, filename)

            # Copiar el archivo a la carpeta aportaciones
            shutil.copy2(file_obj.name, destination_path)

            print(f"✅ Archivo {filename} subido exitosamente a {APORTACIONES_PATH}")
            success_count += 1

        except Exception as e:
            error_message = f"❌ Error al subir {filename}: {str(e)}"
            print(error_message)
            error_messages.append(error_message)
            error_count += 1

    # Crear mensaje de respuesta jovial
    if success_count > 0 and error_count == 0:
        return f"🎉 ¡Genial! Has subido {success_count} archivo(s) exitosamente a la carpeta 'aportaciones'. ¡Tu conocimiento ahora forma parte del sistema! 🚀"
    elif success_count > 0 and error_count > 0:
        return f"⚠️ {success_count} archivo(s) subido(s) correctamente, pero {error_count} archivo(s) tuvieron problemas:\n" + "\n".join(error_messages)
    else:
        return f"😞 ¡Vaya! Hubo problemas al subir los archivos:\n" + "\n".join(error_messages)

def process_user_documents():
    """
    Función que procesa los documentos subidos por usuarios
    """
    global documents_processed, db, embedding_model

    try:
        print("🔄 Procesando documentos de usuarios...")

        # Procesar documentos de ambas carpetas
        processed_documents = chunk_all_documents()

        if not processed_documents:
            return "😅 No se encontraron documentos para procesar. ¡Sube algunos archivos primero!"

        print("🔗 Inicializando modelo de embeddings...")
        embedding_model = OllamaEmbeddings(
            model="nomic-embed-text"
        )

        print("💾 Guardando documentos en la base de datos...")
        db = save_to_chroma_db(processed_documents, embedding_model)

        documents_processed = True

        return f"🎊 ¡Perfecto! Se procesaron {len(processed_documents)} documentos exitosamente. ¡Ya puedes hacer preguntas sobre tu nuevo contenido! 📚✨"

    except Exception as e:
        return f"❌ Error al procesar documentos: {str(e)}. Asegúrate de que todos los servicios estén funcionando correctamente."

# Create Gradio interface
with gr.Blocks(
    title="Sistema RAG - Consulta de Documentos",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 800px;
        margin: auto;
    }
    .title {
        text-align: center;
        color: #2563eb;
        font-size: 2.5em;
        margin-bottom: 1em;
    }
    .subtitle {
        text-align: center;
        color: #64748b;
        font-size: 1.1em;
        margin-bottom: 2em;
    }
    """
) as demo:
    gr.HTML("<h1 class='title'>🤖 Sistema RAG - Consulta de Documentos</h1>")
    gr.HTML("<p class='subtitle'>Haz preguntas sobre el contenido de tus documentos usando IA con ChromaDB Docker</p>")

    gr.HTML("""
    <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
               padding: 20px;
               border-radius: 15px;
               margin: 20px 0;
               text-align: center;
               color: white;">
        <h3 style="margin: 0 0 10px 0;">🚀 ¡Comparte tu conocimiento!</h3>
        <p style="margin: 0; font-size: 1.1em;">
            ¿Tienes documentos interesantes que quieres añadir al sistema?
            ¡Súbelos aquí y forma parte de esta aventura del conocimiento! 📚✨
        </p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=2):
            file_upload = gr.File(
                label="📎 Subir documentos",
                file_count="multiple",
                file_types=[".pdf", ".txt", ".md"],
                elem_id="file_upload"
            )

        with gr.Column(scale=1):
            upload_btn = gr.Button(
                "⬆️ Subir archivos",
                variant="secondary",
                size="lg"
            )

    upload_output = gr.Markdown(
        label="Estado de subida",
        elem_id="upload_status"
    )

    with gr.Row():
        process_btn = gr.Button(
            "🔄 Procesar documentos",
            variant="primary",
            size="lg"
        )

    process_output = gr.Markdown(
        label="Estado de procesamiento",
        elem_id="process_status"
    )

    question_input = gr.Textbox(
        label="Tu pregunta",
        placeholder="Ej: ¿Cuáles son los pasos recomendados para fertilizar un jardín de vegetales?",
        lines=3,
        max_lines=10
    )

    submit_btn = gr.Button(
        "🔍 Consultar",
        variant="primary",
        size="lg"
    )

    answer_output = gr.Markdown(
        label="Respuesta",
        show_copy_button=True
    )

    # Examples
    gr.Examples(
        examples=[
            "¿Cuál es el orgigen étnico de los habitantes de Gaza?",
            "¿Qué documentos históricos están disponibles?",
            "¿Qué ocurrió el 7 de octubre de 2023?",
        ],
        inputs=question_input,
        label="Ejemplos de preguntas"
    )

    # Event handlers
    submit_btn.click(
        fn=answer_question,
        inputs=[question_input],
        outputs=[answer_output]
    )

    question_input.submit(
        fn=answer_question,
        inputs=[question_input],
        outputs=[answer_output]
    )

    # Event handlers para subida de archivos
    upload_btn.click(
        fn=handle_file_upload,
        inputs=[file_upload],
        outputs=[upload_output]
    )

    process_btn.click(
        fn=process_user_documents,
        inputs=[],
        outputs=[process_output]
    )

    gr.HTML("""
    <div style="text-align: center; margin-top: 2em; color: #64748b; font-size: 0.9em;">
        <p>Sistema RAG con LangChain, Ollama y ChromaDB Docker</p>
        <p style="font-size: 0.8em; margin-top: 0.5em;">🌐 ChromaDB corriendo en contenedor Docker (puerto 8000)</p>
    </div>
    """)

if __name__ == "__main__":
    print("🚀 Sistema RAG - Consulta de Documentos")
    print("=" * 50)

    # Usar siempre la opción de base de datos existente
    print("💡 ¡Novedad! Los usuarios ahora pueden subir documentos a la carpeta 'aportaciones' desde la interfaz web")
    print("   ¡Comparte tu conocimiento y enriquecer el sistema! 📚✨")

    print("\n🚀 Usando base de datos existente directamente...")
    process_documents = False

    # Inicializar sistema usando base de datos existente
    print("\nInicializando sistema...")
    db, embedding_model = initialize_system(process_documents)

    # Verificar estado de inicialización
    documents_processed = (db is not None)
    if documents_processed:
        print("✅ Sistema inicializado con documentos existentes")
    else:
        print("⚠️ No se pudo conectar a documentos existentes")
        print("💡 Asegúrate de que la base de datos ChromaDB esté disponible en 'chroma/'")

    print("\n🚀 Iniciando interfaz web...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7862,
        share=True,
        debug=False
    )