FabledScribe/tests/test_integration_pgvector_search.py

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


def _emb(note_id, user_id, chunk_index, vec):
    """A chunk row at the current chunker version (#280, migration 0077)."""
    from scribe.services.embeddings import CHUNKER_VERSION

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


@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).
        # near gets a SECOND, weaker chunk (sim≈0.6) — the collapse to
        # best-chunk-per-note (#280) 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 (#280): 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)