bvandeusenandClaude Opus 5.5
79de8938b9
Module and test docstrings state what is pinned without the record number or step that introduced it. Incident accounts become the failure shape in the present tense: the pipefail head, the code_shapes import, the export order, the facet dialects, the cross-kind repeat, the write-once kind. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
117 lines
4.4 KiB
Python
117 lines
4.4 KiB
Python
"""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
|
|
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, pytest.mark.usefixtures("_dispose_engine")]
|
|
|
|
|
|
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]
|
|
|
|
|
|
def _emb(note_id, user_id, chunk_index, vec):
|
|
"""A chunk row at the current chunker version."""
|
|
from scribe.services.embeddings import CHUNKER_VERSION, EMBEDDING_MODEL
|
|
|
|
return NoteEmbedding(
|
|
note_id=note_id,
|
|
chunk_index=chunk_index,
|
|
user_id=user_id,
|
|
embedding=vec,
|
|
chunk_text=f"chunk {chunk_index} of note {note_id}",
|
|
chunker_version=CHUNKER_VERSION,
|
|
embedding_model=EMBEDDING_MODEL,
|
|
)
|
|
|
|
|
|
@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).
|
|
# near gets a SECOND, weaker chunk (sim≈0.6) — the collapse to
|
|
# best-chunk-per-note is under test: near must come back once,
|
|
# at its best chunk's score, not twice.
|
|
s.add(_emb(near.id, user.id, 0, _vec(1.0)))
|
|
s.add(_emb(near.id, user.id, 1, _vec(0.6, 0.8)))
|
|
s.add(_emb(far.id, user.id, 0, _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
|
|
# Chunk collapse: near has TWO chunk rows above the floor (sim≈1.0
|
|
# and ≈0.6) and must appear exactly once, at its best chunk's score.
|
|
assert ids.count(near_id) == 1
|
|
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)
|