RAG_Simple.py

"""
RAG Multimodal Profesional - Versión Simplificada
Basado en el enfoque directo del notebook docling-basic.ipynb
Procesamiento eficiente con Docling + OpenAI para análisis visual
"""

import os
import streamlit as st
from pathlib import Path
from dotenv import load_dotenv

# Docling - Configuración simple del notebook
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, PictureDescriptionApiOptions

# LangChain
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain.text_splitter import RecursiveCharacterTextSplitter

# UI
from st_social_media_links import SocialMediaIcons

# Configuración
load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")

class SimpleDoclingRAG:
    """RAG simplificado con Docling - Enfoque directo del notebook"""
    
    def __init__(self):
        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-3-large",
            api_key=openai_api_key
        )
        self.llm = ChatOpenAI(
            model="gpt-4o",
            temperature=0.1,
            api_key=openai_api_key
        )
        
        # Configuración Docling siguiendo el notebook
        self.setup_docling()
        
        # Text splitter simple
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=150
        )
        
        self.vector_store = None
    
    def setup_docling(self):
        """Configuración Docling simple basada en el notebook funcional"""
        try:
            # Configuración de análisis de imágenes con OpenAI (del notebook)
            picture_options = PictureDescriptionApiOptions(
                url="https://api.openai.com/v1/chat/completions",
                prompt="Describe esta imagen en detalle, incluyendo gráficos, tablas, diagramas y texto visible.",
                params={"model": "gpt-4o-mini"},
                headers={"Authorization": f"Bearer {openai_api_key}"},
                timeout=60
            )
            
            # Pipeline options del notebook
            pdf_options = PdfPipelineOptions(
                do_picture_description=True,
                picture_description_options=picture_options,
                enable_remote_services=True,
                generate_picture_images=True,
                images_scale=2
            )
            
            # Converter con opciones multimodales
            self.converter = DocumentConverter(
                format_options={
                    InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_options)
                }
            )
            
            st.success("✅ Docling configurado con análisis de imágenes")
            
        except Exception as e:
            st.warning(f"⚠️ Fallback a configuración básica: {e}")
            # Fallback simple
            self.converter = DocumentConverter()
    
    def process_file(self, file_path: str) -> str:
        """Procesa archivo con Docling - Método directo del notebook"""
        try:
            # Conversión directa como en el notebook
            result = self.converter.convert(file_path)
            
            # Exportar markdown con anotaciones (método del notebook)
            markdown = result.document.export_to_markdown(
                mark_annotations=True,
                include_annotations=True
            )
            
            if not markdown or not markdown.strip():
                raise Exception("No se pudo extraer contenido del documento")
            
            return markdown
            
        except Exception as e:
            st.error(f"Error procesando {Path(file_path).name}: {str(e)}")
            raise e
    
    def create_vector_store(self, uploaded_files):
        """Crea vector store de archivos subidos"""
        all_documents = []
        
        for uploaded_file in uploaded_files:
            # Guardar temporalmente
            temp_path = f"temp_{uploaded_file.name}"
            
            try:
                with open(temp_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())
                
                # Procesar con Docling
                st.info(f"📄 Procesando: {uploaded_file.name}")
                markdown_content = self.process_file(temp_path)
                
                # Dividir en chunks
                chunks = self.text_splitter.split_text(markdown_content)
                
                # Crear documentos
                for i, chunk in enumerate(chunks):
                    if chunk.strip():
                        doc = Document(
                            page_content=chunk,
                            metadata={
                                'source': uploaded_file.name,
                                'chunk_id': i,
                                'processed_by': 'docling'
                            }
                        )
                        all_documents.append(doc)
                
                st.success(f"✅ {uploaded_file.name}: {len(chunks)} fragmentos")
                
            finally:
                # Limpiar archivo temporal
                if os.path.exists(temp_path):
                    os.remove(temp_path)
        
        if not all_documents:
            raise Exception("No se pudieron procesar los documentos")
        
        # Crear vector store
        st.info(f"🧠 Creando embeddings para {len(all_documents)} fragmentos...")
        self.vector_store = Chroma.from_documents(
            documents=all_documents,
            embedding=self.embeddings,
            persist_directory="./chroma_db"
        )
        
        st.success(f"✅ Vector store creado con {len(all_documents)} fragmentos")
    
    def query(self, question: str) -> str:
        """Consulta al RAG"""
        if not self.vector_store:
            return "Por favor, carga documentos primero."
        
        # Buscar documentos relevantes
        retriever = self.vector_store.as_retriever(
            search_type="similarity",
            search_kwargs={"k": 5}
        )
        
        docs = retriever.invoke(question)
        
        if not docs:
            return "No se encontró información relevante."
        
        # Construir contexto
        context = "\n\n".join([
            f"**Fuente**: {doc.metadata['source']}\n{doc.page_content}"
            for doc in docs
        ])
        
        # Generar respuesta
        system_prompt = f"""Eres un asistente especializado en análisis de documentos procesados con Docling.

Docling utiliza:
- IA para análisis de layout y estructura
- Reconocimiento avanzado de tablas
- Descripción automática de imágenes y gráficos
- Conversión inteligente a markdown

Responde basándote en el contexto proporcionado. Si la información incluye descripciones de imágenes, gráficos o tablas, incorpórala en tu respuesta.

Contexto:
{context}

Pregunta: {question}"""

        response = self.llm.invoke([SystemMessage(system_prompt)]).content
        return response

def main():
    st.set_page_config(
        page_title="RAG Docling Profesional",
        page_icon="🚀",
        layout="wide"
    )
    
    st.title("🚀 RAG Multimodal con Docling")
    st.markdown("*Procesamiento profesional de documentos con IA - Texto + Imágenes*")
    
    # Información sobre Docling
    with st.expander("ℹ️ Sobre Docling"):
        st.markdown("""
        **Docling** es una biblioteca de IBM Research para conversión avanzada de documentos:
        
        ✅ **Análisis de Layout**: Identificación inteligente de estructura  
        ✅ **Tablas Precisas**: Reconocimiento y preservación de formato  
        ✅ **Análisis Visual**: Descripción automática de gráficos e imágenes  
        ✅ **OCR Avanzado**: Extracción de texto de calidad  
        ✅ **Markdown Enriquecido**: Salida estructurada y semántica
        
        📁 **Formatos**: PDF, DOCX, PPTX, HTML, imágenes
        """)
    
    # Inicializar sistema
    if 'rag_system' not in st.session_state:
        with st.spinner("🔄 Inicializando sistema..."):
            st.session_state.rag_system = SimpleDoclingRAG()
    
    # Sidebar - Carga de archivos
    with st.sidebar:
        st.header("📁 Cargar Documentos")
        
        uploaded_files = st.file_uploader(
            "Selecciona archivos",
            type=['pdf', 'docx', 'pptx', 'html', 'htm', 'png', 'jpg', 'jpeg'],
            accept_multiple_files=True
        )
        
        if st.button("🚀 Procesar con Docling"):
            if uploaded_files:
                try:
                    with st.spinner("📄 Procesando documentos..."):
                        st.session_state.rag_system.create_vector_store(uploaded_files)
                    st.balloons()
                    
                except Exception as e:
                    st.error(f"❌ Error: {str(e)}")
            else:
                st.warning("Selecciona archivos primero")
        
        # Estado del sistema
        st.markdown("---")
        if st.session_state.rag_system.vector_store:
            st.success("✅ Sistema listo")
        else:
            st.warning("⚠️ Sin documentos cargados")
    
    # Chat principal
    if 'messages' not in st.session_state:
        st.session_state.messages = []
    
    # Mostrar historial
    for message in st.session_state.messages:
        if isinstance(message, HumanMessage):
            with st.chat_message("user"):
                st.markdown(message.content)
        elif isinstance(message, AIMessage):
            with st.chat_message("assistant"):
                st.markdown(message.content)
    
    # Input del usuario
    if prompt := st.chat_input("¿Qué necesitas saber de tus documentos?"):
        # Mostrar mensaje del usuario
        st.session_state.messages.append(HumanMessage(prompt))
        with st.chat_message("user"):
            st.markdown(prompt)
        
        # Generar respuesta
        with st.chat_message("assistant"):
            with st.spinner("🔍 Analizando documentos..."):
                response = st.session_state.rag_system.query(prompt)
                st.markdown(response)
                st.session_state.messages.append(AIMessage(response))
    
    # Footer
    st.markdown("---")
    st.markdown("""
    **Desarrollador**: Edwin Quintero Alzate | **Email**: egqa1975@gmail.com
    
    *Powered by Docling (IBM Research) + OpenAI*
    """)
    
    # Redes sociales
    social_media_links = [
        "https://www.facebook.com/edwin.quinteroalzate",
        "https://www.linkedin.com/in/edwinquintero0329/",
        "https://github.com/Edwin1719"
    ]
    
    social_media_icons = SocialMediaIcons(social_media_links)
    social_media_icons.render()

if __name__ == "__main__":
    main()