app.py

import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain_groq import ChatGroq
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv
import re

load_dotenv()

# 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("""
<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; }
</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"

def get_pdf_text(pdf_docs):
    return "".join([page.extract_text() for pdf in pdf_docs for page in PdfReader(pdf).pages])

def get_vector_store(text_chunks):
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    return FAISS.from_texts(text_chunks, embedding=embeddings)

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".
    
    Contexto:
    {context}
    
    Pregunta:
    {question}
    
    Respuesta:
    """
    model = ChatGroq(temperature=0.2, model_name="deepseek-r1-distill-llama-70b", groq_api_key=GROQ_API_KEY)
    return load_qa_chain(model, chain_type="stuff", prompt=PromptTemplate(template=prompt_template, input_variables=["context", "question"]))

def procesar_consulta(pregunta):
    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)
    
    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)

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)

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.title("PDF Consultor 🔍")

    # 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..."):
                raw_text = get_pdf_text(pdf_docs)
                text_chunks = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=500).split_text(raw_text)
                vector_store = get_vector_store(text_chunks)
                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:
        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)

if __name__ == "__main__":
    main()