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
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user