Update app.py

app.py

@@ -1,7 +1,7 @@
-import os
 import streamlit as st
 from PyPDF2 import PdfReader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+import os
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain_groq import ChatGroq
@@ -11,27 +11,36 @@ from dotenv import load_dotenv
 import re
 
 load_dotenv()
+os.getenv("GROQ_API_KEY")
 
-# Configuración inicial
-st.set_page_config(page_title="PDF Consultor 🔍", page_icon="🔍", layout="wide")
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-
-# CSS personalizado
-st.markdown("""
+css_style = """
 <style>
-    .response-box { padding: 20px; background-color: #f8f9fa; border-radius: 10px; border-left: 5px solid #252850; margin: 20px 0; }
-    .metadata-box { padding: 20px; background-color: #f0f2f6; border-radius: 10px; margin-bottom: 20px; }
     .step-number { font-size: 24px; font-weight: bold; }
+    .response-box { padding: 20px; background-color: #f8f9fa; border-radius: 10px; border-left: 5px solid #252850; margin: 20px 0; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }
+    .metadata-box { padding: 20px; background-color: #f0f2f6; border-radius: 10px; margin-bottom: 20px; }
+    .custom-input { font-size: 16px; padding: 10px; border-radius: 5px; border: 1px solid #ccc; }
+    .suggestion-container { border: 1px solid #e0e0e0; border-radius: 8px; padding: 15px; margin: 10px 0; background: #f8f9fa; }
+    .suggestion-btn { width: 100%; margin: 3px 0; padding: 8px; border-radius: 5px; border: 1px solid #252850; background: white; cursor: pointer; transition: all 0.2s; }
+    .suggestion-btn:hover { background: #252850; color: white; }
 </style>
-""", unsafe_allow_html=True)
+"""
 
-# Funciones auxiliares
 def eliminar_proceso_pensamiento(texto):
-    limpio = re.sub(r'<think>.*?</think>', '', texto, flags=re.DOTALL)
-    return limpio.strip(), re.search(r'<think>(.*?)</think>', texto, re.DOTALL).group(1) if "<think>" in texto else "No disponible"
+    texto_limpio = re.sub(r'<.*?>', '', texto, flags=re.DOTALL)
+    lineas = [line.strip() for line in texto_limpio.split('\n') if line.strip()]
+    return lineas[-1] if lineas else "Respuesta no disponible"
 
 def get_pdf_text(pdf_docs):
-    return "".join([page.extract_text() for pdf in pdf_docs for page in PdfReader(pdf).pages])
+    text = ""
+    for pdf in pdf_docs:
+        pdf_reader = PdfReader(pdf)
+        for page in pdf_reader.pages:
+            text += page.extract_text()
+    return text
+
+def get_text_chunks(text):
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=500)
+    return text_splitter.split_text(text)
 
 def get_vector_store(text_chunks):
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
@@ -39,100 +48,189 @@ def get_vector_store(text_chunks):
 
 def get_conversational_chain():
     prompt_template = """
-    Responde en español exclusivamente con la información solicitada usando el contexto.
-    Si no hay información, di "No disponible".
-    
+    Responde en español exclusivamente con la información solicitada usando el contexto, además sé lo más extenso y detallado posible
+    siempre que se pueda desarollar, como explicando el contenido de referencias nombradas.
+    Formato: Respuesta directa sin prefijos. Si no hay información, di "No disponible".
+
     Contexto:
     {context}
-    
+
     Pregunta:
     {question}
-    
+
     Respuesta:
     """
-    model = ChatGroq(temperature=0.2, model_name="deepseek-r1-distill-llama-70b", groq_api_key=GROQ_API_KEY)
+    model = ChatGroq(
+        temperature=0.2,
+        model_name="deepseek-r1-distill-llama-70b",
+        groq_api_key=os.getenv("GROQ_API_KEY")
+    )
     return load_qa_chain(model, chain_type="stuff", prompt=PromptTemplate(template=prompt_template, input_variables=["context", "question"]))
 
-def procesar_consulta(pregunta):
+def extract_metadata(vector_store):
+    metadata_questions = {
+        "title": "¿Cuál es el título principal del documento? Formato: Respuesta simple con algunas letras en mayúscula si hiciera falta",
+        "entity": "¿A qué organización pertenece este documento?. Formato: Respuesta directa con el nombre de la entidad.",
+        "date": "¿A qué fecha corresponde el documento? Si existen indicios indica la fecha, sino di 'No disponible'"
+    }
+
+    metadata = {}
+    chain = get_conversational_chain()
+
+    for key, question in metadata_questions.items():
+        docs = vector_store.similarity_search(question, k=2)
+        response = chain({"input_documents": docs, "question": question}, return_only_outputs=True)
+        clean_response = eliminar_proceso_pensamiento(response['output_text'])
+        metadata[key] = clean_response if clean_response else "No disponible"
+
+    return metadata
+
+def mostrar_respuesta(texto):
+    st.markdown(f'<div class="response-box">{texto}</div>', unsafe_allow_html=True)
+
+def generar_sugerencias():
+    if 'vector_store' not in st.session_state:
+        return
+
+    try:
+        docs = st.session_state.vector_store.similarity_search("", k=3)
+        context = "\n".join([doc.page_content for doc in docs])
+
+        prompt_template = """
+        Genera exactamente 3 preguntas en español basadas en el contexto.
+        Las preguntas deben ser en español, simples y sencillas de máximo 10 palabras.
+        Formato de respuesta:
+        1. [Pregunta completa en español]
+        2. [Pregunta completa en español]
+        3. [Pregunta completa en español]
+
+        Contexto:
+        {context}
+        """
+
+        model = ChatGroq(
+            temperature=0.4,
+            model_name="deepseek-r1-distill-llama-70b",
+            groq_api_key=os.getenv("GROQ_API_KEY")
+        )
+
+        response = model.invoke(prompt_template.format(context=context))
+
+        preguntas = []
+        for line in response.content.split("\n"):
+            line = line.strip()
+            if line and line[0].isdigit():
+                pregunta = line.split('. ', 1)[1] if '. ' in line else line[2:]
+                if pregunta:
+                    preguntas.append(pregunta)
+
+        return preguntas[:3]
+
+    except Exception as e:
+        st.error(f"Error generando sugerencias: {str(e)}")
+        return
+
+def procesar_consulta(user_question):
     if 'vector_store' not in st.session_state:
         st.error("Por favor carga un documento primero")
         return
-    
+
     chain = get_conversational_chain()
-    docs = st.session_state.vector_store.similarity_search(pregunta)
-    
+    docs = st.session_state.vector_store.similarity_search(user_question)
+
     with st.spinner("Analizando documento..."):
-        response = chain({"input_documents": docs, "question": pregunta}, return_only_outputs=True)
-    
-    respuesta_final, pensamiento = eliminar_proceso_pensamiento(response['output_text'])
-    mostrar_respuesta(respuesta_final, pensamiento)
+        response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
 
-def mostrar_respuesta(respuesta, pensamiento):
-    st.markdown(f'<div class="response-box">{respuesta}</div>', unsafe_allow_html=True)
-    with st.expander("💭 Pensamiento del modelo"):
-        st.write(pensamiento)
+    respuesta_final = eliminar_proceso_pensamiento(response['output_text'])
+    mostrar_respuesta(respuesta_final)
 
-def generar_sugerencias():
-    if 'vector_store' not in st.session_state:
-        return []
-    
-    docs = st.session_state.vector_store.similarity_search("", k=3)
-    context = "\n".join([doc.page_content for doc in docs])
-    
-    prompt_template = """
-    Genera exactamente 3 preguntas simples en español basadas en este contexto.
-    
-    Contexto:
-    {context}
-    
-    Preguntas sugeridas:
-    """
-    
-    model = ChatGroq(temperature=0.4, model_name="deepseek-r1-distill-llama-70b", groq_api_key=GROQ_API_KEY)
-    response = model.invoke(prompt_template.format(context=context))
-    
-    preguntas = [line.split('. ', 1)[1] for line in response.content.split("\n") if line.strip() and line[0].isdigit()]
-    return preguntas[:3]
-
-# Aplicación principal
 def main():
+    st.set_page_config(page_title="PDF Consultor 🔍", page_icon="🔍", layout="wide")
     st.title("PDF Consultor 🔍")
+    st.markdown(css_style, unsafe_allow_html=True)
+
+    estados = {
+        'documento_cargado': False,
+        'sugerencias': [],
+        'pregunta_actual': "",
+        'respuestas': []
+    }
+    for key, value in estados.items():
+        if key not in st.session_state:
+            st.session_state[key] = value
 
-    # Estados de sesión
-    if 'documento_cargado' not in st.session_state:
-        st.session_state.documento_cargado = False
-        st.session_state.sugerencias = []
-        st.session_state.pregunta_actual = ""
-
-    # Sidebar de carga de documentos
     with st.sidebar:
-        st.markdown('<p class="step-number">1. Subir archivos</p>', unsafe_allow_html=True)
-        pdf_docs = st.file_uploader("Subir PDF(s)", accept_multiple_files=True, type=["pdf"])
-    
-        if pdf_docs and not st.session_state.documento_cargado:
-            with st.spinner("Procesando documento..."):
+        st.markdown('<p class="step-number">1 Subir archivos</p>', unsafe_allow_html=True)
+        pdf_docs = st.file_uploader("Subir PDF(s)", accept_multiple_files=True, type=["pdf"], label_visibility="collapsed")
+
+    if pdf_docs and not st.session_state.documento_cargado:
+        with st.spinner("Analizando documento..."):
+            try:
                 raw_text = get_pdf_text(pdf_docs)
-                text_chunks = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=500).split_text(raw_text)
+                text_chunks = get_text_chunks(raw_text)
                 vector_store = get_vector_store(text_chunks)
+
+                st.session_state.metadata = extract_metadata(vector_store)
                 st.session_state.vector_store = vector_store
                 st.session_state.documento_cargado = True
                 st.session_state.sugerencias = generar_sugerencias()
-                st.success("Documento procesado exitosamente.")
-                st.experimental_rerun()
 
-    # Mostrar sugerencias y formulario principal
-    if st.session_state.documento_cargado:
+                st.rerun()
+
+            except Exception as e:
+                st.error(f"Error procesando documento: {str(e)}")
+
+    if 'metadata' in st.session_state:
+        st.markdown("---")
+        cols = st.columns(3)
+        campos_metadata = [
+            ("📄 Título", "title"),
+            ("🏛️ Entidad", "entity"),
+            ("📅 Fecha", "date")
+        ]
+
+        for col, (icono, key) in zip(cols, campos_metadata):
+            with col:
+                st.markdown(f"""
+                <div class="metadata-box">
+                    <div style="font-size:16px; margin-bottom:10px;">{icono}</div>
+                    {st.session_state.metadata[key]}
+                </div>
+                """, unsafe_allow_html=True)
+
         if st.session_state.sugerencias:
-            st.subheader("💡 Preguntas sugeridas:")
-            for pregunta in st.session_state.sugerencias:
-                if st.button(pregunta):
-                    procesar_consulta(pregunta)
-
-        with st.form("consulta_form"):
-            pregunta_usuario = st.text_input("Escribe tu pregunta:", placeholder="Ej: ¿Qué normativa regula este proceso?")
-            enviar = st.form_submit_button("Enviar ▶")
-            if enviar and pregunta_usuario:
-                procesar_consulta(pregunta_usuario)
+            st.markdown("---")
+            with st.container():
+                st.markdown("""
+                <div class="suggestion-container">
+                    <div style="font-size:14px; color:#666; margin-bottom:8px;">💡 ¿Necesitas ideas?</div>
+                """, unsafe_allow_html=True)
+
+                cols_sugerencias = st.columns(3)
+                for i, (col, pregunta) in enumerate(zip(cols_sugerencias, st.session_state.sugerencias)):
+                    with col:
+                        if st.button(pregunta, key=f"sug_{i}", help="Haz clic para usar esta pregunta", use_container_width=True):
+                            st.session_state.pregunta_actual = pregunta
+
+                st.markdown("</div>", unsafe_allow_html=True)
+
+    if st.session_state.documento_cargado:
+        with st.form(key="consulta_form"):
+            col1, col2 = st.columns([5, 1])
+            with col1:
+                pregunta_usuario = st.text_input("Escribe tu pregunta:", value=st.session_state.get('pregunta_actual', ''), placeholder="Ej: ¿De qué trata este documento?", label_visibility="collapsed")
+            with col2:
+                st.markdown("<br>", unsafe_allow_html=True)
+                enviar = st.form_submit_button("Enviar ▶")
+
+            if enviar or st.session_state.pregunta_actual:
+                pregunta_final = pregunta_usuario or st.session_state.pregunta_actual
+                procesar_consulta(pregunta_final)
+                if 'pregunta_actual' in st.session_state:
+                    del st.session_state.pregunta_actual
+
+    elif not st.session_state.documento_cargado:
+        st.info("Por favor, sube un documento PDF para comenzar.")
 
 if __name__ == "__main__":
     main()