feat(search): pgvector substrate — vector(384) + HNSW for semantic search

Move semantic_search_notes off the full-table Python cosine scan onto a native
pgvector column: indexed ORDER BY embedding <=> :q LIMIT k (HNSW, cosine).
Migration 0067 enables the extension, converts the JSONB embedding column to
vector(384) (stale-dim rows dropped and regenerated by the startup backfill),
and builds the HNSW cosine index. Postgres image moves postgres:16-alpine ->
pgvector/pgvector:pg17 across prod, quickstart, and CI.

Scribe: project 2, milestone 93, task 1031.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Xz4j1H7pjYSjKsEpgcNH5E

.forgejo/workflows/ci.yml

@@ -165,7 +165,9 @@ jobs:
       SECRET_KEY: ci_integration_placeholder
     services:
       postgres:
-        image: postgres:16-alpine
+        # pgvector image so `alembic upgrade head` can run migration 0067
+        # (CREATE EXTENSION vector). PG17 — matches the prod/quickstart image.
+        image: pgvector/pgvector:pg17
         env:
           POSTGRES_USER: scribe
           POSTGRES_PASSWORD: ci_integration
@@ -189,7 +191,7 @@ jobs:
           set -eux
           echo "=== container landscape (diagnostic for the name filter) ==="
           docker ps -a --format '{{.ID}} {{.Image}} -> {{.Names}}'
-          PG=$(docker ps --filter "name=integration" --filter "ancestor=postgres:16-alpine" -q | head -n1)
+          PG=$(docker ps --filter "name=integration" --filter "ancestor=pgvector/pgvector:pg17" -q | head -n1)
           test -n "$PG"
           PG_IP=$(docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$PG")
           test -n "$PG_IP"

alembic/versions/0067_pgvector_note_embeddings.py

@@ -0,0 +1,73 @@
+"""pgvector: note_embeddings.embedding JSONB -> vector(384) + HNSW index
+
+Revision ID: 0067
+Revises: 0066
+Create Date: 2026-06-22
+
+Moves semantic search off the full-table Python cosine scan onto a native
+pgvector column so ranking + top-k run as an indexed `ORDER BY embedding <=> :q
+LIMIT k` in Postgres (see services/embeddings.semantic_search_notes).
+
+Requires a Postgres image that bundles the `vector` extension — the stack moved
+from postgres:16-alpine to pgvector/pgvector:pg16 in the same change (compose +
+CI). `CREATE EXTENSION IF NOT EXISTS vector` below is the in-db half.
+
+Embeddings are DERIVED data (regenerated from note text by
+backfill_note_embeddings at startup), so this migration is free to drop any row
+it can't cleanly convert: only rows whose stored JSONB array is exactly 384-dim
+are carried over (guarding against stale vectors from an earlier model — the
+same mixed-dim hazard _cosine_similarity defended against). Dropped rows are
+re-embedded on next boot.
+"""
+from alembic import op
+
+
+revision = "0067"
+down_revision = "0066"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.execute("CREATE EXTENSION IF NOT EXISTS vector")
+
+    # New native-vector column, populated only from cleanly-convertible rows.
+    # A JSONB array like [0.1, 0.2, ...] renders to text that is exactly
+    # pgvector's input literal, so (embedding::text)::vector is a direct cast.
+    op.execute("ALTER TABLE note_embeddings ADD COLUMN embedding_vec vector(384)")
+    op.execute(
+        """
+        UPDATE note_embeddings
+           SET embedding_vec = (embedding::text)::vector
+         WHERE jsonb_array_length(embedding) = 384
+        """
+    )
+    # Stale-dim rows (couldn't convert) are derived data — drop and let the
+    # startup backfill regenerate them at the current dimension.
+    op.execute("DELETE FROM note_embeddings WHERE embedding_vec IS NULL")
+
+    op.execute("ALTER TABLE note_embeddings ALTER COLUMN embedding_vec SET NOT NULL")
+    op.execute("ALTER TABLE note_embeddings DROP COLUMN embedding")
+    op.execute("ALTER TABLE note_embeddings RENAME COLUMN embedding_vec TO embedding")
+
+    # HNSW index for cosine distance — matches Vector.cosine_distance (`<=>`).
+    op.execute(
+        """
+        CREATE INDEX ix_note_embeddings_embedding_hnsw
+            ON note_embeddings
+         USING hnsw (embedding vector_cosine_ops)
+        """
+    )
+
+
+def downgrade() -> None:
+    # Back to JSONB. pgvector renders a vector to a text literal that is a valid
+    # JSON array, so the reverse cast is symmetric. The `vector` extension is
+    # intentionally left installed (other objects may depend on it; dropping an
+    # extension is the riskier, rarely-wanted direction).
+    op.execute("DROP INDEX IF EXISTS ix_note_embeddings_embedding_hnsw")
+    op.execute("ALTER TABLE note_embeddings ADD COLUMN embedding_json jsonb")
+    op.execute("UPDATE note_embeddings SET embedding_json = (embedding::text)::jsonb")
+    op.execute("ALTER TABLE note_embeddings ALTER COLUMN embedding_json SET NOT NULL")
+    op.execute("ALTER TABLE note_embeddings DROP COLUMN embedding")
+    op.execute("ALTER TABLE note_embeddings RENAME COLUMN embedding_json TO embedding")

docker-compose.prod.yml

@@ -21,7 +21,11 @@ services:
         max_attempts: 5
 
   db:
-    image: postgres:16-alpine
+    # pgvector image (Debian/glibc, PG17) — bundles the `vector` extension that
+    # migration 0067 enables. Moved off postgres:16-alpine via logical
+    # dump/restore (which doubles as the PG16->PG17 major upgrade); see the
+    # TRANSITION runbook in the PR.
+    image: pgvector/pgvector:pg17
     stop_grace_period: 120s
     volumes:
       - pgdata:/var/lib/postgresql/data

docker-compose.quickstart.yml

@@ -35,7 +35,8 @@ services:
       start_period: 30s
 
   db:
-    image: postgres:16-alpine
+    # pgvector image (PG17) — bundles the `vector` extension (migration 0067).
+    image: pgvector/pgvector:pg17
     stop_grace_period: 120s
     volumes:
       - pgdata:/var/lib/postgresql/data

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "APScheduler>=3.10,<4.0",
     "mcp[cli]>=1.0",
     "fastembed>=0.4",
+    "pgvector>=0.3",
 ]
 
 [project.optional-dependencies]

src/scribe/models/embedding.py

@@ -1,11 +1,17 @@
 from datetime import datetime, timezone
 
+from pgvector.sqlalchemy import Vector
 from sqlalchemy import DateTime, ForeignKey, Integer
-from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import Mapped, mapped_column
 
 from scribe.models import Base
 
+# bge-small-en-v1.5 produces 384-dim unit-normalized vectors. The column is a
+# native pgvector `vector(384)` (see migration 0067) so similarity search runs
+# as an indexed `ORDER BY embedding <=> :q LIMIT k` in Postgres rather than a
+# full-table Python cosine scan.
+EMBEDDING_DIM = 384
+
 
 class NoteEmbedding(Base):
     """Stores the embedding vector for a note, used for semantic search."""
@@ -18,7 +24,7 @@ class NoteEmbedding(Base):
         primary_key=True,
     )
     user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
-    embedding: Mapped[list] = mapped_column(JSONB, nullable=False)
+    embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
     updated_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         default=lambda: datetime.now(timezone.utc),

src/scribe/services/embeddings.py

@@ -28,6 +28,10 @@ logger = logging.getLogger(__name__)
 # loosely-related results that pad the sidebar without adding real value.
 _SIMILARITY_THRESHOLD = 0.45
 
+# Public alias so callers (and telemetry) can record the effective default
+# threshold without reaching for the underscored name.
+DEFAULT_SIMILARITY_THRESHOLD = _SIMILARITY_THRESHOLD
+
 _MODEL_NAME = "BAAI/bge-small-en-v1.5"
 _CACHE_DIR = os.environ.get("FASTEMBED_CACHE_DIR", "/data/fastembed-cache")
 
@@ -115,6 +119,14 @@ async def semantic_search_notes(
 
     Scores are cosine similarities in [-1, 1]; only notes at or above
     *threshold* are returned, sorted highest-first.
+
+    Ranking and the top-k cut happen in Postgres via pgvector's cosine-distance
+    operator (`<=>`, exposed as ``Vector.cosine_distance``) backed by the HNSW
+    index from migration 0067 — so this is an indexed ``ORDER BY ... LIMIT k``
+    rather than a full-table scan. Cosine distance is ``1 - cosine_similarity``,
+    so a similarity floor of *threshold* is a distance ceiling of
+    ``1 - threshold`` and similarity is recovered as ``1 - distance``.
+
     Returns an empty list if the embedder is unavailable or on any error.
     """
     if not query or not query.strip():
@@ -125,10 +137,17 @@ async def semantic_search_notes(
         logger.debug("Semantic search skipped — embedder unavailable")
         return []
 
+    # Distance ceiling equivalent to the similarity floor. Clamp to the valid
+    # cosine-distance range [0, 2] so a threshold of, say, -1 doesn't produce a
+    # nonsensical ceiling.
+    max_distance = min(2.0, max(0.0, 1.0 - threshold))
+    distance = NoteEmbedding.embedding.cosine_distance(query_vec)
+
     try:
         async with async_session() as session:
             stmt = (
-                select(NoteEmbedding, Note)
+                select(Note, distance.label("distance"))
+                .select_from(NoteEmbedding)
                 .join(Note, NoteEmbedding.note_id == Note.id)
                 .where(NoteEmbedding.user_id == user_id, Note.deleted_at.is_(None))
             )
@@ -142,30 +161,14 @@ async def semantic_search_notes(
                 stmt = stmt.where(Note.status.is_(None))
             if exclude_ids:
                 stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids))
+            stmt = stmt.where(distance <= max_distance).order_by(distance.asc()).limit(limit)
             rows = list((await session.execute(stmt)).all())
     except Exception:
         logger.warning("Failed to query note embeddings", exc_info=True)
         return []
 
-    if not rows:
-        return []
-
-    def _score() -> list[tuple[float, Note]]:
-        out: list[tuple[float, Note]] = []
-        for ne, note in rows:
-            try:
-                sim = _cosine_similarity(query_vec, ne.embedding)
-            except Exception:
-                continue
-            if sim >= threshold:
-                out.append((sim, note))
-        out.sort(key=lambda x: x[0], reverse=True)
-        return out[:limit]
-
-    # Offload the O(rows) cosine scoring off the event loop so a large corpus
-    # doesn't stall other requests while ranking. Results are unchanged; the
-    # real scaling fix (ORDER BY / LIMIT in pgvector) is a separate effort.
-    return await asyncio.to_thread(_score)
+    # Recover similarity (1 - distance) and preserve the highest-first contract.
+    return [(1.0 - float(dist), note) for note, dist in rows]
 
 
 async def backfill_note_embeddings() -> None:

tests/test_integration_pgvector_search.py

@@ -0,0 +1,101 @@
+"""Real-Postgres integration test for pgvector semantic search.
+
+Runs only in the CI integration lane (real Postgres + `vector` extension +
+schema built by `alembic upgrade head`, which includes migration 0067). This
+exercises what the unit mocks cannot: the native `vector(384)` column, the
+`<=>` cosine-distance operator behind `Vector.cosine_distance`, the HNSW index,
+and the distance->similarity recovery in `semantic_search_notes`.
+
+The embedder itself is stubbed (get_embedding is patched) so the test does not
+depend on downloading the fastembed model — only the Postgres/pgvector path is
+under test.
+"""
+from unittest.mock import AsyncMock, patch
+
+import pytest
+import pytest_asyncio
+from sqlalchemy import delete
+
+from scribe.models import async_session, engine
+from scribe.models.embedding import EMBEDDING_DIM, NoteEmbedding
+from scribe.models.note import Note
+from scribe.models.user import User
+from scribe.services.embeddings import semantic_search_notes
+
+pytestmark = pytest.mark.integration
+
+
+def _vec(*nonzero_first):
+    """A 384-dim vector with the given leading values, zero-padded."""
+    v = list(nonzero_first) + [0.0] * (EMBEDDING_DIM - len(nonzero_first))
+    return v[:EMBEDDING_DIM]
+
+
+@pytest_asyncio.fixture(autouse=True)
+async def _dispose_engine():
+    # Per-loop pool: dispose after each test (see test_integration_db_maintenance).
+    yield
+    await engine.dispose()
+
+
+@pytest_asyncio.fixture
+async def seeded():
+    """Insert a user + a near and a far note with hand-crafted embeddings.
+
+    Returns (user_id, near_note_id, far_note_id). Cleaned up after the test.
+    """
+    async with async_session() as s:
+        user = User(username="pgvec_itest")
+        s.add(user)
+        await s.flush()
+        near = Note(user_id=user.id, title="near", body="near body")
+        far = Note(user_id=user.id, title="far", body="far body")
+        s.add_all([near, far])
+        await s.flush()
+        # query vector will be [1,0,0,...]; near ~ identical (sim≈1.0),
+        # far is orthogonal (sim≈0.0 -> filtered by the default threshold).
+        s.add(NoteEmbedding(note_id=near.id, user_id=user.id, embedding=_vec(1.0)))
+        s.add(NoteEmbedding(note_id=far.id, user_id=user.id, embedding=_vec(0.0, 1.0)))
+        await s.commit()
+        ids = (user.id, near.id, far.id)
+    yield ids
+    user_id = ids[0]
+    async with async_session() as s:
+        await s.execute(delete(NoteEmbedding).where(NoteEmbedding.user_id == user_id))
+        await s.execute(delete(Note).where(Note.user_id == user_id))
+        await s.execute(delete(User).where(User.id == user_id))
+        await s.commit()
+
+
+@pytest.mark.asyncio
+async def test_semantic_search_ranks_and_thresholds_via_pgvector(seeded):
+    user_id, near_id, far_id = seeded
+    with patch(
+        "scribe.services.embeddings.get_embedding",
+        AsyncMock(return_value=_vec(1.0)),
+    ):
+        results = await semantic_search_notes(user_id=user_id, query="anything", limit=10)
+
+    ids = [note.id for _score, note in results]
+    # Near note returned and ranked first; far (orthogonal, sim≈0) excluded by
+    # the default 0.45 similarity threshold.
+    assert near_id in ids
+    assert far_id not in ids
+    assert ids[0] == near_id
+    top_score = results[0][0]
+    assert top_score == pytest.approx(1.0, abs=1e-3)
+
+
+@pytest.mark.asyncio
+async def test_low_threshold_lets_orthogonal_through(seeded):
+    user_id, near_id, far_id = seeded
+    with patch(
+        "scribe.services.embeddings.get_embedding",
+        AsyncMock(return_value=_vec(1.0)),
+    ):
+        results = await semantic_search_notes(
+            user_id=user_id, query="anything", limit=10, threshold=-1.0,
+        )
+    ids = [note.id for _score, note in results]
+    # With the floor dropped, both come back and near still ranks above far.
+    assert ids.index(near_id) < ids.index(far_id)