wiki-kb/wikikb/indexer.py

"""Build the search index: SQLite (chunks + FTS5) and a row-aligned embeddings file."""

from __future__ import annotations

import sqlite3
import sys
from collections.abc import Callable
from pathlib import Path

import numpy as np

from . import config
from .chunker import Chunk, chunk_document
from .embed import embed_passages

_SCHEMA = """
CREATE TABLE chunks (
    id      INTEGER PRIMARY KEY,
    title   TEXT NOT NULL,
    url     TEXT NOT NULL,
    source  TEXT NOT NULL,
    variant TEXT NOT NULL,
    section TEXT NOT NULL,
    text    TEXT NOT NULL
);
CREATE INDEX chunks_source ON chunks(source);
CREATE INDEX chunks_variant ON chunks(variant);
CREATE VIRTUAL TABLE chunks_fts USING fts5(
    text, title, section,
    content='chunks', content_rowid='id'
);
"""

# Passages are stored at rowid = position+1; embeddings row `i` aligns with rowid i+1.
EmbedFn = Callable[[list[str]], np.ndarray]


def _embed_text(chunk: Chunk) -> str:
    """Context-augmented text for embedding: page title + section + body.

    The stored/displayed chunk text stays clean, but the embedding carries which
    page and section a chunk is from, so trait/topic queries rank the right page up
    instead of any page that merely mentions the query's generic words.
    """
    return f"{chunk.title}. {chunk.section}. {chunk.text}"


def embed_efficiently(texts: list[str], embed_fn: EmbedFn) -> np.ndarray:
    """Embed passages, avoiding two large sources of wasted compute.

    1. **Length bucketing.** fastembed pads every batch to its longest member, so one
       2000-char passage makes the other 31 in its batch pay that cost. Wiki chunks
       are wildly uneven (on the KCD wiki: 184-char median, 2000-char max, 53% under
       200 chars), so most of the compute was padding. Sorting by length before
       batching measured **2.7x faster** end to end.
    2. **Deduplication.** Identical passages are embedded once and the vector reused.
       Note this is usually a small win: ``_embed_text`` prefixes each passage with
       its page title, which makes per-page boilerplate distinct again (16% of raw
       chunk texts were duplicates on the KCD wiki, but only 0.1% after prefixing).

    Vectors are returned in the caller's ORIGINAL order. That is load-bearing: the
    index's one alignment invariant is that embedding row ``i`` is chunk rowid
    ``i+1``, so the sort must be undone here rather than leaking to the caller.
    """
    if not texts:
        return embed_fn([])
    first_index: dict[str, int] = {}
    for text in texts:
        first_index.setdefault(text, len(first_index))
    unique = list(first_index)
    by_length = sorted(range(len(unique)), key=lambda i: len(unique[i]))
    sorted_vectors = embed_fn([unique[i] for i in by_length])
    unique_vectors = np.empty_like(sorted_vectors)
    unique_vectors[by_length] = sorted_vectors  # undo the length sort
    return unique_vectors[[first_index[t] for t in texts]]  # scatter duplicates back


def collect_chunks(corpus_dir: Path) -> list[Chunk]:
    """Chunk every markdown document in the corpus, in stable filename order."""
    chunks: list[Chunk] = []
    for doc in sorted(corpus_dir.glob("*.md")):
        chunks.extend(chunk_document(doc.read_text(encoding="utf-8")))
    return chunks


def build_index(
    *,
    corpus_dir: Path | None = None,
    db_path: Path | None = None,
    embeddings_path: Path | None = None,
    embed_fn: EmbedFn = embed_passages,
    log: Callable[[str], None] = print,
) -> int:
    """(Re)build the index from the corpus. Returns the chunk count.

    Paths default to the live ``config`` values, resolved at call time so tests
    can redirect them via monkeypatch.
    """
    corpus_dir = corpus_dir or config.CORPUS_DIR
    db_path = db_path or config.DB_PATH
    embeddings_path = embeddings_path or config.EMBEDDINGS_PATH
    chunks = collect_chunks(corpus_dir)
    log(f"Chunked {len(chunks)} passages from {corpus_dir}")

    vectors = embed_efficiently([_embed_text(c) for c in chunks], embed_fn)
    log(f"Embedded {vectors.shape[0]} passages (dim={vectors.shape[1] if vectors.size else 0})")

    db_path.parent.mkdir(parents=True, exist_ok=True)
    db_path.unlink(missing_ok=True)
    con = sqlite3.connect(db_path)
    try:
        con.executescript(_SCHEMA)
        con.executemany(
            "INSERT INTO chunks (id, title, url, source, variant, section, text) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            [
                (i + 1, c.title, c.url, c.source, c.variant, c.section, c.text)
                for i, c in enumerate(chunks)
            ],
        )
        con.execute(
            "INSERT INTO chunks_fts (rowid, text, title, section) "
            "SELECT id, text, title, section FROM chunks"
        )
        con.commit()
    finally:
        con.close()

    np.save(embeddings_path, vectors)
    log(f"Wrote {db_path} and {embeddings_path}")
    return len(chunks)


def main(argv: list[str] | None = None) -> int:
    count = build_index()
    print(f"\nIndex built: {count} chunks.")
    return 0


if __name__ == "__main__":  # pragma: no cover
    from . import cpu

    pinned = cpu.isolate()
    print(f"CPU-isolated to {len(pinned)} cores (reserved {sorted(config.RESERVED_CORES)}).")
    sys.exit(main())