"""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)