RAG_Professional.py

"""
RAG Multimodal Profesional - Versión Optimizada y Limpia
Procesamiento eficiente con Docling + URLs + Caché inteligente
"""

import os, hashlib, pickle, time
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
import streamlit as st
from dotenv import load_dotenv
import base64

from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, PictureDescriptionApiOptions

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from openai import OpenAI
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
from st_social_media_links import SocialMediaIcons

load_dotenv()

class OptimizedRAG:
    def __init__(self):
        api_key = os.getenv("OPENAI_API_KEY")
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-large", api_key=api_key)
        self.llm = ChatOpenAI(model="gpt-4o", temperature=0.1, api_key=api_key)
        self.openai_client = OpenAI(api_key=api_key)
        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
        self.cache_dir = Path("./rag_cache")
        self.cache_dir.mkdir(exist_ok=True)
        self.audio_dir = Path("./audio_cache")
        self.audio_dir.mkdir(exist_ok=True)
        self.vector_store = None
        self._setup_docling(api_key)
    
    def _setup_docling(self, api_key):
        try:
            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 {api_key}"},
                timeout=30
            )
            
            pdf_options = PdfPipelineOptions(
                do_picture_description=True,
                picture_description_options=picture_options,
                enable_remote_services=True,
                generate_picture_images=True,
                images_scale=2
            )
            
            self.converter = DocumentConverter(
                format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_options)}
            )
            st.success("✅ Docling configurado con análisis multimodal")
        except Exception as e:
            st.warning(f"⚠️ Fallback básico: {e}")
            self.converter = DocumentConverter()
    
    def _cache_operations(self, file_path=None, content=None, operation='load'):
        """Operaciones de caché unificadas"""
        if operation == 'load' and file_path:
            file_hash = hashlib.md5(open(file_path, 'rb').read()).hexdigest()
            cache_file = self.cache_dir / f"{file_hash}.pkl"
            if cache_file.exists():
                try:
                    with open(cache_file, 'rb') as f:
                        return pickle.load(f).get('markdown')
                except: pass
        
        elif operation == 'save' and file_path and content:
            file_hash = hashlib.md5(open(file_path, 'rb').read()).hexdigest()
            cache_file = self.cache_dir / f"{file_hash}.pkl"
            try:
                with open(cache_file, 'wb') as f:
                    pickle.dump({'markdown': content, 'timestamp': time.time()}, f)
            except: pass
        
        return None
    
    def _process_content(self, source, is_url=False):
        """Procesamiento unificado de contenido"""
        start_time = time.time()
        
        # Caché solo para archivos locales
        if not is_url:
            cached = self._cache_operations(source, operation='load')
            if cached:
                st.success(f"⚡ Caché: {time.time() - start_time:.1f}s")
                return cached
        
        # Procesamiento principal
        if is_url:
            st.info(f"🌐 Procesando URL: {source}")
            try:
                result = self.converter.convert(source)
                markdown = result.document.export_to_markdown(mark_annotations=True, include_annotations=True)
                if markdown and markdown.strip():
                    st.success(f"✅ Docling: {time.time() - start_time:.1f}s")
                    return markdown
            except:
                st.info("🔄 Usando FireCrawl...")
                from firecrawl import Firecrawl
                app = Firecrawl(api_key=os.getenv("FIRECRAWL_API_KEY"))
                result = app.scrape(source, formats=["markdown"])
                
                # Manejar diferentes estructuras de respuesta
                if hasattr(result, 'markdown'):
                    markdown = result.markdown
                elif hasattr(result, 'data') and isinstance(result.data, dict):
                    markdown = result.data.get('markdown', '')
                elif isinstance(result, dict):
                    markdown = result.get('markdown', result.get('data', {}).get('markdown', ''))
                else:
                    markdown = str(result)
                
                if not markdown:
                    raise Exception("FireCrawl no devolvió contenido markdown")
                
                st.success(f"✅ FireCrawl: {time.time() - start_time:.1f}s")
                return markdown
        else:
            st.info(f"📄 Procesando: {Path(source).name}")
            result = self.converter.convert(source)
            markdown = result.document.export_to_markdown(mark_annotations=True, include_annotations=True)
            
            if markdown and markdown.strip():
                self._cache_operations(source, markdown, 'save')
                st.success(f"✅ Procesado: {time.time() - start_time:.1f}s")
                return markdown
        
        raise Exception("No se pudo extraer contenido")
    
    def _create_documents(self, sources, source_type='file'):
        """Creación unificada de documentos"""
        all_documents = []
        
        # Procesamiento paralelo para múltiples fuentes
        if len(sources) > 1 and source_type == 'file':
            st.info(f"🚀 Procesando {len(sources)} archivos en paralelo...")
            with ThreadPoolExecutor(max_workers=3) as executor:
                # Preparar archivos temporales
                temp_files = []
                futures = {}
                
                for source in sources:
                    if source_type == 'file':
                        temp_path = f"temp_{source.name}"
                        with open(temp_path, "wb") as f:
                            f.write(source.getbuffer())
                        temp_files.append(temp_path)
                        futures[executor.submit(self._process_content, temp_path)] = (source.name, temp_path)
                    
                # Recopilar resultados
                progress_bar = st.progress(0)
                for i, future in enumerate(as_completed(futures)):
                    try:
                        file_name, temp_path = futures[future]
                        markdown = future.result()
                        chunks = self.text_splitter.split_text(markdown)
                        
                        for j, chunk in enumerate(chunks):
                            if chunk.strip():
                                all_documents.append(Document(
                                    page_content=chunk,
                                    metadata={
                                        'source': file_name,
                                        'chunk_id': j,
                                        'processed_by': f'docling_{source_type}'
                                    }
                                ))
                        
                        st.success(f"✅ {file_name}: {len(chunks)} fragmentos")
                        progress_bar.progress((i + 1) / len(futures))
                    except Exception as e:
                        st.error(f"❌ Error: {str(e)}")
                
                # Limpiar temporales
                for temp_path in temp_files:
                    if os.path.exists(temp_path):
                        os.remove(temp_path)
        
        else:
            # Procesamiento individual
            for source in sources:
                try:
                    if source_type == 'file':
                        temp_path = f"temp_{source.name}"
                        with open(temp_path, "wb") as f:
                            f.write(source.getbuffer())
                        markdown = self._process_content(temp_path)
                        source_name = source.name
                        if os.path.exists(temp_path):
                            os.remove(temp_path)
                    else:  # URL
                        markdown = self._process_content(source, is_url=True)
                        source_name = source
                    
                    chunks = self.text_splitter.split_text(markdown)
                    for i, chunk in enumerate(chunks):
                        if chunk.strip():
                            all_documents.append(Document(
                                page_content=chunk,
                                metadata={
                                    'source': source_name,
                                    'source_type': 'web_url' if source_type == 'url' else 'file',
                                    'chunk_id': i,
                                    'processed_by': f'docling_{source_type}'
                                }
                            ))
                    
                    st.success(f"✅ {source_name}: {len(chunks)} fragmentos")
                
                except Exception as e:
                    st.error(f"❌ Error con {source}: {str(e)}")
                    continue
        
        return all_documents
    
    def create_vector_store(self, sources, source_type='file'):
        """Creación optimizada de vector store"""
        start_time = time.time()
        documents = self._create_documents(sources, source_type)
        
        if not documents:
            raise Exception("No se pudieron procesar los documentos")
        
        # Embeddings en paralelo
        st.info(f"🧠 Creando embeddings para {len(documents)} fragmentos...")
        batch_size = 25
        batches = [documents[i:i + batch_size] for i in range(0, len(documents), batch_size)]
        
        progress_bar = st.progress(0)
        
        # Primera batch
        self.vector_store = Chroma.from_documents(
            documents=batches[0],
            embedding=self.embeddings,
            persist_directory="./chroma_db"
        )
        
        # Batches restantes en paralelo
        if len(batches) > 1:
            with ThreadPoolExecutor(max_workers=3) as executor:
                futures = [executor.submit(self.vector_store.add_documents, batch) 
                          for batch in batches[1:]]
                
                for i, future in enumerate(as_completed(futures), 2):
                    future.result()
                    progress_bar.progress(i / len(batches))
        
        progress_bar.progress(1.0)
        total_time = time.time() - start_time
        st.success(f"✅ Vector store completo: {len(documents)} fragmentos en {total_time:.1f}s")
        
        if source_type == 'file':
            st.balloons()
    
    def add_web_sources(self, urls):
        """Agregar URLs al RAG existente"""
        web_documents = self._create_documents(urls, 'url')
        if web_documents:
            if self.vector_store:
                self.vector_store.add_documents(web_documents)
                st.success(f"🎉 {len(web_documents)} fragmentos web agregados")
            else:
                self.create_vector_store(urls, 'url')
        else:
            st.warning("⚠️ No se procesaron URLs")
    
    def query(self, question: str) -> str:
        """Consulta optimizada al RAG"""
        if not self.vector_store:
            return "Por favor, carga documentos primero."
        
        docs = self.vector_store.as_retriever(search_kwargs={"k": 5}).invoke(question)
        if not docs:
            return "No se encontró información relevante."
        
        context = "\n\n".join([f"**Fuente**: {doc.metadata['source']}\n{doc.page_content}" for doc in docs])
        
        system_prompt = f"""Eres un asistente especializado en análisis de documentos procesados con Docling.

Docling utiliza IA para análisis de layout, reconocimiento de tablas, descripción de imágenes y conversión inteligente a markdown.

Responde basándote en el contexto proporcionado, incorporando información de imágenes, gráficos y tablas cuando sea relevante.

Contexto:
{context}

Pregunta: {question}"""
        
        return self.llm.invoke([SystemMessage(system_prompt)]).content
    
    def get_summary(self) -> str:
        """Resumen ejecutivo optimizado para múltiples fuentes - Genérico como NotebookLM"""
        if not self.vector_store:
            return "No hay documentos para resumir."
        
        # Obtener fragmentos representativos usando queries genéricas
        queries = [
            "resumen contenido principal temas importantes",
            "conceptos clave información relevante",
            "introducción conclusiones puntos principales"
        ]
        
        all_docs = []
        for query in queries:
            docs = self.vector_store.as_retriever(search_kwargs={"k": 8}).invoke(query)
            all_docs.extend(docs)
        
        # Remover duplicados manteniendo orden
        seen_content = set()
        unique_docs = []
        for doc in all_docs:
            if doc.page_content not in seen_content:
                seen_content.add(doc.page_content)
                unique_docs.append(doc)
        
        # Organizar contenido por fuentes
        sources_content = {}
        for doc in unique_docs:
            source = doc.metadata.get('source', 'Desconocido')
            if source not in sources_content:
                sources_content[source] = []
            sources_content[source].append(doc.page_content)
        
        # Preparar contenido estructurado por fuentes (límite de 15000 caracteres)
        structured_content = ""
        total_chars = 0
        
        for source, contents in sources_content.items():
            structured_content += f"\n## FUENTE: {source}\n"
            for content in contents:
                if total_chars + len(content) < 15000:
                    structured_content += f"{content}\n\n"
                    total_chars += len(content)
                else:
                    # Truncar último contenido si excede el límite
                    remaining_chars = 15000 - total_chars
                    if remaining_chars > 100:  # Solo si queda espacio significativo
                        structured_content += f"{content[:remaining_chars]}...\n\n"
                    break
            if total_chars >= 15000:
                break
        
        prompt = f"""Genera un resumen ejecutivo completo y estructurado basado en el contenido de las fuentes procesadas.

CONTENIDO DE FUENTES:
{structured_content}

INSTRUCCIONES:
1. Analiza TODO el contenido de las fuentes disponibles
2. Identifica automáticamente los temas principales y conceptos clave
3. Crea un resumen que integre información de TODAS las fuentes
4. Adapta la estructura según el tipo de contenido encontrado
5. Usa formato markdown profesional con encabezados claros

REQUISITOS:
- Resume los puntos más importantes de cada fuente
- Identifica conexiones y patrones entre fuentes diferentes
- Organiza la información de manera lógica y coherente
- Incluye datos, cifras o información específica relevante
- Mantén un tono profesional y objetivo

Genera el resumen sin asumir el tema - deja que el contenido determine la estructura y enfoque."""
        
        return self.llm.invoke([SystemMessage(prompt)]).content
    
    def generate_audio_summary(self, summary_text: str, voice: str = "alloy", speed: float = 1.0) -> str:
        """Genera audio del resumen usando OpenAI TTS"""
        if not summary_text or summary_text == "No hay documentos para resumir.":
            return None
        
        try:
            # Crear hash del contenido para caché
            content_hash = hashlib.md5(f"{summary_text}_{voice}_{speed}".encode()).hexdigest()
            audio_file = self.audio_dir / f"summary_{content_hash}.mp3"
            
            # Verificar caché
            if audio_file.exists():
                return str(audio_file)
            
            # Limpiar texto para TTS (remover markdown excesivo)
            clean_text = self._clean_text_for_tts(summary_text)
            
            # Generar audio
            response = self.openai_client.audio.speech.create(
                model="tts-1",
                voice=voice,
                input=clean_text,
                speed=speed
            )
            
            # Guardar archivo
            response.stream_to_file(audio_file)
            return str(audio_file)
            
        except Exception as e:
            st.error(f"Error generando audio: {str(e)}")
            return None
    
    def _clean_text_for_tts(self, text: str) -> str:
        """Limpia el texto markdown para mejor síntesis de voz"""
        import re
        
        # Remover markdown headers pero mantener el texto
        text = re.sub(r'^#{1,6}\s+', '', text, flags=re.MULTILINE)
        
        # Remover markdown bold/italic pero mantener el texto
        text = re.sub(r'\*{1,2}([^\*]+)\*{1,2}', r'\1', text)
        text = re.sub(r'_{1,2}([^_]+)_{1,2}', r'\1', text)
        
        # Remover links markdown pero mantener el texto
        text = re.sub(r'\[([^\]]+)\]\([^\)]+\)', r'\1', text)
        
        # Limpiar bullet points
        text = re.sub(r'^[-\*\+]\s+', '', text, flags=re.MULTILINE)
        
        # Limpiar espacios excesivos
        text = re.sub(r'\n{3,}', '\n\n', text)
        text = re.sub(r'\s{2,}', ' ', text)
        
        return text.strip()

def main():
    st.set_page_config(page_title="RAG Docling Optimizado", page_icon="🚀", layout="wide")
    st.title("🚀 RAG Multimodal Profesional")
    st.markdown("*Procesamiento eficiente con Docling + URLs + Caché inteligente*")
    
    # Inicializar sistema
    if 'rag_system' not in st.session_state:
        with st.spinner("🔄 Inicializando..."):
            st.session_state.rag_system = OptimizedRAG()
    
    # Sidebar
    with st.sidebar:
        st.header("📁 Documentos")
        uploaded_files = st.file_uploader(
            "Archivos",
            type=['pdf', 'docx', 'pptx', 'html', 'htm', 'png', 'jpg', 'jpeg'],
            accept_multiple_files=True,
            help="Procesamiento paralelo automático"
        )
        
        if uploaded_files:
            file_info = f"📊 {len(uploaded_files)} archivo(s) - {sum(len(f.getbuffer()) for f in uploaded_files) / (1024*1024):.1f} MB"
            st.info(file_info)
            if len(uploaded_files) > 1:
                st.success(f"🚀 Paralelo: {len(uploaded_files)} archivos")
        
        if st.button("🚀 Procesar Archivos") and uploaded_files:
            try:
                with st.spinner("⚡ Procesando..."):
                    st.session_state.rag_system.create_vector_store(uploaded_files)
            except Exception as e:
                st.error(f"❌ Error: {str(e)}")
        
        # URLs
        st.markdown("---")
        st.header("🌐 URLs Web")
        web_url = st.text_input("URL individual:", placeholder="https://ejemplo.com/doc.pdf")
        web_urls = st.text_area("URLs múltiples:", height=80, placeholder="Una URL por línea")
        
        if st.button("🌐 Procesar URLs"):
            sources = []
            if web_url.strip():
                sources.append(web_url.strip())
            if web_urls.strip():
                sources.extend([url.strip() for url in web_urls.split('\n') if url.strip()])
            
            if sources:
                try:
                    with st.spinner("🌐 Procesando..."):
                        st.session_state.rag_system.add_web_sources(sources)
                except Exception as e:
                    st.error(f"❌ Error: {str(e)}")
            else:
                st.warning("⚠️ Ingresa URLs")
        
        # Estado
        st.markdown("---")
        if st.session_state.rag_system.vector_store:
            st.success("✅ Sistema listo")
        else:
            st.warning("⚠️ Sin documentos")
        
        # Caché
        st.markdown("**🗄️ Caché:**")
        cache_files = list(st.session_state.rag_system.cache_dir.glob("*.pkl"))
        st.info(f"📁 {len(cache_files)} archivos")
        
        if st.button("🗑️ Limpiar") and cache_files:
            import shutil
            shutil.rmtree(st.session_state.rag_system.cache_dir)
            st.session_state.rag_system.cache_dir.mkdir(exist_ok=True)
            st.success("✅ Caché limpio")
    
    # Pestañas principales
    tab_chat, tab_summary, tab_audio = st.tabs(["💬 Chat", "📋 Resumen", "🎧 Audio"])
    
    with tab_chat:
        # Chat existente
        if 'messages' not in st.session_state:
            st.session_state.messages = []
        
        for message in st.session_state.messages:
            role = "user" if isinstance(message, HumanMessage) else "assistant"
            with st.chat_message(role):
                st.markdown(message.content)
        
        if prompt := st.chat_input("¿Qué necesitas saber?"):
            st.session_state.messages.append(HumanMessage(prompt))
            with st.chat_message("user"):
                st.markdown(prompt)
            
            with st.chat_message("assistant"):
                with st.spinner("🔍 Analizando..."):
                    response = st.session_state.rag_system.query(prompt)
                    st.markdown(response)
                    st.session_state.messages.append(AIMessage(response))
    
    with tab_summary:
        st.header("📋 Resumen Ejecutivo")
        
        if st.session_state.rag_system.vector_store:
            # Generar resumen inicial automáticamente
            if 'summary_content' not in st.session_state:
                with st.spinner("📝 Generando resumen ejecutivo inicial..."):
                    st.session_state.summary_content = st.session_state.rag_system.get_summary()
            
            # Botón para regenerar resumen
            if st.button("🔄 Generar Resumen"):
                with st.spinner("📝 Regenerando resumen ejecutivo..."):
                    st.session_state.summary_content = st.session_state.rag_system.get_summary()
            
            # Mostrar el resumen actual
            if 'summary_content' in st.session_state:
                st.markdown(st.session_state.summary_content)
        else:
            st.info("💡 Carga documentos primero para generar un resumen")
    
    with tab_audio:
        st.header("🎧 Narración del Resumen")
        
        if st.session_state.rag_system.vector_store and 'summary_content' in st.session_state:
            # Configuración de voz
            col1, col2 = st.columns(2)
            
            with col1:
                voice_options = {
                    "Alloy (Neutro)": "alloy",
                    "Echo (Masculino)": "echo", 
                    "Fable (Británico)": "fable",
                    "Onyx (Masculino Profundo)": "onyx",
                    "Nova (Femenino)": "nova",
                    "Shimmer (Femenino Suave)": "shimmer"
                }
                selected_voice = st.selectbox(
                    "🎤 Seleccionar Voz:",
                    options=list(voice_options.keys()),
                    index=0
                )
            
            with col2:
                speed = st.slider(
                    "⚡ Velocidad:",
                    min_value=0.5,
                    max_value=2.0,
                    value=1.0,
                    step=0.1,
                    help="Velocidad de narración"
                )
            
            # Botón para generar audio
            if st.button("🎵 Generar Audio del Resumen"):
                with st.spinner("🎧 Generando narración..."):
                    voice_key = voice_options[selected_voice]
                    audio_file = st.session_state.rag_system.generate_audio_summary(
                        st.session_state.summary_content,
                        voice=voice_key,
                        speed=speed
                    )
                    
                    if audio_file:
                        st.session_state.current_audio = audio_file
                        st.success("✅ Audio generado exitosamente")
                    else:
                        st.error("❌ Error al generar el audio")
            
            # Reproducir audio si existe
            if 'current_audio' in st.session_state and st.session_state.current_audio:
                st.markdown("### 🔊 Reproducir Narración")
                
                # Leer archivo de audio y convertir a base64 para Streamlit
                try:
                    with open(st.session_state.current_audio, "rb") as audio_file:
                        audio_bytes = audio_file.read()
                    
                    st.audio(audio_bytes, format="audio/mp3")
                    
                    # Información adicional
                    file_size = len(audio_bytes) / (1024 * 1024)  # MB
                    st.info(f"📁 Tamaño: {file_size:.2f} MB | 🎤 Voz: {selected_voice} | ⚡ Velocidad: {speed}x")
                    
                except Exception as e:
                    st.error(f"Error cargando audio: {str(e)}")
            
            # Caché de audio
            st.markdown("---")
            st.markdown("**🗄️ Caché de Audio:**")
            audio_files = list(st.session_state.rag_system.audio_dir.glob("*.mp3"))
            if audio_files:
                total_size = sum(f.stat().st_size for f in audio_files) / (1024 * 1024)
                st.info(f"🎵 {len(audio_files)} archivos de audio ({total_size:.2f} MB)")
                
                if st.button("🗑️ Limpiar Audio Caché"):
                    import shutil
                    shutil.rmtree(st.session_state.rag_system.audio_dir)
                    st.session_state.rag_system.audio_dir.mkdir(exist_ok=True)
                    if 'current_audio' in st.session_state:
                        del st.session_state.current_audio
                    st.success("✅ Caché de audio limpio")
                    st.rerun()
            else:
                st.info("📭 Sin archivos de audio")
                
        else:
            st.info("💡 Necesitas generar un resumen primero para crear la narración")
            if not st.session_state.rag_system.vector_store:
                st.warning("⚠️ Carga documentos primero")
    
    # Footer
    st.markdown("---")
    st.markdown("<div style='text-align: center;'><strong>Edwin Quintero Alzate</strong> | egqa1975@gmail.com | <em>Powered by Docling + OpenAI</em></div>", unsafe_allow_html=True)
    
    social_links = [
        "https://www.facebook.com/edwin.quinteroalzate",
        "https://www.linkedin.com/in/edwinquintero0329/",
        "https://github.com/Edwin1719"
    ]
    SocialMediaIcons(social_links).render()

if __name__ == "__main__":
    main()