pagepiper/app/deps.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

84 lines
2.5 KiB
Python

# app/deps.py
"""FastAPI dependency providers."""
from __future__ import annotations
import sqlite3
from dataclasses import dataclass
from pathlib import Path
from typing import Generator
from fastapi import Depends, Request
from app.config import DATA_DIR, LOCAL_USER_ID
from app.services.bm25_index import BM25Index
from app.services.colbert_index import ColBERTIndex
@dataclass
class UserCtx:
"""Per-request context routing DB paths, BM25, and ColBERT to the right user."""
user_id: str
db_path: str
vec_db_path: str
data_dir: Path
watch_dir: Path
bm25: BM25Index
colbert: ColBERTIndex
_user_startup_done: set[str] = set()
def _run_user_startup(user_id: str, user_dir: Path) -> None:
"""Run migrations and vec schema check once per process lifetime per user."""
if user_id in _user_startup_done:
return
_user_startup_done.add(user_id)
from app.config import VEC_DIMENSIONS
from app.startup import apply_migrations, check_and_rebuild_vec_schema
apply_migrations(str(user_dir / "pagepiper.db"))
check_and_rebuild_vec_schema(
str(user_dir / "pagepiper_vecs.db"), VEC_DIMENSIONS, str(user_dir / "pagepiper.db")
)
def get_user_ctx(request: Request) -> UserCtx:
"""Resolve the per-user data directory, DB paths, and BM25 instance for this request."""
import app.main as _main
from app.cloud_session import CLOUD_MODE
if CLOUD_MODE:
from app.cloud_session import resolve_authenticated_user
from app.config import user_data_dir
user_id = resolve_authenticated_user(request)
user_dir = user_data_dir(user_id)
_run_user_startup(user_id, user_dir)
watch_dir = user_dir / "books"
watch_dir.mkdir(parents=True, exist_ok=True)
else:
from app.config import WATCH_DIR
user_id = LOCAL_USER_ID
user_dir = DATA_DIR
watch_dir = WATCH_DIR
return UserCtx(
user_id=user_id,
db_path=str(user_dir / "pagepiper.db"),
vec_db_path=str(user_dir / "pagepiper_vecs.db"),
data_dir=user_dir,
watch_dir=watch_dir,
bm25=_main._get_bm25_for(user_id),
colbert=_main._get_colbert_for(user_id, user_dir),
)
def get_db(ctx: UserCtx = Depends(get_user_ctx)) -> Generator[sqlite3.Connection, None, None]:
conn = sqlite3.connect(ctx.db_path, check_same_thread=False)
conn.execute("PRAGMA foreign_keys = ON")
conn.execute("PRAGMA journal_mode = WAL")
conn.row_factory = sqlite3.Row
try:
yield conn
finally:
conn.close()