atareao / rag_schema.py

"#!/usr/bin/env python3\n\"\"\"\nrag_schema.py — Esquema SQLite para la base de conocimiento RAG.\n\nCrea y gestiona las tablas:\n - documentos → metadatos de cada archivo markdown\n - chunks → fragmentos de texto indexados por documento\n - chunks_fts → índice FTS5 (contenido externo sobre chunks)\n - embeddings → vectores numéricos por chunk\n\"\"\"\n\nfrom __future__ import annotations\n\nimport sqlite3\nfrom typing import Final\n\n# ── SQL DDL ──────────────────────────────────────────────────────────────────\n\nSQL_DOCUMENTOS: Final[str] = \"\"\"\nCREATE TABLE IF NOT EXISTS documentos (\n id INTEGER PRIMARY KEY AUTOINCREMENT,\n path TEXT UNIQUE NOT NULL,\n md5_hash TEXT NOT NULL,\n last_indexed TEXT DEFAULT (datetime('now'))\n);\n\"\"\"\n\nSQL_CHUNKS: Final[str] = \"\"\"\nCREATE TABLE IF NOT EXISTS chunks (\n id INTEGER PRIMARY KEY AUTOINCREMENT,\n doc_id INTEGER NOT NULL,\n chunk_index INTEGER NOT NULL,\n content TEXT NOT NULL,\n tokens INTEGER,\n doc_path TEXT NOT NULL,\n title TEXT,\n tags TEXT,\n created_at TEXT DEFAULT (datetime('now')),\n FOREIGN KEY (doc_id) REFERENCES documentos(id) ON DELETE CASCADE\n);\n\"\"\"\n\nSQL_CHUNKS_FTS: Final[str] = \"\"\"\nCREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(\n content,\n tokenize='porter unicode61 remove_diacritics 1',\n content='chunks',\n content_rowid='id'\n);\n\"\"\"\n\nSQL_EMBEDDINGS: Final[str] = \"\"\"\nCREATE TABLE IF NOT EXISTS embeddings (\n chunk_id INTEGER PRIMARY KEY,\n vector BLOB NOT NULL,\n model TEXT DEFAULT 'bge-m3',\n dimensions INTEGER DEFAULT 1024,\n FOREIGN KEY (chunk_id) REFERENCES chunks(id) ON DELETE CASCADE\n);\n\"\"\"\n\nSQL_INDEX_DOC_PATH: Final[str] = (\n \"CREATE INDEX IF NOT EXISTS idx_chunks_doc_path ON chunks(doc_path);\"\n)\nSQL_INDEX_DOC_ID: Final[str] = (\n \"CREATE INDEX IF NOT EXISTS idx_chunks_doc_id ON chunks(doc_id);\"\n)\n\nSQL_TRIGGER_FTS_INSERT: Final[str] = \"\"\"\nCREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks\nBEGIN\n INSERT INTO chunks_fts(rowid, content) VALUES (new.id, new.content);\nEND;\n\"\"\"\n\nSQL_TRIGGER_FTS_DELETE: Final[str] = \"\"\"\nCREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks\nBEGIN\n INSERT INTO chunks_fts(chunks_fts, rowid, content) VALUES('delete', old.id, old.content);\nEND;\n\"\"\"\n\nSQL_TRIGGER_FTS_UPDATE: Final[str] = \"\"\"\nCREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks\nBEGIN\n INSERT INTO chunks_fts(chunks_fts, rowid, content) VALUES('delete', old.id, old.content);\n INSERT INTO chunks_fts(rowid, content) VALUES (new.id, new.content);\nEND;\n\"\"\"\n\nSQL_PRAGMAS: Final[str] = \"\"\"\nPRAGMA journal_mode = WAL;\nPRAGMA foreign_keys = ON;\nPRAGMA synchronous = NORMAL;\n\"\"\"\n\n\n# ── Funciones públicas ───────────────────────────────────────────────────────\n\n\ndef crear_esquema(conn: sqlite3.Connection) -> None:\n \"\"\"Crea todas las tablas, índices y triggers si no existen.\"\"\"\n conn.executescript(SQL_PRAGMAS)\n conn.execute(SQL_DOCUMENTOS)\n conn.execute(SQL_CHUNKS)\n conn.execute(SQL_CHUNKS_FTS)\n conn.execute(SQL_EMBEDDINGS)\n conn.execute(SQL_INDEX_DOC_PATH)\n conn.execute(SQL_INDEX_DOC_ID)\n conn.execute(SQL_TRIGGER_FTS_INSERT)\n conn.execute(SQL_TRIGGER_FTS_DELETE)\n conn.execute(SQL_TRIGGER_FTS_UPDATE)\n conn.commit()\n\n\ndef rebuild_fts5(conn: sqlite3.Connection) -> None:\n \"\"\"Reconstruye el índice FTS5 desde los datos actuales de la tabla chunks.\n\n Necesario después de inserciones/actualizaciones masivas porque\n la tabla es de contenido externo.\n \"\"\"\n conn.executescript(\"INSERT INTO chunks_fts(chunks_fts) VALUES('rebuild');\")\n\n\ndef limpiar_chunks_por_doc_id(conn: sqlite3.Connection, doc_id: int) -> None:\n \"\"\"Elimina todos los chunks y embeddings de un documento.\n\n Los triggers FTS5 se encargan de limpiar el índice textual.\n El CASCADE de embeddings se activa por la FK.\n \"\"\"\n conn.execute(\n \"DELETE FROM embeddings WHERE chunk_id IN (SELECT id FROM chunks WHERE doc_id = ?)\",\n (doc_id,),\n )\n conn.execute(\"DELETE FROM chunks WHERE doc_id = ?\", (doc_id,))\n conn.commit()\n\n\ndef main() -> None:\n \"\"\"Demo: crea el esquema en la base de datos configurada.\"\"\"\n from rag_config import DB_PATH\n\n print(f\"Creando esquema en: {DB_PATH}\")\n conn = sqlite3.connect(DB_PATH)\n try:\n crear_esquema(conn)\n print(\"✔ Esquema creado correctamente.\")\n\n # Listar las tablas para verificar\n cur = conn.execute(\n \"SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;\"\n )\n tablas = [r[0] for r in cur.fetchall()]\n print(f\" Tablas: {', '.join(tablas)}\")\n\n cur = conn.execute(\n \"SELECT name FROM sqlite_master WHERE type='trigger' ORDER BY name;\"\n )\n triggers = [r[0] for r in cur.fetchall()]\n print(f\" Triggers: {', '.join(triggers)}\")\n finally:\n conn.close()\n\n\nif __name__ == \"__main__\":\n main()\n"