mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 22:15:10 +00:00
feat: optimise embeds
This commit is contained in:
Alex committed
1 parent
e0aff39a1b
commit
8380f9bb47
28 files changed
+1992
-349
No files matched your search
@@ -0,0 +1,124 @@
|
||||
"""Re-embed script: CLI contract, orchestration, and the FAISS rebuild path."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from application.scripts import reembed
|
||||
|
||||
|
||||
class TestCLI:
|
||||
def test_defaults(self):
|
||||
args = reembed.build_parser().parse_args([])
|
||||
assert args.dry_run is False
|
||||
assert args.sources is None
|
||||
assert args.batch_size == reembed.DEFAULT_BATCH_SIZE
|
||||
|
||||
def test_unsupported_store_exits_without_touching_anything(self):
|
||||
with patch.object(reembed.settings, "VECTOR_STORE", "qdrant", create=True):
|
||||
with patch.object(reembed, "run") as run:
|
||||
assert reembed.main([]) == 2
|
||||
run.assert_not_called()
|
||||
|
||||
def test_supported_stores_are_pgvector_and_faiss(self):
|
||||
assert set(reembed.SUPPORTED_STORES) == {"pgvector", "faiss"}
|
||||
|
||||
def test_sources_are_split_and_trimmed(self):
|
||||
with patch.object(reembed.settings, "VECTOR_STORE", "faiss", create=True):
|
||||
with patch.object(reembed, "run", return_value=0) as run:
|
||||
reembed.main(["--sources", " a , b ,, c "])
|
||||
assert run.call_args.args[1] == ["a", "b", "c"]
|
||||
|
||||
def test_batch_size_is_clamped_to_at_least_one(self):
|
||||
with patch.object(reembed.settings, "VECTOR_STORE", "faiss", create=True):
|
||||
with patch.object(reembed, "run", return_value=0) as run:
|
||||
reembed.main(["--batch-size", "0"])
|
||||
assert run.call_args.args[2] == 1
|
||||
|
||||
|
||||
class TestRun:
|
||||
def test_no_sources_is_a_clean_exit(self):
|
||||
with patch.object(reembed, "list_source_ids", return_value=[]):
|
||||
assert reembed.run("faiss", None, 8, False) == 0
|
||||
|
||||
def test_processes_every_discovered_source(self):
|
||||
with patch.object(reembed, "list_source_ids", return_value=["a", "b"]):
|
||||
with patch.object(reembed, "reembed_faiss", return_value=(3, 3)) as handler:
|
||||
assert reembed.run("faiss", None, 8, False) == 0
|
||||
assert [call.args[0] for call in handler.call_args_list] == ["a", "b"]
|
||||
|
||||
def test_explicit_sources_skip_discovery(self):
|
||||
with patch.object(reembed, "list_source_ids") as discover:
|
||||
with patch.object(reembed, "reembed_faiss", return_value=(1, 1)):
|
||||
reembed.run("faiss", ["only-this"], 8, False)
|
||||
discover.assert_not_called()
|
||||
|
||||
def test_one_failing_source_does_not_stop_the_others(self):
|
||||
def handler(source_id, batch_size, dry_run):
|
||||
if source_id == "bad":
|
||||
raise RuntimeError("boom")
|
||||
return (1, 1)
|
||||
|
||||
with patch.object(reembed, "reembed_faiss", side_effect=handler) as spy:
|
||||
code = reembed.run("faiss", ["good", "bad", "also-good"], 8, False)
|
||||
assert code == 1, "a failure must be reported in the exit code"
|
||||
assert spy.call_count == 3, "later sources must still be attempted"
|
||||
|
||||
def test_dry_run_is_reported_as_success(self):
|
||||
with patch.object(reembed, "reembed_faiss", return_value=(5, 0)):
|
||||
assert reembed.run("faiss", ["a"], 8, True) == 0
|
||||
|
||||
def test_pgvector_uses_the_pgvector_handler(self):
|
||||
with patch.object(reembed, "reembed_pgvector", return_value=(1, 1)) as handler:
|
||||
reembed.run("pgvector", ["a"], 8, False)
|
||||
handler.assert_called_once()
|
||||
|
||||
|
||||
class TestFaissRebuild:
|
||||
@pytest.fixture
|
||||
def stores(self):
|
||||
"""An existing store to read from and the rebuilt one written back."""
|
||||
existing = MagicMock()
|
||||
existing.get_chunks.return_value = [
|
||||
{"doc_id": "1", "text": "alpha", "metadata": {"i": 0}},
|
||||
{"doc_id": "2", "text": "beta", "metadata": {"i": 1}},
|
||||
]
|
||||
rebuilt = MagicMock()
|
||||
with patch.object(
|
||||
reembed.VectorCreator, "create_vectorstore", side_effect=[existing, rebuilt]
|
||||
) as factory:
|
||||
yield existing, rebuilt, factory
|
||||
|
||||
def test_rebuilds_from_stored_text_and_saves(self, stores):
|
||||
existing, rebuilt, factory = stores
|
||||
seen, written = reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
|
||||
|
||||
assert (seen, written) == (2, 2)
|
||||
rebuilt.save_local.assert_called_once()
|
||||
docs = factory.call_args_list[1].kwargs["docs_init"]
|
||||
assert [d.page_content for d in docs] == ["alpha", "beta"]
|
||||
assert [d.metadata for d in docs] == [{"i": 0}, {"i": 1}]
|
||||
|
||||
def test_existing_index_is_not_deleted(self, stores):
|
||||
"""The rebuild must not destroy the old index before the new one exists."""
|
||||
existing, rebuilt, _ = stores
|
||||
reembed.reembed_faiss("s1", batch_size=8, dry_run=False)
|
||||
existing.delete_index.assert_not_called()
|
||||
|
||||
def test_dry_run_reads_but_never_rebuilds(self):
|
||||
existing = MagicMock()
|
||||
existing.get_chunks.return_value = [{"doc_id": "1", "text": "a", "metadata": {}}]
|
||||
with patch.object(
|
||||
reembed.VectorCreator, "create_vectorstore", return_value=existing
|
||||
) as factory:
|
||||
seen, written = reembed.reembed_faiss("s1", batch_size=8, dry_run=True)
|
||||
assert (seen, written) == (1, 0)
|
||||
assert factory.call_count == 1, "no rebuild store may be constructed"
|
||||
|
||||
def test_empty_index_is_a_no_op(self):
|
||||
existing = MagicMock()
|
||||
existing.get_chunks.return_value = []
|
||||
with patch.object(
|
||||
reembed.VectorCreator, "create_vectorstore", return_value=existing
|
||||
):
|
||||
assert reembed.reembed_faiss("s1", batch_size=8, dry_run=False) == (0, 0)
|
||||
@@ -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"]
|
||||
Reference in new issue
Block a user