tests/rag/test_pgvector.py

"""Guarded pgvector integration test.

Runs only when TEST_DATABASE_URL points at a Postgres instance with the
pgvector extension available. Verifies the pgvector ordering path matches
the in-Python cosine ranking.
"""

from __future__ import annotations

import os

import pytest

TEST_DB = os.environ.get("TEST_DATABASE_URL", "")
pytestmark = pytest.mark.skipif(
    not TEST_DB.startswith("postgresql"),
    reason="set TEST_DATABASE_URL to a postgres URL to run pgvector tests",
)


async def test_pgvector_search_orders_like_python(tmp_path, monkeypatch):
    from sqlalchemy import text

    from app import db
    from app import embeddings as emb_mod
    from app.models import KnowledgeSource
    from app.rag import index as index_mod

    def _vec(t: str) -> list[float]:
        t = t.lower()
        base = [0.0] * 384
        base[0] = 1.0 if "envio" in t or "envío" in t else 0.0
        base[1] = 1.0 if "garant" in t else 0.0
        return base

    async def fake_embed_texts(texts, *, kind="passage"):
        return [_vec(t) for t in texts]

    async def fake_embed_query(q):
        return _vec(q)

    monkeypatch.setattr(emb_mod, "embed_texts", fake_embed_texts)
    monkeypatch.setattr(emb_mod, "embed_query", fake_embed_query)

    db.init_engine(TEST_DB)
    async with db.get_engine().begin() as conn:
        await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
    await db.create_all()

    async with db.get_sessionmaker()() as session:
        for name, body in [("envios", "Envíos gratis"), ("garantia", "La garantía dura")]:
            p = tmp_path / f"{name}.txt"
            p.write_text(body, encoding="utf-8")
            src = KnowledgeSource(kind="file", name=name, location=str(p))
            session.add(src)
            await session.flush()
            await index_mod.index_source(session, src)
        await session.commit()

        results = await index_mod.search(session, "mi envío", k=1)
        assert results
        top, _ = results[0]
        assert top.meta["source_name"] == "envios"

    async with db.get_engine().begin() as conn:
        await conn.run_sync(db.Base.metadata.drop_all)
    await db.get_engine().dispose()