pagepiper/app/services/colbert_index.py

# app/services/colbert_index.py
"""
ColBERT late-interaction index (pagepiper#8).

Replaces nomic-embed-text bi-encoder + cosine similarity for the semantic
half of hybrid search with Agent-ModernColBERT (lightonai/Agent-ModernColBERT),
a late-interaction retriever that keeps per-token embeddings and scores via
MaxSim at query time — better suited to multi-part rulebook questions than a
single collapsed query vector.

BSL 1.1 — same BYOK gate as the rest of hybrid search (Retriever only reaches
this index when an LLM is configured). The model itself runs locally via
`pylate`, no Ollama call required — the gate matches product tiering, not a
technical dependency on Ollama.

Mirrors BM25Index's dirty-flag, rebuild-from-SQLite pattern: no separate
per-shelve indexing step is needed. `mark_dirty()` is called by the same
callback that already marks BM25 dirty on shelve completion; the next query
triggers a full rebuild from `page_chunks`.
"""
from __future__ import annotations

import logging
import os
import sqlite3
import threading

logger = logging.getLogger(__name__)

_DEFAULT_MODEL = "lightonai/Agent-ModernColBERT"


class ColBERTIndex:
    def __init__(self, index_dir: str, model_name: str | None = None) -> None:
        self._index_dir = index_dir
        self._model_name = model_name or os.environ.get("PAGEPIPER_COLBERT_MODEL", _DEFAULT_MODEL)
        self._model = None
        self._index = None
        self._chunks: dict[str, dict] = {}
        self._dirty = True
        self._lock = threading.Lock()

    def mark_dirty(self) -> None:
        """Signal that the index needs rebuilding (call after any document is shelved)."""
        self._dirty = True

    def _get_model(self):
        if self._model is None:
            from pylate import models

            logger.info("Loading ColBERT model %s", self._model_name)
            self._model = models.ColBERT(model_name_or_path=self._model_name)
        return self._model

    def ensure_fresh(self, db_path: str) -> None:
        """Rebuild from SQLite if dirty."""
        if not self._dirty:
            return
        with self._lock:
            if not self._dirty:
                return
            try:
                conn = sqlite3.connect(db_path)
                conn.row_factory = sqlite3.Row
                try:
                    rows = conn.execute(
                        "SELECT id, doc_id, page_number, text FROM page_chunks ORDER BY doc_id, page_number"
                    ).fetchall()
                finally:
                    conn.close()
            except sqlite3.Error as exc:
                logger.error("ColBERT index rebuild failed: %s", exc)
                return

            self._chunks = {str(r["id"]): dict(r) for r in rows}

            if not rows:
                self._index = None
                self._dirty = False
                return

            from pylate import indexes

            model = self._get_model()
            ids = [str(r["id"]) for r in rows]
            texts = [r["text"] for r in rows]
            embeddings = model.encode(texts, is_query=False, show_progress_bar=False)

            os.makedirs(self._index_dir, exist_ok=True)
            index = indexes.Voyager(
                index_folder=self._index_dir, index_name="colbert", override=True
            )
            index.add_documents(documents_ids=ids, documents_embeddings=embeddings)
            self._index = index
            self._dirty = False
            logger.info("ColBERT index rebuilt: %d chunks", len(rows))

    def query(
        self,
        query_text: str,
        top_k: int = 10,
        doc_ids: list[str] | None = None,
    ) -> list[dict]:
        """Search the corpus. Returns results sorted by descending MaxSim score."""
        if self._index is None:
            return []

        from pylate import retrieve

        model = self._get_model()
        query_embeddings = model.encode([query_text], is_query=True, show_progress_bar=False)
        retriever = retrieve.ColBERT(index=self._index)

        # Oversample when filtering to a doc subset — same pattern as the
        # sqlite-vec path this replaces (see app/services/retriever.py).
        k = top_k * 20 if doc_ids else top_k * 2
        results = retriever.retrieve(queries_embeddings=query_embeddings, k=k)[0]

        hits: list[dict] = []
        for r in results:
            chunk = self._chunks.get(str(r["id"]))
            if not chunk:
                continue
            if doc_ids is not None and chunk["doc_id"] not in doc_ids:
                continue
            hits.append(
                {
                    "chunk_id": chunk["id"],
                    "doc_id": chunk["doc_id"],
                    "page_number": chunk["page_number"],
                    "text": chunk["text"],
                    "score": r["score"],
                }
            )
            if len(hits) >= top_k:
                break
        return hits