diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 17ce38e782..5af1da3300 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -32,7 +32,7 @@ from api.utils.api_utils import deep_merge, get_parser_config, remap_dictionary_ from common import settings from common.constants import PAGERANK_FLD, FileSource, LLMType, RetCode, StatusEnum from common.misc_utils import thread_pool_exec, thread_pool_exec_long_time -from rag.advanced_rag.knowlege_compile.wiki import WIKI_PAGE_COMPILE_KWD, _chunk_hash +from rag.advanced_rag.knowlege_compile.wiki import WIKI_PAGE_COMPILE_KWD # KB-wide structure-graph merge index types. Each (re)builds the ``dataset_graph`` # rows for one structure kind via ``rebuild_dataset_structure_graph_json``; the @@ -2228,117 +2228,49 @@ def _alteration_result(current_doc_ids: set, involved_doc_ids: set, eligible_doc async def _wiki_chunk_alteration( - index_nm, + tenant_id: str, dataset_id: str, - active_doc_ids: set[str], - disabled_doc_ids: set[str], + eligible_doc_ids: set[str], + involved_doc_ids: set[str], ) -> dict: """Compare active source chunks with the hashes used by Wiki MAP. Document-level provenance cannot detect an edited, added, or removed chunk. The MAP resume rows already contain the exact chunk hash used by compilation, so they are the authoritative compiled-side snapshot. - A document is changed when a chunk is added, removed, or has a different - hash. A disabled document with retained MAP rows is also changed because - those rows must no longer participate in the next Wiki build. + A previously involved document is changed when a chunk is added, removed, + or has a different hash. Documents removed or disabled at document level + remain represented by the existing ``removed`` fields rather than being + duplicated in ``changed``. """ - from common.doc_store.doc_store_base import OrderByExpr - empty = { "changed": 0, "changed_doc_ids": [], } - if not active_doc_ids and not disabled_doc_ids: + if not eligible_doc_ids and not involved_doc_ids: return empty - current: dict[str, dict] = {} - for doc_id in active_doc_ids: - offset = 0 - while True: - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - ["id", "doc_id", "content_with_weight"], - [], - {"doc_id": [doc_id], "available_int": 1, "must_not": {"exists": "compile_kwd"}}, - [], - OrderByExpr(), - offset, - 1000, - index_nm, - [dataset_id], - ) - rows = settings.docStoreConn.get_fields(res, ["id", "doc_id", "content_with_weight"]) or {} - except Exception as exc: - logging.exception("alteration: failed to load source chunks for doc=%s", doc_id) - raise RuntimeError(f"Failed to load source chunks for Wiki alteration (kb={dataset_id}, doc={doc_id})") from exc - if not rows: - break - for row in rows.values(): - chunk_id = str(row.get("id") or "") - if not chunk_id: - continue - current[chunk_id] = { - "id": chunk_id, - "doc_id": doc_id, - "hash": _chunk_hash(row.get("content_with_weight") or ""), - } - if len(rows) < 1000: - break - offset += 1000 + from rag.advanced_rag.knowlege_compile.wiki import ( + _wiki_compare_chunk_states, + _wiki_load_active_map_state, + _wiki_scan_current_chunk_state, + ) - compiled: dict[str, dict] = {} - offset = 0 - while True: - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - ["id", "doc_id", "source_chunk_ids", "chunk_hash_kwd"], - [], - {"compile_kwd": ["wiki_map_extract"]}, - [], - OrderByExpr(), - offset, - 1000, - index_nm, - [dataset_id], - ) - rows = settings.docStoreConn.get_fields(res, ["id", "doc_id", "source_chunk_ids", "chunk_hash_kwd"]) or {} - except Exception as exc: - logging.exception("alteration: failed to load Wiki MAP hashes for kb=%s", dataset_id) - raise RuntimeError(f"Failed to load Wiki MAP hashes for alteration (kb={dataset_id})") from exc - if not rows: - break - for row in rows.values(): - row_doc_ids = _flatten_provenance_doc_ids(row.get("doc_id")) - matching_doc_ids = row_doc_ids & (active_doc_ids | disabled_doc_ids) - if not matching_doc_ids: - continue - doc_id = next(iter(matching_doc_ids)) - chunk_ids = row.get("source_chunk_ids") or [] - if isinstance(chunk_ids, str): - chunk_ids = [chunk_ids] - saved_hash = row.get("chunk_hash_kwd") - saved_hash = saved_hash if isinstance(saved_hash, str) else "" - for chunk_id in chunk_ids: - chunk_id = str(chunk_id or "") - if chunk_id: - compiled.setdefault(chunk_id, {"id": chunk_id, "doc_id": doc_id, "hash": saved_hash}) - if len(rows) < 1000: - break - offset += 1000 + try: + current = await _wiki_scan_current_chunk_state(tenant_id, dataset_id, eligible_doc_ids) + previous = await _wiki_load_active_map_state(tenant_id, dataset_id) + except Exception as exc: + logging.exception("alteration: failed to compare Wiki chunk state for kb=%s", dataset_id) + raise RuntimeError(f"Failed to compare Wiki chunk state for alteration (kb={dataset_id})") from exc - changed_doc_ids: set[str] = set() - for chunk_id, item in current.items(): - old = compiled.get(chunk_id) - if old is None or old["hash"] != item["hash"]: - changed_doc_ids.add(item["doc_id"]) - - # A missing current chunk means it was removed from an active document. - # For a disabled document, all retained MAP chunks are stale by definition. - for chunk_id, item in compiled.items(): - if item["doc_id"] in disabled_doc_ids or (item["doc_id"] in active_doc_ids and chunk_id not in current): - changed_doc_ids.add(item["doc_id"]) + delta = _wiki_compare_chunk_states(previous, current) + changed_doc_ids = { + str((current.get(chunk_id) or previous.get(chunk_id) or {}).get("doc_id") or "") for chunk_id in delta["new_chunk_ids"] | delta["changed_chunk_ids"] | delta["deleted_chunk_ids"] + } + # ``newly_uploaded`` owns eligible documents which have not contributed to + # the current Wiki; ``removed`` owns previously involved documents which + # are no longer eligible. Keep ``changed`` disjoint from both categories. + changed_doc_ids &= eligible_doc_ids & involved_doc_ids return { "changed": len(changed_doc_ids), @@ -2496,15 +2428,18 @@ async def _get_alteration(dataset_id: str, tenant_id: str, kind: str): index_nm, _ = pack involved_doc_ids = await _involved_doc_ids_for_kind(index_nm, dataset_id, kind) if kind == "wiki": - disabled_doc_ids = await _disabled_dataset_doc_ids(dataset_id) chunk_changes = await _wiki_chunk_alteration( - index_nm, + kb.tenant_id, dataset_id, eligible_doc_ids, - disabled_doc_ids, + involved_doc_ids, ) - result = _alteration_result(current_doc_ids, involved_doc_ids, eligible_doc_ids) + # Wiki membership follows compilation eligibility. Disabling a document or + # removing its Wiki template is therefore a removal; enabling it again or + # restoring the template after a rebuild makes it newly uploaded. + alteration_current_doc_ids = eligible_doc_ids if kind == "wiki" else current_doc_ids + result = _alteration_result(alteration_current_doc_ids, involved_doc_ids, eligible_doc_ids) if chunk_changes is not None: result.update(chunk_changes) return True, result @@ -4162,6 +4097,8 @@ async def update_wiki_page( # the dataset Artifact tab's graph view reads exactly this row. _WIKI_COMPILE_KWDS = ( "wiki_map_extract", + "wiki_map_state", + "wiki_map_state_meta", "wiki_reduce_result", "wiki_compilation_plan", "wiki_page_draft", diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index 556c0c1f2a..3a71d98fbc 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -24,11 +24,12 @@ - Citation anchor is the source chunk id (``source_chunk_ids`` list per item), not a byte position. The LLM is prompted to tag each extracted item with the ``[CHUNK_ID …]`` of the chunk it came from. - - Resume: per-chunk extracts are persisted to ES under + - Cache: per-chunk extraction versions are persisted to ES under ``compile_kwd="wiki_map_extract"`` with ``available_int=0`` and no vector / token-list fields, so retrievers ignore them but downstream phases can - fetch them by ``doc_id`` + ``source_chunk_ids``. Re-running MAP for the same - ``doc_id`` skips chunks that already have an extract row. + fetch them by ``doc_id`` + ``source_chunk_ids`` + ``chunk_hash_kwd``. + Historical versions are retained so reverted chunk content can reuse its + earlier extraction. Public entry: ``wiki_map_from_chunks``. """ @@ -37,6 +38,7 @@ import asyncio import json import logging import re +import uuid from typing import Callable, Optional from urllib.parse import urlsplit from common.misc_utils import thread_pool_exec @@ -86,6 +88,8 @@ from .structure import ( # --------------------------------------------------------------------------- WIKI_MAP_COMPILE_KWD = "wiki_map_extract" +WIKI_MAP_STATE_COMPILE_KWD = "wiki_map_state" +WIKI_MAP_STATE_META_COMPILE_KWD = "wiki_map_state_meta" DEFAULT_WIKI_MAP_WORKERS = 20 DEFAULT_WIKI_MAP_TIMEOUT = 600 @@ -114,6 +118,19 @@ def _wiki_doc_ids(value) -> set[str]: return {value} if value else set() +def _wiki_compare_chunk_states(previous: dict[str, dict], current: dict[str, dict]) -> dict[str, set[str]]: + """Return the chunk-level delta between two successful Wiki states.""" + previous_ids = set(previous) + current_ids = set(current) + common_ids = previous_ids & current_ids + return { + "new_chunk_ids": current_ids - previous_ids, + "changed_chunk_ids": {chunk_id for chunk_id in common_ids if previous[chunk_id].get("hash") != current[chunk_id].get("hash")}, + "deleted_chunk_ids": previous_ids - current_ids, + "unchanged_chunk_ids": {chunk_id for chunk_id in common_ids if previous[chunk_id].get("hash") == current[chunk_id].get("hash")}, + } + + WIKI_MAP_SYSTEM = ( "You are a knowledge extraction engine. Extract structured knowledge from the " "provided document section. Return ONLY valid JSON matching the schema exactly. " @@ -615,6 +632,11 @@ def _wiki_resolve_chunk_ids( per_chunk: dict[str, dict] = {real_id: _wiki_empty_extract() for real_id in label_to_id.values()} merged = _wiki_empty_extract() merged["topics"] = list(extract.get("topics") or []) + # MAP topics are batch-level strings and carry no source_chunk_id. Store + # them with every chunk version in the batch so a later cache hit or state + # reload preserves the same topic pool; downstream deduplicates them. + for chunk_extract in per_chunk.values(): + chunk_extract["topics"] = list(merged["topics"]) dropped = 0 dropped_identifier = 0 @@ -684,7 +706,10 @@ def _wiki_build_resume_doc( content_with_weight = json.dumps(per_chunk_extract, ensure_ascii=False) doc_id_str = str(doc_id) return { - "id": _stable_row_id(content_with_weight, doc_id_str, chunk_id), + # A MAP row is an immutable cache version. Keeping the hash in the + # identity lets A -> B -> A reuse the first extraction instead of + # replacing it when B is compiled. + "id": _stable_row_id(WIKI_MAP_COMPILE_KWD, doc_id_str, chunk_id, chunk_hash), "doc_id": doc_id_str, "compile_kwd": WIKI_MAP_COMPILE_KWD, "source_chunk_ids": [chunk_id], @@ -694,104 +719,71 @@ def _wiki_build_resume_doc( } -async def _wiki_load_resume_map( - doc_id: str, +async def _wiki_load_map_versions( + doc_ids: str | set[str], tenant_id: str, kb_id: str, -) -> dict[str, str]: - """Query ES for chunks that already have a wiki_map_extract row for - this doc. Returns ``{chunk_id → chunk_hash}``. - - ``chunk_hash`` may be empty for legacy rows that predate the field — - callers treat empty as "definitely re-MAP" (no hash to compare). - """ + requested_versions: Optional[dict[str, str]] = None, +) -> dict[str, dict[str, dict]]: + """Load historical MAP versions as ``chunk_id -> hash -> extract``.""" from common import settings from common.doc_store.doc_store_base import OrderByExpr from rag.nlp import search as _rag_search index = _rag_search.index_name(tenant_id) - condition = { - "compile_kwd": [WIKI_MAP_COMPILE_KWD], - "doc_id": [str(doc_id)], - } - select_fields = ["id", "source_chunk_ids", "chunk_hash_kwd"] - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - select_fields, - [], - condition, - [], - OrderByExpr(), - 0, - 10000, - index, - [kb_id], - ) - field_map = settings.docStoreConn.get_fields(res, select_fields) - except Exception: - logging.exception("wiki_map: failed to query resume map; will re-extract all chunks") - return {} - - seen: dict[str, str] = {} - for row in field_map.values(): - src = row.get("source_chunk_ids") or [] - hh = row.get("chunk_hash_kwd") - if not isinstance(hh, str): - hh = "" - if isinstance(src, list): - for cid in src: - if isinstance(cid, str) and cid: - # First-write-wins is fine: if a doc has two rows for - # the same chunk_id (legacy / dirty state), we treat - # the first as the canonical and let the changed-hash - # path or the deletion sweep clean it up later. - seen.setdefault(cid, hh) - return seen - - -async def _wiki_delete_map_rows( - doc_id: str, - chunk_ids: list[str], - tenant_id: str, - kb_id: str, -) -> int: - """Delete ``wiki_map_extract`` rows for ``(doc_id, chunk_id)`` pairs. - - Used by the incremental MAP path: - * stale rows whose chunk content has changed → re-extracted next. - * rows whose chunk_id is gone from the doc (chunk deleted upstream). - - Returns the number of distinct ``chunk_ids`` we attempted to drop; - the backend may delete more (e.g. duplicate rows) — we don't try to - track that precisely. - """ - if not chunk_ids: - return 0 - from common import settings - from rag.nlp import search as _rag_search - - index = _rag_search.index_name(tenant_id) - condition = { - "compile_kwd": [WIKI_MAP_COMPILE_KWD], - "doc_id": [str(doc_id)], - "source_chunk_ids": list(chunk_ids), - } - try: - await thread_pool_exec( - settings.docStoreConn.delete, - condition, - index, - kb_id, - ) - except Exception: - logging.exception( - "wiki_map: failed to delete %d stale extract row(s) for doc %s", - len(chunk_ids), - doc_id, - ) - return 0 - return len(chunk_ids) + select_fields = ["source_chunk_ids", "chunk_hash_kwd", "content_with_weight"] + offset = 0 + page_size = 1000 + versions: dict[str, dict[str, dict]] = {} + requested_chunk_ids = set(requested_versions or {}) + requested_hashes = {chunk_hash for chunk_hash in (requested_versions or {}).values() if chunk_hash} + normalized_doc_ids = {str(doc_id) for doc_id in ({doc_ids} if isinstance(doc_ids, str) else doc_ids) if doc_id} + condition = {"compile_kwd": [WIKI_MAP_COMPILE_KWD], "doc_id": sorted(normalized_doc_ids)} + if requested_chunk_ids: + condition["source_chunk_ids"] = sorted(requested_chunk_ids) + if requested_hashes: + condition["chunk_hash_kwd"] = sorted(requested_hashes) + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + condition, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + except Exception: + logging.exception("wiki_map: failed to load historical versions for docs %s", sorted(normalized_doc_ids)) + return versions + for row in field_map.values(): + chunk_ids = _wiki_doc_ids(row.get("source_chunk_ids")) + chunk_hash = row.get("chunk_hash_kwd") + if not isinstance(chunk_hash, str) or not chunk_hash: + continue + if requested_hashes and chunk_hash not in requested_hashes: + continue + try: + extract = json.loads(row.get("content_with_weight") or "{}") + except (TypeError, json.JSONDecodeError): + continue + if not isinstance(extract, dict): + continue + for chunk_id in chunk_ids: + if requested_chunk_ids and chunk_id not in requested_chunk_ids: + continue + if requested_versions is not None and requested_versions.get(chunk_id) != chunk_hash: + continue + versions.setdefault(chunk_id, {}).setdefault(chunk_hash, extract) + if len(field_map) < page_size: + break + offset += page_size + return versions async def _wiki_persist_extracts( @@ -833,6 +825,231 @@ async def _wiki_persist_extracts( logging.exception("wiki_map: failed to persist %d resume docs", len(docs)) +async def _wiki_scan_current_chunk_state( + tenant_id: str, + kb_id: str, + doc_ids: set[str], +) -> dict[str, dict]: + """Scan enabled source chunks and return their current MAP input hashes.""" + if not doc_ids: + return {} + from common import settings + from common.doc_store.doc_store_base import OrderByExpr + from rag.nlp import search as _rag_search + + index = _rag_search.index_name(tenant_id) + state: dict[str, dict] = {} + for doc_id in sorted(doc_ids): + offset = 0 + while True: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["id", "doc_id", "content_with_weight"], + [], + { + "doc_id": [doc_id], + "available_int": 1, + "must_not": {"exists": "compile_kwd"}, + }, + [], + OrderByExpr(), + offset, + 1000, + index, + [kb_id], + ) + rows = settings.docStoreConn.get_fields(res, ["id", "doc_id", "content_with_weight"]) or {} + for row_id, row in rows.items(): + chunk_id = str(row.get("id") or row_id or "") + if chunk_id: + state[chunk_id] = { + "doc_id": str(row.get("doc_id") or doc_id), + "hash": _chunk_hash(row.get("content_with_weight") or ""), + } + if len(rows) < 1000: + break + offset += 1000 + return state + + +async def _wiki_load_active_map_state( + tenant_id: str, + kb_id: str, +) -> dict[str, dict]: + """Load the chunk versions used by the last successful Wiki build.""" + from common import settings + from common.doc_store.doc_store_base import OrderByExpr + from rag.nlp import search as _rag_search + + index = _rag_search.index_name(tenant_id) + generation = await _wiki_load_active_map_generation(tenant_id, kb_id) + if not generation: + return {} + + state: dict[str, dict] = {} + offset = 0 + while True: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["doc_id", "source_chunk_ids", "chunk_hash_kwd"], + [], + {"compile_kwd": [WIKI_MAP_STATE_COMPILE_KWD], "type_kwd": [generation]}, + [], + OrderByExpr(), + offset, + 1000, + index, + [kb_id], + ) + rows = settings.docStoreConn.get_fields(res, ["doc_id", "source_chunk_ids", "chunk_hash_kwd"]) or {} + for row in rows.values(): + chunk_hash = row.get("chunk_hash_kwd") + if not isinstance(chunk_hash, str) or not chunk_hash: + continue + doc_id = next(iter(_wiki_doc_ids(row.get("doc_id"))), "") + for chunk_id in _wiki_doc_ids(row.get("source_chunk_ids")): + state[chunk_id] = {"doc_id": doc_id, "hash": chunk_hash} + if len(rows) < 1000: + break + offset += 1000 + return state + + +async def _wiki_load_active_map_generation(tenant_id: str, kb_id: str) -> str: + from common import settings + from common.doc_store.doc_store_base import OrderByExpr + from rag.nlp import search as _rag_search + + res = await thread_pool_exec( + settings.docStoreConn.search, + ["type_kwd"], + [], + { + "compile_kwd": [WIKI_MAP_STATE_META_COMPILE_KWD], + "id": [_stable_row_id(WIKI_MAP_STATE_META_COMPILE_KWD, kb_id)], + }, + [], + OrderByExpr(), + 0, + 1, + _rag_search.index_name(tenant_id), + [kb_id], + ) + marker_rows = settings.docStoreConn.get_fields(res, ["type_kwd"]) or {} + for row in marker_rows.values(): + values = _wiki_doc_ids(row.get("type_kwd")) + if values: + return next(iter(values)) + return "" + + +async def _wiki_load_map_extracts_for_state( + tenant_id: str, + kb_id: str, + state: dict[str, dict], + chunk_ids: Optional[set[str]] = None, +) -> list[dict]: + """Load only MAP versions selected by a source-state snapshot.""" + selected_ids = set(state) + if chunk_ids is not None: + selected_ids &= set(chunk_ids) + if not selected_ids: + return [] + + by_doc: dict[str, set[str]] = {} + for chunk_id in selected_ids: + doc_id = str(state[chunk_id].get("doc_id") or "") + if doc_id: + by_doc.setdefault(doc_id, set()).add(chunk_id) + + requested_versions = {chunk_id: str(state[chunk_id].get("hash") or "") for chunk_ids_for_doc in by_doc.values() for chunk_id in chunk_ids_for_doc} + versions = await _wiki_load_map_versions(set(by_doc), tenant_id, kb_id, requested_versions) + extracts: list[dict] = [] + for doc_id, doc_chunk_ids in by_doc.items(): + for chunk_id in doc_chunk_ids: + chunk_hash = str(state[chunk_id].get("hash") or "") + extract = versions.get(chunk_id, {}).get(chunk_hash) + if not isinstance(extract, dict): + continue + item = dict(extract) + item["doc_id"] = doc_id + item["_map_version"] = { + "chunk_id": chunk_id, + "hash": chunk_hash, + } + extracts.append(item) + return extracts + + +async def _wiki_commit_active_map_state( + tenant_id: str, + kb_id: str, + state: dict[str, dict], +) -> None: + """Commit the source snapshot only after Wiki compilation succeeds.""" + from common import settings + from rag.nlp import search as _rag_search + + index = _rag_search.index_name(tenant_id) + try: + previous_generation = await _wiki_load_active_map_generation(tenant_id, kb_id) + except Exception: + logging.exception("wiki_map: failed to read the previous active-state generation") + raise + + generation = uuid.uuid4().hex + rows = [] + for chunk_id, item in state.items(): + doc_id = str(item.get("doc_id") or "") + chunk_hash = str(item.get("hash") or "") + if not doc_id or not chunk_hash: + continue + rows.append( + { + "id": _stable_row_id(WIKI_MAP_STATE_COMPILE_KWD, doc_id, chunk_id), + "doc_id": doc_id, + "compile_kwd": WIKI_MAP_STATE_COMPILE_KWD, + "type_kwd": generation, + "source_chunk_ids": [chunk_id], + "chunk_hash_kwd": chunk_hash, + "content_with_weight": "{}", + "available_int": 0, + } + ) + # State rows use generation-qualified IDs so an interrupted write cannot + # overwrite the active snapshot. The single marker is switched only after + # every row in the new generation has been persisted. + for row in rows: + row["id"] = _stable_row_id(WIKI_MAP_STATE_COMPILE_KWD, generation, row["doc_id"], row["source_chunk_ids"][0]) + if rows: + await thread_pool_exec(settings.docStoreConn.insert, rows, index, kb_id) + marker = { + "id": _stable_row_id(WIKI_MAP_STATE_META_COMPILE_KWD, kb_id), + "doc_id": "", + "compile_kwd": WIKI_MAP_STATE_META_COMPILE_KWD, + "type_kwd": generation, + "source_chunk_ids": ["__wiki_map_state__"], + "chunk_hash_kwd": "committed", + "content_with_weight": "{}", + "available_int": 0, + } + await thread_pool_exec(settings.docStoreConn.insert, [marker], index, kb_id) + if previous_generation and previous_generation != generation: + try: + await thread_pool_exec( + settings.docStoreConn.delete, + {"compile_kwd": [WIKI_MAP_STATE_COMPILE_KWD], "type_kwd": [previous_generation]}, + index, + kb_id, + ) + except Exception: + logging.warning( + "wiki_map: failed to remove inactive state generation %s", + previous_generation, + exc_info=True, + ) + + # --------------------------------------------------------------------------- # Per-batch extraction # --------------------------------------------------------------------------- @@ -975,6 +1192,7 @@ async def wiki_map_from_chunks( parser_config: Optional[dict] = None, batch_size_cap: Optional[int] = None, window_fraction: Optional[float] = None, + target_chunk_ids: Optional[set[str]] = None, ) -> dict: """Phase 1 (MAP) of the artifact compilation pipeline. @@ -1012,46 +1230,15 @@ async def wiki_map_from_chunks( """ _ = embd_mdl # noqa: F841 — accepted for symmetry with downstream phases if not chunks: - # Even with zero chunks we still want to sweep any orphaned MAP - # rows that point at chunks the doc no longer has — otherwise - # deletions never propagate. - prior_resume_map = await _wiki_load_resume_map(doc_id, tenant_id, kb_id) - if prior_resume_map: - await _wiki_delete_map_rows( - doc_id, - list(prior_resume_map.keys()), - tenant_id, - kb_id, - ) - logging.info( - "wiki_map: doc %s now has zero chunks; swept %d stale extract row(s)", - doc_id, - len(prior_resume_map), - ) out = _wiki_empty_extract() out["_meta"] = { "doc_id": str(doc_id), - "new": 0, - "changed": 0, - "deleted": len(prior_resume_map), - "unchanged": 0, - "had_delta": bool(prior_resume_map), + "requested": 0, + "cache_hits": 0, + "extracted": 0, } return out - # Incremental decision per current chunk: - # - # * Compute the fresh chunk hash for every chunk in this call. - # * Load the prior resume map (chunk_id → hash from the last MAP). - # * NEW — chunk_id not in prior → MAP this chunk. - # * UNCHANGED — chunk_id in prior, hash matches → skip (resume). - # * CHANGED — chunk_id in prior, hash differs → delete prior - # row, then MAP this chunk. - # * DELETED — chunk_id only in prior → delete prior row - # (chunk was removed upstream). - # - # The "resume set" handed to ``_build_chunk_batches`` is just the - # UNCHANGED ids — those are the only ones the packer should skip. current_chunk_hashes: dict[str, str] = {} for chunk in chunks: cid = chunk.get("id") or chunk.get("chunk_id") @@ -1060,43 +1247,23 @@ async def wiki_map_from_chunks( text = _wiki_pick_chunk_text(chunk) or "" current_chunk_hashes[cid] = _chunk_hash(text) - prior_resume_map = await _wiki_load_resume_map(doc_id, tenant_id, kb_id) - unchanged_ids: set[str] = set() - changed_ids: list[str] = [] - new_ids: list[str] = [] - for cid, h in current_chunk_hashes.items(): - prior_h = prior_resume_map.get(cid) - if prior_h is None: - new_ids.append(cid) - elif prior_h and prior_h == h: - unchanged_ids.add(cid) - else: - # Empty stored hash = legacy row written before chunk_hash_kwd - # existed → re-MAP. Differing hash = content changed → re-MAP. - changed_ids.append(cid) - deleted_ids = [cid for cid in prior_resume_map if cid not in current_chunk_hashes] + requested_ids = set(current_chunk_hashes) + if target_chunk_ids is not None: + requested_ids &= set(target_chunk_ids) - if changed_ids or deleted_ids: - await _wiki_delete_map_rows( - doc_id, - list(set(changed_ids) | set(deleted_ids)), - tenant_id, - kb_id, - ) + requested_versions = {chunk_id: current_chunk_hashes[chunk_id] for chunk_id in requested_ids} + historical_versions = await _wiki_load_map_versions(doc_id, tenant_id, kb_id, requested_versions) + cache_hits: list[dict] = [] + cache_hit_ids: set[str] = set() + for chunk_id in requested_ids: + extract = historical_versions.get(chunk_id, {}).get(current_chunk_hashes[chunk_id]) + if extract is not None: + cache_hit_ids.add(chunk_id) + cache_hits.append(extract) - if unchanged_ids or changed_ids or deleted_ids or new_ids: - logging.info( - "wiki_map: doc %s — new=%d changed=%d unchanged=%d deleted=%d", - doc_id, - len(new_ids), - len(changed_ids), - len(unchanged_ids), - len(deleted_ids), - ) - - # The packer's "resume" set is the UNCHANGED ids only — NEW and - # CHANGED both need re-extraction. - resume_set = unchanged_ids + extract_ids = requested_ids - cache_hit_ids + # Skip chunks outside this run's delta as well as historical cache hits. + resume_set = set(current_chunk_hashes) - extract_ids # Defensive scrub: chunkers sometimes embed the chunk_id / doc_id into # the body (e.g. as a header). Without this the extraction LLM tends to @@ -1120,8 +1287,15 @@ async def wiki_map_from_chunks( batch_size_cap=batch_size_cap, window_fraction=window_fraction, ) + cached_merged = _wiki_merge_extracts(cache_hits) if not packed_batches: - return _wiki_empty_extract() + cached_merged["_meta"] = { + "doc_id": str(doc_id), + "requested": len(requested_ids), + "cache_hits": len(cache_hit_ids), + "extracted": 0, + } + return cached_merged async def _process_one(batch: list[dict], bi: int, total: int) -> dict: # ``_run_chunked_pipeline`` already wraps each task in the engine's @@ -1142,7 +1316,7 @@ async def wiki_map_from_chunks( chunk_hashes=current_chunk_hashes, ) - merged = await _run_chunked_pipeline( + extracted = await _run_chunked_pipeline( packed_batches, process_batch=_process_one, aggregate=_wiki_merge_extracts, @@ -1150,25 +1324,24 @@ async def wiki_map_from_chunks( callback=callback, log_prefix="wiki_map", ) + merged = _wiki_merge_extracts([cached_merged, extracted]) logging.info( - "wiki_map: doc %s — entities=%d concepts=%d claims=%d relations=%d topics=%d", + "wiki_map: doc %s — requested=%d cache_hits=%d extracted=%d entities=%d concepts=%d claims=%d relations=%d topics=%d", doc_id, + len(requested_ids), + len(cache_hit_ids), + len(extract_ids), len(merged["entities"]), len(merged["concepts"]), len(merged["claims"]), len(merged["relations"]), len(merged["topics"]), ) - # Surface the incremental decisions to the orchestrator. ``had_delta`` - # is the most useful summary: REDUCE/PLAN/REFINE can short-circuit - # KB-wide when no doc's MAP touched any rows on this run. merged["_meta"] = { "doc_id": str(doc_id), - "new": len(new_ids), - "changed": len(changed_ids), - "unchanged": len(unchanged_ids), - "deleted": len(deleted_ids), - "had_delta": bool(new_ids or changed_ids or deleted_ids), + "requested": len(requested_ids), + "cache_hits": len(cache_hit_ids), + "extracted": len(extract_ids), } return merged diff --git a/rag/advanced_rag/knowlege_compile/wiki_incremental.py b/rag/advanced_rag/knowlege_compile/wiki_incremental.py index 775e77b658..3333b135e4 100644 --- a/rag/advanced_rag/knowlege_compile/wiki_incremental.py +++ b/rag/advanced_rag/knowlege_compile/wiki_incremental.py @@ -172,7 +172,7 @@ def _wiki_should_re_synthesize( existing_sources = set(page.get("source_doc_ids", [])) total_sources = existing_sources | new_source_doc_ids claim_count = len(page.get("claims", [])) - last_synth_ver = page.get("synthesis_version_int", 1) + last_synth_ver = _as_int(page.get("synthesis_version_int"), 1) versions_since = next_version - last_synth_ver return ( @@ -998,7 +998,12 @@ async def _search_existing_pages( return results -async def _load_map_relations(tenant_id: str, kb_id: str, excluded_doc_ids: set[str] | None = None) -> list[dict]: +async def _load_map_relations( + tenant_id: str, + kb_id: str, + excluded_doc_ids: set[str] | None = None, + chunk_state: dict[str, dict] | None = None, +) -> list[dict]: """Load all extracted (from, to, type) relations from wiki_map_extract rows. These are the semantic edges the LLM extracted during MAP. When both @@ -1006,52 +1011,24 @@ async def _load_map_relations(tenant_id: str, kb_id: str, excluded_doc_ids: set[ connections (outlinks), which is far more reliable than hoping REFINE sprinkled [[wikilinks]] into prose. """ - from common.doc_store.doc_store_base import OrderByExpr + if chunk_state is None: + from rag.advanced_rag.knowlege_compile.wiki import _wiki_load_active_map_state - index = search.index_name(tenant_id) - if not settings.docStoreConn.index_exist(index, kb_id): - return [] + chunk_state = await _wiki_load_active_map_state(tenant_id, kb_id) + from rag.advanced_rag.knowlege_compile.wiki import _wiki_load_map_extracts_for_state + + extracts = await _wiki_load_map_extracts_for_state(tenant_id, kb_id, chunk_state) relations: list[dict] = [] - offset, page_size = 0, 1000 - while True: - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - ["content_with_weight", "doc_id"], - [], - {"compile_kwd": ["wiki_map_extract"]}, - [], - OrderByExpr(), - offset, - page_size, - index, - [kb_id], - ) - field_map = settings.docStoreConn.get_fields(res, ["content_with_weight", "doc_id"]) or {} - except Exception: - logging.exception("wiki: failed to load map relations for kb=%s", kb_id) - return relations - for row in field_map.values(): - row_doc_ids = _as_str_list(row.get("doc_id")) - if excluded_doc_ids and any(doc_id in excluded_doc_ids for doc_id in row_doc_ids): + for extract in extracts: + if excluded_doc_ids and str(extract.get("doc_id") or "") in excluded_doc_ids: + continue + for relation in extract.get("relations") or []: + if not isinstance(relation, dict): continue - raw = row.get("content_with_weight") - if isinstance(raw, str): - try: - raw = json.loads(raw) - except (json.JSONDecodeError, TypeError): - raw = None - if not isinstance(raw, dict): - continue - for r in raw.get("relations") or []: - if isinstance(r, dict): - frm = r.get("from") - to = r.get("to") - if isinstance(frm, str) and isinstance(to, str): - relations.append({"from": frm, "to": to, "type": r.get("type", "related")}) - if len(field_map) < page_size: - break - offset += page_size + source = relation.get("from") + target = relation.get("to") + if isinstance(source, str) and isinstance(target, str): + relations.append({"from": source, "to": target, "type": relation.get("type", "related")}) return relations @@ -1059,6 +1036,7 @@ async def _wiki_load_pages_for_graph( tenant_id: str, kb_id: str, excluded_doc_ids: set[str] | None = None, + chunk_state: dict[str, dict] | None = None, ) -> list[dict]: """Reload compiled wiki_page rows and project them onto the canvas-graph shape expected by ``dataset_wiki_generator.build_wiki_page_graph``. @@ -1157,7 +1135,12 @@ async def _wiki_load_pages_for_graph( from api.db.services.document_service import DocumentService excluded_doc_ids = await thread_pool_exec(DocumentService.get_disabled_doc_ids_by_kb_id, kb_id) - map_relations = await _load_map_relations(tenant_id, kb_id, excluded_doc_ids=excluded_doc_ids) + map_relations = await _load_map_relations( + tenant_id, + kb_id, + excluded_doc_ids=excluded_doc_ids, + chunk_state=chunk_state, + ) except Exception: logging.exception("wiki: failed to load MAP relations for graph fallback kb=%s", kb_id) map_relations = [] @@ -1328,6 +1311,14 @@ def _as_str_list(raw) -> list[str]: return [] +def _as_int(raw, default: int = 0) -> int: + """Coerce numeric doc-store fields, which some backends return as strings.""" + try: + return int(raw) + except (TypeError, ValueError): + return default + + def _wiki_claim_chunk_ids(claim: dict) -> list[str]: """Return the source chunk id(s) a MAP claim is attributed to. @@ -1524,6 +1515,7 @@ async def _wiki_reduce_entity( new_claims: list[dict], existing_page: dict | None, deleted_doc_ids: set[str], + invalidated_chunk_ids: set[str] | None = None, entity_type: str = "entity", aliases: list[str] | None = None, source_doc_ids: list[str] | None = None, @@ -1576,10 +1568,15 @@ async def _wiki_reduce_entity( else: existing_claims = [] deleted_set = deleted_doc_ids or set() + invalidated_set = invalidated_chunk_ids or set() + existing_chunk_ids = set(_as_str_list(existing_page.get("source_chunk_ids"))) + all_page_evidence_invalidated = bool(existing_chunk_ids) and existing_chunk_ids <= invalidated_set - retractions = [c for c in existing_claims if c.get("source_doc_id") in deleted_set] + retractions = [ + c for c in existing_claims if c.get("source_doc_id") in deleted_set or bool(set(_wiki_claim_chunk_ids(c)) & invalidated_set) or (all_page_evidence_invalidated and not _wiki_claim_chunk_ids(c)) + ] - retained_claims = [c for c in existing_claims if c.get("source_doc_id") not in deleted_set] + retained_claims = [c for c in existing_claims if c not in retractions] retained_texts = {c.get("statement", c.get("text", "")) for c in retained_claims} additions = [c for c in new_claims if c.get("statement", c.get("text", "")) not in retained_texts] @@ -1587,8 +1584,8 @@ async def _wiki_reduce_entity( all_doc_ids = ( {c.get("source_doc_id") for c in retained_claims if c.get("source_doc_id")} | {c.get("source_doc_id") for c in additions if c.get("source_doc_id")} | (set(source_doc_ids or []) - deleted_set) ) - current_chunk_ids = sorted(set(source_chunk_ids or [])) - evidence_changed = bool(set(current_chunk_ids) - set(_as_str_list(existing_page.get("source_chunk_ids")))) + current_chunk_ids = sorted(set(source_chunk_ids or []) if source_chunk_ids else existing_chunk_ids - invalidated_set) + evidence_changed = set(current_chunk_ids) != existing_chunk_ids if not all_doc_ids: return { @@ -1627,6 +1624,7 @@ async def _wiki_reduce_batch( affected_names: set[str], existing_pages: dict[str, dict], deleted_doc_ids: set[str], + invalidated_chunk_ids: set[str] | None = None, canonical_claims: dict[str, list[dict]] | None = None, canonical_map: dict[str, dict] | None = None, name_resolution: dict[str, str] | None = None, @@ -1685,6 +1683,7 @@ async def _wiki_reduce_batch( new_claims=claims, existing_page=name_to_page.get(name, existing_pages.get(name)), deleted_doc_ids=deleted_doc_ids, + invalidated_chunk_ids=invalidated_chunk_ids, ) ) if not tasks: @@ -1864,6 +1863,7 @@ async def _wiki_refine_page( # Blank page_id would produce `slug_kwd: [""]` queries → Infinity 3052. if not page_id or not str(page_id).strip(): return existing_page + page_version = _as_int(page_version) topic_candidates = await _wiki_rank_topic_candidates( page_title, @@ -2462,7 +2462,11 @@ def _wiki_build_contextual_hints( rp = existing_page.get("related_kb_pages_kwd") if rp: if isinstance(rp, str): - related = json.loads(rp) + try: + parsed = json.loads(rp) + related = parsed if isinstance(parsed, list) else [rp] + except (json.JSONDecodeError, TypeError): + related = [rp] elif isinstance(rp, list): related = rp if not related: @@ -2473,8 +2477,14 @@ def _wiki_build_contextual_hints( lines = ["## Context: Related Entities & Concepts", "Reference them in the opening paragraph and relevant sections:"] for r in related[:10]: - entity_name = r.get("entity_name") or r.get("name", "") - relation = r.get("relation") or r.get("type", "related") + if isinstance(r, dict): + entity_name = r.get("entity_name") or r.get("name") or r.get("slug", "") + relation = r.get("relation") or r.get("type", "related") + else: + entity_name = str(r or "").strip() + relation = "related" + if not entity_name: + continue lines.append(f"- [[{entity_name}]] — {relation}") return "\n".join(lines) @@ -3071,6 +3081,7 @@ async def _wiki_finalize( kb_id: str, embd_mdl, page_ids: list[str] | None = None, + chunk_state: dict[str, dict] | None = None, ) -> None: """Post-REFINE cleanup: dead wikilinks + cross-reference update. @@ -3142,7 +3153,12 @@ async def _wiki_finalize( from api.db.services.document_service import DocumentService disabled_doc_ids = await thread_pool_exec(DocumentService.get_disabled_doc_ids_by_kb_id, kb_id) - map_relations = await _load_map_relations(tenant_id, kb_id, excluded_doc_ids=disabled_doc_ids) + map_relations = await _load_map_relations( + tenant_id, + kb_id, + excluded_doc_ids=disabled_doc_ids, + chunk_state=chunk_state, + ) relation_edges: dict[str, set[str]] = {} # pid → {target slug} if map_relations: for rel in map_relations: @@ -3322,8 +3338,10 @@ async def wiki_compile_incremental( tenant_id: str, kb_id: str, mode: str, + chunk_delta: dict[str, set[str]], + previous_chunk_state: dict[str, dict], + current_chunk_state: dict[str, dict], incremental: bool = False, # True = incremental run - map_results: list[dict] | None = None, # from MAP phase deleted_doc_ids: set[str] | None = None, callback: Callable | None = None, ) -> dict: @@ -3332,7 +3350,6 @@ async def wiki_compile_incremental( Args: mode: ``entity`` for Mode A, or ``topic`` for Mode B. incremental: True=incremental update, False=full build - map_results: MAP outputs. If None, loads from ES. deleted_doc_ids: Documents that were removed. callback: Progress callback. @@ -3340,7 +3357,6 @@ async def wiki_compile_incremental( """ from common.misc_utils import thread_pool_exec from rag.nlp import search - from common.doc_store.doc_store_base import OrderByExpr summary = {"pages_created": 0, "pages_modified": 0, "pages_deleted": 0, "errors": []} @@ -3351,64 +3367,33 @@ async def wiki_compile_incremental( except Exception: pass - # ----- Phase 1: Load MAP results if not provided ----- - if not map_results: - _progress("Loading MAP results from doc store ...") - map_results = [] - index = search.index_name(tenant_id) - # Each wiki_map_extract row stores its per-chunk extract as a JSON blob in - # ``content_with_weight`` (see _wiki_build_resume_doc) — the entity / - # concept / claim / relation / topic lists are NOT separate columns, so - # they must be parsed out of that blob to rebuild the map_result shape - # that _extract_raw_entities and REDUCE expect. - from api.db.services.document_service import DocumentService - from rag.advanced_rag.knowlege_compile.wiki import _wiki_doc_ids + invalidated_chunk_ids = set(chunk_delta.get("changed_chunk_ids") or set()) | set(chunk_delta.get("deleted_chunk_ids") or set()) + delta_current_chunk_ids = set(chunk_delta.get("new_chunk_ids") or set()) | set(chunk_delta.get("changed_chunk_ids") or set()) + delta_before_results: list[dict] = [] + delta_after_results: list[dict] = [] - disabled_doc_ids = _wiki_doc_ids(await thread_pool_exec(DocumentService.get_disabled_doc_ids_by_kb_id, kb_id)) - select_fields = ["content_with_weight", "doc_id"] - offset = 0 - page_size = 1000 - while True: - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - select_fields, - [], - {"compile_kwd": ["wiki_map_extract"]}, - [], - OrderByExpr(), - offset, - page_size, - index, - [kb_id], - ) - field_map = settings.docStoreConn.get_fields(res, select_fields) or {} - except Exception: - logging.exception("wiki: failed to load MAP results for kb=%s", kb_id) - break - for row in field_map.values(): - row_doc_ids = _wiki_doc_ids(row.get("doc_id")) - if row_doc_ids & disabled_doc_ids: - continue - raw = row.get("content_with_weight") - if isinstance(raw, str) and raw: - try: - extract = json.loads(raw) - except Exception: - extract = None - elif isinstance(raw, dict): - extract = raw - else: - extract = None - if not isinstance(extract, dict): - continue - extract["doc_id"] = next(iter(row_doc_ids), "") - map_results.append(extract) - if len(field_map) < page_size: - break - offset += page_size + # Versioned MAP storage contains historical rows. Incremental compilation + # must select exactly the versions referenced by the candidate current + # state, never every historical row in the index. + from rag.advanced_rag.knowlege_compile.wiki import _wiki_load_map_extracts_for_state - if not map_results: + map_results = await _wiki_load_map_extracts_for_state(tenant_id, kb_id, current_chunk_state) + if delta_current_chunk_ids: + delta_after_results = await _wiki_load_map_extracts_for_state( + tenant_id, + kb_id, + current_chunk_state, + delta_current_chunk_ids, + ) + if invalidated_chunk_ids: + delta_before_results = await _wiki_load_map_extracts_for_state( + tenant_id, + kb_id, + previous_chunk_state, + invalidated_chunk_ids, + ) + + if not map_results and not delta_before_results: _progress("No MAP results found. Skipping wiki compilation.") return summary @@ -3436,7 +3421,12 @@ async def wiki_compile_incremental( # claim loading. Keeps Entity Matching operating on small metadata only # (mirrors old-mode dedup); full claim text is loaded per-affected-name # after matching, so peak memory stays bounded. + map_results = map_results or [] raw_entities, claim_index = _extract_raw_entities(map_results) + before_raw_entities, _before_claim_index = _extract_raw_entities(delta_before_results) + before_raw_names = {entry.get("name") for entry in before_raw_entities if entry.get("name")} + after_raw_entities, _after_claim_index = _extract_raw_entities(delta_after_results) + after_raw_names = {entry.get("name") for entry in after_raw_entities if entry.get("name")} # Preserve MAP topic provenance so each page's writer chooses among topics # extracted from that page's own source documents, rather than from an @@ -3496,16 +3486,6 @@ async def wiki_compile_incremental( kb_id=kb_id, incremental=incremental, ) - if mode == "topic" and incremental: - try: - from api.db.services.document_service import DocumentService - - disabled_doc_ids = await thread_pool_exec(DocumentService.get_disabled_doc_ids_by_kb_id, kb_id) - historical_relations = await _load_map_relations(tenant_id, kb_id, excluded_doc_ids=disabled_doc_ids) - raw_relations.extend(historical_relations) - except Exception: - logging.exception("wiki: failed to load historical relations for incremental routing") - canonical_resolution = dict(name_resolution) for canonical_name, canonical_entry in canonical_entities.items(): canonical_resolution.setdefault(canonical_name, canonical_name) @@ -3529,10 +3509,28 @@ async def wiki_compile_incremental( del raw_relations del canonical_resolution + # ``map_results`` is the complete current snapshot when chunk deltas are + # supplied. Replace provenance accumulated in the canonical cache with + # that snapshot so changed/deleted chunk ids do not survive indefinitely. + current_evidence: dict[str, dict[str, set[str] | int]] = {} + for entry in raw_entities: + cname = name_resolution.get(entry["name"], entry["name"]) + evidence = current_evidence.setdefault(cname, {"docs": set(), "chunks": set(), "claims": 0}) + evidence["docs"].update(entry.get("source_doc_ids") or []) + evidence["chunks"].update(entry.get("source_chunk_ids") or []) + evidence["claims"] += int(entry.get("claim_count") or 0) + for cname, centry in canonical_map.items(): + evidence = current_evidence.get(cname) + if evidence is None: + continue + centry["source_doc_ids"] = sorted(evidence["docs"]) + centry["source_chunk_ids"] = sorted(evidence["chunks"]) + centry["claim_count"] = evidence["claims"] + # raw_entities (lightweight) no longer needed after matching. del raw_entities - if not canonical_map: + if not canonical_map and not before_raw_names: _progress("Entity Matching: no canonical entities found. Skipping.") return summary @@ -3622,6 +3620,28 @@ async def wiki_compile_incremental( kb_id, ) + if invalidated_chunk_ids: + for cname, existing in canonical_entities.items(): + if cname in canonical_map: + continue + old_chunks = set(existing.get("source_chunk_ids") or []) + if not old_chunks & invalidated_chunk_ids: + continue + remaining_chunks = old_chunks - invalidated_chunk_ids + if not remaining_chunks: + await _delete_canonical_entity(tenant_id, kb_id, cname) + else: + await _update_canonical_entity( + tenant_id, + kb_id, + cname, + existing.get("entity_type_kwd", "entity"), + existing.get("aliases", []), + existing.get("source_doc_ids", []), + existing.get("mention_count_int", 0), + source_chunk_ids=sorted(remaining_chunks), + ) + # Clean up deleted canonical entities (from doc deletion) if deleted_doc_ids: for cname, centry in list(canonical_map.items()): @@ -3637,27 +3657,17 @@ async def wiki_compile_incremental( canonical_names: set[str] = set(canonical_map.keys()) if incremental: - # Affected doc ids are the docs contributing to this batch's canonical - # entities, plus any deleted docs. Derived from canonical_map (which - # carries source_doc_ids) — no need to re-read the released map_results. - affected_doc_ids = set() - for centry in canonical_map.values(): - affected_doc_ids.update(centry.get("source_doc_ids", [])) - affected_doc_ids = affected_doc_ids | (deleted_doc_ids or set()) - - # Map doc_page_source entity_names (raw) through name_resolution -> canonical - affected_names: set[str] = set() - if affected_doc_ids: - dps_tasks = [_wiki_load_doc_page_source(tenant_id, kb_id, did) for did in affected_doc_ids] - dps_results = await asyncio.gather(*dps_tasks) - for dps in dps_results: - if dps: - for raw_name in dps.get("entity_names", []): - cname = name_resolution.get(raw_name, raw_name) - if cname in canonical_names: - affected_names.add(cname) - if not affected_names: - affected_names = canonical_names + # Resolve names from both sides of the chunk delta. Deleted names may + # no longer exist in the current canonical map, but their old pages + # still need retraction/deletion. + existing_aliases: dict[str, str] = {} + for cname, centry in canonical_entities.items(): + existing_aliases[_normalize_key(cname)] = cname + for alias in centry.get("aliases") or []: + if isinstance(alias, str) and alias: + existing_aliases[_normalize_key(alias)] = cname + affected_names = {name_resolution.get(raw_name) or existing_aliases.get(_normalize_key(raw_name)) or raw_name for raw_name in before_raw_names | after_raw_names} + affected_names.discard("") else: affected_names = canonical_names @@ -3710,6 +3720,7 @@ async def wiki_compile_incremental( affected_names=affected_names, existing_pages=existing_pages, deleted_doc_ids=deleted_doc_ids or set(), + invalidated_chunk_ids=invalidated_chunk_ids, canonical_claims=canonical_claims, canonical_map=canonical_map, name_resolution=name_resolution, @@ -3780,7 +3791,7 @@ async def wiki_compile_incremental( # (doc_page_source page_ids is already handled in mode_run) _progress("FINALIZE: updating cross-references ...") try: - await _wiki_finalize(tenant_id, kb_id, embd_mdl) + await _wiki_finalize(tenant_id, kb_id, embd_mdl, chunk_state=current_chunk_state) except Exception: logging.exception("wiki: FINALIZE failed for kb=%s", kb_id) summary["errors"].append("FAILED_FINALIZE") @@ -3961,7 +3972,7 @@ async def _wiki_mode_a_run( summary["pages_deleted"] += 1 return - next_version = (existing.get("page_version_int", 0) if existing else 0) + 1 + next_version = _as_int(existing.get("page_version_int")) + 1 if existing else 1 new_doc_ids = {c.get("source_doc_id") for c in entry["additions"] if c.get("source_doc_id")} if existing and _wiki_should_re_synthesize(existing, new_doc_ids, next_version): refine_mode = "re-synthesize" @@ -4296,7 +4307,7 @@ async def _wiki_mode_b_run( if existing and _wiki_should_re_synthesize( existing, {c.get("source_doc_id") for c in additions if c.get("source_doc_id")}, - existing.get("page_version_int", 0) + 1, + _as_int(existing.get("page_version_int")) + 1, ): refine_mode = "re-synthesize" diff --git a/rag/svr/task_executor_refactor/dataset_wiki_generator.py b/rag/svr/task_executor_refactor/dataset_wiki_generator.py index 5a47f3c2d9..739a54cba1 100644 --- a/rag/svr/task_executor_refactor/dataset_wiki_generator.py +++ b/rag/svr/task_executor_refactor/dataset_wiki_generator.py @@ -57,7 +57,13 @@ from common.misc_utils import thread_pool_exec from rag.nlp import search from rag.advanced_rag.knowlege_compile.structure import LLMCallPool from rag.advanced_rag.knowlege_compile.wiki import ( - _wiki_doc_ids, + WIKI_MAP_STATE_COMPILE_KWD, + WIKI_MAP_STATE_META_COMPILE_KWD, + _wiki_commit_active_map_state, + _wiki_compare_chunk_states, + _wiki_load_active_map_state, + _wiki_load_map_extracts_for_state, + _wiki_scan_current_chunk_state, wiki_map_from_chunks, wiki_plan_from_reduction, wiki_reduce_from_extracts, @@ -333,71 +339,8 @@ def _wiki_eligible_docs(all_docs, tenant_id: str, skip_doc_ids=None) -> list[tup async def _wiki_existing_map_doc_ids(tenant_id: str, kb_id: str) -> set[str]: - from common.doc_store.doc_store_base import OrderByExpr - from api.db.services.document_service import DocumentService - - index = search.index_name(tenant_id) - if not settings.docStoreConn.index_exist(index, kb_id): - return set() - - doc_ids: set[str] = set() - disabled_doc_ids = _wiki_doc_ids(await thread_pool_exec(DocumentService.get_disabled_doc_ids_by_kb_id, kb_id)) - select_fields = ["id", "doc_id"] - offset = 0 - page_size = 1000 - while True: - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - select_fields, - [], - {"compile_kwd": [WIKI_MAP_COMPILE_KWD]}, - [], - OrderByExpr(), - offset, - page_size, - index, - [kb_id], - ) - field_map = settings.docStoreConn.get_fields(res, select_fields) or {} - except Exception: - logging.exception("wiki: failed to scan MAP doc ids for kb=%s", kb_id) - return doc_ids - if not field_map: - break - for row in field_map.values(): - doc_ids.update(_wiki_doc_ids(row.get("doc_id")) - disabled_doc_ids) - if len(field_map) < page_size: - break - offset += page_size - return doc_ids - - -async def _wiki_delete_map_rows_for_docs( - tenant_id: str, - kb_id: str, - doc_ids: set[str], -) -> None: - """Remove MAP resume rows so the next run must re-extract the documents.""" - if not doc_ids: - return - try: - await thread_pool_exec( - settings.docStoreConn.delete, - { - "compile_kwd": [WIKI_MAP_COMPILE_KWD], - "doc_id": sorted(str(doc_id) for doc_id in doc_ids), - }, - search.index_name(tenant_id), - kb_id, - ) - except Exception: - logging.exception( - "wiki: failed to invalidate MAP resume rows for kb=%s docs=%s", - kb_id, - sorted(doc_ids), - ) - raise + state = await _wiki_load_active_map_state(tenant_id, kb_id) + return {str(item.get("doc_id") or "") for item in state.values() if item.get("doc_id")} async def _wiki_has_compiled_pages(tenant_id: str, kb_id: str) -> bool | None: @@ -444,26 +387,11 @@ async def _wiki_delete_deleted_doc_state( if not settings.docStoreConn.index_exist(index, kb_id): return - # 1. MAP resume rows are keyed by the real doc_id — delete outright. - try: - await thread_pool_exec( - settings.docStoreConn.delete, - { - "compile_kwd": [WIKI_MAP_COMPILE_KWD], - "doc_id": sorted(deleted_doc_ids), - }, - index, - kb_id, - ) - except Exception: - logging.exception( - "wiki: failed to delete MAP rows for removed docs in kb=%s docs=%s", - kb_id, - sorted(deleted_doc_ids), - ) - return + # MAP extraction versions are historical cache entries and deliberately + # survive document deletion. The committed active-state snapshot decides + # which versions are allowed to participate in the current Wiki. - # 1b. doc_page_source rows are keyed by doc_id too — delete outright. + # doc_page_source rows are current derived state and are deleted outright. try: await thread_pool_exec( settings.docStoreConn.delete, @@ -717,7 +645,8 @@ async def _wiki_reset_all_wiki_state(tenant_id: str, kb_id: str) -> None: "wiki_reduce_result", "wiki_page_draft", "wiki_doc_page_source", - "wiki_map_extract", + WIKI_MAP_STATE_COMPILE_KWD, + WIKI_MAP_STATE_META_COMPILE_KWD, "wiki_mode_meta", ] # Delete in one bulk call using compile_kwd IN filter. @@ -1445,10 +1374,9 @@ async def run_wiki( return _cb - # 4. MAP per eligible doc. Each MAP call's own resume mechanism - # (wiki_map_extract rows keyed by chunk_id) skips chunks that - # were already processed in a prior run — this is the incremental - # behavior the user asked for. + # 4. MAP per eligible doc. Historical extraction versions are keyed by + # doc_id + chunk_id + input hash, so unchanged or reverted content can + # reuse prior LLM output. # # Resolve templates before starting workers so the first eligible # template remains the deterministic source for KB-wide REDUCE/PLAN/ @@ -1491,8 +1419,7 @@ async def run_wiki( "saw_any": False, "status": "ok", "agg": {"entities": 0, "concepts": 0, "claims": 0, "relations": 0}, - "delta": {"new": 0, "changed": 0, "unchanged": 0, "deleted": 0}, - "had_delta": False, + "map": {"requested": 0, "cache_hits": 0, "extracted": 0}, } for i, (doc, _, _) in enumerate(resolved_eligible) } @@ -1551,9 +1478,8 @@ async def run_wiki( stats["agg"][key] += len(phase1.get(key) or []) meta = phase1.get("_meta") or {} if isinstance(meta, dict): - for key in stats["delta"]: - stats["delta"][key] += int(meta.get(key, 0) or 0) - stats["had_delta"] |= bool(meta.get("had_delta")) + for key in stats["map"]: + stats["map"][key] += int(meta.get(key, 0) or 0) except Exception: logging.exception( "wiki: MAP failed for doc %s batch %d", @@ -1582,20 +1508,18 @@ async def run_wiki( stats["status"] = "empty" logging.info("wiki: no chunks for doc %s; skipping", doc_id) agg = stats["agg"] - delta = stats["delta"] + map_stats = stats["map"] logging.info( - "wiki: MAP doc=%s entities=%d concepts=%d claims=%d relations=%d (batches=%d, new=%d changed=%d unchanged=%d deleted=%d, delta=%s)", + "wiki: MAP doc=%s entities=%d concepts=%d claims=%d relations=%d (batches=%d, requested=%d cache_hits=%d extracted=%d)", doc_id, agg["entities"], agg["concepts"], agg["claims"], agg["relations"], stats["batch_count"], - delta["new"], - delta["changed"], - delta["unchanged"], - delta["deleted"], - stats["had_delta"], + map_stats["requested"], + map_stats["cache_hits"], + map_stats["extracted"], ) # 5. REDUCE / PLAN / REFINE KB-wide. Each phase has its own @@ -1679,7 +1603,6 @@ async def run_wiki_incremental( embedding_model, load_chunks_for_doc: Callable[..., AsyncIterator[list[dict]]], mode: str | None = None, - _map_rebuild_retry: bool = False, ) -> None: """Dual-mode wiki compilation with incremental support. @@ -1758,6 +1681,25 @@ async def run_wiki_incremental( return pipeline_chat_llm_ids = _validate_wiki_eligible_docs(eligible) if eligible else {} + eligible_doc_ids = {str(doc.get("id")) for doc, _ in eligible if doc.get("id")} + previous_chunk_state = await _wiki_load_active_map_state(ctx.tenant_id, ctx.kb_id) + current_chunk_state = await _wiki_scan_current_chunk_state( + ctx.tenant_id, + ctx.kb_id, + eligible_doc_ids, + ) + chunk_delta = _wiki_compare_chunk_states(previous_chunk_state, current_chunk_state) + target_chunk_ids = chunk_delta["new_chunk_ids"] | chunk_delta["changed_chunk_ids"] + has_chunk_delta = bool(target_chunk_ids or chunk_delta["deleted_chunk_ids"]) + logging.info( + "wiki chunk delta: kb=%s new=%d changed=%d deleted=%d unchanged=%d", + ctx.kb_id, + len(chunk_delta["new_chunk_ids"]), + len(chunk_delta["changed_chunk_ids"]), + len(chunk_delta["deleted_chunk_ids"]), + len(chunk_delta["unchanged_chunk_ids"]), + ) + # Resolve mode from the eligible documents' templates. Each eligible # doc resolves to a wiki template either via its own parser_config or via # its ingestion pipeline (doc.pipeline_id → pipeline dsl → compiler → @@ -1801,6 +1743,10 @@ async def run_wiki_incremental( is_incremental = False existing_map_doc_ids = set() deleted_doc_ids = set() + previous_chunk_state = {} + chunk_delta = _wiki_compare_chunk_states(previous_chunk_state, current_chunk_state) + target_chunk_ids = set(chunk_delta["new_chunk_ids"]) + has_chunk_delta = bool(target_chunk_ids) await _wiki_save_mode(ctx.tenant_id, ctx.kb_id, mode, current_embedding) # 3. Resolve chat model @@ -1833,7 +1779,6 @@ async def run_wiki_incremental( # 4. MAP per doc (same as run_wiki's MAP phase) map_queue: asyncio.Queue = asyncio.Queue(maxsize=WIKI_MAP_QUEUE_SIZE) n_docs = len(eligible) - all_map_results: list[dict] = [] # Pre-resolve parser_cfg for each eligible doc (avoids sync DB call in worker) doc_configs: dict[str, dict] = {} @@ -1874,7 +1819,7 @@ async def run_wiki_incremental( doc_id = doc["id"] map_llm_id = pipeline_chat_llm_ids.get(str(doc_id)) - result = await wiki_map_from_chunks( + await wiki_map_from_chunks( chunks=batch, chat_mdl=map_llm_pool.wrap( _bundle_for(map_llm_id), @@ -1891,15 +1836,8 @@ async def run_wiki_incremental( batch_size_cap=8, window_fraction=0.5, max_workers=WIKI_MAP_LLM_POOL_SIZE, + target_chunk_ids=target_chunk_ids, ) - # Only forward extracts that actually produced content. An - # all-unchanged doc (MAP fully resumed from prior rows) returns an - # empty extract; appending it would make ``all_map_results`` - # non-empty and suppress wiki_compile_incremental's "load stored - # extracts from ES" fallback, so nothing would ever be compiled. - if result and any(result.get(k) for k in ("entities", "concepts", "claims", "relations", "topics")): - result["doc_id"] = doc_id - all_map_results.append(result) except Exception: logging.exception("wiki: MAP failed for doc %s", doc_id) finally: @@ -1916,7 +1854,25 @@ async def run_wiki_incremental( task.cancel() await asyncio.gather(*producers, *workers, return_exceptions=True) - if not all_map_results and not deleted_doc_ids: + if target_chunk_ids: + resolved_versions = await _wiki_load_map_extracts_for_state( + ctx.tenant_id, + ctx.kb_id, + current_chunk_state, + target_chunk_ids, + ) + resolved_chunk_ids = {str((extract.get("_map_version") or {}).get("chunk_id") or "") for extract in resolved_versions} + missing_chunk_ids = target_chunk_ids - resolved_chunk_ids + if missing_chunk_ids: + logging.error( + "wiki: MAP extraction/cache resolution incomplete kb=%s missing_chunks=%s", + ctx.kb_id, + sorted(missing_chunk_ids), + ) + progress(-1, f"Wiki MAP failed for {len(missing_chunk_ids)} chunk(s).") + return + + if not has_chunk_delta and not deleted_doc_ids: # Nothing fresh, changed, or deleted this run. Skip only when there is # genuinely nothing to build: an existing MAP baseline already has # compiled pages. When the Wiki was explicitly cleared, both @@ -1937,7 +1893,12 @@ async def run_wiki_incremental( # the persisted pages (zero LLM cost) so a re-run backfills graph # edges for pages written before auto-linking existed. try: - await _wiki_finalize(ctx.tenant_id, ctx.kb_id, embedding_model) + await _wiki_finalize( + ctx.tenant_id, + ctx.kb_id, + embedding_model, + chunk_state=current_chunk_state, + ) except Exception: logging.exception("wiki: up-to-date FINALIZE failed for kb=%s", ctx.kb_id) @@ -1945,39 +1906,19 @@ async def run_wiki_incremental( # persistence existed (or a graph lost to an interrupted run) still # render. Reload pages → project → persist wiki_entity/relation. try: - graph_pages = await _wiki_load_pages_for_graph(ctx.tenant_id, ctx.kb_id) + graph_pages = await _wiki_load_pages_for_graph( + ctx.tenant_id, + ctx.kb_id, + chunk_state=current_chunk_state, + ) if graph_pages: await persist_wiki_page_graph(ctx=ctx, pages=graph_pages) except Exception: logging.exception("wiki: up-to-date page-graph persist failed for kb=%s", ctx.kb_id) + await _wiki_commit_active_map_state(ctx.tenant_id, ctx.kb_id, current_chunk_state) progress(1.0, "Wiki is up to date.") return - if existing_map_doc_ids and has_compiled_pages is False and not _map_rebuild_retry: - # A first build can be interrupted after MAP rows are written. If - # those rows are subsequently removed or become unreadable, the - # resume check still suppresses MAP and the restore phase sees no - # payload. Invalidate only the affected documents and retry once; - # the second invocation then treats them as a clean MAP build. - stale_doc_ids = {str(doc.get("id")) for doc, _ in eligible if doc.get("id")} - logging.warning( - "wiki: MAP resume state is unusable with no compiled pages; forcing one MAP rebuild kb=%s docs=%s", - ctx.kb_id, - sorted(stale_doc_ids), - ) - try: - await _wiki_delete_map_rows_for_docs(ctx.tenant_id, ctx.kb_id, stale_doc_ids) - except Exception: - progress(-1, "Failed to reset stale MAP state for wiki rebuild.") - return - await run_wiki_incremental( - ctx=ctx, - embedding_model=embedding_model, - load_chunks_for_doc=load_chunks_for_doc, - mode=mode, - _map_rebuild_retry=True, - ) - return logging.info("wiki: MAP rows exist but no pages found for kb=%s; rebuilding from stored extracts.", ctx.kb_id) # 5. Run incremental wiki compilation (Mode A or Mode B) @@ -1996,8 +1937,10 @@ async def run_wiki_incremental( kb_id=ctx.kb_id, mode=mode, incremental=is_incremental, - map_results=all_map_results or None, deleted_doc_ids=deleted_doc_ids or None, + chunk_delta=chunk_delta, + previous_chunk_state=previous_chunk_state, + current_chunk_state=current_chunk_state, callback=lambda p, msg: progress(p, msg), ) @@ -2010,12 +1953,21 @@ async def run_wiki_incremental( _wiki_load_pages_for_graph, ) - graph_pages = await _wiki_load_pages_for_graph(ctx.tenant_id, ctx.kb_id) + graph_pages = await _wiki_load_pages_for_graph( + ctx.tenant_id, + ctx.kb_id, + chunk_state=current_chunk_state, + ) if graph_pages: await persist_wiki_page_graph(ctx=ctx, pages=graph_pages) except Exception: logging.exception("wiki: page-graph persist failed for kb=%s", ctx.kb_id) - progress(1.0, f"Wiki done: +{summary.get('pages_created', 0)} ~{summary.get('pages_modified', 0)} -{summary.get('pages_deleted', 0)}") + if not summary.get("errors"): + await _wiki_commit_active_map_state(ctx.tenant_id, ctx.kb_id, current_chunk_state) + if summary.get("errors"): - logging.warning("wiki: non-fatal errors: %s", summary["errors"]) + logging.warning("wiki: incomplete compilation errors: %s", summary["errors"]) + progress(-1, f"Wiki incomplete: {len(summary['errors'])} page(s) failed; retry required.") + else: + progress(1.0, f"Wiki done: +{summary.get('pages_created', 0)} ~{summary.get('pages_modified', 0)} -{summary.get('pages_deleted', 0)}") diff --git a/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py b/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py index 5425885d6f..f30d20aac1 100644 --- a/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py +++ b/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py @@ -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"] diff --git a/test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py b/test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py index fe05e83f9d..bbf2996f73 100644 --- a/test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py +++ b/test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py @@ -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 diff --git a/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py b/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py index a7474c75c4..aaffa6a7c1 100644 --- a/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py +++ b/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py @@ -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 diff --git a/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_map_state.py b/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_map_state.py new file mode 100644 index 0000000000..7961e819df --- /dev/null +++ b/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_map_state.py @@ -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"}}