pagepiper/app/main.py
pyr0ball 89a58ec9b0 feat: replace nomic-embed-text retriever with Agent-ModernColBERT
Bi-encoder embeddings collapse a whole query into one vector, losing
multi-part reasoning structure — queries like "the procedure for setting
an IP on an AVC-X" or "what is the action economy for a fighter casting
a spell while prone" lose nuance. Agent-ModernColBERT is a late-interaction
retriever: per-token embeddings, scored via MaxSim at query time, built
specifically for agentic/multi-hop queries.

Implements Option A from the issue (in-process, via `pylate`) rather than
Option B (managed cf-orch service) — cf-orch already has `agent-moderncolbert`
registered in model_registry.yaml with a `pagepiper/retrieve` assignment
in assignments.yaml pointing at it and referencing this issue directly,
someone had already pre-wired that side.

- app/services/colbert_index.py: new ColBERTIndex class, mirrors
  BM25Index's dirty-flag/rebuild-from-SQLite pattern exactly — no
  separate per-shelve indexing step needed, just mark_dirty() on the
  same callback that already marks BM25 dirty.
- app/services/retriever.py: hybrid_search's semantic half now merges
  BM25 with ColBERT MaxSim scores (min-max normalized per-batch, since
  MaxSim is unbounded unlike the old sqlite-vec L2-distance path) instead
  of Ollama-embed + sqlite-vec cosine. BM25 merge/rank/per-doc-cap/
  adjacent-chunk-window logic is unchanged.
- app/main.py / app/deps.py: per-user ColBERTIndex registry, same
  pattern as the existing per-user BM25Index registry.
- Existing BYOK tier gate preserved exactly (llm is None check) — this
  is a retrieval-technology swap, not a tier/licensing change. The
  ColBERT model runs locally via pylate with no Ollama dependency, but
  gating still follows product tiering.
- 12 new tests. pylate is intentionally NOT installed in the dev/test
  env — see the cf-sysadmin skill's "Known Gotchas" for why (installing
  it directly into the shared `cf` conda env broke several other
  services' torch/transformers pins on 2026-07-10). Tests inject fake
  pylate modules via sys.modules instead.

Known follow-up (not addressed here): shelve scripts still compute and
store Ollama embeddings into `page_vecs` at shelve time — that table is
no longer read by search/chat now that retrieval uses the ColBERT index.
Removing the now-redundant embedding step is separate cleanup.

Closes: #8
2026-07-10 19:02:12 -07:00

78 lines
2.7 KiB
Python

# app/main.py
"""FastAPI application factory for pagepiper."""
from __future__ import annotations
import logging
import os
from contextlib import asynccontextmanager
from fastapi import FastAPI
from app.config import DB_PATH, VEC_DB_PATH, VEC_DIMENSIONS
from app.services.bm25_index import BM25Index
from app.services.colbert_index import ColBERTIndex
logger = logging.getLogger("pagepiper")
# Per-user BM25 registry — keyed by user_id; "__local__" for single-user mode
_bm25_map: dict[str, BM25Index] = {}
# Per-user ColBERT registry — keyed by user_id; index files live under
# <user_dir>/colbert_index/ (see _get_colbert_for)
_colbert_map: dict[str, ColBERTIndex] = {}
def _get_bm25_for(user_id: str) -> BM25Index:
if user_id not in _bm25_map:
_bm25_map[user_id] = BM25Index()
return _bm25_map[user_id]
def _get_colbert_for(user_id: str, user_dir) -> ColBERTIndex:
if user_id not in _colbert_map:
index_dir = str(user_dir / "colbert_index")
_colbert_map[user_id] = ColBERTIndex(index_dir=index_dir)
return _colbert_map[user_id]
@asynccontextmanager
async def lifespan(app: FastAPI):
from app.cloud_session import CLOUD_MODE
from app.config import LOCAL_USER_ID
from app.startup import apply_migrations, check_and_rebuild_vec_schema
embed_model = os.environ.get("PAGEPIPER_EMBED_MODEL", "nomic-embed-text")
logger.info("Pagepiper starting — embed model: %s, dims: %d", embed_model, VEC_DIMENSIONS)
if CLOUD_MODE:
from app.startup import warn_if_unencrypted
from app.config import DATA_DIR
warn_if_unencrypted(str(DATA_DIR))
else:
# In cloud mode, per-user migration and vec schema check run on first request (deps.py).
from app.config import DATA_DIR
apply_migrations(DB_PATH)
check_and_rebuild_vec_schema(VEC_DB_PATH, VEC_DIMENSIONS, DB_PATH)
_get_bm25_for(LOCAL_USER_ID).mark_dirty()
_get_colbert_for(LOCAL_USER_ID, DATA_DIR).mark_dirty()
yield
app = FastAPI(title="Pagepiper", lifespan=lifespan)
# Register routers
from app.api.library import router as library_router # noqa: E402
from app.api.ingest import router as ingest_router # noqa: E402
from app.api.search import router as search_router # noqa: E402
from app.api.chat import router as chat_router # noqa: E402
from app.api.feedback import router as feedback_router # noqa: E402
from app.api.feedback_attach import router as feedback_attach_router # noqa: E402
app.include_router(library_router)
app.include_router(ingest_router)
app.include_router(search_router)
app.include_router(chat_router)
app.include_router(feedback_router, prefix="/api/v1/feedback")
app.include_router(feedback_attach_router, prefix="/api/v1/feedback")