backend/app/rag_store.py

"""Lightweight but REAL persistent vector store (the production RAG DB).

This replaces the in-memory keyword stub (kept in prototype/rag_stub.py). It is an
actual on-disk vector index:

  • embeddings persisted as float32 BLOBs in SQLite (collections: vendor_master,
    documents, …);
  • cosine similarity search in NumPy;
  • embedding backend = sentence-transformers if installed, else a deterministic
    hashed-trigram embedding (real vectors, zero extra deps).

Swap-in path for scale (same interface): sqlite-vec / pgvector / Chroma / LanceDB —
see docs/DATABASE.md and the ROADMAP.
"""
from __future__ import annotations

import hashlib
import json
import re
import sqlite3
import struct
import threading
from pathlib import Path

import numpy as np

EMBED_DIM = 256


def _norm(s: str) -> str:
    return re.sub(r"[^a-z0-9 ]", " ", (s or "").lower())


class _HashEmbedder:
    """Deterministic hashed character-trigram embedding, L2-normalized.

    Not a neural model, but a genuine fixed-dim vector representation — good enough
    for demo retrieval and fully offline. Replaced by sentence-transformers when present.
    """

    dim = EMBED_DIM
    name = "hashed-trigram"

    def encode(self, text: str) -> np.ndarray:
        v = np.zeros(self.dim, dtype=np.float32)
        s = f"  {_norm(text)}  "
        for i in range(len(s) - 2):
            tri = s[i:i + 3]
            h = int(hashlib.md5(tri.encode()).hexdigest(), 16)
            v[h % self.dim] += 1.0
        n = np.linalg.norm(v)
        return v / n if n else v


class _STEmbedder:
    name = "sentence-transformers/all-MiniLM-L6-v2"

    def __init__(self, model):
        self._m = model
        self.dim = model.get_sentence_embedding_dimension()

    def encode(self, text: str) -> np.ndarray:
        return np.asarray(self._m.encode(_norm(text), normalize_embeddings=True),
                          dtype=np.float32)


def _load_embedder():
    try:
        from sentence_transformers import SentenceTransformer

        return _STEmbedder(SentenceTransformer("all-MiniLM-L6-v2"))
    except Exception:
        return _HashEmbedder()


class VectorStore:
    def __init__(self, db_path: str | Path) -> None:
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        self._conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._embedder = _load_embedder()
        self._init()

    def _init(self):
        with self._lock:
            self._conn.execute(
                """CREATE TABLE IF NOT EXISTS chunks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    collection TEXT, ref TEXT, text TEXT, metadata TEXT,
                    dim INTEGER, embedding BLOB, ts REAL DEFAULT (strftime('%s','now')))"""
            )
            self._conn.execute("CREATE INDEX IF NOT EXISTS idx_chunks_col ON chunks(collection)")
            self._conn.commit()

    @property
    def backend(self) -> str:
        return self._embedder.name

    def _pack(self, vec: np.ndarray) -> bytes:
        return struct.pack(f"{len(vec)}f", *vec.tolist())

    def _unpack(self, blob: bytes, dim: int) -> np.ndarray:
        return np.array(struct.unpack(f"{dim}f", blob), dtype=np.float32)

    # --- writes ---
    def add(self, collection: str, ref: str, text: str, metadata: dict | None = None) -> int:
        vec = self._embedder.encode(text)
        with self._lock:
            cur = self._conn.execute(
                "INSERT INTO chunks (collection, ref, text, metadata, dim, embedding) "
                "VALUES (?,?,?,?,?,?)",
                (collection, ref, text, json.dumps(metadata or {}), len(vec), self._pack(vec)),
            )
            self._conn.commit()
            return cur.lastrowid

    def index_document(self, doc_id: str, text: str, metadata: dict | None = None,
                       collection: str = "documents", chunk_size: int = 800) -> int:
        chunks = _chunk(text, chunk_size)
        for i, ch in enumerate(chunks):
            self.add(collection, f"{doc_id}#{i}", ch, {**(metadata or {}), "chunk": i})
        return len(chunks)

    def seed(self, collection: str, items: list[dict], text_key, ref_key) -> int:
        """Idempotent seed: only populates if the collection is empty."""
        if self.count(collection) > 0:
            return 0
        for it in items:
            self.add(collection, str(it.get(ref_key, "")), text_key(it), it)
        return len(items)

    # --- reads ---
    def count(self, collection: str | None = None) -> int:
        with self._lock:
            if collection:
                r = self._conn.execute(
                    "SELECT COUNT(*) c FROM chunks WHERE collection=?", (collection,)).fetchone()
            else:
                r = self._conn.execute("SELECT COUNT(*) c FROM chunks").fetchone()
        return r["c"]

    def search(self, query: str, k: int = 3, collection: str | None = None) -> list[dict]:
        qv = self._embedder.encode(query)
        with self._lock:
            if collection:
                rows = self._conn.execute(
                    "SELECT * FROM chunks WHERE collection=?", (collection,)).fetchall()
            else:
                rows = self._conn.execute("SELECT * FROM chunks").fetchall()
        scored = []
        for row in rows:
            vec = self._unpack(row["embedding"], row["dim"])
            score = float(np.dot(qv, vec))
            scored.append((score, row))
        scored.sort(key=lambda x: x[0], reverse=True)
        out = []
        for score, row in scored[:k]:
            md = json.loads(row["metadata"] or "{}")
            out.append({"ref": row["ref"], "text": row["text"], "collection": row["collection"],
                        "score": round(score, 4), "metadata": md})
        return out

    def collections(self) -> dict:
        with self._lock:
            rows = self._conn.execute(
                "SELECT collection, COUNT(*) c FROM chunks GROUP BY collection").fetchall()
        return {r["collection"]: r["c"] for r in rows}

    def info(self) -> dict:
        return {"backend": self.backend, "dim": self._embedder.dim,
                "collections": self.collections(), "path": str(self.db_path)}


def _chunk(text: str, size: int) -> list[str]:
    text = text.strip()
    if len(text) <= size:
        return [text] if text else []
    out, cur = [], []
    n = 0
    for para in text.split("\n"):
        cur.append(para)
        n += len(para)
        if n >= size:
            out.append("\n".join(cur))
            cur, n = [], 0
    if cur:
        out.append("\n".join(cur))
    return out