fix: track wiki chunk changes incrementally (#18285)

This commit is contained in:
buua436
2026-08-14 20:11:55 +08:00
committed by GitHub
parent 34e59c8ec9
commit 024c35f085
8 changed files with 966 additions and 574 deletions

View File

@@ -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"]

View File

@@ -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

View File

@@ -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

View File

@@ -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"}}