feat: optimise embeds

tests/scripts/test_reembed_pgvector_live.py

@@ -0,0 +1,187 @@
+"""Live pgvector run of the re-embed script.
+
+Uses the ephemeral pytest-postgresql cluster with a stub embeddings model, so
+a real ``UPDATE ... ::vector`` round trip is exercised without downloading a
+model. Skips when the cluster has no pgvector build.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+
+from application.scripts import reembed
+from application.vectorstore import pgvector as pgvector_module
+from application.vectorstore.pgvector import PGVectorStore
+
+pytestmark = pytest.mark.integration
+
+DIM = 8
+
+
+class _Embeddings:
+    """Returns a distinct constant per generation, so a rewrite is visible."""
+
+    dimension = DIM
+
+    def __init__(self, seed: float):
+        self.seed = seed
+        self.calls = 0
+
+    def embed_documents(self, texts):
+        self.calls += 1
+        return [[self.seed] + [0.0] * (DIM - 1) for _ in texts]
+
+    def embed_query(self, query):
+        return [self.seed] + [0.0] * (DIM - 1)
+
+
+def _dsn(info) -> str:
+    password = f":{info.password}" if info.password else ""
+    return f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}"
+
+
+@pytest.fixture(autouse=True)
+def _close_pools():
+    yield
+    for dsn, pool in list(pgvector_module._POOLS.items()):
+        try:
+            pool.close()
+        except Exception:
+            pass
+        pgvector_module._POOLS.pop(dsn, None)
+
+
+@pytest.fixture
+def live_dsn(postgresql, monkeypatch):
+    try:
+        with postgresql.cursor() as cursor:
+            cursor.execute("CREATE EXTENSION vector;")
+        postgresql.rollback()
+    except Exception as exc:
+        postgresql.rollback()
+        pytest.skip(f"pgvector extension unavailable: {exc}")
+
+    dsn = _dsn(postgresql.info)
+    from application.core import settings as settings_module
+
+    settings = settings_module.settings
+    monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False)
+    monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_IVFFLAT_PROBES", None, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_POOL_MAX_SIZE", 4, raising=False)
+    monkeypatch.setattr(settings, "EMBEDDINGS_NAME", "granite-311m", raising=False)
+    return dsn
+
+
+def _seed(dsn, source_id, texts, embeddings):
+    """Create the schema and insert ``texts`` embedded by ``embeddings``."""
+    with patch(
+        "application.vectorstore.base.BaseVectorStore._get_embeddings",
+        return_value=embeddings,
+    ):
+        store = PGVectorStore(source_id=source_id, connection_string=dsn)
+        conn = store._get_connection()
+        PGVectorStore.create_schema(conn, dimension=DIM)
+        conn.commit()
+        store.add_texts(list(texts), metadatas=[{"i": i} for i in range(len(texts))])
+        store.close()
+
+
+def _as_list(vector):
+    """Normalise a stored vector to a list of floats.
+
+    Depending on whether pgvector's adapter is registered on the reading
+    connection, the value comes back as a ``Vector`` or as its text form
+    ``'[1,0,...]'``. Both are valid; the assertions should not care.
+    """
+    if vector is None:
+        return []
+    if hasattr(vector, "to_list"):
+        return list(vector.to_list())
+    if isinstance(vector, str):
+        return [float(part) for part in vector.strip("[]").split(",") if part]
+    return list(vector)
+
+
+def _vectors(dsn, source_id):
+    store = PGVectorStore(source_id=source_id, connection_string=dsn)
+    conn = store._get_connection()
+    cursor = conn.cursor()
+    try:
+        cursor.execute(
+            "SELECT text, embedding FROM documents WHERE source_id = %s ORDER BY id",
+            (source_id,),
+        )
+        # pgvector hands back a ``Vector``; normalise to a plain list so the
+        # assertions read the same whichever adapter is registered.
+        return [(text, _as_list(vector)) for text, vector in cursor.fetchall()]
+    finally:
+        cursor.close()
+        store.close()
+
+
+TEXTS = ["alpha document", "beta document", "gamma document"]
+
+
+class TestReembedPgvectorLive:
+    def test_rewrites_vectors_and_preserves_text(self, live_dsn):
+        _seed(live_dsn, "src-a", TEXTS, _Embeddings(1.0))
+        before = _vectors(live_dsn, "src-a")
+        assert [row[0] for row in before] == TEXTS
+        assert all(row[1][0] == pytest.approx(1.0) for row in before)
+
+        new_model = _Embeddings(9.0)
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=new_model,
+        ):
+            seen, written = reembed.reembed_pgvector("src-a", batch_size=2, dry_run=False)
+
+        assert (seen, written) == (3, 3)
+        after = _vectors(live_dsn, "src-a")
+        # Text is untouched; only the vectors moved.
+        assert [row[0] for row in after] == TEXTS
+        assert all(row[1][0] == pytest.approx(9.0) for row in after)
+
+    def test_dry_run_counts_without_writing(self, live_dsn):
+        _seed(live_dsn, "src-b", TEXTS, _Embeddings(1.0))
+        new_model = _Embeddings(9.0)
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=new_model,
+        ):
+            seen, written = reembed.reembed_pgvector("src-b", batch_size=2, dry_run=True)
+
+        assert (seen, written) == (3, 0)
+        assert new_model.calls == 0, "dry run must not embed"
+        assert all(row[1][0] == pytest.approx(1.0) for row in _vectors(live_dsn, "src-b"))
+
+    def test_batches_are_respected(self, live_dsn):
+        _seed(live_dsn, "src-c", TEXTS, _Embeddings(1.0))
+        new_model = _Embeddings(9.0)
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=new_model,
+        ):
+            reembed.reembed_pgvector("src-c", batch_size=2, dry_run=False)
+        assert new_model.calls == 2, "3 chunks at batch 2 is two embed calls"
+
+    def test_only_the_named_source_is_touched(self, live_dsn):
+        _seed(live_dsn, "src-d", TEXTS, _Embeddings(1.0))
+        _seed(live_dsn, "src-e", TEXTS, _Embeddings(1.0))
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=_Embeddings(9.0),
+        ):
+            reembed.reembed_pgvector("src-d", batch_size=64, dry_run=False)
+
+        assert all(row[1][0] == pytest.approx(9.0) for row in _vectors(live_dsn, "src-d"))
+        assert all(row[1][0] == pytest.approx(1.0) for row in _vectors(live_dsn, "src-e"))
+
+    def test_source_discovery_lists_every_source(self, live_dsn):
+        _seed(live_dsn, "src-f", TEXTS, _Embeddings(1.0))
+        _seed(live_dsn, "src-g", TEXTS, _Embeddings(1.0))
+        assert reembed.list_source_ids("pgvector") == ["src-f", "src-g"]