From 513019786e48a57d9b615bcc0217b71cad77552a Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Mon, 22 Jun 2026 20:10:15 -0400 Subject: [PATCH] =?UTF-8?q?feat(search):=20pgvector=20substrate=20?= =?UTF-8?q?=E2=80=94=20vector(384)=20+=20HNSW=20for=20semantic=20search?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_01Xz4j1H7pjYSjKsEpgcNH5E --- .forgejo/workflows/ci.yml | 6 +- .../versions/0067_pgvector_note_embeddings.py | 73 +++++++++++++ docker-compose.prod.yml | 6 +- docker-compose.quickstart.yml | 3 +- pyproject.toml | 1 + src/scribe/models/embedding.py | 10 +- src/scribe/services/embeddings.py | 43 ++++---- tests/test_integration_pgvector_search.py | 101 ++++++++++++++++++ 8 files changed, 217 insertions(+), 26 deletions(-) create mode 100644 alembic/versions/0067_pgvector_note_embeddings.py create mode 100644 tests/test_integration_pgvector_search.py diff --git a/.forgejo/workflows/ci.yml b/.forgejo/workflows/ci.yml index 8a3aa81..9bf7ac0 100644 --- a/.forgejo/workflows/ci.yml +++ b/.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" diff --git a/alembic/versions/0067_pgvector_note_embeddings.py b/alembic/versions/0067_pgvector_note_embeddings.py new file mode 100644 index 0000000..14c20e4 --- /dev/null +++ b/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") diff --git a/docker-compose.prod.yml b/docker-compose.prod.yml index 3d07c21..a540bf6 100644 --- a/docker-compose.prod.yml +++ b/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 diff --git a/docker-compose.quickstart.yml b/docker-compose.quickstart.yml index 6fd53d0..80011b3 100644 --- a/docker-compose.quickstart.yml +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index be0e53f..420c60d 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/scribe/models/embedding.py b/src/scribe/models/embedding.py index de1a282..5de4fd9 100644 --- a/src/scribe/models/embedding.py +++ b/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), diff --git a/src/scribe/services/embeddings.py b/src/scribe/services/embeddings.py index e9eaf2e..b893345 100644 --- a/src/scribe/services/embeddings.py +++ b/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: diff --git a/tests/test_integration_pgvector_search.py b/tests/test_integration_pgvector_search.py new file mode 100644 index 0000000..d9938e6 --- /dev/null +++ b/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)