Tres tablas: documentos (con su hash), trozos (con documento, página y posición) y los vectores. Solo persistencia: este módulo no sabe de PDFs ni de modelos.
src/inventario/rag/store.py"""Almacén de trozos y vectores en SQLite (+ extensión sqlite-vec).
Solo persistencia: no sabe de PDFs ni de modelos de embeddings.
"""
from __future__ import annotations
import sqlite3
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
import sqlite_vec
DIM = 384 # dimensión de los vectores del modelo elegido (ver ingest.MODELO)
@dataclass(frozen=True)
class Trozo:
documento_id: int
pagina: int
pos: int
texto: str
@dataclass(frozen=True)
class Resultado:
trozo_id: int
texto: str
documento: str
ruta: str
pagina: int
pos: int
distancia: float
class Store:
def __init__(self, path: str | Path = "rag.db") -> None:
self.path = str(path)
self.db = sqlite3.connect(self.path)
self.db.row_factory = sqlite3.Row
self.db.enable_load_extension(True)
sqlite_vec.load(self.db)
self.db.enable_load_extension(False)
self._crear_tablas()
def _crear_tablas(self) -> None:
self.db.executescript(
f"""
CREATE TABLE IF NOT EXISTS documentos (
id INTEGER PRIMARY KEY,
nombre TEXT NOT NULL,
ruta TEXT NOT NULL UNIQUE,
hash TEXT NOT NULL,
paginas INTEGER NOT NULL,
ingestado_en TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS trozos (
id INTEGER PRIMARY KEY,
documento_id INTEGER NOT NULL REFERENCES documentos(id) ON DELETE CASCADE,
pagina INTEGER NOT NULL,
pos INTEGER NOT NULL,
texto TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS ix_trozos_doc ON trozos(documento_id);
CREATE VIRTUAL TABLE IF NOT EXISTS vec_trozos USING vec0(
trozo_id INTEGER PRIMARY KEY,
embedding float[{DIM}] distance_metric=cosine
);
"""
)
self.db.execute("PRAGMA foreign_keys = ON")
# --- documentos -------------------------------------------------------
def documento_por_ruta(self, ruta: str) -> sqlite3.Row | None:
return self.db.execute("SELECT * FROM documentos WHERE ruta = ?", (ruta,)).fetchone()
def borrar_documento(self, documento_id: int) -> None:
ids = [
r[0]
for r in self.db.execute(
"SELECT id FROM trozos WHERE documento_id = ?", (documento_id,)
)
]
for tid in ids:
self.db.execute("DELETE FROM vec_trozos WHERE trozo_id = ?", (tid,))
self.db.execute("DELETE FROM trozos WHERE documento_id = ?", (documento_id,))
self.db.execute("DELETE FROM documentos WHERE id = ?", (documento_id,))
self.db.commit()
def guardar_documento(self, nombre: str, ruta: str, hash_: str, paginas: int) -> int:
cur = self.db.execute(
"INSERT INTO documentos (nombre, ruta, hash, paginas, ingestado_en) VALUES (?,?,?,?,?)",
(nombre, ruta, hash_, paginas, datetime.now(UTC).isoformat()),
)
return int(cur.lastrowid)
# --- trozos + vectores ------------------------------------------------
def guardar_trozos(self, trozos: list[Trozo], embeddings: list[list[float]]) -> None:
if len(trozos) != len(embeddings):
raise ValueError("trozos y embeddings deben tener el mismo largo")
with self.db: # transacción: o entran todos o ninguno
for trozo, emb in zip(trozos, embeddings, strict=True):
cur = self.db.execute(
"INSERT INTO trozos (documento_id, pagina, pos, texto) VALUES (?,?,?,?)",
(trozo.documento_id, trozo.pagina, trozo.pos, trozo.texto),
)
self.db.execute(
"INSERT INTO vec_trozos (trozo_id, embedding) VALUES (?, ?)",
(cur.lastrowid, sqlite_vec.serialize_float32(emb)),
)
def buscar_vectores(self, embedding: list[float], k: int = 5) -> list[Resultado]:
filas = self.db.execute(
"""
SELECT t.id, t.texto, d.nombre, d.ruta, t.pagina, t.pos, v.distance
FROM vec_trozos v
JOIN trozos t ON t.id = v.trozo_id
JOIN documentos d ON d.id = t.documento_id
WHERE v.embedding MATCH ? AND k = ?
ORDER BY v.distance
""",
(sqlite_vec.serialize_float32(embedding), k),
).fetchall()
return [Resultado(f[0], f[1], f[2], f[3], f[4], f[5], float(f[6])) for f in filas]
# --- utilidades -------------------------------------------------------
def stats(self) -> dict[str, int]:
docs = self.db.execute("SELECT COUNT(*) FROM documentos").fetchone()[0]
trozos = self.db.execute("SELECT COUNT(*) FROM trozos").fetchone()[0]
return {"documentos": docs, "trozos": trozos}
def close(self) -> None:
self.db.close()