mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 14:15:52 +00:00
Keep the chunker's metadata when re-embedding a wiki page
`reembed_wiki_page_worker` built each chunk's metadata from scratch as
`{source, title, filename}`, discarding everything the chunker attached
to the document -- including the per-chunk `token_count` every other
ingest path records. Wiki chunks therefore reached the source viewer
with no count at all.
Carry `extra_info` over and let the page's own path and title win on top,
so a chunk that inherited a stale source from an earlier conversion is
still keyed and filtered correctly.
This commit is contained in:
1 parent
0aecdc32c9
commit
d787f54c2d
2 files changed
+80
-4
No files matched your search
+6
-4
@@ -2615,10 +2615,12 @@ def reembed_wiki_page_worker(self, source_id, path, content_hash, user):
|
||||
|
||||
added = 0
|
||||
for chunk in chunks:
|
||||
store.add_chunk(
|
||||
chunk.text,
|
||||
metadata={"source": path, "title": title, "filename": path},
|
||||
)
|
||||
# Start from what the chunker produced -- ``token_count`` above
|
||||
# all, which the source viewer reads per chunk -- and let the
|
||||
# page's own identity win over anything stale it inherited.
|
||||
metadata = dict(chunk.extra_info or {})
|
||||
metadata.update({"source": path, "title": title, "filename": path})
|
||||
store.add_chunk(chunk.text, metadata=metadata)
|
||||
added += 1
|
||||
|
||||
with db_session() as conn:
|
||||
|
||||
@@ -100,6 +100,80 @@ class TestReembedWikiPageWorker:
|
||||
)
|
||||
assert result == {"status": "embedded", "added": 2, "deleted": 3}
|
||||
|
||||
def test_chunk_metadata_survives_the_rewrite(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
||||
):
|
||||
"""The chunker's own metadata -- ``token_count`` above all -- has to
|
||||
reach the store, or the source viewer shows "-" for every wiki chunk.
|
||||
"""
|
||||
from docsgpt import worker
|
||||
|
||||
source_id = _seed_source(pg_conn)
|
||||
|
||||
store = MagicMock(name="vector_store")
|
||||
store.delete_chunks_by_source_path.return_value = 0
|
||||
_patch_store(monkeypatch, store)
|
||||
|
||||
_patch_repo(monkeypatch, {"content": "page body", "title": "Page Title"})
|
||||
_patch_chunker(
|
||||
monkeypatch,
|
||||
[
|
||||
Document(text="c1", extra_info={"token_count": 11, "custom": "keep"}),
|
||||
Document(text="c2", extra_info={"token_count": 7}),
|
||||
],
|
||||
)
|
||||
|
||||
worker.reembed_wiki_page_worker(
|
||||
task_self, source_id, "guide/intro.md", "hash-1", "alice"
|
||||
)
|
||||
|
||||
counts = [
|
||||
call.kwargs["metadata"]["token_count"]
|
||||
for call in store.add_chunk.call_args_list
|
||||
]
|
||||
assert counts == [11, 7]
|
||||
assert store.add_chunk.call_args_list[0].kwargs["metadata"]["custom"] == "keep"
|
||||
|
||||
def test_page_identity_wins_over_stale_chunk_metadata(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
||||
):
|
||||
"""Whatever the chunker inherited, the page's own path and title are
|
||||
what the store is keyed and filtered on.
|
||||
"""
|
||||
from docsgpt import worker
|
||||
|
||||
source_id = _seed_source(pg_conn)
|
||||
|
||||
store = MagicMock(name="vector_store")
|
||||
store.delete_chunks_by_source_path.return_value = 0
|
||||
_patch_store(monkeypatch, store)
|
||||
|
||||
_patch_repo(monkeypatch, {"content": "body", "title": "Page Title"})
|
||||
_patch_chunker(
|
||||
monkeypatch,
|
||||
[
|
||||
Document(
|
||||
text="c1",
|
||||
extra_info={
|
||||
"source": "stale/path.md",
|
||||
"title": "Stale",
|
||||
"filename": "stale/path.md",
|
||||
"token_count": 5,
|
||||
},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
worker.reembed_wiki_page_worker(
|
||||
task_self, source_id, "guide/intro.md", "hash-1", "alice"
|
||||
)
|
||||
|
||||
metadata = store.add_chunk.call_args_list[0].kwargs["metadata"]
|
||||
assert metadata["source"] == "guide/intro.md"
|
||||
assert metadata["filename"] == "guide/intro.md"
|
||||
assert metadata["title"] == "Page Title"
|
||||
assert metadata["token_count"] == 5
|
||||
|
||||
def test_a_reembed_stamps_the_model_on_the_source(
|
||||
self, pg_conn, patch_worker_db, task_self, monkeypatch
|
||||
):
|
||||
|
||||
Reference in new issue
Block a user