shell1809/Asistente_UNEFA

1"""2Script de indexación — Asistente de Optimización Administrativa3y Orientación Normativa UNEFA Lara4 5Este script:61. Lee todos los PDFs de la carpeta 'documentos/'72. Extrae el texto83. Lo divide en fragmentos (chunks)94. Genera el embedding de cada fragmento vía Gemini API105. Guarda todo en la carpeta 'indice/' para que rag.py lo use después11 12Se corre UNA SOLA VEZ, y cada vez que agregues, quites o cambies un PDF13en la carpeta 'documentos/'.14"""15 16import os17import io18import json19import time20import base6421import fitz  # PyMuPDF22import requests23import numpy as np24from dotenv import load_dotenv25 26# --- Nota: pytesseract ya no es obligatorio con este método de OCR vía Gemini,27# pero lo dejamos disponible por si quieres volver al OCR local más adelante.28# import pytesseract29# from PIL import Image30 31# Si vuelves a usar Tesseract localmente, descomenta la línea de arriba y esta:32# pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"33 34MIN_CARACTERES_PARA_CONSIDERAR_TEXTO = 2035 36load_dotenv()37GEMINI_KEY = os.getenv("GEMINI_API_KEY")38 39if not GEMINI_KEY:40    raise SystemExit(41        "ERROR: No se encontró GEMINI_API_KEY en el archivo .env. "42        "Revisa que el archivo .env exista y tenga la línea GEMINI_API_KEY=..."43    )44 45DOCS_DIR = "documentos"46OUT_DIR = "indice"47CHUNK_SIZE = 800       # caracteres por fragmento48CHUNK_OVERLAP = 150    # solapamiento entre fragmentos para no cortar ideas a la mitad49 50os.makedirs(OUT_DIR, exist_ok=True)51 52 53def ocr_pagina(pagina, zoom=2, reintentos=3):54    """Convierte una página del PDF en imagen y le pide a Gemini que transcriba55    el texto exactamente (alternativa a Tesseract, más precisa en documentos56    institucionales/legales en español)."""57    matriz = fitz.Matrix(zoom, zoom)58    pix = pagina.get_pixmap(matrix=matriz)59    img_bytes = pix.tobytes("png")60    img_b64 = base64.b64encode(img_bytes).decode("utf-8")61 62    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"63    headers = {"Content-Type": "application/json"}64    params = {"key": GEMINI_KEY}65    body = {66        "contents": [{67            "parts": [68                {"text": "Transcribe EXACTAMENTE todo el texto visible en esta imagen "69                          "(documento institucional en español). No traduzcas, no "70                          "resumas, no agregues comentarios. Si hay tablas, "71                          "transcribe fila por fila. Si no hay texto legible, "72                          "responde solo con: [SIN TEXTO]"},73                {"inline_data": {"mime_type": "image/png", "data": img_b64}}74            ]75        }]76    }77 78    for intento in range(reintentos):79        try:80            r = requests.post(url, headers=headers, params=params, json=body, timeout=60)81            if r.status_code == 429:82                espera = 20 * (intento + 1)83                print(f"      Límite alcanzado en OCR, esperando {espera}s...")84                time.sleep(espera)85                continue86            r.raise_for_status()87            data = r.json()88            texto = data["candidates"][0]["content"]["parts"][0]["text"]89            return "" if "[SIN TEXTO]" in texto else texto90        except (requests.exceptions.RequestException, KeyError, IndexError) as e:91            print(f"      (reintento OCR {intento + 1}/{reintentos} por error: {e})")92            time.sleep(5)93    print("      No se pudo hacer OCR de esta página, se deja vacía.")94    return ""95 96 97def extraer_texto(ruta_pdf):98    """Extrae el texto completo de un PDF, página por página.99    Si una página no tiene texto extraíble (PDF escaneado/imagen),100    usa OCR automáticamente como respaldo."""101    doc = fitz.open(ruta_pdf)102    texto_completo = ""103    paginas_con_ocr = 0104 105    for num_pagina, pagina in enumerate(doc, start=1):106        texto_pagina = pagina.get_text().strip()107 108        if len(texto_pagina) < MIN_CARACTERES_PARA_CONSIDERAR_TEXTO:109            print(f"    Página {num_pagina} sin texto extraíble, aplicando OCR...")110            texto_pagina = ocr_pagina(pagina)111            paginas_con_ocr += 1112 113        texto_completo += f"\n[Página {num_pagina}]\n{texto_pagina}"114 115    doc.close()116 117    if paginas_con_ocr:118        print(f"    ({paginas_con_ocr} página(s) procesada(s) con OCR)")119 120    return texto_completo121 122 123def dividir_en_chunks(texto, tam=CHUNK_SIZE, overlap=CHUNK_OVERLAP):124    """Divide un texto largo en fragmentos con solapamiento."""125    chunks = []126    inicio = 0127    while inicio < len(texto):128        fin = inicio + tam129        fragmento = texto[inicio:fin].strip()130        if fragmento:131            chunks.append(fragmento)132        inicio += tam - overlap133    return chunks134 135 136BATCH_SIZE = 25  # fragmentos por solicitud a la API (menos solicitudes = menos riesgo de 429)137 138 139def obtener_embeddings_lote(textos, reintentos=5):140    """Obtiene los embeddings de una LISTA de textos en una sola llamada a la API,141    usando el endpoint de lote de Gemini. Reduce drásticamente el número de142    solicitudes comparado con pedir un embedding a la vez."""143    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents"144    headers = {"Content-Type": "application/json"}145    params = {"key": GEMINI_KEY}146    body = {147        "requests": [148            {149                "model": "models/gemini-embedding-001",150                "content": {"parts": [{"text": t}]},151                "task_type": "RETRIEVAL_DOCUMENT",152            }153            for t in textos154        ]155    }156 157    for intento in range(reintentos):158        try:159            r = requests.post(url, headers=headers, params=params, json=body, timeout=60)160            if r.status_code == 429:161                espera = 20 * (intento + 1)162                print(f"    Límite de frecuencia alcanzado, esperando {espera}s antes de reintentar...")163                time.sleep(espera)164                continue165            r.raise_for_status()166            data = r.json()167            return [e["values"] for e in data["embeddings"]]168        except requests.exceptions.RequestException as e:169            print(f"    (reintento {intento + 1}/{reintentos} por error: {e})")170            time.sleep(5)171    raise RuntimeError("No se pudo obtener el lote de embeddings tras varios intentos.")172 173 174def main():175    if not os.path.isdir(DOCS_DIR):176        raise SystemExit(f"ERROR: No existe la carpeta '{DOCS_DIR}'. Créala y coloca ahí tus PDFs.")177 178    archivos_pdf = [f for f in os.listdir(DOCS_DIR) if f.lower().endswith(".pdf")]179 180    if not archivos_pdf:181        raise SystemExit(f"No se encontró ningún PDF dentro de '{DOCS_DIR}/'. Copia tus reglamentos ahí.")182 183    print(f"Se encontraron {len(archivos_pdf)} PDF(s): {', '.join(archivos_pdf)}")184    print()185 186    todos_chunks = []187    for archivo in archivos_pdf:188        ruta = os.path.join(DOCS_DIR, archivo)189        print(f"Procesando: {archivo}")190        texto = extraer_texto(ruta)191 192        if len(texto.strip()) < 50:193            print(f"  AVISO: '{archivo}' parece no tener texto extraíble (¿es un PDF escaneado/imagen?). Se omite.")194            continue195 196        chunks = dividir_en_chunks(texto)197        print(f"  -> {len(chunks)} fragmentos generados")198 199        for i, c in enumerate(chunks):200            todos_chunks.append({201                "texto": c,202                "fuente": archivo,203                "chunk_id": i,204            })205 206    if not todos_chunks:207        raise SystemExit("No se generó ningún fragmento de texto. Revisa tus PDFs.")208 209    print(f"\nTotal de fragmentos a indexar: {len(todos_chunks)}")210    print("Generando embeddings (esto usa la API de Gemini, puede tardar unos minutos)...\n")211 212    vectores = []213    total = len(todos_chunks)214    for inicio in range(0, total, BATCH_SIZE):215        lote = todos_chunks[inicio:inicio + BATCH_SIZE]216        textos_lote = [ch["texto"] for ch in lote]217        fin = min(inicio + BATCH_SIZE, total)218        print(f"  Embeddings {inicio + 1}-{fin} de {total}...")219        embs_lote = obtener_embeddings_lote(textos_lote)220        vectores.extend(embs_lote)221        time.sleep(2)  # pausa breve entre lotes222 223    matriz = np.array(vectores, dtype=np.float32)224    np.save(os.path.join(OUT_DIR, "embeddings.npy"), matriz)225 226    with open(os.path.join(OUT_DIR, "chunks.json"), "w", encoding="utf-8") as f:227        json.dump(todos_chunks, f, ensure_ascii=False, indent=2)228 229    print("\n✅ Índice creado correctamente en la carpeta 'indice/'.")230    print(f"   - {len(todos_chunks)} fragmentos indexados de {len(archivos_pdf)} documento(s).")231    print("   Ya puedes correr: streamlit run app.py")232 233 234if __name__ == "__main__":235    main()236