mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-15 21:16:03 +08:00
fix: track wiki chunk changes incrementally (#18285)
This commit is contained in:
@@ -420,3 +420,70 @@ def test_wiki_alteration_treats_wiki_template_as_eligible(monkeypatch):
|
||||
]
|
||||
|
||||
assert module._eligible_doc_ids_for_kind(docs, "tenant-1", "wiki") == {"doc-wiki"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wiki_chunk_alteration_uses_full_successful_state(monkeypatch):
|
||||
module, _, _ = _load_list_datasets_module(
|
||||
monkeypatch,
|
||||
kbs=[],
|
||||
parsing_status_by_kb={},
|
||||
)
|
||||
|
||||
previous = {
|
||||
"chunk-changed": {"doc_id": "doc-existing", "hash": "old"},
|
||||
"chunk-deleted": {"doc_id": "doc-existing", "hash": "same"},
|
||||
"chunk-template-off": {"doc_id": "doc-template-off", "hash": "same"},
|
||||
"chunk-removed-doc": {"doc_id": "doc-removed", "hash": "same"},
|
||||
}
|
||||
current = {
|
||||
"chunk-changed": {"doc_id": "doc-existing", "hash": "new"},
|
||||
"chunk-new-existing": {"doc_id": "doc-existing", "hash": "same"},
|
||||
"chunk-new-document": {"doc_id": "doc-new", "hash": "same"},
|
||||
}
|
||||
|
||||
async def _load_state(tenant_id, dataset_id):
|
||||
return previous
|
||||
|
||||
async def _scan_state(tenant_id, dataset_id, doc_ids):
|
||||
assert doc_ids == {"doc-existing", "doc-new"}
|
||||
return current
|
||||
|
||||
def _compare_states(old, new):
|
||||
old_ids = set(old)
|
||||
new_ids = set(new)
|
||||
common = old_ids & new_ids
|
||||
return {
|
||||
"new_chunk_ids": new_ids - old_ids,
|
||||
"changed_chunk_ids": {chunk_id for chunk_id in common if old[chunk_id]["hash"] != new[chunk_id]["hash"]},
|
||||
"deleted_chunk_ids": old_ids - new_ids,
|
||||
"unchanged_chunk_ids": {chunk_id for chunk_id in common if old[chunk_id]["hash"] == new[chunk_id]["hash"]},
|
||||
}
|
||||
|
||||
_stub(
|
||||
monkeypatch,
|
||||
"rag.advanced_rag.knowlege_compile.wiki",
|
||||
_wiki_compare_chunk_states=_compare_states,
|
||||
_wiki_load_active_map_state=_load_state,
|
||||
_wiki_scan_current_chunk_state=_scan_state,
|
||||
)
|
||||
|
||||
result = await module._wiki_chunk_alteration(
|
||||
"tenant-1",
|
||||
"kb-1",
|
||||
{"doc-existing", "doc-new"},
|
||||
{"doc-existing", "doc-template-off", "doc-removed"},
|
||||
)
|
||||
|
||||
assert result == {
|
||||
"changed": 1,
|
||||
"changed_doc_ids": ["doc-existing"],
|
||||
}
|
||||
|
||||
membership = module._alteration_result(
|
||||
{"doc-existing", "doc-new"},
|
||||
{"doc-existing", "doc-template-off", "doc-removed"},
|
||||
{"doc-existing", "doc-new"},
|
||||
)
|
||||
assert membership["removed_doc_ids"] == ["doc-removed", "doc-template-off"]
|
||||
assert membership["newly_uploaded_doc_ids"] == ["doc-new"]
|
||||
|
||||
@@ -19,7 +19,17 @@ import pytest
|
||||
@pytest.fixture(autouse=True)
|
||||
def _mock_disabled_document_lookup(monkeypatch):
|
||||
"""Keep knowledge-compile unit tests independent of the MySQL database."""
|
||||
from api.db.services.document_service import DocumentService
|
||||
try:
|
||||
from api.db.services.document_service import DocumentService
|
||||
except ModuleNotFoundError:
|
||||
module = types.ModuleType("api.db.services.document_service")
|
||||
module.DocumentService = type(
|
||||
"DocumentService",
|
||||
(),
|
||||
{"get_disabled_doc_ids_by_kb_id": MagicMock(return_value=set())},
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "api.db.services.document_service", module)
|
||||
return
|
||||
|
||||
monkeypatch.setattr(DocumentService, "get_disabled_doc_ids_by_kb_id", MagicMock(return_value=set()))
|
||||
|
||||
@@ -70,6 +80,8 @@ if not hasattr(sys.modules["rag.prompts.generator"], "message_fit_in"):
|
||||
return True
|
||||
|
||||
sys.modules["rag.prompts.generator"].message_fit_in = _message_fit_in
|
||||
if not hasattr(sys.modules["rag.prompts.generator"], "gen_json"):
|
||||
sys.modules["rag.prompts.generator"].gen_json = MagicMock(return_value={})
|
||||
|
||||
# ---- Modules that wiki_incremental.py imports at module level — use
|
||||
# real import when possible to avoid polluting other test suites.
|
||||
@@ -136,9 +148,15 @@ if hasattr(sys.modules["rag.advanced_rag.knowlege_compile"], "__package__"):
|
||||
|
||||
# _common.py symbols used by wiki_incremental at import time
|
||||
_common_mod = sys.modules["rag.advanced_rag.knowlege_compile._common"]
|
||||
_common_mod.build_chunk_batches = lambda *a, **k: ([], {})
|
||||
_common_mod.bulk_dedup_items = lambda items, *a, **k: items
|
||||
_common_mod.ensure_llm_bundle = lambda model: model
|
||||
_common_mod.knowledge_compile_gen_conf = lambda *a, **k: {}
|
||||
_common_mod.run_chunked_pipeline = MagicMock(return_value={})
|
||||
_common_mod.stable_row_id = lambda *a, **k: ""
|
||||
|
||||
# ---- Test helper constants (same values as structure.py) ----
|
||||
sys.modules["rag.advanced_rag.knowlege_compile.structure"].CONCEPT_MIN_CLAIMS = 3
|
||||
sys.modules["rag.advanced_rag.knowlege_compile.structure"].CONCEPT_MIN_SOURCES = 2
|
||||
sys.modules["rag.advanced_rag.knowlege_compile.structure"]._struct_get = lambda *a, **k: None
|
||||
sys.modules["rag.advanced_rag.knowlege_compile.structure"]._struct_localize = lambda value, *a, **k: value
|
||||
|
||||
@@ -83,10 +83,10 @@ def make_doc_store(search_results: list[dict] | None = None):
|
||||
"""Create a mock settings.docStoreConn."""
|
||||
conn = MagicMock()
|
||||
conn.index_exist = MagicMock(return_value=True)
|
||||
conn.search = AsyncMock(return_value={"hits": {"total": {"value": len(search_results or [])}, "hits": search_results or []}})
|
||||
conn.insert = AsyncMock(return_value=None)
|
||||
conn.update = AsyncMock(return_value=None)
|
||||
conn.delete = AsyncMock(return_value=None)
|
||||
conn.search = MagicMock(return_value={"hits": {"total": {"value": len(search_results or [])}, "hits": search_results or []}})
|
||||
conn.insert = MagicMock(return_value=None)
|
||||
conn.update = MagicMock(return_value=None)
|
||||
conn.delete = MagicMock(return_value=None)
|
||||
conn.refresh_idx = MagicMock(return_value=True)
|
||||
|
||||
def _get_fields(res, fields):
|
||||
@@ -1270,6 +1270,8 @@ async def test_finalize_links_via_map_relations():
|
||||
"_source": {
|
||||
"doc_id": "map1",
|
||||
"compile_kwd": "wiki_map_extract",
|
||||
"source_chunk_ids": ["chunk-1"],
|
||||
"chunk_hash_kwd": "hash-1",
|
||||
"content_with_weight": json.dumps(
|
||||
{
|
||||
"entities": [{"name": "肖亮", "type": "person"}, {"name": "肖立", "type": "person"}],
|
||||
@@ -1286,7 +1288,12 @@ async def test_finalize_links_via_map_relations():
|
||||
patch("common.settings.docStoreConn", doc_store),
|
||||
patch(f"{_wiki.__name__}._load_canonical_entities", new_callable=AsyncMock, return_value={}),
|
||||
):
|
||||
await _wiki._wiki_finalize(tenant_id="t1", kb_id="kb1", embd_mdl=None)
|
||||
await _wiki._wiki_finalize(
|
||||
tenant_id="t1",
|
||||
kb_id="kb1",
|
||||
embd_mdl=None,
|
||||
chunk_state={"chunk-1": {"doc_id": "map1", "hash": "hash-1"}},
|
||||
)
|
||||
|
||||
update_calls = doc_store.update.call_args_list
|
||||
xiaoliang_upd = None
|
||||
@@ -1904,3 +1911,63 @@ async def test_local_split_preserves_all_members_and_original_slug():
|
||||
assert any(page_id.startswith("_new_") for page_id in split)
|
||||
assert {entity["entity_name"] for entities in split.values() for entity in entities} == set(names)
|
||||
assert "entity-0" in {entity["entity_name"] for entity in split["entity/original"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reduce_entity_retracts_only_invalidated_chunk_claims():
|
||||
existing_page = {
|
||||
"claims": [
|
||||
{"statement": "old-c1", "source_doc_id": "doc-1", "chunk_ids": ["c1"]},
|
||||
{"statement": "keep-c2", "source_doc_id": "doc-1", "chunk_ids": ["c2"]},
|
||||
],
|
||||
"source_chunk_ids": ["c1", "c2"],
|
||||
}
|
||||
|
||||
delta = await _wiki._wiki_reduce_entity(
|
||||
entity_name="Alpha",
|
||||
new_claims=[{"statement": "new-c1", "source_doc_id": "doc-1", "chunk_ids": ["c1"]}],
|
||||
existing_page=existing_page,
|
||||
deleted_doc_ids=set(),
|
||||
invalidated_chunk_ids={"c1"},
|
||||
source_doc_ids=["doc-1"],
|
||||
source_chunk_ids=["c1", "c2"],
|
||||
)
|
||||
|
||||
assert delta["action"] == "update"
|
||||
assert [claim["statement"] for claim in delta["retractions"]] == ["old-c1"]
|
||||
assert [claim["statement"] for claim in delta["additions"]] == ["new-c1"]
|
||||
assert delta["source_chunk_ids"] == ["c1", "c2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reduce_entity_deletes_page_when_last_chunk_is_removed():
|
||||
delta = await _wiki._wiki_reduce_entity(
|
||||
entity_name="Alpha",
|
||||
new_claims=[],
|
||||
existing_page={
|
||||
"claims": [{"statement": "only claim", "source_doc_id": "doc-1", "chunk_ids": ["c1"]}],
|
||||
"source_chunk_ids": ["c1"],
|
||||
},
|
||||
deleted_doc_ids=set(),
|
||||
invalidated_chunk_ids={"c1"},
|
||||
)
|
||||
|
||||
assert delta["action"] == "delete"
|
||||
assert delta["source_chunk_ids"] == []
|
||||
assert [claim["statement"] for claim in delta["retractions"]] == ["only claim"]
|
||||
|
||||
|
||||
def test_as_int_accepts_string_doc_store_values():
|
||||
assert _wiki._as_int("7") == 7
|
||||
assert _wiki._as_int(None, 3) == 3
|
||||
|
||||
|
||||
def test_contextual_hints_accepts_native_string_relations():
|
||||
hints = _wiki._wiki_build_contextual_hints(
|
||||
"entity/Alpha",
|
||||
{"related_kb_pages_kwd": ["entity/Beta", "concept/Gamma"]},
|
||||
{},
|
||||
)
|
||||
|
||||
assert "[[entity/Beta]] — related" in hints
|
||||
assert "[[concept/Gamma]] — related" in hints
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _load_wiki_module(monkeypatch):
|
||||
token_utils = ModuleType("common.token_utils")
|
||||
token_utils.num_tokens_from_string = lambda text: len(text)
|
||||
monkeypatch.setitem(sys.modules, "common.token_utils", token_utils)
|
||||
xxhash = ModuleType("xxhash")
|
||||
xxhash.xxh64 = lambda value: SimpleNamespace(hexdigest=lambda: "test-hash")
|
||||
monkeypatch.setitem(sys.modules, "xxhash", xxhash)
|
||||
|
||||
generator = sys.modules["rag.prompts.generator"]
|
||||
monkeypatch.setattr(generator, "gen_json", lambda *args, **kwargs: {}, raising=False)
|
||||
|
||||
common = sys.modules["rag.advanced_rag.knowlege_compile._common"]
|
||||
monkeypatch.setattr(common, "build_chunk_batches", lambda *args, **kwargs: ([], {}), raising=False)
|
||||
monkeypatch.setattr(common, "bulk_dedup_items", lambda items, *args, **kwargs: items, raising=False)
|
||||
monkeypatch.setattr(common, "ensure_llm_bundle", lambda model: model, raising=False)
|
||||
monkeypatch.setattr(common, "knowledge_compile_gen_conf", lambda *args, **kwargs: {}, raising=False)
|
||||
monkeypatch.setattr(common, "run_chunked_pipeline", lambda *args, **kwargs: {}, raising=False)
|
||||
monkeypatch.setattr(common, "stable_row_id", lambda *parts: "|".join(str(part) for part in parts), raising=False)
|
||||
|
||||
structure = sys.modules["rag.advanced_rag.knowlege_compile.structure"]
|
||||
monkeypatch.setattr(structure, "_struct_get", lambda *args, **kwargs: None, raising=False)
|
||||
monkeypatch.setattr(structure, "_struct_localize", lambda value, *args, **kwargs: value, raising=False)
|
||||
|
||||
module_name = "rag.advanced_rag.knowlege_compile.wiki"
|
||||
module_path = os.path.normpath(os.path.join(os.path.dirname(__file__), "../../../../../rag/advanced_rag/knowlege_compile/wiki.py"))
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class _StateDocStore:
|
||||
def __init__(self, rows=None):
|
||||
self.rows = {row["id"]: dict(row) for row in rows or []}
|
||||
self.insert_calls = []
|
||||
|
||||
def search(self, fields, _highlights, condition, _matches, _order, offset, limit, _index, _dataset_ids):
|
||||
rows = []
|
||||
for row in self.rows.values():
|
||||
if any(condition.get(key) and row.get(key) not in condition[key] for key in ("id", "compile_kwd", "type_kwd")):
|
||||
continue
|
||||
if condition.get("doc_id") and row.get("doc_id") not in condition["doc_id"]:
|
||||
continue
|
||||
if condition.get("source_chunk_ids") and not set(row.get("source_chunk_ids") or []) & set(condition["source_chunk_ids"]):
|
||||
continue
|
||||
if condition.get("chunk_hash_kwd") and row.get("chunk_hash_kwd") not in condition["chunk_hash_kwd"]:
|
||||
continue
|
||||
rows.append(row)
|
||||
return rows[offset : offset + limit]
|
||||
|
||||
def get_fields(self, result, fields):
|
||||
return {row["id"]: {field: row.get(field) for field in fields} for row in result}
|
||||
|
||||
def insert(self, rows, _index, _dataset_id):
|
||||
self.insert_calls.append([dict(row) for row in rows])
|
||||
for row in rows:
|
||||
self.rows[row["id"]] = dict(row)
|
||||
|
||||
def delete(self, condition, _index, _dataset_id):
|
||||
for row_id, row in list(self.rows.items()):
|
||||
if all(not condition.get(key) or row.get(key) in condition[key] for key in ("compile_kwd", "type_kwd")):
|
||||
del self.rows[row_id]
|
||||
|
||||
|
||||
def test_active_map_state_switches_marker_after_new_generation(monkeypatch):
|
||||
wiki = _load_wiki_module(monkeypatch)
|
||||
marker_id = wiki._stable_row_id(wiki.WIKI_MAP_STATE_META_COMPILE_KWD, "kb-1")
|
||||
store = _StateDocStore(
|
||||
[
|
||||
{
|
||||
"id": marker_id,
|
||||
"compile_kwd": wiki.WIKI_MAP_STATE_META_COMPILE_KWD,
|
||||
"type_kwd": "old",
|
||||
},
|
||||
{
|
||||
"id": "old-state",
|
||||
"doc_id": "doc-old",
|
||||
"compile_kwd": wiki.WIKI_MAP_STATE_COMPILE_KWD,
|
||||
"type_kwd": "old",
|
||||
"source_chunk_ids": ["chunk-old"],
|
||||
"chunk_hash_kwd": "hash-old",
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(sys.modules["common.settings"], "docStoreConn", store)
|
||||
|
||||
asyncio.run(wiki._wiki_commit_active_map_state("tenant-1", "kb-1", {"chunk-new": {"doc_id": "doc-new", "hash": "hash-new"}}))
|
||||
|
||||
assert store.insert_calls[0][0]["compile_kwd"] == wiki.WIKI_MAP_STATE_COMPILE_KWD
|
||||
assert store.insert_calls[1][0]["compile_kwd"] == wiki.WIKI_MAP_STATE_META_COMPILE_KWD
|
||||
assert asyncio.run(wiki._wiki_load_active_map_state("tenant-1", "kb-1")) == {"chunk-new": {"doc_id": "doc-new", "hash": "hash-new"}}
|
||||
assert "old-state" not in store.rows
|
||||
|
||||
|
||||
def test_map_version_query_is_limited_to_requested_chunk_hash(monkeypatch):
|
||||
wiki = _load_wiki_module(monkeypatch)
|
||||
store = _StateDocStore(
|
||||
[
|
||||
{
|
||||
"id": "wanted",
|
||||
"doc_id": "doc-1",
|
||||
"compile_kwd": wiki.WIKI_MAP_COMPILE_KWD,
|
||||
"source_chunk_ids": ["chunk-1"],
|
||||
"chunk_hash_kwd": "hash-b",
|
||||
"content_with_weight": json.dumps({"entities": [{"name": "B"}]}),
|
||||
},
|
||||
{
|
||||
"id": "historical",
|
||||
"doc_id": "doc-1",
|
||||
"compile_kwd": wiki.WIKI_MAP_COMPILE_KWD,
|
||||
"source_chunk_ids": ["chunk-1"],
|
||||
"chunk_hash_kwd": "hash-a",
|
||||
"content_with_weight": json.dumps({"entities": [{"name": "A"}]}),
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(sys.modules["common.settings"], "docStoreConn", store)
|
||||
|
||||
versions = asyncio.run(wiki._wiki_load_map_versions("doc-1", "tenant-1", "kb-1", {"chunk-1": "hash-b"}))
|
||||
|
||||
assert set(versions["chunk-1"]) == {"hash-b"}
|
||||
|
||||
|
||||
def test_failed_state_write_keeps_previous_generation_active(monkeypatch):
|
||||
wiki = _load_wiki_module(monkeypatch)
|
||||
marker_id = wiki._stable_row_id(wiki.WIKI_MAP_STATE_META_COMPILE_KWD, "kb-1")
|
||||
|
||||
class _FailingStore(_StateDocStore):
|
||||
def insert(self, rows, index, dataset_id):
|
||||
if rows[0].get("compile_kwd") == wiki.WIKI_MAP_STATE_COMPILE_KWD:
|
||||
raise RuntimeError("state write failed")
|
||||
super().insert(rows, index, dataset_id)
|
||||
|
||||
store = _FailingStore(
|
||||
[
|
||||
{
|
||||
"id": marker_id,
|
||||
"compile_kwd": wiki.WIKI_MAP_STATE_META_COMPILE_KWD,
|
||||
"type_kwd": "old",
|
||||
},
|
||||
{
|
||||
"id": "old-state",
|
||||
"doc_id": "doc-old",
|
||||
"compile_kwd": wiki.WIKI_MAP_STATE_COMPILE_KWD,
|
||||
"type_kwd": "old",
|
||||
"source_chunk_ids": ["chunk-old"],
|
||||
"chunk_hash_kwd": "hash-old",
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(sys.modules["common.settings"], "docStoreConn", store)
|
||||
|
||||
with pytest.raises(RuntimeError, match="state write failed"):
|
||||
asyncio.run(wiki._wiki_commit_active_map_state("tenant-1", "kb-1", {"chunk-new": {"doc_id": "doc-new", "hash": "hash-new"}}))
|
||||
|
||||
assert asyncio.run(wiki._wiki_load_active_map_state("tenant-1", "kb-1")) == {"chunk-old": {"doc_id": "doc-old", "hash": "hash-old"}}
|
||||
Reference in New Issue
Block a user