atareao / rag_consulta.py

"#!/usr/bin/env python3\n\"\"\"\nrag_consulta.py — Consulta FTS5 sobre la base de conocimiento RAG.\n\nBusca en el índice FTS5 y devuelve los fragmentos más relevantes\ncon resaltado de términos, información del documento y score BM25.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport json\nimport sqlite3\nimport sys\nfrom typing import Any\n\nfrom rag_config import DB_PATH\n\n\n# ── Colores ANSI ─────────────────────────────────────────────────────────────\n\n\nclass ANSI:\n \"\"\"Códigos de color ANSI para la salida por terminal.\"\"\"\n\n RESET = \"\\033[0m\"\n BOLD = \"\\033[1m\"\n DIM = \"\\033[2m\"\n CYAN = \"\\033[36m\"\n GREEN = \"\\033[32m\"\n YELLOW = \"\\033[33m\"\n RED = \"\\033[31m\"\n MAGENTA = \"\\033[35m\"\n BLUE = \"\\033[34m\"\n\n\n# ── Consulta FTS5 ────────────────────────────────────────────────────────────\n\n\ndef consulta_fts5(\n db: sqlite3.Connection,\n query: str,\n limit: int = 5,\n) -> list[dict[str, Any]]:\n \"\"\"Ejecuta una consulta FTS5 y devuelve los resultados ordenados por relevancia.\n\n Args:\n db: Conexión SQLite.\n query: Término(s) de búsqueda en sintaxis FTS5.\n limit: Número máximo de resultados.\n\n Returns:\n Lista de dicts con: id, doc_id, doc_path, title, snippet, rank, tokens.\n \"\"\"\n sql = \"\"\"\n SELECT\n c.id,\n c.doc_id,\n c.doc_path,\n c.title,\n snippet(chunks_fts, 0, '<<<', '>>>', '...', 50) AS fragmento,\n rank,\n c.tokens,\n c.tags,\n c.chunk_index\n FROM chunks_fts\n JOIN chunks c ON chunks_fts.rowid = c.id\n WHERE chunks_fts MATCH ?\n ORDER BY rank\n LIMIT ?\n \"\"\"\n cur = db.execute(sql, (query, limit))\n resultados: list[dict[str, Any]] = []\n for row in cur.fetchall():\n resultados.append(\n {\n \"id\": row[0],\n \"doc_id\": row[1],\n \"doc_path\": row[2],\n \"title\": row[3] or \"(sin título)\",\n \"snippet\": row[4] or \"(sin fragmento disponible)\",\n \"rank\": row[5],\n \"tokens\": row[6],\n \"tags\": row[7],\n \"chunk_index\": row[8],\n }\n )\n return resultados\n\n\n# ── Salida formateada ────────────────────────────────────────────────────────\n\n\ndef mostrar_resultados(resultados: list[dict[str, Any]], query: str) -> None:\n \"\"\"Muestra los resultados de la consulta con colores y formato bonito.\n\n Args:\n resultados: Lista de resultados de consulta_fts5.\n query: La consulta original (para mostrarla en el encabezado).\n \"\"\"\n if not resultados:\n print(\n f\"\\n{ANSI.YELLOW}😕 No se encontraron resultados para: {ANSI.BOLD}{query}{ANSI.RESET}\"\n )\n print(f\" Prueba con términos más generales o revisa la ortografía.\")\n print(\n f' Sintaxis FTS5: \"frase exacta\", termino AND otro, termino OR otro, term*'\n )\n return\n\n print(f\"\\n{ANSI.CYAN}{ANSI.BOLD}🔍 Resultados para: {query}{ANSI.RESET}\")\n print(f\"{ANSI.DIM} {len(resultados)} resultado(s) encontrado(s){ANSI.RESET}\\n\")\n\n for i, r in enumerate(resultados, 1):\n # Encabezado del resultado\n print(f\"{ANSI.CYAN}{ANSI.BOLD}─── Resultado #{i} ───{ANSI.RESET}\")\n\n # Título\n print(f\" {ANSI.GREEN}{ANSI.BOLD}📄 {r['title']}{ANSI.RESET}\")\n\n # Ruta\n print(f\" {ANSI.DIM}📍 {r['doc_path']} (chunk {r['chunk_index']}){ANSI.RESET}\")\n\n # Tags si existen\n if r[\"tags\"] and r[\"tags\"] != \"[]\" and r[\"tags\"] != \"null\":\n try:\n tags = json.loads(r[\"tags\"])\n if tags:\n print(f\" {ANSI.MAGENTA}🏷️ {', '.join(tags)}{ANSI.RESET}\")\n except (json.JSONDecodeError, TypeError):\n pass\n\n # Snippet con resaltado\n snippet = r[\"snippet\"]\n # Reemplazar marcadores por colores ANSI\n snippet_coloreado = snippet.replace(\"<<<\", f\"{ANSI.YELLOW}{ANSI.BOLD}\")\n snippet_coloreado = snippet_coloreado.replace(\">>>\", f\"{ANSI.RESET}\")\n print(f\" {ANSI.DIM}Fragmento:{ANSI.RESET}\")\n print(f\" {snippet_coloreado}\")\n\n # Score y tokens\n print(\n f\" {ANSI.BLUE}Score: {ANSI.BOLD}{r['rank']:.4f}{ANSI.RESET}\"\n f\" {ANSI.DIM}|{ANSI.RESET} \"\n f\"{ANSI.BLUE}Tokens: {r['tokens']}{ANSI.RESET}\"\n )\n print()\n\n\n# ── Estadísticas de la base de datos ─────────────────────────────────────────\n\n\ndef mostrar_estadisticas(db: sqlite3.Connection) -> None:\n \"\"\"Muestra estadísticas básicas de la base de datos.\"\"\"\n try:\n docs = db.execute(\"SELECT COUNT(*) FROM documentos\").fetchone()[0]\n chunks = db.execute(\"SELECT COUNT(*) FROM chunks\").fetchone()[0]\n embeds = db.execute(\"SELECT COUNT(*) FROM embeddings\").fetchone()[0]\n\n # Última indexación\n ultima = db.execute(\"SELECT MAX(last_indexed) FROM documentos\").fetchone()[0]\n\n print(\n f\"\\n{ANSI.CYAN}{ANSI.BOLD}📊 Estadísticas de la base de datos{ANSI.RESET}\"\n )\n print(f\" Documentos: {docs}\")\n print(f\" Chunks: {chunks}\")\n print(f\" Embeddings: {embeds}\")\n if ultima:\n print(f\" Última indexación: {ultima}\")\n print()\n except Exception:\n pass\n\n\n# ── CLI ──────────────────────────────────────────────────────────────────────\n\n\ndef main() -> None:\n \"\"\"Punto de entrada: procesa argumentos y ejecuta la consulta.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Consulta la base de conocimiento RAG con FTS5\",\n formatter_class=argparse.RawDescriptionHelpFormatter,\n epilog=\"\"\"\nEjemplos:\n %(prog)s \"cómo configurar docker\"\n %(prog)s \"systemd AND timers\"\n %(prog)s \"docker compose volumes\"\n %(prog)s \"inteligencia artificial\" -l 3\n %(prog)s --stats\n %(prog)s \"nginx AND proxy AND ssl\" -l 10\n \"\"\",\n )\n parser.add_argument(\n \"query\",\n nargs=\"?\",\n help=\"Término(s) de búsqueda en formato FTS5\",\n )\n parser.add_argument(\n \"-l\",\n \"--limit\",\n type=int,\n default=5,\n help=\"Número de resultados (default: 5)\",\n )\n parser.add_argument(\n \"--db\",\n help=f\"Ruta a la base de datos (default: {DB_PATH})\",\n )\n parser.add_argument(\n \"--stats\",\n action=\"store_true\",\n help=\"Muestra estadísticas de la base de datos\",\n )\n\n args = parser.parse_args()\n\n db_path = args.db or DB_PATH\n\n # Conectar a la base de datos\n try:\n conn = sqlite3.connect(db_path)\n except sqlite3.Error as e:\n print(f\"{ANSI.RED}Error conectando a la base de datos: {e}{ANSI.RESET}\")\n sys.exit(1)\n\n try:\n # Mostrar estadísticas si se pide\n if args.stats:\n mostrar_estadisticas(conn)\n return\n\n # Consulta\n if not args.query:\n parser.print_help()\n return\n\n resultados = consulta_fts5(conn, query=args.query, limit=args.limit)\n mostrar_resultados(resultados, query=args.query)\n\n finally:\n conn.close()\n\n\nif __name__ == \"__main__\":\n main()\n"