mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 04:17:56 +00:00
feat: ingestion optimisations
This commit is contained in:
Alex committed
1 parent
949b8d7a18
commit
e0aff39a1b
8 files changed
+1037
-114
No files matched your search
@@ -89,21 +89,22 @@ class TestEmbedCheckpoint:
|
||||
docs = _make_docs(5)
|
||||
captured_chunks: list[str] = []
|
||||
|
||||
def _fake_add(store, doc, sid):
|
||||
captured_chunks.append(doc.page_content)
|
||||
def _fake_add(store, batch, sid):
|
||||
captured_chunks.extend(d.page_content for d in batch)
|
||||
|
||||
# The retry decorator wraps the real fn. Patch the module-level name
|
||||
# so our loop calls the spy directly.
|
||||
# so our loop calls the spy directly. The loop embeds a batch at a
|
||||
# time, so the spy receives a list and flattens it back out.
|
||||
import application.parser.embedding_pipeline as ep_mod
|
||||
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = _fake_add
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = _fake_add
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
docs, str(tmp_path), source_id, MagicMock()
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
# Faiss seeds the store with docs[0]; the loop picks up at idx=1.
|
||||
assert captured_chunks == ["chunk-1", "chunk-2", "chunk-3", "chunk-4"]
|
||||
@@ -136,17 +137,17 @@ class TestEmbedCheckpoint:
|
||||
|
||||
captured_chunks: list[str] = []
|
||||
|
||||
def _fake_add(store, doc, sid):
|
||||
captured_chunks.append(doc.page_content)
|
||||
def _fake_add(store, batch, sid):
|
||||
captured_chunks.extend(d.page_content for d in batch)
|
||||
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = _fake_add
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = _fake_add
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
docs, str(tmp_path), source_id, MagicMock()
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
# On resume the FAISS store is loaded from storage (no docs_init);
|
||||
# the loop iterates the un-popped docs list starting at resume_index.
|
||||
@@ -187,14 +188,14 @@ class TestEmbedCheckpoint:
|
||||
docs = _make_docs(6)
|
||||
_seed_progress_row(pg_conn, source_id, total=6, last_index=2)
|
||||
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
docs, str(tmp_path), source_id, MagicMock()
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
assert len(captured_kwargs) == 1
|
||||
# No ``docs_init`` on resume — this is what triggers FaissStore to
|
||||
@@ -219,14 +220,14 @@ class TestEmbedCheckpoint:
|
||||
docs = _make_docs(4)
|
||||
_seed_progress_row(pg_conn, source_id, total=4, last_index=1)
|
||||
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
docs, str(tmp_path), source_id, MagicMock()
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
fake_store.delete_index.assert_not_called()
|
||||
|
||||
@@ -245,9 +246,11 @@ class TestEmbedCheckpoint:
|
||||
)
|
||||
|
||||
captured: list[str] = []
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = (
|
||||
lambda store, doc, sid: captured.append(doc.page_content)
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = (
|
||||
lambda store, batch, sid: captured.extend(
|
||||
d.page_content for d in batch
|
||||
)
|
||||
)
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
@@ -255,7 +258,7 @@ class TestEmbedCheckpoint:
|
||||
attempt_id="att-A",
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
# Same attempt → resume past the last persisted index.
|
||||
assert captured == ["chunk-3", "chunk-4", "chunk-5"]
|
||||
@@ -278,9 +281,11 @@ class TestEmbedCheckpoint:
|
||||
docs = _make_docs(5)
|
||||
|
||||
captured: list[str] = []
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = (
|
||||
lambda store, doc, sid: captured.append(doc.page_content)
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = (
|
||||
lambda store, batch, sid: captured.extend(
|
||||
d.page_content for d in batch
|
||||
)
|
||||
)
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
@@ -288,7 +293,7 @@ class TestEmbedCheckpoint:
|
||||
attempt_id="att-new",
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
# Fresh attempt → reset to chunk 0; FAISS branch seeds with
|
||||
# docs[0] and the loop picks up at idx=1.
|
||||
@@ -327,9 +332,11 @@ class TestEmbedCheckpoint:
|
||||
docs = _make_docs(5)
|
||||
|
||||
captured: list[str] = []
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = (
|
||||
lambda store, doc, sid: captured.append(doc.page_content)
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = (
|
||||
lambda store, batch, sid: captured.extend(
|
||||
d.page_content for d in batch
|
||||
)
|
||||
)
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
@@ -337,7 +344,7 @@ class TestEmbedCheckpoint:
|
||||
attempt_id="sync-2",
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
# All non-seed chunks re-embedded under the new attempt.
|
||||
assert captured == ["chunk-1", "chunk-2", "chunk-3", "chunk-4"]
|
||||
@@ -358,16 +365,18 @@ class TestEmbedCheckpoint:
|
||||
docs = _make_docs(4)
|
||||
|
||||
captured: list[str] = []
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = (
|
||||
lambda store, doc, sid: captured.append(doc.page_content)
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = (
|
||||
lambda store, batch, sid: captured.extend(
|
||||
d.page_content for d in batch
|
||||
)
|
||||
)
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
docs, str(tmp_path), source_id, MagicMock(),
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
# Resumed past last_index=1.
|
||||
assert captured == ["chunk-2", "chunk-3"]
|
||||
@@ -387,15 +396,15 @@ class TestEmbedCheckpoint:
|
||||
source_id = str(uuid.uuid4())
|
||||
docs = _make_docs(1)
|
||||
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
docs, str(tmp_path), source_id, MagicMock(),
|
||||
attempt_id="att-single",
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
@@ -426,15 +435,15 @@ class TestEmbedCheckpoint:
|
||||
source_id = str(uuid.uuid4())
|
||||
docs = _make_docs(4)
|
||||
|
||||
original = ep_mod.add_text_to_store_with_retry
|
||||
ep_mod.add_text_to_store_with_retry = lambda store, doc, sid: None
|
||||
original = ep_mod.add_texts_to_store_with_retry
|
||||
ep_mod.add_texts_to_store_with_retry = lambda store, batch, sid: None
|
||||
try:
|
||||
ep_mod.embed_and_store_documents(
|
||||
docs, str(tmp_path), source_id, MagicMock(),
|
||||
attempt_id="att-multi",
|
||||
)
|
||||
finally:
|
||||
ep_mod.add_text_to_store_with_retry = original
|
||||
ep_mod.add_texts_to_store_with_retry = original
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
|
||||
Reference in new issue
Block a user