ai-resume/corpus/store.py

"""SQLite: the keyword index and the incremental state.

Two jobs, one file (`RESUME_CORPUS_DB`, default `data/corpus/corpus.db`):

* **`chunk_fts`** — an FTS5 index over chunk text with the same filter columns
  the vectors carry. This is the keyword half of hybrid retrieval (decision
  D10). Phase 2 indexes *pages* for the site's own search box; this indexes
  *chunks*, because retrieval hands the model a passage, not a page.
* **`manifest`** — one row per source document: its content hash, its chunk
  count and the embedding model/dimension it was embedded with. Incremental mode
  is a hash comparison against this table, and nothing else.

`meta` records the collection contract (provider, model, dimension, chunker
version) so a provider change is detected instead of silently corrupting the
index.

Standard library `sqlite3` only. The same file is opened read-only by the Phase
5 retrieval service.
"""

from __future__ import annotations

import json
import os
import re
import sqlite3
import time
import unicodedata

SCHEMA_VERSION = "1"

SCHEMA = """
CREATE TABLE IF NOT EXISTS chunks (
    id               TEXT PRIMARY KEY,
    source_path      TEXT NOT NULL,
    source_id        TEXT NOT NULL,
    chunk_index      INTEGER NOT NULL,
    area             TEXT NOT NULL,
    level            INTEGER NOT NULL,
    source_type      TEXT NOT NULL,
    confidence       TEXT NOT NULL,
    period           TEXT NOT NULL,
    last_verified    TEXT NOT NULL,
    title            TEXT NOT NULL,
    claim            TEXT NOT NULL,
    url              TEXT NOT NULL,
    content_hash     TEXT NOT NULL,
    embedding_model  TEXT NOT NULL,
    embedding_dim    INTEGER NOT NULL,
    chunker_version  TEXT NOT NULL,
    chunk_tokens     INTEGER NOT NULL,
    public           INTEGER NOT NULL,
    artefacts        TEXT NOT NULL,
    tech             TEXT NOT NULL,
    tags             TEXT NOT NULL,
    text             TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS chunks_source ON chunks(source_path);
CREATE INDEX IF NOT EXISTS chunks_area ON chunks(area);

CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts USING fts5(
    id UNINDEXED,
    source_path UNINDEXED,
    title,
    terms,
    text,
    area UNINDEXED,
    confidence UNINDEXED,
    tokenize = 'porter unicode61 remove_diacritics 2'
);

CREATE TABLE IF NOT EXISTS manifest (
    source_path      TEXT PRIMARY KEY,
    content_hash     TEXT NOT NULL,
    chunk_count      INTEGER NOT NULL,
    embedding_model  TEXT NOT NULL,
    embedding_dim    INTEGER NOT NULL,
    chunker_version  TEXT NOT NULL,
    last_verified    TEXT NOT NULL,
    updated_at       TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS meta (
    key   TEXT PRIMARY KEY,
    value TEXT NOT NULL
);
"""

# Column order used by every INSERT in this module.
_CHUNK_COLUMNS = (
    "id", "source_path", "source_id", "chunk_index", "area", "level", "source_type",
    "confidence", "period", "last_verified", "title", "claim", "url", "content_hash",
    "embedding_model", "embedding_dim", "chunker_version", "chunk_tokens",
    "public", "artefacts", "tech", "tags", "text",
)

_JSON_FIELDS = ("artefacts", "tech", "tags")


class StoreError(RuntimeError):
    """The corpus database is unreadable or inconsistent."""


class Store:
    def __init__(self, path: str) -> None:
        self.path = path
        os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
        self.db = sqlite3.connect(path)
        self.db.row_factory = sqlite3.Row
        self.db.execute("PRAGMA journal_mode = WAL")
        self.db.execute("PRAGMA synchronous = NORMAL")
        self.db.executescript(SCHEMA)
        self.set_meta("schema_version", SCHEMA_VERSION)

    # -- lifecycle --------------------------------------------------------

    def close(self) -> None:
        self.db.close()

    def __enter__(self) -> "Store":
        return self

    def __exit__(self, *_exc) -> None:
        self.close()

    def reset(self) -> None:
        """Empty every table. Used by a from-scratch rebuild."""
        for table in ("chunk_fts", "chunks", "manifest", "meta"):
            self.db.execute(f"DELETE FROM {table}")
        self.db.commit()

    # -- writing ----------------------------------------------------------

    def write_document(self, doc, payloads: list) -> None:
        """Replace all rows for one document: chunks, FTS rows and manifest."""
        self.delete_document(doc.path, keep_manifest=True)
        for payload in payloads:
            values = []
            for column in _CHUNK_COLUMNS:
                value = payload[column]
                if column in _JSON_FIELDS:
                    value = json.dumps(value, sort_keys=True)
                values.append(value)
            cursor = self.db.execute(
                f"INSERT INTO chunks ({', '.join(_CHUNK_COLUMNS)}) "
                f"VALUES ({', '.join('?' for _ in _CHUNK_COLUMNS)})",
                values,
            )
            # `terms` is the keyword surface: the words a visitor or a recruiter
            # would type that live in front-matter rather than the body — tech
            # names, tags, the one-line claim. Without it, a search for "zigbee"
            # misses a page whose body says "mesh". It is indexed, never retrieved.
            terms = " ".join(
                str(value)
                for value in (
                    payload["claim"],
                    payload["title"],
                    " ".join(payload["tech"]),
                    " ".join(payload["tags"]),
                )
                if value
            )
            self.db.execute(
                "INSERT INTO chunk_fts (rowid, id, source_path, title, terms, text, area, confidence) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    cursor.lastrowid,
                    payload["id"],
                    payload["source_path"],
                    payload["title"],
                    terms,
                    payload["text"],
                    payload["area"],
                    payload["confidence"],
                ),
            )
        self.db.execute(
            "INSERT OR REPLACE INTO manifest "
            "(source_path, content_hash, chunk_count, embedding_model, embedding_dim, "
            " chunker_version, last_verified, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                doc.path,
                doc.digest,
                len(payloads),
                payloads[0]["embedding_model"] if payloads else "",
                payloads[0]["embedding_dim"] if payloads else 0,
                payloads[0]["chunker_version"] if payloads else "",
                doc.last_verified,
                time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            ),
        )
        self.db.commit()

    def delete_document(self, source_path: str, keep_manifest: bool = False) -> int:
        rowids = [
            row["rowid"]
            for row in self.db.execute("SELECT rowid FROM chunks WHERE source_path = ?", (source_path,))
        ]
        for rowid in rowids:
            self.db.execute("DELETE FROM chunk_fts WHERE rowid = ?", (rowid,))
        self.db.execute("DELETE FROM chunks WHERE source_path = ?", (source_path,))
        if not keep_manifest:
            self.db.execute("DELETE FROM manifest WHERE source_path = ?", (source_path,))
        self.db.commit()
        return len(rowids)

    # -- reading ----------------------------------------------------------

    def manifest(self) -> dict:
        return {
            row["source_path"]: dict(row)
            for row in self.db.execute("SELECT * FROM manifest")
        }

    def chunk_count(self) -> int:
        return int(self.db.execute("SELECT count(*) FROM chunks").fetchone()[0])

    def source_count(self) -> int:
        return int(self.db.execute("SELECT count(DISTINCT source_path) FROM chunks").fetchone()[0])

    def chunk_ids(self) -> list:
        return [row["id"] for row in self.db.execute("SELECT id FROM chunks ORDER BY id")]

    def chunk_payloads(self, source_path: str = None) -> list:
        """Every stored payload, optionally for one document. Used by verify()."""
        sql = "SELECT * FROM chunks"
        args = ()
        if source_path:
            sql += " WHERE source_path = ?"
            args = (source_path,)
        sql += " ORDER BY source_path, chunk_index"
        out = []
        for row in self.db.execute(sql, args):
            payload = dict(row)
            for field in _JSON_FIELDS:
                payload[field] = json.loads(payload[field])
            out.append(payload)
        return out

    def keyword_search(self, query: str, area: str = None, confidence: str = None, limit: int = 10) -> list:
        """BM25 over chunk text, optionally filtered before ranking."""
        match = match_expression(query)
        if not match:
            return []
        sql = [
            "SELECT c.id, c.url, c.title, c.area, c.confidence, c.source_path,",
            "       c.chunk_index, snippet(chunk_fts, 4, '<mark>', '</mark>', '…', 18) AS snippet,",
            "       bm25(chunk_fts) AS score",
            "FROM chunk_fts JOIN chunks c ON c.rowid = chunk_fts.rowid",
            "WHERE chunk_fts MATCH ?",
        ]
        args = [match]
        if area:
            sql.append("AND c.area = ?")
            args.append(area)
        if confidence:
            sql.append("AND c.confidence = ?")
            args.append(confidence)
        sql.append("ORDER BY bm25(chunk_fts) LIMIT ?")
        args.append(int(limit))
        rows = self.db.execute(" ".join(sql), args)
        return [dict(row) for row in rows]

    # -- meta -------------------------------------------------------------

    def set_meta(self, key: str, value) -> None:
        self.db.execute(
            "INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)",
            (key, json.dumps(value, sort_keys=True)),
        )
        self.db.commit()

    def get_meta(self, key: str, default=None):
        row = self.db.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone()
        if row is None:
            return default
        return json.loads(row["value"])

    def meta(self) -> dict:
        return {row["key"]: json.loads(row["value"]) for row in self.db.execute("SELECT key, value FROM meta")}


def match_expression(query: str) -> str:
    """Free text -> a safe FTS5 expression: bare words, each a prefix, all required.

    Mirrors `MatchExpression` in internal/search/search.go so the site's search
    box and the retrieval path treat a visitor's words identically. FTS5
    operators in the input are never honoured.
    """
    terms = []
    for raw in re.split(r"[^0-9a-zA-Z\-]+", unicodedata.normalize("NFKC", query or "").lower()):
        term = raw.strip("-")
        if term:
            terms.append('"' + term.replace('"', "") + '"*')
    return " AND ".join(terms)