From 3e4c6dfc0a3b927c15989a70b2e7b6d2309310ca Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Thu, 23 Jul 2026 18:39:16 +0800 Subject: [PATCH] Feat: refine the tree navigation during compilations (#17140) --- api/apps/restful_apis/dataset_api.py | 202 ++++++ api/apps/services/dataset_api_service.py | 673 +++++++++++++++++- api/db/__init__.py | 10 + api/db/db_models.py | 19 + .../compilation_templates/empty.yaml | 17 - api/db/services/document_service.py | 121 +++- api/db/services/knowledgebase_service.py | 12 + .../pipeline_operation_log_service.py | 52 +- common/constants.py | 13 + rag/advanced_rag/knowlege_compile/_common.py | 28 +- .../knowlege_compile/dataset_nav.py | 104 ++- rag/advanced_rag/knowlege_compile/raptor.py | 4 +- rag/advanced_rag/knowlege_compile/runner.py | 210 +++++- .../knowlege_compile/structure.py | 265 +++++-- rag/advanced_rag/knowlege_compile/wiki.py | 79 +- rag/flow/compiler/compiler.py | 2 + rag/nlp/search.py | 2 +- rag/svr/task_executor.py | 24 +- .../dataset_structure_merger.py | 214 ++++++ .../dataset_wiki_generator.py | 337 ++++++++- .../task_executor_refactor/task_handler.py | 16 +- rag/utils/infinity_conn.py | 6 +- 22 files changed, 2244 insertions(+), 166 deletions(-) delete mode 100644 api/db/init_data/compilation_templates/empty.yaml create mode 100644 rag/svr/task_executor_refactor/dataset_structure_merger.py diff --git a/api/apps/restful_apis/dataset_api.py b/api/apps/restful_apis/dataset_api.py index de0615f1da..f062fceafa 100644 --- a/api/apps/restful_apis/dataset_api.py +++ b/api/apps/restful_apis/dataset_api.py @@ -672,6 +672,73 @@ async def get_wiki_graph(tenant_id, dataset_id): return get_error_data_result(message="Internal server error") +@manager.route("/datasets//artifacts_structure", methods=["GET"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def get_dataset_structure(tenant_id, dataset_id): + """Return the dataset-scope (KB-wide) structure graph for one kind. + + GET /api/v1/datasets//artifacts_structure?kind= + where ``kind`` is one of: + graph | mindmap | timeline | session_essence | session_graph + + These are the non-tree structured artifacts merged across the whole + dataset (written when a template has ``dataset_merge`` enabled). Response + mirrors the per-document structure graph so the frontend reuses its view:: + + {"code": 0, "data": {"kind": "", "templates": [ + {"template_id", "template_name", "kind", "entities", "relations"} + ]}} + """ + try: + kind = request.args.get("kind", "") + if isinstance(kind, str): + kind = kind.strip() + if not kind: + return get_error_data_result( + message="`kind` is required (one of: graph, mindmap, timeline, session_essence, session_graph).", + code=RetCode.ARGUMENT_ERROR, + ) + if dataset_api_service._resolve_dataset_structure_kind(kind) is None: + return get_error_data_result( + message=f"Unsupported structure kind: {kind!r}. Expected one of: graph, mindmap, timeline, session_essence, session_graph.", + code=RetCode.ARGUMENT_ERROR, + ) + success, result = await dataset_api_service.get_dataset_structure( + dataset_id, + tenant_id, + kind, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + +@manager.route("/datasets//artifacts/alteration", methods=["GET"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def get_wiki_alteration(tenant_id, dataset_id): + """Return document drift for the dataset Artifact wiki. + + GET /api/v1/datasets//artifacts/alteration + Success: {"code": 0, "data": {"removed": int, "newly_uploaded": int, ...}} + """ + try: + success, result = await dataset_api_service.get_wiki_alteration( + dataset_id, + tenant_id, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + @manager.route("/datasets//artifacts", methods=["DELETE"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs @@ -764,6 +831,28 @@ async def get_skill_tree(tenant_id, dataset_id): return get_error_data_result(message="Internal server error") +@manager.route("/datasets//skills", methods=["DELETE"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def delete_all_skills(tenant_id, dataset_id): + """Delete every compiled skill for this dataset. + + DELETE /api/v1/datasets//skills + Success: {"code": 0, "data": {"deleted": }} + """ + try: + success, result = await dataset_api_service.delete_skills( + dataset_id, + tenant_id, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + @manager.route("/datasets//skills/", methods=["GET"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs @@ -787,6 +876,119 @@ async def get_skill_page(tenant_id, dataset_id, skill_kwd): return get_error_data_result(message="Internal server error") +@manager.route("/datasets//nav", methods=["GET"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def list_dataset_nav(tenant_id, dataset_id): + """First level of the dataset navigation tree — the top-level clusters. + + GET /api/v1/datasets//nav + Success: {"code": 0, "data": {"total": , "items": [{name, description, doc_count, type, has_children}, ...]}} + """ + try: + success, result = await dataset_api_service.list_nav_clusters( + dataset_id, + tenant_id, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + +@manager.route("/datasets//nav//children", methods=["GET"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def list_dataset_nav_children(tenant_id, dataset_id, name): + """Direct children of a navigation node (hierarchical, one level per call). + + GET /api/v1/datasets//nav//children + Success: {"code": 0, "data": {"total": , "items": [{name, description, doc_count, type, doc_id, has_children}, ...]}} + """ + try: + success, result = await dataset_api_service.list_nav_children( + dataset_id, + tenant_id, + name, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + +@manager.route("/datasets//nav", methods=["DELETE"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def delete_dataset_nav(tenant_id, dataset_id): + """Delete the entire dataset navigation tree. + + DELETE /api/v1/datasets//nav + Success: {"code": 0, "data": {"deleted": }} + """ + try: + success, result = await dataset_api_service.delete_nav( + dataset_id, + tenant_id, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + +@manager.route("/datasets//nav/", methods=["DELETE"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def delete_dataset_nav_node(tenant_id, dataset_id, name): + """Delete one navigation node and its whole subtree. + + DELETE /api/v1/datasets//nav/ + Success: {"code": 0, "data": {"deleted": }} + """ + try: + success, result = await dataset_api_service.delete_nav_node( + dataset_id, + tenant_id, + name, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + +@manager.route("/datasets//skills/", methods=["DELETE"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def delete_skill_page(tenant_id, dataset_id, skill_kwd): + """Delete one compiled skill node by skill_kwd. + + DELETE /api/v1/datasets//skills/ + Success: {"code": 0, "data": {"deleted": }} + """ + try: + success, result = await dataset_api_service.delete_skill( + dataset_id, + tenant_id, + skill_kwd, + ) + if success: + return get_result(data=result) + return get_result(data=False, message=result, code=RetCode.AUTHENTICATION_ERROR) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + # The two artifact-commit endpoints # GET /datasets//artifacts///commits # GET /datasets//artifacts/commits/ diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index e09a18b61d..06538ceccd 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -35,7 +35,22 @@ from api.utils.api_utils import deep_merge, get_parser_config, remap_dictionary_ from common.misc_utils import thread_pool_exec from rag.advanced_rag.knowlege_compile.wiki import WIKI_PAGE_COMPILE_KWD -_VALID_INDEX_TYPES = {"graph", "raptor", "mindmap", "artifact", "skill"} +# 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 +# task_type equals the index_type and the KB task-id column is ``_task_id``. +# The value is the friendly kind resolvable through ``_resolve_dataset_structure_kind`` +# (defined below), except ``structure`` which is the merge-all variant. +_STRUCTURE_INDEX_TYPE_TO_KIND = { + "structure_graph": "graph", + "structure_mindmap": "mindmap", + "timeline": "timeline", + "session_graph": "session_graph", + "session_essence": "session_essence", + "structure": None, # merge-all: rebuild every dataset-merge kind +} +_STRUCTURE_INDEX_TYPES = frozenset(_STRUCTURE_INDEX_TYPE_TO_KIND) + +_VALID_INDEX_TYPES = {"graph", "raptor", "mindmap", "artifact", "skill"} | set(_STRUCTURE_INDEX_TYPES) _INDEX_TYPE_TO_TASK_TYPE = { "graph": "graphrag", @@ -43,6 +58,9 @@ _INDEX_TYPE_TO_TASK_TYPE = { "mindmap": "mindmap", "artifact": "artifact", "skill": "skill", + # Structure merge types carry their own task_type (== index_type) so the + # executor can resolve which kind to merge from the task body. + **{t: t for t in _STRUCTURE_INDEX_TYPES}, } _INDEX_TYPE_TO_TASK_ID_FIELD = { @@ -51,6 +69,7 @@ _INDEX_TYPE_TO_TASK_ID_FIELD = { "mindmap": "mindmap_task_id", "artifact": "artifact_task_id", "skill": "skill_task_id", + **{t: f"{t}_task_id" for t in _STRUCTURE_INDEX_TYPES}, } _INDEX_TYPE_TO_DISPLAY_NAME = { @@ -59,6 +78,12 @@ _INDEX_TYPE_TO_DISPLAY_NAME = { "mindmap": "Mindmap", "artifact": "Artifact", "skill": "Skill", + "structure_graph": "Structure Graph", + "structure_mindmap": "Structure Mindmap", + "timeline": "Timeline", + "session_graph": "Session Graph", + "session_essence": "Session Essence", + "structure": "Structure", } @@ -884,6 +909,18 @@ def delete_index(dataset_id: str, tenant_id: str, index_type: str, wipe: bool = from rag.nlp import search settings.docStoreConn.delete({"compile_kwd": ["skill", "skill_all"]}, search.index_name(kb.tenant_id), dataset_id) + elif wipe and index_type in _STRUCTURE_INDEX_TYPES: + from rag.nlp import search + + # Wipe the merged KB-wide dataset_graph rows for the requested kind + # (all kinds for the merge-all "structure" type). The per-document + # entity/relation rows the merge reads from are left intact. + friendly = _STRUCTURE_INDEX_TYPE_TO_KIND.get(index_type) + resolved_kind = _resolve_dataset_structure_kind(friendly) if friendly else None + condition: dict = {"knowledge_graph_kwd": ["dataset_graph"]} + if resolved_kind: + condition["compilation_template_kind_kwd"] = [resolved_kind] + settings.docStoreConn.delete(condition, search.index_name(kb.tenant_id), dataset_id) KnowledgebaseService.update_by_id(kb.id, {task_id_field: "", task_finish_at_field: None}) return True, {} @@ -1502,6 +1539,124 @@ def _wiki_index_or_none(tenant_id: str, kb_id: str): return _compiled_index_or_none(tenant_id, kb_id) +def _compilation_template_kind(kind) -> str: + if not isinstance(kind, str): + return "" + normalized = kind.strip().lower().replace("-", "_") + if normalized in {"pageindex", "page_index", "knowledge_graph"}: + return "timeline" + return normalized + + +def _normalize_compilation_template_group_ids(raw) -> list[str]: + if isinstance(raw, str): + raw = [raw] + if not isinstance(raw, list): + return [] + ids: list[str] = [] + seen: set[str] = set() + for group_id in raw: + if not isinstance(group_id, str): + continue + group_id = group_id.strip() + if group_id and group_id not in seen: + seen.add(group_id) + ids.append(group_id) + return ids + + +def _extract_pipeline_compiler_group_ids(dsl) -> list[str]: + if isinstance(dsl, str): + try: + dsl = json.loads(dsl) + except Exception: + return [] + if not isinstance(dsl, dict): + return [] + components = dsl.get("components") + if not isinstance(components, dict): + return [] + + group_ids: list[str] = [] + seen: set[str] = set() + for component in components.values(): + if not isinstance(component, dict): + continue + obj = component.get("obj") if isinstance(component.get("obj"), dict) else {} + component_name = obj.get("component_name") or component.get("component_name") or component.get("name") + if not isinstance(component_name, str) or component_name.lower() != "compiler": + continue + candidates = [ + obj.get("params") if isinstance(obj.get("params"), dict) else {}, + obj, + component.get("params") if isinstance(component.get("params"), dict) else {}, + component, + ] + for candidate in candidates: + for key in ("compilation_template_group_ids", "compilation_template_group_id"): + for group_id in _normalize_compilation_template_group_ids(candidate.get(key)): + if group_id not in seen: + seen.add(group_id) + group_ids.append(group_id) + return group_ids + + +def _template_is_wiki(template: dict | None) -> bool: + if not isinstance(template, dict): + return False + config = template.get("config") if isinstance(template.get("config"), dict) else {} + raw_kind = config.get("kind") or template.get("kind") or "" + return _compilation_template_kind(raw_kind) == "artifacts" + + +def _group_has_wiki_template(group_id: str, tenant_id: str, group_cache: dict[str, bool]) -> bool: + if group_id in group_cache: + return group_cache[group_id] + from api.db.services.compilation_template_group_service import CompilationTemplateGroupService + + group = CompilationTemplateGroupService.get_saved(group_id, tenant_id) + has_wiki = any(_template_is_wiki(template) for template in (group or {}).get("templates") or []) + group_cache[group_id] = has_wiki + return has_wiki + + +def _parser_config_has_wiki_template(parser_config, tenant_id: str, template_cache: dict[str, bool]) -> bool: + from api.db.services.compilation_template_service import CompilationTemplateService + from rag.svr.task_executor_refactor.chunk_post_processor import _parser_config_compilation_template_ids + + for template_id in _parser_config_compilation_template_ids(parser_config, tenant_id): + if template_id not in template_cache: + template_cache[template_id] = _template_is_wiki(CompilationTemplateService.get_saved(template_id, tenant_id)) + if template_cache[template_id]: + return True + return False + + +def _pipeline_has_wiki_compiler( + pipeline_id: str, + tenant_id: str, + pipeline_cache: dict[str, bool], + group_cache: dict[str, bool], +) -> bool: + pipeline_id = (pipeline_id or "").strip() + if not pipeline_id: + return False + if pipeline_id in pipeline_cache: + return pipeline_cache[pipeline_id] + + from api.db.services.canvas_service import UserCanvasService + + ok, canvas = UserCanvasService.get_by_id(pipeline_id) + if not ok or not canvas: + pipeline_cache[pipeline_id] = False + return False + + group_ids = _extract_pipeline_compiler_group_ids(getattr(canvas, "dsl", None)) + has_wiki = any(_group_has_wiki_template(group_id, tenant_id, group_cache) for group_id in group_ids) + pipeline_cache[pipeline_id] = has_wiki + return has_wiki + + def _skill_index_or_none(tenant_id: str, kb_id: str): return _compiled_index_or_none(tenant_id, kb_id) @@ -1543,6 +1698,253 @@ async def has_any_wiki(dataset_id: str, tenant_id: str): return True, {"has": bool(total)} +# Dataset-scope structure kinds the artifacts_structure API serves. Keys are the +# friendly names the frontend passes; values are the template's *top-level* kind +# as stamped on ``dataset_graph`` rows. The canonical names map to themselves so +# a caller can pass either form. Note this deliberately does NOT reuse +# ``_compilation_template_kind`` — that helper folds ``knowledge_graph`` into +# ``timeline`` and would merge distinct kinds here. +_DATASET_STRUCTURE_ROW_KWD = "dataset_graph" +_DATASET_STRUCTURE_KIND_ALIASES = { + "graph": "knowledge_graph", + "knowledge_graph": "knowledge_graph", + "mindmap": "mind_map", + "mind_map": "mind_map", + "timeline": "timeline", + "session_essence": "session_essence", + "session_graph": "session_graph", +} + + +def _resolve_dataset_structure_kind(kind) -> str | None: + """Map a friendly/canonical kind string to the stored top-level kind.""" + if not isinstance(kind, str): + return None + return _DATASET_STRUCTURE_KIND_ALIASES.get(kind.strip().lower().replace("-", "_")) + + +async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str): + """Load the dataset-scope (KB-wide) structure graph for one ``kind``. + + ``kind`` is one of ``graph`` / ``mindmap`` / ``timeline`` / + ``session_essence`` / ``session_graph``. Reads the + ``knowledge_graph_kwd="dataset_graph"`` rows written by + ``rebuild_dataset_structure_graph_json`` (one per template), filters to the + requested kind, and returns them grouped by template — mirroring the + per-document ``structure/graph`` response so the frontend graph view is + reused unchanged. + + Returns ``(True, {"kind": , "templates": [...]})`` or + ``(False, message)`` on auth/validation failure. + """ + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + + resolved_kind = _resolve_dataset_structure_kind(kind) + if not resolved_kind: + return False, f"Unsupported structure kind: {kind!r}. Expected one of: graph, mindmap, timeline, session_essence, session_graph." + + _, kb = KnowledgebaseService.get_by_id(dataset_id) + empty = {"kind": kind, "templates": []} + + pack = _compiled_index_or_none(kb.tenant_id, dataset_id) + if pack is None: + return True, empty + index_nm, _ = pack + + from common.doc_store.doc_store_base import OrderByExpr + from api.db.services.compilation_template_service import CompilationTemplateService + + select_fields = [ + "content_with_weight", + "compile_kwd", + "compilation_template_ids", + "compilation_template_kind_kwd", + ] + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"knowledge_graph_kwd": [_DATASET_STRUCTURE_ROW_KWD]}, + [], + OrderByExpr(), + 0, + 1000, + index_nm, + [dataset_id], + ) + rows = settings.docStoreConn.get_fields(res, select_fields) or {} + except Exception: + logging.exception("get_dataset_structure: docStore search failed for kb=%s", dataset_id) + return True, empty + + def _row_template_id(row: dict) -> str | None: + raw = row.get("compilation_template_ids") + if isinstance(raw, list): + for v in raw: + if isinstance(v, str) and v.strip(): + return v.strip() + if isinstance(raw, str) and raw.strip(): + return raw.strip() + return None + + # Resolve a template's top-level kind + display name, memoized. Used both to + # label buckets and as a fallback for rows written before the kind stamp + # (they carry ``compilation_template_ids`` but no ``compilation_template_kind_kwd``). + template_kind_cache: dict[str, str | None] = {} + template_name_cache: dict[str, str] = {} + + def _template_meta(tid: str | None) -> str | None: + if not tid: + return None + if tid in template_kind_cache: + return template_kind_cache[tid] + top_kind = None + try: + saved = CompilationTemplateService.get_saved(tid, tenant_id) + if saved: + top_kind = (saved.get("kind") or "").strip() or None + template_name_cache[tid] = saved.get("name") or tid + except Exception: + logging.exception("get_dataset_structure: template lookup failed for %s", tid) + template_kind_cache[tid] = top_kind + return top_kind + + grouped: dict[str, dict] = {} + for row in rows.values(): + tid = _row_template_id(row) + stamped_kind = (row.get("compilation_template_kind_kwd") or "").strip() + row_kind = stamped_kind or _template_meta(tid) or "" + if _resolve_dataset_structure_kind(row_kind) != resolved_kind: + continue + + try: + graph = json.loads(row.get("content_with_weight") or "{}") + except Exception: + continue + if not isinstance(graph, dict): + continue + entities = graph.get("entities") or [] + relations = graph.get("relations") or [] + if not entities and not relations: + continue + + bucket_id = tid or f"kind:{resolved_kind}" + if bucket_id not in grouped: + if tid and tid not in template_name_cache: + _template_meta(tid) + grouped[bucket_id] = { + "template_id": bucket_id, + "template_name": template_name_cache.get(tid or "", bucket_id), + "kind": row_kind or resolved_kind, + "entities": [], + "relations": [], + } + grouped[bucket_id]["entities"].extend(entities) + grouped[bucket_id]["relations"].extend(relations) + + templates_out = [g for g in grouped.values() if g["entities"] or g["relations"]] + return True, {"kind": kind, "templates": templates_out} + + +async def get_wiki_alteration(dataset_id: str, tenant_id: str): + """Return doc-level drift between current dataset docs and compiled wiki provenance.""" + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + ok, kb = KnowledgebaseService.get_by_id(dataset_id) + if not ok: + return False, "Invalid Dataset ID" + + docs, _ = await thread_pool_exec( + DocumentService.get_by_kb_id, + kb_id=dataset_id, + page_number=0, + items_per_page=0, + orderby="create_time", + desc=False, + keywords="", + run_status=[], + types=[], + suffix=[], + ) + current_doc_ids = {str(doc.get("id")) for doc in docs or [] if doc.get("id")} + + wiki_involved_doc_ids: set[str] = set() + pack = _wiki_index_or_none(kb.tenant_id, dataset_id) + if pack is not None: + index_nm, _ = pack + from common.doc_store.doc_store_base import OrderByExpr + + select_fields = ["id", "source_doc_ids"] + offset = 0 + page_size = 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields=select_fields, + highlight_fields=[], + condition={"compile_kwd": [WIKI_PAGE_COMPILE_KWD]}, + match_expressions=[], + order_by=OrderByExpr(), + offset=offset, + limit=page_size, + index_names=index_nm, + knowledgebase_ids=[dataset_id], + ) + rows = settings.docStoreConn.get_fields(res, select_fields) or {} + except Exception: + logging.exception("get_wiki_alteration: docStore search failed for kb=%s", dataset_id) + rows = {} + + if not rows: + break + for row in rows.values(): + source_doc_ids = row.get("source_doc_ids") + if isinstance(source_doc_ids, str): + source_doc_ids = [source_doc_ids] + if not isinstance(source_doc_ids, list): + continue + wiki_involved_doc_ids.update(str(doc_id) for doc_id in source_doc_ids if doc_id) + + offset += page_size + total = settings.docStoreConn.get_total(res) + if not total or offset >= int(total): + break + + template_cache: dict[str, bool] = {} + group_cache: dict[str, bool] = {} + pipeline_cache: dict[str, bool] = {} + eligible_wiki_doc_ids: set[str] = set() + for doc in docs or []: + doc_id = str(doc.get("id") or "") + if not doc_id: + continue + parser_config = doc.get("parser_config") or {} + if _parser_config_has_wiki_template(parser_config, kb.tenant_id, template_cache): + eligible_wiki_doc_ids.add(doc_id) + continue + if _pipeline_has_wiki_compiler( + doc.get("pipeline_id") or "", + kb.tenant_id, + pipeline_cache, + group_cache, + ): + eligible_wiki_doc_ids.add(doc_id) + + removed_doc_ids = sorted(wiki_involved_doc_ids - current_doc_ids) + newly_uploaded_doc_ids = sorted(eligible_wiki_doc_ids - wiki_involved_doc_ids) + return True, { + "removed": len(removed_doc_ids), + "newly_uploaded": len(newly_uploaded_doc_ids), + "removed_doc_ids": removed_doc_ids, + "newly_uploaded_doc_ids": newly_uploaded_doc_ids, + "involved_doc_ids": sorted(wiki_involved_doc_ids), + "eligible_doc_ids": sorted(eligible_wiki_doc_ids), + } + + async def list_wiki_pages( dataset_id: str, tenant_id: str, @@ -1892,6 +2294,70 @@ async def get_skill_tree(dataset_id: str, tenant_id: str): } +async def delete_skills(dataset_id: str, tenant_id: str): + """Delete every compiled skill row (``skill`` + ``skill_all``) for a dataset. + + Returns ``(True, {"deleted": })`` on success. When the tenant index does + not exist yet there is nothing to delete, so it succeeds with ``0``. + """ + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + _, kb = KnowledgebaseService.get_by_id(dataset_id) + + pack = _skill_index_or_none(kb.tenant_id, dataset_id) + if pack is None: + return True, {"deleted": 0} + index_nm, _ = pack + + try: + deleted = settings.docStoreConn.delete( + {"compile_kwd": [_SKILL_COMPILE_KWD, _SKILL_ALL_COMPILE_KWD]}, + index_nm, + dataset_id, + ) + except Exception: + logging.exception("delete_skills: docStore delete failed for kb=%s", dataset_id) + return False, "Failed to delete skills." + + # Clear the skill compilation markers so the dataset reflects "no skill" + # (and a later re-compile isn't short-circuited by a stale task id). + try: + KnowledgebaseService.update_by_id(kb.id, {"skill_task_id": "", "skill_task_finish_at": None}) + except Exception: + logging.exception("delete_skills: failed clearing skill task markers for kb=%s", dataset_id) + + return True, {"deleted": int(deleted or 0)} + + +async def delete_skill(dataset_id: str, tenant_id: str, skill_kwd: str): + """Delete a single compiled skill node identified by ``skill_kwd``. + + Removes the per-node ``skill`` row(s) matching ``skill_kwd`` (the same + identity ``get_skill_page`` reads). The aggregate ``skill_all`` tree row is + left untouched. Returns ``(True, {"deleted": })``. + """ + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + _, kb = KnowledgebaseService.get_by_id(dataset_id) + + pack = _skill_index_or_none(kb.tenant_id, dataset_id) + if pack is None: + return True, {"deleted": 0} + index_nm, _ = pack + + try: + deleted = settings.docStoreConn.delete( + {"compile_kwd": [_SKILL_COMPILE_KWD], "skill_kwd": [skill_kwd]}, + index_nm, + dataset_id, + ) + except Exception: + logging.exception("delete_skill: docStore delete failed for kb=%s skill=%s", dataset_id, skill_kwd) + return False, "Failed to delete skill." + + return True, {"deleted": int(deleted or 0)} + + async def get_skill_page(dataset_id: str, tenant_id: str, skill_kwd: str): """Fetch the full markdown body for a single skill node.""" if not KnowledgebaseService.accessible(dataset_id, tenant_id): @@ -1957,6 +2423,211 @@ async def get_skill_page(dataset_id: str, tenant_id: str, skill_kwd: str): } +# --------------------------------------------------------------------------- +# Dataset navigation tree (written by rag/advanced_rag/knowlege_compile/ +# dataset_nav.py as nav_cluster / nav_doc rows). Loaded hierarchically: the +# first call returns the top clusters (parent = "root"); clicking a cluster +# returns its direct children (sub-clusters + document leaves). These rows are +# ``available_int=0`` so they're read via docStoreConn.search directly. +# --------------------------------------------------------------------------- + +_NAV_COMPILE_KWD = "dataset_nav" +_NAV_ROOT_PARENT = "root" +_NAV_FIELDS = [ + "id", + "name", + "type_kwd", + "content_with_weight", + "doc_count_int", + "doc_ids_kwd", + "doc_id", + "depth_int", + "parent_kwd", +] + + +def _nav_item(row: dict) -> dict: + """Shape one nav row into a UI node: name, description, doc count, type.""" + try: + payload = json.loads(row.get("content_with_weight") or "{}") + except Exception: + payload = {} + is_cluster = (row.get("type_kwd") or payload.get("type")) == "nav_cluster" + return { + "name": row.get("name") or "", + "description": payload.get("description") or "", + # doc_id count under this node: the cluster's tally, or 1 for a leaf. + "doc_count": int(row.get("doc_count_int") or 0) if is_cluster else 1, + "type": "cluster" if is_cluster else "doc", + "doc_id": None if is_cluster else (row.get("doc_id") or row.get("name")), + "has_children": is_cluster, + } + + +async def _nav_search(dataset_id: str, tenant_id: str, condition: dict, page: int, page_size: int): + """Run one nav-tree search and shape the hits into UI nodes.""" + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + _, kb = KnowledgebaseService.get_by_id(dataset_id) + + pack = _compiled_index_or_none(kb.tenant_id, dataset_id) + if pack is None: + return True, {"total": 0, "items": []} + index_nm, _ = pack + + from common.doc_store.doc_store_base import OrderByExpr + + page = max(1, int(page or 1)) + page_size = max(1, min(int(page_size or 1000), 2000)) + offset = (page - 1) * page_size + + order_by = OrderByExpr() + try: + # Biggest clusters first; leaves (no doc_count_int) fall to the end. + order_by.desc("doc_count_int") + except Exception: + order_by = OrderByExpr() + + try: + res = settings.docStoreConn.search( + select_fields=_NAV_FIELDS, + highlight_fields=[], + condition=condition, + match_expressions=[], + order_by=order_by, + offset=offset, + limit=page_size, + index_names=index_nm, + knowledgebase_ids=[dataset_id], + ) + field_map = settings.docStoreConn.get_fields(res, _NAV_FIELDS) + except Exception: + logging.exception("dataset_nav: docStore search failed for kb=%s", dataset_id) + return True, {"total": 0, "items": []} + + total = settings.docStoreConn.get_total(res) + items = [_nav_item(row) for row in (field_map or {}).values()] + return True, {"total": int(total or 0), "items": items} + + +async def list_nav_clusters(dataset_id: str, tenant_id: str, page: int = 1, page_size: int = 1000): + """First level of the nav tree: the clusters with no parent.""" + condition = { + "compile_kwd": [_NAV_COMPILE_KWD], + "type_kwd": ["nav_cluster"], + "parent_kwd": [_NAV_ROOT_PARENT], + } + return await _nav_search(dataset_id, tenant_id, condition, page, page_size) + + +async def list_nav_children(dataset_id: str, tenant_id: str, name: str, page: int = 1, page_size: int = 1000): + """Direct children of the node ``name`` — sub-clusters and document leaves. + + One level at a time (lazy) so the tree loads hierarchically as the user + expands each node. + """ + if not isinstance(name, str) or not name.strip(): + return True, {"total": 0, "items": []} + condition = { + "compile_kwd": [_NAV_COMPILE_KWD], + "parent_kwd": [name.strip()], + } + return await _nav_search(dataset_id, tenant_id, condition, page, page_size) + + +async def delete_nav(dataset_id: str, tenant_id: str): + """Delete the entire dataset navigation tree for a dataset. + + Returns ``(True, {"deleted": })``; succeeds with ``0`` when there is no + index yet. + """ + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + _, kb = KnowledgebaseService.get_by_id(dataset_id) + + pack = _compiled_index_or_none(kb.tenant_id, dataset_id) + if pack is None: + return True, {"deleted": 0} + index_nm, _ = pack + + try: + deleted = settings.docStoreConn.delete( + {"compile_kwd": [_NAV_COMPILE_KWD]}, + index_nm, + dataset_id, + ) + except Exception: + logging.exception("delete_nav: docStore delete failed for kb=%s", dataset_id) + return False, "Failed to delete the navigation tree." + + return True, {"deleted": int(deleted or 0)} + + +async def delete_nav_node(dataset_id: str, tenant_id: str, name: str): + """Delete one navigation node (identified by ``name``) and its whole subtree. + + Children reference their parent by ``name`` (``parent_kwd``), so removing a + cluster without its descendants would leave them orphaned in the tree view. + We therefore walk the subtree top-down and delete every node in it. + """ + if not isinstance(name, str) or not name.strip(): + return True, {"deleted": 0} + name = name.strip() + + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + _, kb = KnowledgebaseService.get_by_id(dataset_id) + + pack = _compiled_index_or_none(kb.tenant_id, dataset_id) + if pack is None: + return True, {"deleted": 0} + index_nm, _ = pack + + from common.doc_store.doc_store_base import OrderByExpr + + # Collect the node plus every descendant, level by level via parent_kwd. + names: set[str] = {name} + frontier: list[str] = [name] + for _ in range(64): # depth guard against a malformed (cyclic) tree + if not frontier: + break + try: + res = settings.docStoreConn.search( + select_fields=["name"], + highlight_fields=[], + condition={"compile_kwd": [_NAV_COMPILE_KWD], "parent_kwd": frontier}, + match_expressions=[], + order_by=OrderByExpr(), + offset=0, + limit=10000, + index_names=index_nm, + knowledgebase_ids=[dataset_id], + ) + rows = settings.docStoreConn.get_fields(res, ["name"]) or {} + except Exception: + logging.exception("delete_nav_node: subtree scan failed for kb=%s name=%s", dataset_id, name) + break + nxt: list[str] = [] + for row in rows.values(): + child = row.get("name") + if isinstance(child, str) and child and child not in names: + names.add(child) + nxt.append(child) + frontier = nxt + + try: + deleted = settings.docStoreConn.delete( + {"compile_kwd": [_NAV_COMPILE_KWD], "name": list(names)}, + index_nm, + dataset_id, + ) + except Exception: + logging.exception("delete_nav_node: docStore delete failed for kb=%s name=%s", dataset_id, name) + return False, "Failed to delete the navigation node." + + return True, {"deleted": int(deleted or 0)} + + async def update_wiki_page( dataset_id: str, tenant_id: str, diff --git a/api/db/__init__.py b/api/db/__init__.py index 2468d59473..c3e071faf0 100644 --- a/api/db/__init__.py +++ b/api/db/__init__.py @@ -85,6 +85,16 @@ PIPELINE_SPECIAL_PROGRESS_FREEZE_TASK_TYPES = { PipelineTaskType.MINDMAP.lower(), PipelineTaskType.ARTIFACT.lower(), PipelineTaskType.SKILL.lower(), + # Structure-graph merge fan-out task types. These are the raw task_type + # strings (== the index type), which — unlike the types above — do not equal + # their PipelineTaskType value lowercased (e.g. "structure_graph" vs + # "structuregraph"), so they are listed literally. + "structure_graph", + "structure_mindmap", + "timeline", + "session_graph", + "session_essence", + "structure", } diff --git a/api/db/db_models.py b/api/db/db_models.py index b79dddc4cd..bbe693a6fb 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -866,6 +866,21 @@ class Knowledgebase(DataBaseModel): artifact_task_finish_at = DateTimeField(null=True) skill_task_id = CharField(max_length=32, null=True, help_text="Skill generation task ID", index=True) skill_task_finish_at = DateTimeField(null=True) + # KB-wide structure-graph merge tasks, one traceable task id per merged kind + # (rebuild_dataset_structure_graph_json). ``structure_task_id`` is the + # merge-all variant that rebuilds every dataset-merge kind at once. + structure_graph_task_id = CharField(max_length=32, null=True, help_text="Structure graph merge task ID", index=True) + structure_graph_task_finish_at = DateTimeField(null=True) + structure_mindmap_task_id = CharField(max_length=32, null=True, help_text="Structure mindmap merge task ID", index=True) + structure_mindmap_task_finish_at = DateTimeField(null=True) + timeline_task_id = CharField(max_length=32, null=True, help_text="Timeline merge task ID", index=True) + timeline_task_finish_at = DateTimeField(null=True) + session_graph_task_id = CharField(max_length=32, null=True, help_text="Session graph merge task ID", index=True) + session_graph_task_finish_at = DateTimeField(null=True) + session_essence_task_id = CharField(max_length=32, null=True, help_text="Session essence merge task ID", index=True) + session_essence_task_finish_at = DateTimeField(null=True) + structure_task_id = CharField(max_length=32, null=True, help_text="Structure merge-all task ID", index=True) + structure_task_finish_at = DateTimeField(null=True) status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True) @@ -1464,6 +1479,7 @@ def alter_db_drop_index(migrator, table_name, index_name): # logging.critical(f"Failed to rename {settings.DATABASE_TYPE.upper()}.{table_name} column {old_column_name} to {new_column_name}, error: {ex}") pass + def ensure_model_indexes(migrator): """Create indexes declared by the Peewee models when they are missing.""" members = inspect.getmembers(sys.modules[__name__], inspect.isclass) @@ -1756,6 +1772,9 @@ def migrate_db(): alter_db_add_column(migrator, "knowledgebase", "artifact_task_finish_at", DateTimeField(null=True)) alter_db_add_column(migrator, "knowledgebase", "skill_task_id", CharField(max_length=32, null=True, help_text="Skill generation task ID", index=True)) alter_db_add_column(migrator, "knowledgebase", "skill_task_finish_at", DateTimeField(null=True)) + for _structure_type in ("structure_graph", "structure_mindmap", "timeline", "session_graph", "session_essence", "structure"): + alter_db_add_column(migrator, "knowledgebase", f"{_structure_type}_task_id", CharField(max_length=32, null=True, help_text=f"{_structure_type} merge task ID", index=True)) + alter_db_add_column(migrator, "knowledgebase", f"{_structure_type}_task_finish_at", DateTimeField(null=True)) alter_db_column_type(migrator, "tenant_llm", "api_key", TextField(null=True, help_text="API KEY")) alter_db_add_column(migrator, "tenant_llm", "status", CharField(max_length=1, null=False, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True)) alter_db_add_column(migrator, "connector2kb", "auto_parse", CharField(max_length=1, null=False, default="1", index=False)) diff --git a/api/db/init_data/compilation_templates/empty.yaml b/api/db/init_data/compilation_templates/empty.yaml deleted file mode 100644 index 4be03bac6b..0000000000 --- a/api/db/init_data/compilation_templates/empty.yaml +++ /dev/null @@ -1,17 +0,0 @@ -kind: empty -display_name: Empty -config: - kind: empty - entity: - description: '' - fields: - - type: '' - description: '' - rule: '' - relation: - description: '' - fields: - - type: '' - description: '' - rule: '' - global_rules: '' diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index 187ade0a7a..bb4c9acf1c 100644 --- a/api/db/services/document_service.py +++ b/api/db/services/document_service.py @@ -498,6 +498,15 @@ class DocumentService(CommonService): except Exception as e: logging.error(f"Failed to delete chunks from doc store for document {doc.id}: {e}") + # Ref-counted cleanup of wiki/artifact products this doc fed into + # (non-critical, log and continue). A product shared by other docs + # survives; one this doc solely owned is removed. + try: + if chunk_index_exists: + cls.remove_artifact_products(doc, tenant_id) + except Exception as e: + logging.warning(f"Failed to clean up artifact products for document {doc.id}: {e}") + # Prune this doc's line from the KB's tree-kind navigation # markdown (best-effort — the markdown is a downstream artifact, # and failure here must not block the document delete). @@ -558,6 +567,103 @@ class DocumentService(CommonService): settings.STORAGE_IMPL.rm(doc.kb_id, cid) page += 1 + @classmethod + def remove_artifact_products(cls, doc, tenant_id): + """Reference-counted cleanup of KB-scoped wiki/artifact products + in the doc store when a document is deleted. + + Every derived artifact row (pages, entities, relations, drafts, + topics, reduce/plan aggregates) carries a ``source_doc_ids`` list + of the documents that contributed to it. On delete we detach + ``doc.id`` from that list and drop the row only when this document + was its sole contributor — a product shared by other docs + survives. ``artifact_map_extract`` resume rows are 1:1 with a + document and are removed directly by ``doc_id``. + + The compile_kwd set is pulled from the wiki generator so new + artifact row types are covered automatically (single source of + truth). Deletion is not a hot path, so the module import cost is + acceptable here. + """ + from rag.svr.task_executor_refactor.dataset_wiki_generator import ( + WIKI_MAP_COMPILE_KWD, + WIKI_DERIVED_COMPILE_KWDS, + ) + + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, doc.kb_id): + return + + # 1. Per-doc MAP resume rows are keyed by the real doc_id. + settings.docStoreConn.delete( + {"compile_kwd": [WIKI_MAP_COMPILE_KWD], "doc_id": doc.id}, + index, + doc.kb_id, + ) + + # 2. Derived KB-scoped rows: reference-counted via source_doc_ids. + # Read every row this doc contributed to, partitioning into rows it + # solely owned (delete by id) vs. rows shared with other docs + # (detach this doc). Reading first — rather than a blanket + # ``must_not exists`` sweep — avoids deleting rows that legitimately + # carry no source_doc_ids. + derived_kwds = list(WIKI_DERIVED_COMPILE_KWDS) + select_fields = ["id", "source_doc_ids"] + sole_owner_ids: list[str] = [] + shared_seen = False + offset = 0 + page_size = 1000 + while True: + res = settings.docStoreConn.search( + select_fields, + [], + {"compile_kwd": derived_kwds, "source_doc_ids": [doc.id]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [doc.kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + if not field_map: + break + for row_id, row in field_map.items(): + raw = row.get("source_doc_ids") + if isinstance(raw, str): + owners = [raw] if raw else [] + elif isinstance(raw, list): + owners = [d for d in raw if isinstance(d, str) and d] + else: + owners = [] + if any(d != doc.id for d in owners): + shared_seen = True + else: + sole_owner_ids.append(row_id) + if len(field_map) < page_size: + break + offset += page_size + + # Drop rows this document solely owned (delete by id in batches). + for i in range(0, len(sole_owner_ids), page_size): + settings.docStoreConn.delete( + {"id": sole_owner_ids[i : i + page_size]}, + index, + doc.kb_id, + ) + + # Detach this document from rows still owned by others. The filter + # guarantees source_doc_ids contains doc.id, so the store's + # list-remove is safe; any sole-owner rows already deleted above are + # simply not matched. + if shared_seen: + settings.docStoreConn.update( + {"compile_kwd": derived_kwds, "source_doc_ids": doc.id}, + {"remove": {"source_doc_ids": doc.id}}, + index, + doc.kb_id, + ) + @classmethod @DB.connection_context() def get_newly_uploaded(cls): @@ -1101,7 +1207,20 @@ def queue_raptor_o_graphrag_tasks(sample_doc, ty, priority, fake_doc_id="", doc_ """ if doc_ids is None: doc_ids = [] - assert ty in ["graphrag", "raptor", "mindmap", "artifact", "skill"], "type should be graphrag, raptor, mindmap, artifact or skill" + assert ty in [ + "graphrag", + "raptor", + "mindmap", + "artifact", + "skill", + # KB-wide structure-graph merge task types (rebuild dataset_graph rows). + "structure_graph", + "structure_mindmap", + "timeline", + "session_graph", + "session_essence", + "structure", + ], f"unsupported task type '{ty}'" chunking_config = DocumentService.get_chunking_config(sample_doc["id"]) hasher = xxhash.xxh64() diff --git a/api/db/services/knowledgebase_service.py b/api/db/services/knowledgebase_service.py index c05d46466e..5e21a02ade 100644 --- a/api/db/services/knowledgebase_service.py +++ b/api/db/services/knowledgebase_service.py @@ -310,6 +310,18 @@ class KnowledgebaseService(CommonService): cls.model.artifact_task_finish_at, cls.model.skill_task_id, cls.model.skill_task_finish_at, + cls.model.structure_graph_task_id, + cls.model.structure_graph_task_finish_at, + cls.model.structure_mindmap_task_id, + cls.model.structure_mindmap_task_finish_at, + cls.model.timeline_task_id, + cls.model.timeline_task_finish_at, + cls.model.session_graph_task_id, + cls.model.session_graph_task_finish_at, + cls.model.session_essence_task_id, + cls.model.session_essence_task_finish_at, + cls.model.structure_task_id, + cls.model.structure_task_finish_at, cls.model.create_time, cls.model.update_time, ] diff --git a/api/db/services/pipeline_operation_log_service.py b/api/db/services/pipeline_operation_log_service.py index 12fbef5201..850f88b613 100644 --- a/api/db/services/pipeline_operation_log_service.py +++ b/api/db/services/pipeline_operation_log_service.py @@ -32,6 +32,25 @@ from common.misc_utils import get_uuid from common.time_utils import current_timestamp, datetime_format +# KB-level fan-out pipeline task types (task row carries a fake doc_id; the real +# participants live in task["doc_ids"]) → the KB ``_task_finish_at`` column +# stamped when the task completes. Membership also marks a task as KB-scoped so +# the per-document progress update is skipped. +_PIPELINE_TASK_TYPE_TO_FINISH_FIELD = { + PipelineTaskType.GRAPH_RAG: "graphrag_task_finish_at", + PipelineTaskType.RAPTOR: "raptor_task_finish_at", + PipelineTaskType.MINDMAP: "mindmap_task_finish_at", + PipelineTaskType.ARTIFACT: "artifact_task_finish_at", + PipelineTaskType.SKILL: "skill_task_finish_at", + PipelineTaskType.STRUCTURE_GRAPH: "structure_graph_task_finish_at", + PipelineTaskType.STRUCTURE_MINDMAP: "structure_mindmap_task_finish_at", + PipelineTaskType.TIMELINE: "timeline_task_finish_at", + PipelineTaskType.SESSION_GRAPH: "session_graph_task_finish_at", + PipelineTaskType.SESSION_ESSENCE: "session_essence_task_finish_at", + PipelineTaskType.STRUCTURE: "structure_task_finish_at", +} + + class PipelineOperationLogService(CommonService): model = PipelineOperationLog @@ -99,7 +118,7 @@ class PipelineOperationLogService(CommonService): referred_document_id = document_id # no need to update document for KB-level fan-out tasks - if task_type not in [PipelineTaskType.GRAPH_RAG, PipelineTaskType.RAPTOR, PipelineTaskType.MINDMAP, PipelineTaskType.ARTIFACT, PipelineTaskType.SKILL]: + if task_type not in _PIPELINE_TASK_TYPE_TO_FINISH_FIELD: ok, document = DocumentService.get_by_id(referred_document_id) if not ok: logging.warning(f"Document for referred_document_id {referred_document_id} not found") @@ -137,7 +156,7 @@ class PipelineOperationLogService(CommonService): if task_type not in VALID_PIPELINE_TASK_TYPES: raise ValueError(f"Invalid task type: {task_type}") - if task_type in [PipelineTaskType.GRAPH_RAG, PipelineTaskType.RAPTOR, PipelineTaskType.MINDMAP, PipelineTaskType.ARTIFACT, PipelineTaskType.SKILL]: + if task_type in _PIPELINE_TASK_TYPE_TO_FINISH_FIELD: # query task to get progress information from task ok, task = TaskService.get_by_id(task_id) if not ok: @@ -151,31 +170,10 @@ class PipelineOperationLogService(CommonService): process_duration = task.process_duration finish_at = process_begin_at + timedelta(seconds=process_duration) - if task_type == PipelineTaskType.GRAPH_RAG: - KnowledgebaseService.update_by_id( - document.kb_id, - {"graphrag_task_finish_at": finish_at}, - ) - elif task_type == PipelineTaskType.RAPTOR: - KnowledgebaseService.update_by_id( - document.kb_id, - {"raptor_task_finish_at": finish_at}, - ) - elif task_type == PipelineTaskType.MINDMAP: - KnowledgebaseService.update_by_id( - document.kb_id, - {"mindmap_task_finish_at": finish_at}, - ) - elif task_type == PipelineTaskType.ARTIFACT: - KnowledgebaseService.update_by_id( - document.kb_id, - {"artifact_task_finish_at": finish_at}, - ) - elif task_type == PipelineTaskType.SKILL: - KnowledgebaseService.update_by_id( - document.kb_id, - {"skill_task_finish_at": finish_at}, - ) + KnowledgebaseService.update_by_id( + document.kb_id, + {_PIPELINE_TASK_TYPE_TO_FINISH_FIELD[task_type]: finish_at}, + ) log = dict( id=get_uuid(), diff --git a/common/constants.py b/common/constants.py index 9ccef14186..42d9c9f8a9 100644 --- a/common/constants.py +++ b/common/constants.py @@ -187,6 +187,13 @@ class PipelineTaskType(StrEnum): MEMORY = "Memory" ARTIFACT = "Artifact" SKILL = "Skill" + # KB-wide structure-graph merge tasks (rebuild_dataset_structure_graph_json). + STRUCTURE_GRAPH = "StructureGraph" + STRUCTURE_MINDMAP = "StructureMindmap" + TIMELINE = "Timeline" + SESSION_GRAPH = "SessionGraph" + SESSION_ESSENCE = "SessionEssence" + STRUCTURE = "Structure" VALID_PIPELINE_TASK_TYPES = { @@ -197,6 +204,12 @@ VALID_PIPELINE_TASK_TYPES = { PipelineTaskType.MINDMAP, PipelineTaskType.ARTIFACT, PipelineTaskType.SKILL, + PipelineTaskType.STRUCTURE_GRAPH, + PipelineTaskType.STRUCTURE_MINDMAP, + PipelineTaskType.TIMELINE, + PipelineTaskType.SESSION_GRAPH, + PipelineTaskType.SESSION_ESSENCE, + PipelineTaskType.STRUCTURE, } diff --git a/rag/advanced_rag/knowlege_compile/_common.py b/rag/advanced_rag/knowlege_compile/_common.py index b180d5779e..1fceea9408 100644 --- a/rag/advanced_rag/knowlege_compile/_common.py +++ b/rag/advanced_rag/knowlege_compile/_common.py @@ -214,7 +214,7 @@ def ensure_llm_bundle(mdl, method: str, *, label: str = "model"): # --------------------------------------------------------------------------- -async def es_search( +async def doc_storage_search( select_fields: list[str], condition: dict, *, @@ -223,7 +223,7 @@ async def es_search( match_expressions: list | None = None, offset: int = 0, limit: int = 1000, - label: str = "es_search", + label: str = "doc_storage_search", ) -> dict: """Thin wrapper around ``docStoreConn.search`` + ``get_fields``. @@ -255,12 +255,12 @@ async def es_search( return {} -async def es_insert( +async def doc_storage_insert( rows: list[dict], tenant_id: str, kb_id: str, *, - label: str = "es_insert", + label: str = "doc_storage_insert", ) -> None: """Bulk insert wrapped in ``thread_pool_exec``. Logs on failure.""" if not rows: @@ -275,12 +275,12 @@ async def es_insert( logging.exception("%s failed (%d row(s))", label, len(rows)) -async def es_delete( +async def doc_storage_delete( condition: dict, tenant_id: str, kb_id: str, *, - label: str = "es_delete", + label: str = "doc_storage_delete", ) -> None: """Bulk delete wrapped in ``thread_pool_exec``. Best-effort; logs on failure (some callers rely on id-based upsert as a fallback).""" @@ -294,13 +294,13 @@ async def es_delete( logging.debug("%s failed (condition=%r); caller may rely on id-upsert", label, condition) -async def es_upsert_one( +async def doc_storage_upsert_one( filter_condition: dict, row: dict, tenant_id: str, kb_id: str, *, - label: str = "es_upsert_one", + label: str = "doc_storage_upsert_one", ) -> None: """Delete-by-filter then insert. Used when an in-place update would require knowing the existing row's id and we'd rather drop+re-create. @@ -310,8 +310,8 @@ async def es_upsert_one( (:func:`stable_row_id`) so id-based dedup at the connector catches any race that bypasses the delete. """ - await es_delete(filter_condition, tenant_id, kb_id, label=f"{label}.delete") - await es_insert([row], tenant_id, kb_id, label=f"{label}.insert") + await doc_storage_delete(filter_condition, tenant_id, kb_id, label=f"{label}.delete") + await doc_storage_insert([row], tenant_id, kb_id, label=f"{label}.insert") # --------------------------------------------------------------------------- @@ -977,10 +977,10 @@ __all__ = [ "union_ordered", "make_input_budget", "ensure_llm_bundle", - "es_search", - "es_insert", - "es_delete", - "es_upsert_one", + "doc_storage_search", + "doc_storage_insert", + "doc_storage_delete", + "doc_storage_upsert_one", "find_vec_field", # New engines "normalize_key", diff --git a/rag/advanced_rag/knowlege_compile/dataset_nav.py b/rag/advanced_rag/knowlege_compile/dataset_nav.py index b2dade9e5d..d618c8ad27 100644 --- a/rag/advanced_rag/knowlege_compile/dataset_nav.py +++ b/rag/advanced_rag/knowlege_compile/dataset_nav.py @@ -219,7 +219,7 @@ async def _store_knn( scoring = [] for r in rows: stored = r.get(vf) - if stored and len(stored) == len(vec): + if _vector_len(stored) > 0 and _vector_len(stored) == _vector_len(vec): sim = sum(a * b for a, b in zip(stored, vec)) scoring.append((sim, r)) scoring.sort(key=lambda x: -x[0]) @@ -281,7 +281,7 @@ async def _embed(embd_mdl, text: str) -> list[float]: """Encode a single text string and return its embedding vector.""" global _EMBED_DIM vecs = await _encode(embd_mdl, [text]) - if vecs and len(vecs[0]) > 0: + if vecs and _vector_len(vecs[0]) > 0: dim = len(vecs[0]) if _EMBED_DIM is None: _EMBED_DIM = dim @@ -289,9 +289,20 @@ async def _embed(embd_mdl, text: str) -> list[float]: return [] -def _cosine_sim(a: list[float], b: list[float]) -> float: +def _vector_len(vec) -> int: + if vec is None: + return 0 + try: + return len(vec) + except TypeError: + return 0 + + +def _cosine_sim(a, b) -> float: """Compute cosine similarity between two vectors.""" - if not a or not b or len(a) != len(b): + a_len = _vector_len(a) + b_len = _vector_len(b) + if a_len == 0 or b_len == 0 or a_len != b_len: return 0.0 dot = sum(x * y for x, y in zip(a, b)) na = sum(x * x for x in a) ** 0.5 @@ -323,7 +334,7 @@ def _make_nav_doc_row( "compile_kwd": _COMPILE_KWD, "knowledge_graph_kwd": "entity", "type_kwd": "nav_doc", - "name": doc_id, + "name": f"{parent_kwd}_{xxhash.xxh64(summary.encode()).hexdigest()[:12]}", "parent_kwd": parent_kwd, "depth_int": depth_int, "available_int": 0, @@ -333,7 +344,7 @@ def _make_nav_doc_row( ltks = _tokenize(summary) row["content_ltks"] = ltks row["content_sm_ltks"] = _fine_tokenize(ltks) - if embedding: + if _vector_len(embedding) > 0: dim = len(embedding) row[_vec_field(dim)] = embedding return row @@ -369,7 +380,7 @@ def _make_nav_cluster_row( ltks = _tokenize(description) row["content_ltks"] = ltks row["content_sm_ltks"] = _fine_tokenize(ltks) - if embedding: + if _vector_len(embedding) > 0: dim = len(embedding) row[_vec_field(dim)] = embedding return row @@ -426,7 +437,7 @@ async def _find_best_cluster( best_parent = best.get("parent_kwd", "") # compute actual similarity to root stored = best.get(_vec_field(vec_dim)) - sim = _cosine_sim(doc_embedding, stored) if stored else 0.0 + sim = _cosine_sim(doc_embedding, stored) # Step 2: recursively descend into children while sim >= _RECURSE_THRESHOLD: @@ -441,7 +452,7 @@ async def _find_best_cluster( break child = children[0] stored = child.get(_vec_field(vec_dim)) - child_sim = _cosine_sim(doc_embedding, stored) if stored else 0.0 + child_sim = _cosine_sim(doc_embedding, stored) if child_sim < _RECURSE_THRESHOLD: break best_name = child.get("name", best_name) @@ -476,23 +487,59 @@ async def _llm_merge(chat_mdl, cluster_desc: str, doc_summary: str) -> str: return cluster_desc -async def _llm_create_summary(chat_mdl, doc_summaries: list[str]) -> str: - """LLM create a cluster summary from one or more doc summaries.""" +def _clean_title(title: str) -> str: + """Normalize an LLM title into a one-line, length-capped display name.""" + return " ".join((title or "").split())[:48].strip() + + +def _fallback_title(summary: str) -> str: + """Derive a short readable title from a summary when the LLM gives none.""" + words = " ".join((summary or "").split()).split(" ") + return " ".join(words[:6]).strip() or "Cluster" + + +def _readable_cluster_name(title: str, seed: str) -> str: + """A readable yet unique nav-cluster key: ``" <8-hex>"``. + + The title makes the node name human-readable; the short hash of ``seed`` + (the cluster description) preserves the per-KB uniqueness the tree keying + relies on (``_nav_cluster_id`` / ``parent_kwd``). + """ + suffix = xxhash.xxh64((seed or "").encode("utf-8")).hexdigest()[:8] + return f"{_clean_title(title) or 'Cluster'} {suffix}" + + +async def _llm_create_summary(chat_mdl, doc_summaries: list[str]) -> tuple[str, str]: + """LLM-derive a cluster's readable ``(name, summary)`` from doc summaries. + + ``name`` is a short human-readable topic title used to make the nav node + name readable (the caller still appends a short hash to keep the tree key + unique); ``summary`` is the 1-3 sentence description. + """ + fallback_summary = doc_summaries[0] if doc_summaries else "" if not chat_mdl: - return doc_summaries[0] if doc_summaries else "" + return _fallback_title(fallback_summary), fallback_summary + from rag.prompts.generator import gen_json texts = "\n---\n".join(doc_summaries) - prompt = f"Summarize the common topic of the following document excerpts in 1-3 concise sentences:\n\n{texts}\n\nReturn ONLY the summary text, no commentary." + prompt = ( + "Given the document excerpts below, produce a short human-readable topic " + "name and a concise description of their common topic.\n\n" + f"{texts}\n\n" + 'Return ONLY JSON: {"name": "<2-6 word topic title>", "summary": "<1-3 sentence description>"}' + ) try: resp = await gen_json("", prompt, chat_mdl, gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1})) if isinstance(resp, dict): - return str(resp.get("summary", resp.get("result", doc_summaries[0]))) + summary = str(resp.get("summary") or resp.get("result") or fallback_summary).strip() + name = _clean_title(str(resp.get("name") or "")) or _fallback_title(summary) + return name, (summary or fallback_summary) if isinstance(resp, str) and resp.strip(): - return resp.strip() + return _fallback_title(resp), resp.strip() except Exception: logging.exception("dataset_nav: LLM summary failed") - return doc_summaries[0] if doc_summaries else "" + return _fallback_title(fallback_summary), fallback_summary # --------------------------------------------------------------------------- @@ -583,7 +630,7 @@ async def upsert_dataset_nav_doc( # Re-compute embedding for the new summary if embd_mdl and new_desc != old_desc: new_emb = await _embed(embd_mdl, new_desc) - if new_emb: + if _vector_len(new_emb) > 0: cluster_row[_vec_field(len(new_emb))] = new_emb await _store_upsert(tenant_id, kb_id, cluster_row) @@ -621,8 +668,8 @@ async def upsert_dataset_nav_doc( if parent_row: depth_of_parent = parent_row.get("depth_int", 1) new_depth = depth_of_parent + 1 - new_name = f"navc_{xxhash.xxh64(summary.encode()).hexdigest()[:12]}" - new_desc = await _llm_create_summary(chat_mdl, [summary]) + new_title, new_desc = await _llm_create_summary(chat_mdl, [summary]) + new_name = _readable_cluster_name(new_title, summary) new_cluster = _make_nav_cluster_row( kb_id, new_name, @@ -632,7 +679,7 @@ async def upsert_dataset_nav_doc( [doc_id], doc_embedding, ) - if embd_mdl and doc_embedding: + if embd_mdl and _vector_len(doc_embedding) > 0: new_cluster[_vec_field(len(doc_embedding))] = doc_embedding await _store_upsert(tenant_id, kb_id, new_cluster) @@ -648,8 +695,8 @@ async def upsert_dataset_nav_doc( await _store_upsert(tenant_id, kb_id, nav_doc_row) else: # ── Create root-level new cluster ── - new_name = f"navc_{xxhash.xxh64(summary.encode()).hexdigest()[:12]}" - new_desc = await _llm_create_summary(chat_mdl, [summary]) + root_title, new_desc = await _llm_create_summary(chat_mdl, [summary]) + new_name = _readable_cluster_name(root_title, summary) new_cluster = _make_nav_cluster_row( kb_id, new_name, @@ -659,7 +706,7 @@ async def upsert_dataset_nav_doc( [doc_id], doc_embedding, ) - if embd_mdl and doc_embedding: + if embd_mdl and _vector_len(doc_embedding) > 0: new_cluster[_vec_field(len(doc_embedding))] = doc_embedding await _store_upsert(tenant_id, kb_id, new_cluster) @@ -905,8 +952,11 @@ async def _maybe_split_cluster( for d in dids: if d not in doc_ids: doc_ids.append(d) - group_desc = await _llm_create_summary(chat_mdl, descs) if descs else f"Group {gi + 1}" - group_name = f"navc_split_{xxhash.xxh64(group_desc.encode()).hexdigest()[:12]}" + if descs: + group_title, group_desc = await _llm_create_summary(chat_mdl, descs) + else: + group_title = group_desc = f"Group {gi + 1}" + group_name = _readable_cluster_name(group_title, group_desc) group_emb = await _embed(embd_mdl, group_desc) if embd_mdl else [] new_cluster = _make_nav_cluster_row( kb_id, @@ -964,11 +1014,11 @@ async def search_dataset_nav( except Exception: logging.exception("search_dataset_nav: embed failed for kb=%s", kb_id) vec = [] - if vec: + if _vector_len(vec) > 0: try: rows = await _store_knn(tenant_id, kb_id, vec, len(vec), condition, top_k=top_k) vf = _vec_field(len(vec)) - rows_with_scores = [(r, _cosine_sim(vec, r.get(vf) or [])) for r in rows] + rows_with_scores = [(r, _cosine_sim(vec, r.get(vf))) for r in rows] except Exception: logging.exception("search_dataset_nav: knn failed for kb=%s", kb_id) rows_with_scores = [] diff --git a/rag/advanced_rag/knowlege_compile/raptor.py b/rag/advanced_rag/knowlege_compile/raptor.py index 92618dfa49..d68b265af3 100644 --- a/rag/advanced_rag/knowlege_compile/raptor.py +++ b/rag/advanced_rag/knowlege_compile/raptor.py @@ -778,7 +778,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: to materialize. """ if len(chunks) <= 1: - return None if is_tree else ([], []) + return (None, None) if is_tree else ([], []) # Normalize input to the 3-tuple shape. Reject empties / bad # vectors at the same time the legacy path used to. @@ -800,7 +800,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval: normalized = [t for t in (_normalize(c) for c in chunks) if t is not None] if len(normalized) <= 1: - return None if is_tree else (normalized, [(0, len(normalized))]) + return (None, None) if is_tree else (normalized, [(0, len(normalized))]) chunks = normalized if self._tree_builder == PSI_TREE_BUILDER: diff --git a/rag/advanced_rag/knowlege_compile/runner.py b/rag/advanced_rag/knowlege_compile/runner.py index 65f16addfb..58c04bbc78 100644 --- a/rag/advanced_rag/knowlege_compile/runner.py +++ b/rag/advanced_rag/knowlege_compile/runner.py @@ -44,9 +44,13 @@ from api.db.services.llm_service import LLMBundle from common.exceptions import TaskCanceledException from rag.advanced_rag.knowlege_compile.structure import ( LLMCallPool, + MERGE_SCOPE_DATASET, + MERGE_SCOPE_DOC, compile_structure_from_text, cleanup_timeline_isolated_entities, merge_compiled_structures, + rebuild_dataset_structure_graph_json, + rebuild_structure_graph_json, ) @@ -146,6 +150,99 @@ def split_tree_templates( return tree_templates, non_tree_templates +def _is_page_index_template(parser_cfg: dict) -> bool: + kind = (parser_cfg or {}).get("kind") + if not isinstance(kind, str): + return False + return kind.strip().lower().replace("-", "_") in {"page_index", "pageindex"} + + +def _page_index_graph_summary(graph: dict, limit: int = 80) -> str: + entities = graph.get("entities") if isinstance(graph, dict) else None + if not isinstance(entities, list): + return "" + + lines: list[str] = [] + for entity in entities: + if not isinstance(entity, dict): + continue + name = str(entity.get("name") or "").strip() + description = str(entity.get("discription") or entity.get("description") or "").strip() + text = f"{name}: {description}".strip(": ").strip() + if text: + lines.append(text) + if len(lines) >= limit: + break + return "\n".join(lines) + + +async def _upsert_dataset_nav_from_page_index( + *, + active_templates: list[tuple[str, dict]], + chat_mdl_by_tid: dict[str, LLMBundle], + embedding_model: LLMBundle, + tenant_id: str, + kb_id: str, + doc_id: str, + progress_cb: Callable[..., None], + cancel_check: Callable[[], bool], +) -> None: + page_index_templates = [(template_id, parser_cfg) for template_id, parser_cfg in active_templates if _is_page_index_template(parser_cfg)] + if not page_index_templates: + return + + summaries: list[str] = [] + chat_mdl = None + for template_id, _ in page_index_templates: + if cancel_check(): + raise TaskCanceledException("Task was cancelled before dataset navigation update") + try: + graph = await rebuild_structure_graph_json( + tenant_id, + kb_id, + doc_id, + "timeline", + compilation_template_id=template_id, + ) + except Exception: + logging.exception( + "page_index: failed to rebuild graph summary for dataset_nav doc %s template %s", + doc_id, + template_id, + ) + continue + + summary = _page_index_graph_summary(graph) + if summary: + summaries.append(summary) + chat_mdl = chat_mdl or chat_mdl_by_tid.get(template_id) + + if not summaries: + logging.info("page_index: no dataset_nav summary for doc %s", doc_id) + return + + if cancel_check(): + raise TaskCanceledException("Task was cancelled before dataset navigation upsert") + try: + from rag.advanced_rag.knowlege_compile.dataset_nav import ( + upsert_dataset_nav_doc, + ) + + progress_cb(msg=f"page_index: updating dataset navigation for doc {doc_id} ...") + await upsert_dataset_nav_doc( + tenant_id, + kb_id, + doc_id, + "\n\n".join(summaries), + embd_mdl=embedding_model, + chat_mdl=chat_mdl, + ) + except TaskCanceledException: + raise + except Exception: + logging.exception("page_index: dataset_nav upsert failed for doc %s", doc_id) + + # ----- non-tree compilation core ------------------------------------- @@ -185,12 +282,20 @@ async def run_structure_compile_over_batches( accumulators: dict[str, list[dict]] = {tid: [] for tid, _ in active_templates} template_kinds: dict[str, str] = {tid: _compilation_template_kind((cfg or {}).get("kind")) for tid, cfg in active_templates} + # ``dataset_merge`` hyper-parameter (per template config): when truthy the + # merge dedups entities/relations across the whole dataset (KB), collapsing + # cross-document duplicates onto one canonical row, instead of merging only + # within the current document. + merge_scope_by_tid: dict[str, str] = {tid: (MERGE_SCOPE_DATASET if bool((cfg or {}).get("dataset_merge")) else MERGE_SCOPE_DOC) for tid, cfg in active_templates} + # compile_kwd(s) each template actually wrote, harvested from flush results + # so a dataset-scope template can rebuild its dataset graph once at the end. + compile_kwds_by_tid: dict[str, set[str]] = {tid: set() for tid, _ in active_templates} agg_infos: dict[str, dict] = {tid: {"inserted": 0, "updated": 0, "duplicates_dropped": 0} for tid, _ in active_templates} chunks_by_id: dict[str, str] = {} flush_sequence = 0 flush_tasks: set[asyncio.Task[None]] = set() - es_condition = asyncio.Condition() - next_es_sequence = 0 + doc_storage_condition = asyncio.Condition() + next_doc_storage_sequence = 0 async def _flush(template_id: str) -> None: nonlocal flush_sequence @@ -204,24 +309,24 @@ async def run_structure_compile_over_batches( timing_context = f"{doc_id}:{template_id}:flush-{flush_sequence}" async def _run_flush() -> None: - nonlocal next_es_sequence - es_acquired = False - es_released = False + nonlocal next_doc_storage_sequence + doc_storage_acquired = False + doc_storage_released = False - async def _wait_for_es() -> None: - nonlocal es_acquired - async with es_condition: - await es_condition.wait_for(lambda: next_es_sequence == sequence) - es_acquired = True + async def _wait_for_doc_storage() -> None: + nonlocal doc_storage_acquired + async with doc_storage_condition: + await doc_storage_condition.wait_for(lambda: next_doc_storage_sequence == sequence) + doc_storage_acquired = True - async def _release_es() -> None: - nonlocal next_es_sequence, es_released - async with es_condition: - if next_es_sequence != sequence: - raise RuntimeError(f"ES sequence mismatch: expected {next_es_sequence}, releasing {sequence}") - next_es_sequence += 1 - es_released = True - es_condition.notify_all() + async def _release_doc_storage() -> None: + nonlocal next_doc_storage_sequence, doc_storage_released + async with doc_storage_condition: + if next_doc_storage_sequence != sequence: + raise RuntimeError(f"ES sequence mismatch: expected {next_doc_storage_sequence}, releasing {sequence}") + next_doc_storage_sequence += 1 + doc_storage_released = True + doc_storage_condition.notify_all() kind = template_kinds.get(template_id, "") merge_chat_mdl = llm_pool.wrap( @@ -244,18 +349,22 @@ async def run_structure_compile_over_batches( chain_kind=kind, chain_callback=progress_cb, chain_timeout_seconds=STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S, - es_waiter=_wait_for_es, - es_releaser=_release_es, + doc_storage_waiter=_wait_for_doc_storage, + doc_storage_releaser=_release_doc_storage, + merge_scope=merge_scope_by_tid[template_id], ) finally: - if not es_released: - if not es_acquired: - await _wait_for_es() - await _release_es() + if not doc_storage_released: + if not doc_storage_acquired: + await _wait_for_doc_storage() + await _release_doc_storage() if isinstance(info, dict): agg = agg_infos[template_id] for k in ("inserted", "updated", "duplicates_dropped"): agg[k] = agg.get(k, 0) + int(info.get(k, 0) or 0) + for compile_kwd in info.get("compile_kwds") or []: + if compile_kwd: + compile_kwds_by_tid[template_id].add(str(compile_kwd)) flush_tasks.add(asyncio.create_task(_run_flush())) @@ -375,6 +484,59 @@ async def run_structure_compile_over_batches( finally: flush_tasks.clear() + # ── Dataset structure graph ────────────────────────────────────────── + # For dataset-scope templates the entity/relation rows are now merged + # across documents, so (re)project them into a single KB-wide graph. This + # runs once per document-parse completion; the last document to finish + # produces the complete dataset graph. Best-effort — a failure here must + # not fail the parse. + for template_id, _ in active_templates: + if merge_scope_by_tid[template_id] != MERGE_SCOPE_DATASET: + continue + # The row is filtered by the template's *top-level* kind (e.g. + # ``knowledge_graph``, ``session_graph``), which — unlike ``config.kind`` + # — distinguishes the knowledge_graph family. Resolve it from the + # template record; on failure leave it unstamped (the read side falls + # back to resolving kind from the template id). + structure_kind = None + try: + saved_template = CompilationTemplateService.get_saved(template_id, tenant_id) + if saved_template: + structure_kind = (saved_template.get("kind") or "").strip() or None + except Exception: + logging.exception("dataset structure graph: failed to resolve top-level kind for template %s", template_id) + for compile_kwd in sorted(compile_kwds_by_tid[template_id]): + if cancel_check(): + raise TaskCanceledException("Task was cancelled before dataset structure graph rebuild") + try: + progress_cb(msg=f"Rebuilding dataset structure graph (compile_kwd={compile_kwd}) ...") + await rebuild_dataset_structure_graph_json( + tenant_id, + kb_id, + compile_kwd, + compilation_template_id=template_id, + structure_kind=structure_kind, + ) + except TaskCanceledException: + raise + except Exception: + logging.exception( + "dataset structure graph rebuild failed for kb=%s compile_kwd=%s template=%s", + kb_id, + compile_kwd, + template_id, + ) + + await _upsert_dataset_nav_from_page_index( + active_templates=active_templates, + chat_mdl_by_tid=chat_mdl_by_tid, + embedding_model=embedding_model, + tenant_id=tenant_id, + kb_id=kb_id, + doc_id=doc_id, + progress_cb=progress_cb, + cancel_check=cancel_check, + ) # Timeline entity cleanup must happen after every flush has completed; # otherwise an entity can look isolated in one flush and be referenced by # a relation from a later flush. Keep this scoped to timeline templates. diff --git a/rag/advanced_rag/knowlege_compile/structure.py b/rag/advanced_rag/knowlege_compile/structure.py index c8171a4ea4..87ef5e189f 100644 --- a/rag/advanced_rag/knowlege_compile/structure.py +++ b/rag/advanced_rag/knowlege_compile/structure.py @@ -50,6 +50,27 @@ _ES_DEDUP_LLM_CONCURRENCY = 16 _ES_DEDUP_LLM_BATCH_SIZE = 16 _ES_DEDUP_EMBED_BATCH_SIZE = 64 _ES_DEDUP_INSERT_BATCH_SIZE = 256 +_STRUCT_INVALID_SENTINELS = {"-1"} + +# Merge scopes. ``doc`` (default) dedups an incoming entity/relation only +# against rows already stored for the *same* document; ``dataset`` widens the +# candidate lookup to the whole knowledge base so the same logical entity found +# in different documents collapses onto one canonical row. +MERGE_SCOPE_DOC = "doc" +MERGE_SCOPE_DATASET = "dataset" + +# Dataset-scope merges read-modify-write shared KB-wide rows, so concurrent +# per-document parse tasks (possibly in different worker processes) must be +# serialized to avoid inserting duplicate canonical rows before either sees the +# other. Mirrors the per-kb lock ``dataset_nav`` already uses for the same +# cross-document upsert hazard. +_STRUCT_MERGE_LOCK_TIMEOUT_S = 60 +_STRUCT_MERGE_LOCK_BLOCKING_TIMEOUT_S = 5 + + +def _struct_merge_lock_key(kb_id: str, compilation_template_id: str | None) -> str: + """Per-(kb, template) lock so different templates can merge in parallel.""" + return f"struct_merge:{kb_id}:{compilation_template_id or ''}" class LLMCallPool: @@ -310,7 +331,7 @@ def _struct_hypergraph_prompts(parser_config: dict, language: str = "en") -> Tup if rel_desc: edge_parts.append(f"## Relation Description:\n{rel_desc}") edge_parts.append(f"## Relation Fields:\n{rel_fields_text}") - edge_parts.append("## Known Entities:\n{known_nodes}") + edge_parts.append("## Known Entities:\nSee the source text below.") edge_parts.append( "## Response Format:\n" "Reply with a single JSON object of the form: " @@ -337,6 +358,10 @@ def _struct_entity_id_field(parser_config: dict) -> str: return "name" +def _struct_is_invalid_sentinel(value) -> bool: + return isinstance(value, str) and value.strip() in _STRUCT_INVALID_SENTINELS + + def _struct_unwrap_items(res) -> list: if res is None: return [] @@ -378,9 +403,6 @@ async def _struct_extract_hypergraph(text: str, parser_config: dict, chat_mdl, l return nodes, edges -# Backwards-compat alias for the shared helper. New code should use -# ``_common.encode`` directly; kept here so existing references inside this -# module keep working without a wider rename. _struct_embed = _encode @@ -413,7 +435,7 @@ def _struct_load_payload(doc: dict) -> dict: def _struct_graph_entity(payload: dict, source_chunk_ids: list | None = None) -> dict | None: name = payload.get("name") or payload.get("text") or payload.get("term") or payload.get("title") name = str(name).strip() if name is not None else "" - if not name: + if not name or _struct_is_invalid_sentinel(name): return None typ = payload.get("type") or "other" typ = str(typ).strip() if typ is not None else "other" @@ -442,7 +464,7 @@ def _struct_graph_relation(payload: dict) -> dict | None: tgt = payload.get("target") or payload.get("tgt") or payload.get("to") src = str(src).strip() if src is not None else "" tgt = str(tgt).strip() if tgt is not None else "" - if not src or not tgt: + if not src or not tgt or _struct_is_invalid_sentinel(src) or _struct_is_invalid_sentinel(tgt): return None typ = payload.get("type") or "related" return { @@ -509,7 +531,7 @@ def _struct_relation_member_fields(parser_config: dict) -> Tuple: return None, None -def _struct_to_es_doc( +def _struct_to_doc_storage_doc( payload: dict, compile_kwd: str, doc_id: str, @@ -648,7 +670,7 @@ async def _struct_process_batch( return [] docs = [ - _struct_to_es_doc( + _struct_to_doc_storage_doc( payload, autotype, doc_id, @@ -917,7 +939,7 @@ async def _struct_rewrite_relation_doc(doc: dict, aliases: dict[str, str], embd_ base["content_with_weight"] = json.dumps(payload, ensure_ascii=False) base["from_entity_kwd"] = _struct_resolve_entity_alias(base.get("from_entity_kwd", ""), aliases) base["to_entity_kwd"] = _struct_resolve_entity_alias(base.get("to_entity_kwd", ""), aliases) - return _struct_rebuild_es_doc(payload, base, vecs[0], doc.get("source_chunk_ids") or [], preserve_id=True) + return _struct_rebuild_doc_storage_doc(payload, base, vecs[0], doc.get("source_chunk_ids") or [], preserve_id=True) async def _struct_merge_pair(existing: dict, incoming: dict, chat_mdl) -> dict | None: @@ -972,14 +994,14 @@ def _struct_apply_merge_invariants(existing: dict, merged_payload: dict) -> dict return merged_payload -def _struct_rebuild_es_doc( +def _struct_rebuild_doc_storage_doc( payload: dict, base_doc: dict, vec, chunk_ids: list, preserve_id: bool = True, ) -> dict: - """Rebuild an ES doc from a merged payload using _struct_to_es_doc, then + """Rebuild an ES doc from a merged payload using _struct_to_doc_storage_doc, then overlay identity fields (id, from_entity_kwd, to_entity_kwd) from base_doc. """ kind = base_doc.get("knowledge_graph_kwd") or "entity" @@ -994,7 +1016,7 @@ def _struct_rebuild_es_doc( except Exception: pass - new_doc = _struct_to_es_doc( + new_doc = _struct_to_doc_storage_doc( payload=payload, compile_kwd=base_doc.get("compile_kwd"), doc_id=base_doc.get("doc_id"), @@ -1022,11 +1044,16 @@ async def _struct_reembed_payload(payload: dict, embd_mdl): return vecs[0] if vecs else None -def _struct_es_dedup_condition(doc: dict) -> dict: +def _struct_doc_storage_dedup_condition(doc: dict, merge_scope: str = MERGE_SCOPE_DOC) -> dict: condition = { "compile_kwd": [doc["compile_kwd"]], - "doc_id": [doc["doc_id"]], } + # Doc scope: only an entity/relation already stored for the same document is + # a merge candidate. Dataset scope: widen to the whole KB (``kb_id`` is + # already the search index scope) so the same entity across documents + # collapses onto one canonical row. + if merge_scope != MERGE_SCOPE_DATASET: + condition["doc_id"] = [doc["doc_id"]] if doc.get("knowledge_graph_kwd"): condition["knowledge_graph_kwd"] = [doc["knowledge_graph_kwd"]] if doc.get("from_entity_kwd"): @@ -1039,7 +1066,7 @@ def _struct_es_dedup_condition(doc: dict) -> dict: return condition -async def _struct_es_knn_candidate( +async def _struct_doc_storage_knn_candidate( doc: dict, tenant_id: str, kb_id: str, @@ -1048,6 +1075,7 @@ async def _struct_es_knn_candidate( select_fields: list[str], timing_context: str | None, item_index: int, + merge_scope: str = MERGE_SCOPE_DOC, ) -> dict | None: """Run one KNN lookup; the caller controls concurrency.""" from common import settings @@ -1069,7 +1097,7 @@ async def _struct_es_knn_candidate( settings.docStoreConn.search, select_fields, [], - _struct_es_dedup_condition(doc), + _struct_doc_storage_dedup_condition(doc, merge_scope), [match_expr], OrderByExpr(), 0, @@ -1157,7 +1185,7 @@ Groups: """ -async def _struct_judge_es_group_batch(group_specs: list[dict], chat_mdl) -> dict[str, set[int]]: +async def _struct_judge_doc_storage_group_batch(group_specs: list[dict], chat_mdl) -> dict[str, set[int]]: """Judge every incoming item independently without generating a merge.""" prompt_groups = [] for spec in group_specs: @@ -1206,7 +1234,7 @@ async def _struct_judge_es_group_batch(group_specs: list[dict], chat_mdl) -> dic return result -async def _struct_merge_es_group_batch(group_specs: list[dict], chat_mdl) -> dict[str, tuple[list[dict], dict | None]]: +async def _struct_merge_doc_storage_group_batch(group_specs: list[dict], chat_mdl) -> dict[str, tuple[list[dict], dict | None]]: """Judge multiple old_id groups in one LLM request.""" prompt_groups = [] for spec in group_specs: @@ -1261,7 +1289,7 @@ async def _struct_merge_es_group_batch(group_specs: list[dict], chat_mdl) -> dic return result -async def _struct_merge_es_group(old_doc: dict, incoming_docs: list[dict], chat_mdl) -> tuple[list[dict], dict | None]: +async def _struct_merge_doc_storage_group(old_doc: dict, incoming_docs: list[dict], chat_mdl) -> tuple[list[dict], dict | None]: """Judge one ES candidate group with one LLM request. Returns ``(non_duplicate_docs, merged_payload)``. The existing ES row is @@ -1302,7 +1330,7 @@ async def _struct_merge_es_group(old_doc: dict, incoming_docs: list[dict], chat_ return separate, merged -async def _struct_es_dedup_batch( +async def _struct_doc_storage_dedup_batch( docs: list[dict], chat_mdl, embd_mdl, @@ -1311,8 +1339,14 @@ async def _struct_es_dedup_batch( similarity_threshold: float, timing_context: str | None = None, cancel_check: Callable[[], bool] | None = None, + merge_scope: str = MERGE_SCOPE_DOC, ) -> tuple[int, int]: - """Batch ES dedup: concurrent KNN, parallel decisions, then grouped merges.""" + """Batch ES dedup: concurrent KNN, parallel decisions, then grouped merges. + + ``merge_scope`` controls the candidate horizon: ``doc`` restricts KNN and + the relation-alias rewrite to the incoming rows' own ``doc_id``; ``dataset`` + widens both to the whole knowledge base so cross-document duplicates merge. + """ from common import settings from rag.nlp import search as _rag_search @@ -1342,7 +1376,7 @@ async def _struct_es_dedup_batch( async def run_knn_shared(item_index: int, doc: dict): _raise_if_canceled() async with knn_semaphore: - return doc, await _struct_es_knn_candidate( + return doc, await _struct_doc_storage_knn_candidate( doc, tenant_id, kb_id, @@ -1351,6 +1385,7 @@ async def _struct_es_dedup_batch( select_fields, timing_context, item_index, + merge_scope, ) _raise_if_canceled() @@ -1416,7 +1451,7 @@ async def _struct_es_dedup_batch( _raise_if_canceled() async with llm_semaphore: try: - result = await _struct_judge_es_group_batch(batch_specs, chat_mdl) + result = await _struct_judge_doc_storage_group_batch(batch_specs, chat_mdl) except Exception: logging.exception("merge_compiled_structures: ES decision batch failed") result = {spec["request_group_id"]: set() for spec in batch_specs} @@ -1450,7 +1485,7 @@ async def _struct_es_dedup_batch( for start in range(0, len(duplicate_docs), _ES_DEDUP_LLM_BATCH_SIZE): _raise_if_canceled() candidate_docs = duplicate_docs[start : start + _ES_DEDUP_LLM_BATCH_SIZE] - separate, candidate_merged = await _struct_merge_es_group(current_doc, candidate_docs, chat_mdl) + separate, candidate_merged = await _struct_merge_doc_storage_group(current_doc, candidate_docs, chat_mdl) state["separate"].extend(separate) if candidate_merged is None: continue @@ -1508,7 +1543,7 @@ async def _struct_es_dedup_batch( logging.exception("merge_compiled_structures: grouped embedding failed for %d docs", len(batch)) vectors = [] for job, vec in zip(batch, vectors): - job["rebuilt"] = _struct_rebuild_es_doc( + job["rebuilt"] = _struct_rebuild_doc_storage_doc( job["payload"], job["old_doc"], vec, @@ -1556,9 +1591,14 @@ async def _struct_es_dedup_batch( ] from common.doc_store.doc_store_base import OrderByExpr + # In doc scope a renamed entity only affects relations inside the same + # document; in dataset scope the canonical entity is shared, so every + # relation in the KB that references an alias must be rewritten. Drop + # ``doc_id`` from the scope key (and the search condition) accordingly. + dataset_scope = merge_scope == MERGE_SCOPE_DATASET scopes = { ( - state["old_doc"].get("doc_id"), + None if dataset_scope else state["old_doc"].get("doc_id"), state["old_doc"].get("compile_kwd"), _struct_doc_template_id(state["old_doc"]), ) @@ -1567,10 +1607,11 @@ async def _struct_es_dedup_batch( } for doc_id, compile_kwd, template_id in scopes: condition = { - "doc_id": [doc_id], "compile_kwd": [compile_kwd], "knowledge_graph_kwd": ["relation"], } + if doc_id is not None: + condition["doc_id"] = [doc_id] if template_id: condition["compilation_template_ids"] = [template_id] try: @@ -1604,7 +1645,7 @@ async def _struct_es_dedup_batch( for start in range(0, len(rewrite_batch), _ES_DEDUP_EMBED_BATCH_SIZE): batch = rewrite_batch[start : start + _ES_DEDUP_EMBED_BATCH_SIZE] vectors = await _struct_embed(embd_mdl, [_struct_payload_description(payload) for _, payload in batch]) - rewritten = [_struct_rebuild_es_doc(payload, base, vector, base.get("source_chunk_ids") or [], preserve_id=True) for (base, payload), vector in zip(batch, vectors)] + rewritten = [_struct_rebuild_doc_storage_doc(payload, base, vector, base.get("source_chunk_ids") or [], preserve_id=True) for (base, payload), vector in zip(batch, vectors)] if rewritten: await thread_pool_exec(settings.docStoreConn.insert, rewritten, index, kb_id) existing_relation_updates += len(rewritten) @@ -1618,7 +1659,7 @@ async def _struct_es_dedup_batch( continue vector = await _struct_reembed_payload(job["payload"], embd_mdl) if vector is not None: - rewritten = _struct_rebuild_es_doc(job["payload"], job["old_doc"], vector, job["chunk_ids"], preserve_id=True) + rewritten = _struct_rebuild_doc_storage_doc(job["payload"], job["old_doc"], vector, job["chunk_ids"], preserve_id=True) await thread_pool_exec(settings.docStoreConn.insert, [rewritten], index, kb_id) return inserted, updated + existing_relation_updates @@ -1692,7 +1733,7 @@ async def _struct_local_dedup( # Re-embed failed: keep existing, drop incoming silently. dropped += 1 continue - rebuilt = _struct_rebuild_es_doc( + rebuilt = _struct_rebuild_doc_storage_doc( merged_payload, existing, new_vec, @@ -1858,7 +1899,7 @@ def _struct_graph_row_id( async def _struct_rebuild_graph_json( tenant_id: str, kb_id: str, - doc_id: str, + doc_id: str | None, compile_kwd: str, compilation_template_id: str | None = None, ) -> dict: @@ -1868,11 +1909,14 @@ async def _struct_rebuild_graph_json( index = _rag_search.index_name(tenant_id) fields = ["content_with_weight", "knowledge_graph_kwd", "source_chunk_ids"] + # ``doc_id is None`` collects every document's entities/relations in the KB + # for the dataset-level graph; a concrete id keeps it document-scoped. condition: dict = { - "doc_id": [doc_id], "compile_kwd": [compile_kwd], "knowledge_graph_kwd": ["entity", "relation"], } + if doc_id is not None: + condition["doc_id"] = [doc_id] if compilation_template_id: condition["compilation_template_ids"] = [compilation_template_id] res = await thread_pool_exec( @@ -2052,6 +2096,107 @@ async def rebuild_structure_graph_json( return graph +def _dataset_struct_graph_row_id( + kb_id: str, + compile_kwd: str, + compilation_template_id: str | None = None, +) -> str: + """Stable id for the KB-wide (dataset) structure graph row, keyed by + (kb, compile_kwd, template). Distinct namespace from the per-doc row id so + the dataset graph never collides with any document's graph.""" + tpl_part = compilation_template_id or "" + return xxhash.xxh64( + f"{kb_id}:dataset_structure_graph:{compile_kwd}:{tpl_part}".encode( + "utf-8", + "surrogatepass", + ), + ).hexdigest() + + +async def _struct_upsert_dataset_graph_json( + graph: dict, + tenant_id: str, + kb_id: str, + compile_kwd: str, + compilation_template_id: str | None = None, + structure_kind: str | None = None, +) -> None: + from common import settings + from rag.nlp import search as _rag_search + + index = _rag_search.index_name(tenant_id) + row_id = _dataset_struct_graph_row_id(kb_id, compile_kwd, compilation_template_id) + # The dataset graph is not per-document generated data, but the store schema + # requires ``doc_id``; carry ``kb_id`` there so the row is self-describing + # and per-doc graph queries (filtered by a real doc_id) never pick it up. + row = { + "id": row_id, + "content_with_weight": json.dumps(graph, ensure_ascii=False), + "compile_kwd": compile_kwd, + "knowledge_graph_kwd": "dataset_graph", + "doc_id": kb_id, + "kb_id": kb_id, + "available_int": 0, + } + if compilation_template_id: + row["compilation_template_ids"] = [compilation_template_id] + # On ``dataset_graph`` rows ``compilation_template_kind_kwd`` carries the + # template's *top-level* kind (graph/mind_map/timeline/session_essence/ + # session_graph) so the artifacts_structure endpoint can filter by it — + # unlike entity/relation rows where the field holds the ``config.kind`` + # (which collapses the knowledge_graph family together). + if structure_kind: + row["compilation_template_kind_kwd"] = str(structure_kind) + old = await thread_pool_exec(settings.docStoreConn.get, row_id, index, [kb_id]) + if old: + await thread_pool_exec( + settings.docStoreConn.update, + {"id": row_id}, + {k: v for k, v in row.items() if k != "id"}, + index, + kb_id, + ) + else: + await thread_pool_exec(settings.docStoreConn.insert, [row], index, kb_id) + + +async def rebuild_dataset_structure_graph_json( + tenant_id: str, + kb_id: str, + compile_kwd: str, + compilation_template_id: str | None = None, + structure_kind: str | None = None, +) -> dict: + """Rebuild and persist the KB-wide (dataset) structure graph for one + (compile_kwd, template_id) pair. + + Reads every document's merged entities/relations in the KB (no ``doc_id`` + filter) and writes a single ``knowledge_graph_kwd="dataset_graph"`` row. + Meant to run once per template after all per-document flushes finish; when + ``dataset`` merge scope is on the entity/relation rows are already + deduplicated across documents, so this simply projects them into a graph. + + ``structure_kind`` is the template's top-level kind (e.g. ``knowledge_graph``, + ``session_graph``); it is stamped on the row so the dataset structure API + can filter by kind.""" + graph = await _struct_rebuild_graph_json( + tenant_id, + kb_id, + None, + compile_kwd, + compilation_template_id, + ) + await _struct_upsert_dataset_graph_json( + graph, + tenant_id, + kb_id, + compile_kwd, + compilation_template_id, + structure_kind=structure_kind, + ) + return graph + + # --------------------------------------------------------------------------- # Chain-shape validation for ``list`` / ``timeline`` kinds. # @@ -2387,8 +2532,9 @@ async def merge_compiled_structures( chain_kind: str = "", chain_callback=None, chain_timeout_seconds: float = 120.0, - es_waiter: Callable[[], Awaitable[None]] | None = None, - es_releaser: Callable[[], Awaitable[None]] | None = None, + doc_storage_waiter: Callable[[], Awaitable[None]] | None = None, + doc_storage_releaser: Callable[[], Awaitable[None]] | None = None, + merge_scope: str = MERGE_SCOPE_DOC, ) -> dict: """Merge ``docs`` (the output of ``compile_structure_from_text``) before inserting them into ES. @@ -2419,9 +2565,18 @@ async def merge_compiled_structures( cancel_check: optional callable returning True when the owning parse task has been canceled. Checked between ES-dedup iterations so a long merge can stop promptly. + merge_scope: ``"doc"`` (default) dedups only against rows already + stored for the incoming ``doc_id``; ``"dataset"`` dedups against + the whole KB so cross-document duplicates collapse. Dataset-scope + ES writes are serialized with a per-(kb, template) Redis lock so + concurrent per-document parses don't insert duplicate canonical + rows. Surviving rows keep the existing row's ``doc_id``. Returns: - {"inserted": N, "updated": M, "duplicates_dropped": K} summary. + {"inserted": N, "updated": M, "duplicates_dropped": K, + "compile_kwds": [...]} summary. ``compile_kwds`` lists the compile + keywords touched so a dataset-scope caller can rebuild the dataset + structure graph once all flushes finish. """ if not docs: return {"inserted": 0, "updated": 0, "duplicates_dropped": 0} @@ -2471,11 +2626,30 @@ async def merge_compiled_structures( if callable(cancel_check) and cancel_check(): raise TaskCanceledException("Task was cancelled during structure ES dedup merge") - if es_waiter is not None: - await es_waiter() + if doc_storage_waiter is not None: + await doc_storage_waiter() _raise_if_canceled() + # Dataset scope: hold a per-(kb, template) lock across the read-modify-write + # KNN merge so concurrent per-document parses can't both KNN-miss the same + # canonical entity and insert duplicates. Acquired *after* the in-doc + # ``doc_storage_waiter`` gate (never before) to avoid a deadlock where a later flush + # holds the KB lock while waiting on an earlier flush that also needs it. + merge_lock = None + if merge_scope == MERGE_SCOPE_DATASET: + from rag.utils.redis_conn import RedisDistributedLock + + merge_lock = RedisDistributedLock( + _struct_merge_lock_key(kb_id, compilation_template_id), + timeout=_STRUCT_MERGE_LOCK_TIMEOUT_S, + blocking_timeout=_STRUCT_MERGE_LOCK_BLOCKING_TIMEOUT_S, + ) + try: + await merge_lock.spin_acquire() + except Exception: + logging.exception("merge_compiled_structures: dataset merge lock acquire failed for kb=%s", kb_id) + merge_lock = None try: - inserted, updated = await _struct_es_dedup_batch( + inserted, updated = await _struct_doc_storage_dedup_batch( deduped, chat_mdl, embd_mdl, @@ -2484,12 +2658,19 @@ async def merge_compiled_structures( similarity_threshold, timing_context=timing_context, cancel_check=cancel_check, + merge_scope=merge_scope, ) except Exception: logging.exception("merge_compiled_structures: batched ES dedup failed") inserted = updated = 0 - if es_releaser is not None: - await es_releaser() + finally: + if merge_lock is not None: + try: + merge_lock.release() + except Exception: + logging.exception("merge_compiled_structures: dataset merge lock release failed for kb=%s", kb_id) + if doc_storage_releaser is not None: + await doc_storage_releaser() graphs = 0 for graph_index, (doc_id, compile_kwd, template_id) in enumerate(graph_keys): @@ -2516,6 +2697,7 @@ async def merge_compiled_structures( "updated": updated, "duplicates_dropped": dropped, "graphs": graphs, + "compile_kwds": sorted({compile_kwd for _, compile_kwd, _ in graph_keys}), } return info @@ -2525,4 +2707,7 @@ __all__ = [ "merge_compiled_structures", "cleanup_timeline_isolated_entities", "rebuild_structure_graph_json", + "rebuild_dataset_structure_graph_json", + "MERGE_SCOPE_DOC", + "MERGE_SCOPE_DATASET", ] diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index 99d88bcd8b..713da1bcb3 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -1237,6 +1237,59 @@ async def _wiki_load_all_map_extracts(tenant_id: str, kb_id: str) -> dict: return merged +async def _wiki_all_map_doc_ids(tenant_id: str, kb_id: str) -> list[str]: + """Distinct ``doc_id`` across every ``artifact_map_extract`` row in this KB. + + These are the documents that fed the current compilation. Stamped onto + the KB-wide aggregate rows (REDUCE / PLAN) as ``source_doc_ids`` so a + document delete can reference-count them — the aggregate is dropped only + once its last contributing document is gone. + """ + 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]} + select_fields = ["id", "doc_id"] + + PAGE_SIZE = 1000 + offset = 0 + doc_ids: list[str] = [] + seen: set[str] = set() + 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) + except Exception: + logging.exception("wiki: failed to scan MAP doc ids for kb=%s (offset=%d)", kb_id, offset) + break + if not field_map: + break + for row in field_map.values(): + raw = row.get("doc_id") + candidates = raw if isinstance(raw, list) else [raw] + for d in candidates: + if isinstance(d, str) and d and d not in seen: + seen.add(d) + doc_ids.append(d) + if len(field_map) < PAGE_SIZE: + break + offset += PAGE_SIZE + return doc_ids + + async def _wiki_compute_map_input_hash(tenant_id: str, kb_id: str) -> str: """xxh64 fingerprint of the **current** ``artifact_map_extract`` rows for this KB — used by REDUCE / PLAN to cache-bust when MAP changed. @@ -1365,11 +1418,14 @@ async def _wiki_persist_reduce( tenant_id: str, kb_id: str, input_hash: str = "", + source_doc_ids: Optional[list[str]] = None, ) -> None: """Upsert the single non-searchable artifact_reduce_result row for this KB. ``input_hash`` records the MAP-state fingerprint this reduction was computed from; the next call compares it before re-running. + ``source_doc_ids`` is the set of documents that fed this reduction, used + for delete-time reference counting. """ from common import settings from rag.nlp import search as _rag_search @@ -1384,6 +1440,7 @@ async def _wiki_persist_reduce( "doc_id": kb_id_str, # sentinel — KB-scoped row, not a real document "compile_kwd": WIKI_REDUCE_COMPILE_KWD, "source_id": [kb_id_str], + "source_doc_ids": list(source_doc_ids or []), "input_hash_kwd": input_hash, "content_with_weight": content_with_weight, "available_int": 0, @@ -1466,6 +1523,7 @@ async def wiki_reduce_from_extracts( # correct. ``force_rerun=True`` bypasses both checks for the # legacy / admin "rebuild from scratch" path. current_input_hash = await _wiki_compute_map_input_hash(tenant_id, kb_id) + reduce_source_doc_ids = await _wiki_all_map_doc_ids(tenant_id, kb_id) if not force_rerun: cached_pair = await _wiki_load_reduce_resume(tenant_id, kb_id) if cached_pair is not None: @@ -1501,7 +1559,7 @@ async def wiki_reduce_from_extracts( if not raw_entities and not raw_concepts: # Nothing to reduce; persist an empty result so resume can short-circuit. empty = _wiki_empty_extract() - await _wiki_persist_reduce(empty, tenant_id, kb_id, input_hash=current_input_hash) + await _wiki_persist_reduce(empty, tenant_id, kb_id, input_hash=current_input_hash, source_doc_ids=reduce_source_doc_ids) return empty if callback: @@ -1560,7 +1618,7 @@ async def wiki_reduce_from_extracts( callback(0.9, "wiki REDUCE: persisting result") except Exception: pass - await _wiki_persist_reduce(reduced, tenant_id, kb_id, input_hash=current_input_hash) + await _wiki_persist_reduce(reduced, tenant_id, kb_id, input_hash=current_input_hash, source_doc_ids=reduce_source_doc_ids) logging.info( "wiki_reduce: kb=%s done — entities=%d concepts=%d claims=%d relations=%d topics=%d", @@ -2167,11 +2225,14 @@ async def _wiki_persist_plan( tenant_id: str, kb_id: str, input_hash: str = "", + source_doc_ids: Optional[list[str]] = None, ) -> None: """Upsert the single non-searchable artifact_compilation_plan row for this KB. ``input_hash`` records the REDUCE-state fingerprint this plan was derived from; the next call compares it before re-planning. + ``source_doc_ids`` is the set of documents that fed this plan, used for + delete-time reference counting. """ from common import settings from rag.nlp import search as _rag_search @@ -2185,6 +2246,7 @@ async def _wiki_persist_plan( "doc_id": kb_id_str, # sentinel — KB-scoped row, not a real document "compile_kwd": WIKI_PLAN_COMPILE_KWD, "source_id": [kb_id_str], + "source_doc_ids": list(source_doc_ids or []), "input_hash_kwd": input_hash, "content_with_weight": content_with_weight, "available_int": 0, @@ -2262,6 +2324,7 @@ async def wiki_plan_from_reduction( # plan was stamped with the same hash REDUCE is currently exposing, # nothing upstream has changed and the plan is still valid. current_reduce_hash = await _wiki_load_reduce_input_hash(tenant_id, kb_id) + plan_source_doc_ids = await _wiki_all_map_doc_ids(tenant_id, kb_id) if not force_rerun: cached_pair = await _wiki_load_plan_resume(tenant_id, kb_id) if cached_pair is not None: @@ -2295,7 +2358,7 @@ async def wiki_plan_from_reduction( "_topics": [], "_reconciliation": {}, } - await _wiki_persist_plan(empty, tenant_id, kb_id, input_hash=current_reduce_hash) + await _wiki_persist_plan(empty, tenant_id, kb_id, input_hash=current_reduce_hash, source_doc_ids=plan_source_doc_ids) return empty canonical_entities = reduced.get("entities") or [] @@ -2326,7 +2389,7 @@ async def wiki_plan_from_reduction( "_topics": raw_topics, "_reconciliation": {}, } - await _wiki_persist_plan(empty, tenant_id, kb_id, input_hash=current_reduce_hash) + await _wiki_persist_plan(empty, tenant_id, kb_id, input_hash=current_reduce_hash, source_doc_ids=plan_source_doc_ids) return empty if callback: @@ -2391,7 +2454,7 @@ async def wiki_plan_from_reduction( callback(0.9, "wiki PLAN: persisting plan") except Exception: pass - await _wiki_persist_plan(plan, tenant_id, kb_id, input_hash=current_reduce_hash) + await _wiki_persist_plan(plan, tenant_id, kb_id, input_hash=current_reduce_hash, source_doc_ids=plan_source_doc_ids) logging.info( "wiki_plan: kb=%s done — pages=%d (target=%d) updates=%d creates=%d", @@ -3227,12 +3290,14 @@ async def _wiki_persist_draft( return index = _rag_search.index_name(tenant_id) content_with_weight = json.dumps(page, ensure_ascii=False) + draft_doc_ids = [d for d in (page.get("source_doc_ids") or []) if isinstance(d, str) and d] row = { "id": _wiki_draft_row_id(kb_id, slug), "doc_id": str(kb_id), "compile_kwd": WIKI_DRAFT_COMPILE_KWD, "artifact_slug_kwd": slug, "source_id": [str(kb_id)], + "source_doc_ids": draft_doc_ids, "input_hash_kwd": plan_input_hash, "content_with_weight": content_with_weight, "available_int": 0, # non-searchable @@ -3570,8 +3635,8 @@ async def wiki_refine_from_plan( return None # Searchable artifact_page persistence has moved to the task - # handler (TaskHandler._persist_wiki_pages_to_es) so the ES - # schema can be controlled in one place at the ingest layer. + # handler so the doc-storage schema can be controlled in one + # place at the ingest layer. # REFINE now just builds the page dict and resume cache. try: await _wiki_persist_draft( diff --git a/rag/flow/compiler/compiler.py b/rag/flow/compiler/compiler.py index 82bf990359..5ccd750f27 100644 --- a/rag/flow/compiler/compiler.py +++ b/rag/flow/compiler/compiler.py @@ -53,6 +53,8 @@ class CompilerParam(ProcessParamBase, LLMParam): def check(self): super().check() self.check_empty(self.compilation_template_group_ids, "Compilation Template Groups") + if isinstance(self.compilation_template_group_ids, str): + self.compilation_template_group_ids = [self.compilation_template_group_ids] class Compiler(ProcessBase, LLM): diff --git a/rag/nlp/search.py b/rag/nlp/search.py index 0f97a1a537..6f69f83e10 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -207,7 +207,7 @@ class Dealer: if settings.DOC_ENGINE_OCEANBASE: src.append(f"q_{len(q_vec)}_vec") - fusionExpr = FusionExpr("weighted_sum", topk, {"weights": "0.05,0.95"}) + fusionExpr = FusionExpr("weighted_sum", topk, {"weights": "0.001,1"}) matchExprs = [matchText, matchDense, fusionExpr] res = await thread_pool_exec(self.dataStore.search, src, highlightFields, filters, matchExprs, orderBy, offset, limit, idx_names, kb_ids, rank_feature=rank_feature) diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index 2f785cce96..626070375c 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -138,8 +138,30 @@ TASK_TYPE_TO_PIPELINE_TASK_TYPE = { "memory": PipelineTaskType.MEMORY, "artifact": PipelineTaskType.ARTIFACT, "skill": PipelineTaskType.SKILL, + "structure_graph": PipelineTaskType.STRUCTURE_GRAPH, + "structure_mindmap": PipelineTaskType.STRUCTURE_MINDMAP, + "timeline": PipelineTaskType.TIMELINE, + "session_graph": PipelineTaskType.SESSION_GRAPH, + "session_essence": PipelineTaskType.SESSION_ESSENCE, + "structure": PipelineTaskType.STRUCTURE, } +# KB-wide fan-out task types: their task row's ``doc_id`` is a fake sentinel and +# the participating documents live in ``task["doc_ids"]``. +_KB_FANOUT_TASK_TYPES = [ + "graphrag", + "raptor", + "mindmap", + "artifact", + "skill", + "structure_graph", + "structure_mindmap", + "timeline", + "session_graph", + "session_essence", + "structure", +] + UNACKED_ITERATOR = None # Task type and executor index (consistent with SAAS version) TASK_TYPE = "common" @@ -1771,7 +1793,7 @@ async def handle_task(): finally: if not task.get("dataflow_id", ""): referred_document_id = None - if task_type in ["graphrag", "raptor", "mindmap", "artifact", "skill"]: + if task_type in _KB_FANOUT_TASK_TYPES: # KB-level fan-out tasks store the participating doc list in # task["doc_ids"]; the first entry is used as a referent so # the pipeline operation log has something to anchor to. diff --git a/rag/svr/task_executor_refactor/dataset_structure_merger.py b/rag/svr/task_executor_refactor/dataset_structure_merger.py new file mode 100644 index 0000000000..e05977f708 --- /dev/null +++ b/rag/svr/task_executor_refactor/dataset_structure_merger.py @@ -0,0 +1,214 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""KB-wide structure-graph merge task. + +Runs when the user POSTs to ``/datasets/<id>/index`` with a structure index +type (``structure_graph`` / ``structure_mindmap`` / ``timeline`` / +``session_graph`` / ``session_essence``, or ``structure`` for merge-all). It +re-projects every document's already-merged ``entity`` / ``relation`` rows into +the KB-wide ``knowledge_graph_kwd="dataset_graph"`` rows via +:func:`rag.advanced_rag.knowlege_compile.structure.rebuild_dataset_structure_graph_json` +— the same merge the per-document parse runner performs at flush time +(``rag.advanced_rag.knowlege_compile.runner``), exposed here as an on-demand +re-merge. + +The task carries the target kind in ``task_type``; the driver enumerates the +``(compile_kwd, template_id)`` pairs actually present in the store, keeps only +dataset-merge templates of the requested kind (all kinds for merge-all), and +rebuilds one dataset graph per pair. It performs no LLM work. +""" + +from __future__ import annotations + +import logging +from typing import Optional + +from common import settings +from common.misc_utils import thread_pool_exec +from rag.nlp import search +from rag.advanced_rag.knowlege_compile.structure import ( + rebuild_dataset_structure_graph_json, +) +from rag.svr.task_executor_refactor.task_context import TaskContext + + +# Structure merge task_type -> the template's top-level ``kind`` as stored on +# ``dataset_graph`` rows. ``None`` (the merge-all ``structure`` type) rebuilds +# every dataset-merge kind regardless of its top-level kind. +_STRUCTURE_TASK_TYPE_TO_KIND: dict[str, Optional[str]] = { + "structure_graph": "knowledge_graph", + "structure_mindmap": "mind_map", + "timeline": "timeline", + "session_graph": "session_graph", + "session_essence": "session_essence", + "structure": None, # merge-all +} + +STRUCTURE_MERGE_TASK_TYPES = frozenset(_STRUCTURE_TASK_TYPE_TO_KIND) + + +def is_structure_merge_task(task_type: str) -> bool: + return (task_type or "").lower() in STRUCTURE_MERGE_TASK_TYPES + + +async def _collect_structure_pairs(tenant_id: str, kb_id: str) -> set[tuple[str, str]]: + """Distinct ``(compile_kwd, template_id)`` pairs across the KB's structure + ``entity`` / ``relation`` rows — the inputs each dataset graph is rebuilt + from. Rows without a ``compilation_template_ids`` are skipped: the merge is + always template-scoped (an untemplated rebuild would over-merge unrelated + kinds sharing a ``compile_kwd``).""" + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return set() + + select_fields = ["id", "compile_kwd", "compilation_template_ids"] + pairs: set[tuple[str, str]] = set() + offset = 0 + page_size = 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"knowledge_graph_kwd": ["entity", "relation"]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + except Exception: + logging.exception("structure_merge: failed to scan entity/relation rows for kb=%s", kb_id) + break + if not field_map: + break + for row in field_map.values(): + compile_kwd = row.get("compile_kwd") + if not isinstance(compile_kwd, str) or not compile_kwd: + continue + raw_tids = row.get("compilation_template_ids") + if isinstance(raw_tids, str): + tids = [raw_tids] if raw_tids else [] + elif isinstance(raw_tids, list): + tids = [t for t in raw_tids if isinstance(t, str) and t] + else: + tids = [] + for tid in tids: + pairs.add((compile_kwd, tid)) + if len(field_map) < page_size: + break + offset += page_size + return pairs + + +async def run_structure_merge(ctx: TaskContext) -> None: + """Rebuild the KB-wide dataset structure graph(s) for the task's kind. + + Enumerates the ``(compile_kwd, template_id)`` pairs present in the store, + filters to dataset-merge templates of the requested kind (all kinds for the + merge-all ``structure`` type), and rebuilds one ``dataset_graph`` row per + pair. Per-pair failures are logged and skipped so one bad template does not + abort the whole merge. + """ + from api.db.services.compilation_template_service import CompilationTemplateService + + progress = ctx.progress_cb + task_type = (ctx.task_type or "").lower() + target_kind = _STRUCTURE_TASK_TYPE_TO_KIND.get(task_type) + merge_all = task_type == "structure" + + def _canceled() -> bool: + try: + return bool(ctx.has_canceled_func(ctx.id)) + except Exception: + return False + + progress(0.0, "Collecting structure graphs to merge...") + pairs = await _collect_structure_pairs(ctx.tenant_id, ctx.kb_id) + if not pairs: + progress(1.0, "No structure graphs to merge.") + return + + # Resolve each template once: keep only dataset-merge templates, and (unless + # merge-all) only those whose top-level kind matches the requested kind. + # ``template_meta`` maps template_id -> (keep: bool, structure_kind: str|None). + template_meta: dict[str, tuple[bool, Optional[str]]] = {} + + def _resolve_template(template_id: str) -> tuple[bool, Optional[str]]: + cached = template_meta.get(template_id) + if cached is not None: + return cached + keep = False + structure_kind: Optional[str] = None + try: + saved = CompilationTemplateService.get_saved(template_id, ctx.tenant_id) + if saved: + structure_kind = (saved.get("kind") or "").strip() or None + config = saved.get("config") or {} + dataset_merge = bool(config.get("dataset_merge")) if isinstance(config, dict) else False + kind_ok = merge_all or (structure_kind == target_kind) + keep = dataset_merge and kind_ok + except Exception: + logging.exception("structure_merge: failed to resolve template %s for kb=%s", template_id, ctx.kb_id) + result = (keep, structure_kind) + template_meta[template_id] = result + return result + + eligible = [] + for compile_kwd, template_id in sorted(pairs): + keep, structure_kind = _resolve_template(template_id) + if keep: + eligible.append((compile_kwd, template_id, structure_kind)) + + if not eligible: + kind_label = "any kind" if merge_all else (target_kind or task_type) + progress(1.0, f"No dataset-merge structure templates to merge for {kind_label}.") + return + + total = len(eligible) + rebuilt = 0 + for i, (compile_kwd, template_id, structure_kind) in enumerate(eligible): + if _canceled(): + progress(-1, "Task has been canceled.") + return + progress( + 0.05 + 0.9 * (i / total), + f"Merging structure graph {i + 1}/{total} (compile_kwd={compile_kwd}) ...", + ) + try: + await rebuild_dataset_structure_graph_json( + ctx.tenant_id, + ctx.kb_id, + compile_kwd, + compilation_template_id=template_id, + structure_kind=structure_kind, + ) + rebuilt += 1 + except Exception: + logging.exception( + "structure_merge: rebuild failed for kb=%s compile_kwd=%s template=%s", + ctx.kb_id, + compile_kwd, + template_id, + ) + + progress(1.0, f"Merged {rebuilt}/{total} structure graph(s).") diff --git a/rag/svr/task_executor_refactor/dataset_wiki_generator.py b/rag/svr/task_executor_refactor/dataset_wiki_generator.py index 26d817558d..cce39ae7af 100644 --- a/rag/svr/task_executor_refactor/dataset_wiki_generator.py +++ b/rag/svr/task_executor_refactor/dataset_wiki_generator.py @@ -101,8 +101,22 @@ WIKI_GRAPH_MAX_CHUNK_IDS_PER_NODE = 64 # path, which carries the user's own title/comments). WIKI_REGEN_COMMIT_TITLE = "Regenerated by artifact compilation" WIKI_REGEN_COMMIT_COMMENTS_TEMPLATE = "Auto-update via run_wiki (action={action})" +WIKI_MAP_COMPILE_KWD = "artifact_map_extract" +WIKI_REDUCE_COMPILE_KWD = "artifact_reduce_result" +WIKI_PLAN_COMPILE_KWD = "artifact_compilation_plan" +WIKI_DRAFT_COMPILE_KWD = "artifact_page_draft" WIKI_PAGE_COMPILE_KWD = "artifact_page" WIKI_PAGE_TOPIC_COMPILE_KWD = "artifact_page_topic" +WIKI_DERIVED_COMPILE_KWDS = ( + WIKI_REDUCE_COMPILE_KWD, + WIKI_PLAN_COMPILE_KWD, + WIKI_DRAFT_COMPILE_KWD, + WIKI_PAGE_COMPILE_KWD, + WIKI_PAGE_TOPIC_COMPILE_KWD, + "artifact_entity", + "artifact_relation", + "artifact_page_graph", +) # ----- helpers ------------------------------------------------------- @@ -129,6 +143,240 @@ def _parser_config_compilation_template_ids(parser_config, tenant_id: str) -> li return template_ids +def _normalize_compilation_template_group_ids(raw) -> list[str]: + if isinstance(raw, str): + raw = [raw] + if not isinstance(raw, list): + return [] + ids: list[str] = [] + seen: set[str] = set() + for group_id in raw: + if not isinstance(group_id, str): + continue + group_id = group_id.strip() + if group_id and group_id not in seen: + seen.add(group_id) + ids.append(group_id) + return ids + + +def _extract_pipeline_compiler_group_ids(dsl) -> list[str]: + if isinstance(dsl, str): + try: + dsl = json.loads(dsl) + except Exception: + return [] + if not isinstance(dsl, dict): + return [] + components = dsl.get("components") + if not isinstance(components, dict): + return [] + + group_ids: list[str] = [] + seen: set[str] = set() + for component in components.values(): + if not isinstance(component, dict): + continue + obj = component.get("obj") if isinstance(component.get("obj"), dict) else {} + component_name = obj.get("component_name") or component.get("component_name") or component.get("name") + if not isinstance(component_name, str) or component_name.lower() != "compiler": + continue + candidates = [ + obj.get("params") if isinstance(obj.get("params"), dict) else {}, + obj, + component.get("params") if isinstance(component.get("params"), dict) else {}, + component, + ] + for candidate in candidates: + for key in ("compilation_template_group_ids", "compilation_template_group_id"): + for group_id in _normalize_compilation_template_group_ids(candidate.get(key)): + if group_id not in seen: + seen.add(group_id) + group_ids.append(group_id) + return group_ids + + +def _pipeline_compilation_template_ids(pipeline_id: str, tenant_id: str) -> list[str]: + pipeline_id = (pipeline_id or "").strip() + if not pipeline_id: + return [] + from api.db.services.canvas_service import UserCanvasService + from api.db.services.compilation_template_group_service import ( + CompilationTemplateGroupService, + ) + + ok, canvas = UserCanvasService.get_by_id(pipeline_id) + if not ok or not canvas: + return [] + template_ids: list[str] = [] + seen: set[str] = set() + for group_id in _extract_pipeline_compiler_group_ids(getattr(canvas, "dsl", None)): + for template_id in CompilationTemplateGroupService.resolve_template_ids(group_id, tenant_id): + if template_id in seen: + continue + seen.add(template_id) + template_ids.append(template_id) + return template_ids + + +async def _wiki_existing_map_doc_ids(tenant_id: str, kb_id: str) -> set[str]: + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return set() + + doc_ids: set[str] = set() + 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_id = row.get("doc_id") + if isinstance(doc_id, str) and doc_id: + doc_ids.add(doc_id) + if len(field_map) < page_size: + break + offset += page_size + return doc_ids + + +async def _wiki_delete_deleted_doc_state( + tenant_id: str, + kb_id: str, + deleted_doc_ids: set[str], +) -> None: + if not deleted_doc_ids: + return + + index = search.index_name(tenant_id) + 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 + + # 2. Derived KB-scoped rows: reference-counted self-healing backstop for + # the eager delete-time cleanup (DocumentService.remove_artifact_products). + # Read every row referencing any deleted doc, drop the ones left with no + # surviving owner, and shrink the rest to their surviving doc set. This + # replaces the former blunt "delete every derived row" wipe so products + # shared with still-present docs survive a peer's deletion. + from common.doc_store.doc_store_base import OrderByExpr + + deleted = set(deleted_doc_ids) + derived_kwds = list(WIKI_DERIVED_COMPILE_KWDS) + select_fields = ["id", "source_doc_ids"] + to_delete: list[str] = [] + to_shrink: list[tuple[str, list[str]]] = [] + offset = 0 + page_size = 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"compile_kwd": derived_kwds, "source_doc_ids": sorted(deleted)}, + [], + 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 derived rows for removed docs in kb=%s", kb_id) + return + if not field_map: + break + for row_id, row in field_map.items(): + raw = row.get("source_doc_ids") + if isinstance(raw, str): + owners = [raw] if raw else [] + elif isinstance(raw, list): + owners = [d for d in raw if isinstance(d, str) and d] + else: + owners = [] + remaining = [d for d in owners if d not in deleted] + if remaining: + to_shrink.append((row_id, remaining)) + else: + to_delete.append(row_id) + if len(field_map) < page_size: + break + offset += page_size + + # Drop rows with no surviving owner (delete by id in batches). + for i in range(0, len(to_delete), page_size): + try: + await thread_pool_exec( + settings.docStoreConn.delete, + {"id": to_delete[i : i + page_size]}, + index, + kb_id, + ) + except Exception: + logging.exception("wiki: failed to drop orphaned derived rows in kb=%s", kb_id) + + # Shrink rows still owned by surviving docs to just those docs. + for row_id, remaining in to_shrink: + try: + await thread_pool_exec( + settings.docStoreConn.update, + {"id": row_id}, + {"source_doc_ids": remaining}, + index, + kb_id, + ) + except Exception: + logging.exception("wiki: failed to shrink source_doc_ids for row=%s in kb=%s", row_id, kb_id) + + logging.info( + "wiki: ref-counted cleanup for %d deleted doc(s) in kb=%s (dropped=%d, shrunk=%d)", + len(deleted), + kb_id, + len(to_delete), + len(to_shrink), + ) + + def _wiki_topic_from_page(page: Dict, fallback: str = "") -> str: for key in ("topic", "title", "page_type"): value = page.get(key) @@ -149,6 +397,7 @@ async def _ensure_wiki_topic_rows( index: str, kb_id_str: str, topics_by_name: dict[str, str], + topic_doc_ids: dict[str, list[str]] | None = None, ) -> None: if not topics_by_name: return @@ -156,6 +405,7 @@ async def _ensure_wiki_topic_rows( from common.doc_store.doc_store_base import OrderByExpr from rag.nlp import rag_tokenizer + topic_doc_ids = topic_doc_ids or {} topics = list(topics_by_name.keys()) existing: set[str] = set() try: @@ -190,6 +440,24 @@ async def _ensure_wiki_topic_rows( rows: list[dict] = [] for topic, slug in topics_by_name.items(): if topic in existing: + # Topic row already present — refresh only its doc provenance so + # delete-time ref-counting stays accurate as new docs contribute + # to (or stop contributing to) the topic across recompiles. + doc_ids = topic_doc_ids.get(topic) or [] + try: + await thread_pool_exec( + settings.docStoreConn.update, + {"compile_kwd": WIKI_PAGE_TOPIC_COMPILE_KWD, "topic_kwd": topic}, + {"source_doc_ids": list(doc_ids)}, + index, + ctx.kb_id, + ) + except Exception: + logging.exception( + "wiki_persist: topic provenance update failed for kb=%s topic=%s", + kb_id_str, + topic, + ) continue topic_id = xxhash.xxh64( f"{kb_id_str}:{WIKI_PAGE_TOPIC_COMPILE_KWD}:{topic}".encode( @@ -207,6 +475,7 @@ async def _ensure_wiki_topic_rows( "topic_kwd": topic, "title_kwd": topic, "slug_kwd": slug, + "source_doc_ids": list(topic_doc_ids.get(topic) or []), "content_with_weight": topic, "content_ltks": content_ltks, "content_sm_ltks": rag_tokenizer.fine_grained_tokenize(content_ltks), @@ -324,6 +593,9 @@ async def persist_wiki_pages_to_es( rows: List[Dict] = [] topics_by_name: dict[str, str] = {} + # Per-topic union of contributing doc ids, so topic rows can be + # reference-counted at document-delete time like every other product. + topic_doc_ids: dict[str, list[str]] = {} for page, vec in zip(pages, embeddings): slug = page.get("slug") or "" if not slug: @@ -378,6 +650,12 @@ async def persist_wiki_pages_to_es( ) if topic: topics_by_name.setdefault(topic, _wiki_topic_slug(topic)) + bucket = topic_doc_ids.setdefault(topic, []) + seen_bucket = set(bucket) + for d in page.get("source_doc_ids") or []: + if isinstance(d, str) and d and d not in seen_bucket: + seen_bucket.add(d) + bucket.append(d) if not rows: return @@ -394,7 +672,7 @@ async def persist_wiki_pages_to_es( if topics_by_name: try: - await _ensure_wiki_topic_rows(ctx, index, kb_id_str, topics_by_name) + await _ensure_wiki_topic_rows(ctx, index, kb_id_str, topics_by_name, topic_doc_ids) except Exception: logging.exception( "wiki_persist: topic row insert failed for kb=%s", @@ -484,6 +762,11 @@ def build_wiki_page_graph( name = p.get("title") or slug aliases = list(p.get("entity_names") or []) + # Per-node doc provenance: the documents that fed this page. Stamped + # onto the entity row so document deletion can reference-count it + # (drop the entity only when its last source doc is removed). + page_doc_ids = [d for d in (p.get("source_doc_ids") or []) if isinstance(d, str) and d] + by_slug[slug] = { "slug": slug, "name": name, @@ -492,6 +775,7 @@ def build_wiki_page_graph( "type": page_type, "weight": weight, "source_chunk_ids": capped_chunk_ids, + "source_doc_ids": page_doc_ids, } # Per-entity ES row. content_ltks is built from slug + summary @@ -518,6 +802,7 @@ def build_wiki_page_graph( "slug_kwd": slug, "weight_int": int(weight), "source_chunk_ids": capped_chunk_ids, + "source_doc_ids": page_doc_ids, "content_ltks": rag_tokenizer.tokenize(content_text) if content_text else "", "content_with_weight": json.dumps(entity_payload, ensure_ascii=False), } @@ -537,6 +822,16 @@ def build_wiki_page_graph( tgt = "" if not tgt or tgt == src or tgt not in by_slug: continue + # Edge provenance is the union of both endpoints' source docs, so + # the relation is dropped only once neither endpoint traces to a + # surviving document. + edge_doc_ids: list[str] = [] + edge_seen: set[str] = set() + for endpoint in (src, tgt): + for d in by_slug[endpoint].get("source_doc_ids") or []: + if d not in edge_seen: + edge_seen.add(d) + edge_doc_ids.append(d) relation_payload = {"from": src, "to": tgt} relation_rows.append( { @@ -550,6 +845,7 @@ def build_wiki_page_graph( "type_kwd": "artifact_relation", "from_kwd": src, "to_kwd": tgt, + "source_doc_ids": edge_doc_ids, "content_with_weight": json.dumps(relation_payload, ensure_ascii=False), } ) @@ -693,10 +989,42 @@ async def run_wiki( types=[], suffix=[], ) + current_doc_ids = {str(d.get("id")) for d in all_docs or [] if d.get("id")} + existing_map_doc_ids = await _wiki_existing_map_doc_ids(ctx.tenant_id, ctx.kb_id) + deleted_doc_ids = existing_map_doc_ids - current_doc_ids + if deleted_doc_ids: + progress( + 0.02, + f"Removing stale wiki state for {len(deleted_doc_ids)} deleted document(s)...", + ) + await _wiki_delete_deleted_doc_state( + ctx.tenant_id, + ctx.kb_id, + deleted_doc_ids, + ) + eligible = [] + pipeline_template_ids_cache: dict[str, list[str]] = {} for d in all_docs or []: pc = d.get("parser_config") or {} + template_ids: list[str] = [] + seen_template_ids: set[str] = set() for template_id in _parser_config_compilation_template_ids(pc, ctx.tenant_id): + if template_id in seen_template_ids: + continue + seen_template_ids.add(template_id) + template_ids.append(template_id) + pipeline_id = (d.get("pipeline_id") or "").strip() + if pipeline_id: + if pipeline_id not in pipeline_template_ids_cache: + pipeline_template_ids_cache[pipeline_id] = _pipeline_compilation_template_ids(pipeline_id, ctx.tenant_id) + for template_id in pipeline_template_ids_cache[pipeline_id]: + if template_id in seen_template_ids: + continue + seen_template_ids.add(template_id) + template_ids.append(template_id) + + for template_id in template_ids: template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id) config = template.get("config") if template else {} kind = _compilation_template_kind(config.get("kind") if isinstance(config, dict) else "") @@ -864,7 +1192,11 @@ async def run_wiki( stats["delta"][key] += int(meta.get(key, 0) or 0) stats["had_delta"] |= bool(meta.get("had_delta")) except Exception: - logging.exception("wiki: MAP failed for doc %s batch %d", doc_id, batch_no) + logging.exception( + "wiki: MAP failed for doc %s batch %d", + doc_id, + batch_no, + ) stats["status"] = "error" finally: map_queue.task_done() @@ -902,6 +1234,7 @@ async def run_wiki( delta["deleted"], stats["had_delta"], ) + # 5. REDUCE / PLAN / REFINE KB-wide. Each phase has its own # input_hash gate (REDUCE keys off the MAP-state hash, PLAN off # REDUCE's hash, REFINE off PLAN's hash) so re-runs without an diff --git a/rag/svr/task_executor_refactor/task_handler.py b/rag/svr/task_executor_refactor/task_handler.py index 3d6b507f4e..5bb0cb0cf4 100644 --- a/rag/svr/task_executor_refactor/task_handler.py +++ b/rag/svr/task_executor_refactor/task_handler.py @@ -154,6 +154,10 @@ class TaskHandler: @staticmethod def _is_standard_chunking_task(task_type: str) -> bool: + from rag.svr.task_executor_refactor.dataset_structure_merger import ( + STRUCTURE_MERGE_TASK_TYPES, + ) + task_type = (task_type or "").lower() return task_type not in { "memory", @@ -165,7 +169,7 @@ class TaskHandler: "evaluation", "reembedding", "clone", - } and not task_type.startswith("dataflow") + } | STRUCTURE_MERGE_TASK_TYPES and not task_type.startswith("dataflow") async def handle_task(self) -> None: try: @@ -241,6 +245,10 @@ class TaskHandler: return # Route to appropriate handler + from rag.svr.task_executor_refactor.dataset_structure_merger import ( + is_structure_merge_task, + ) + if task_type == "raptor": await self._run_raptor(embedding_model, vector_size) elif task_type == "graphrag": @@ -267,6 +275,12 @@ class TaskHandler: embedding_model, self._load_chunks_for_doc, ) + elif is_structure_merge_task(task_type): + from rag.svr.task_executor_refactor.dataset_structure_merger import ( + run_structure_merge, + ) + + await run_structure_merge(self._task_context) elif task_type == "evaluation": await self._run_evaluation() elif task_type == "reembedding": diff --git a/rag/utils/infinity_conn.py b/rag/utils/infinity_conn.py index 7cb571fa86..aac207e608 100644 --- a/rag/utils/infinity_conn.py +++ b/rag/utils/infinity_conn.py @@ -35,7 +35,11 @@ class InfinityConnection(InfinityConnectionBase): @staticmethod def field_keyword(field_name: str): # Treat "*_kwd" tag-like columns as keyword lists except knowledge_graph_kwd; source_id is also keyword-like. - if field_name == "source_id" or (field_name.endswith("_kwd") and field_name not in ["knowledge_graph_kwd", "docnm_kwd", "important_kwd", "question_kwd"]): + # source_doc_ids / source_chunk_ids are multi-valued provenance lists (artifact/wiki rows) and must be + # stored/read/updated as keyword lists so the delete-time ref-count (remove one id, drop row when empty) works. + if field_name in ("source_id", "source_doc_ids", "source_chunk_ids") or ( + field_name.endswith("_kwd") and field_name not in ["knowledge_graph_kwd", "docnm_kwd", "important_kwd", "question_kwd"] + ): return True return False