Refactor: merge dataset scope graph. (#17526)

### Summary

merge dataset scope graph.

---------

Co-authored-by: Yingfeng Zhang <yingfeng.zhang@gmail.com>

api/apps/services/dataset_api_service.py

@@ -1833,45 +1833,86 @@ async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str, keyw
             "kind": row_kind or template_kind_cache.get(tid) or resolved_kind,
         }
 
-    # ── Discovery: dataset_graph blob rows, metadata only (no huge content). ──
-    # Keep only rows whose TOP-LEVEL kind matches the request. Raw entity/relation
-    # rows stamp ``compilation_template_kind_kwd`` with config.kind (which folds the
-    # knowledge_graph family), so we scope raw-row queries by template id — resolved
-    # here from the blobs, whose stamp is the top-level kind — not by kind directly.
-    meta_fields = ["compile_kwd", "compilation_template_ids", "compilation_template_kind_kwd"]
-    try:
-        res = await thread_pool_exec(
-            settings.docStoreConn.search,
-            meta_fields,
-            [],
-            {"knowledge_graph_kwd": [_DATASET_STRUCTURE_ROW_KWD]},
-            [],
-            OrderByExpr(),
-            0,
-            1000,
-            index_nm,
-            [dataset_id],
-        )
-        meta_rows = settings.docStoreConn.get_fields(res, meta_fields) or {}
-    except Exception:
-        logging.exception("get_dataset_structure: docStore discovery failed for kb=%s", dataset_id)
-        return True, empty
-
+    # ── Discovery: the dataset-scoped rows the structure merge writes. ──
+    # ``run_structure_merge`` (rag.svr.task_executor_refactor.dataset_structure_merger)
+    # writes merged ``knowledge_graph_kwd="entity"/"relation"`` rows with
+    # ``scope_kwd="dataset"``, stamped with ``compilation_template_ids`` and the
+    # top-level ``compilation_template_kind_kwd``. Page through them (metadata
+    # fields only) to enumerate the distinct template ids whose top-level kind
+    # matches the request; ``build_bucket`` below then reads each template's rows.
+    meta_fields = ["id", "compile_kwd", "compilation_template_ids", "compilation_template_kind_kwd"]
     kind_template_ids: list[str] = []
     seen_tid: set[str] = set()
     has_templateless = False
-    for row in meta_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
-        if tid:
-            if tid not in seen_tid:
-                seen_tid.add(tid)
-                kind_template_ids.append(tid)
-        else:
-            has_templateless = True
+    offset = 0
+    page_size = 1000
+    pages = 0
+    while True:
+        pages += 1
+        try:
+            res = await thread_pool_exec(
+                settings.docStoreConn.search,
+                meta_fields,
+                [],
+                {"knowledge_graph_kwd": ["entity", "relation"], "scope_kwd": ["dataset"]},
+                [],
+                OrderByExpr(),
+                offset,
+                page_size,
+                index_nm,
+                [dataset_id],
+            )
+            meta_rows = settings.docStoreConn.get_fields(res, meta_fields) or {}
+        except Exception:
+            logging.exception("get_dataset_structure: docStore discovery failed for kb=%s", dataset_id)
+            return True, empty
+        if not meta_rows:
+            break
+        for row in meta_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
+            if tid:
+                if tid not in seen_tid:
+                    seen_tid.add(tid)
+                    kind_template_ids.append(tid)
+            else:
+                has_templateless = True
+        if len(meta_rows) < page_size:
+            break
+        offset += page_size
+
+    # Detect datasets that have ONLY the legacy dataset_graph blob (no
+    # entity/relation rows yet) so the fallback path below handles them.
+    if not kind_template_ids and not has_templateless:
+        try:
+            legacy_check = await thread_pool_exec(
+                settings.docStoreConn.search,
+                ["id"],
+                [],
+                {"knowledge_graph_kwd": [_DATASET_STRUCTURE_ROW_KWD]},
+                [],
+                OrderByExpr(),
+                0,
+                1,
+                index_nm,
+                [dataset_id],
+            )
+            legacy_fm = settings.docStoreConn.get_fields(legacy_check, ["id"]) or {}
+            if legacy_fm:
+                has_templateless = True
+        except Exception:
+            pass
+
+    logging.debug(
+        "get_dataset_structure: discovered %d template(s) in %d page(s) for kb=%s kind=%s",
+        len(kind_template_ids),
+        pages,
+        dataset_id,
+        resolved_kind,
+    )
 
     # ── keywords mode: global KNN across the kind → top-1's focused subgraph. ──
     if keywords:
@@ -1888,18 +1929,20 @@ async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str, keyw
             tid = _row_template_id(row) or ""
             stamped = (row.get("compilation_template_kind_kwd") or "").strip()
             meta = _bucket_meta_for(tid, stamped) if tid else {"template_id": f"kind:{resolved_kind}", "template_name": f"kind:{resolved_kind}", "kind": resolved_kind}
-            return meta, {"compilation_template_ids": [tid]} if tid else {"compilation_template_kind_kwd": [stamped]}
+            if tid:
+                return meta, {"compilation_template_ids": [tid], "scope_kwd": ["dataset"]}
+            return meta, {"compilation_template_kind_kwd": [stamped], "scope_kwd": ["dataset"]}
 
         bucket_meta, kw_entities, kw_relations = await sgc.keyword_subgraph(
             index_nm,
             dataset_id,
             embd_mdl,
-            {"compilation_template_ids": kind_template_ids, "knowledge_graph_kwd": ["entity"]},
+            {"compilation_template_ids": kind_template_ids, "knowledge_graph_kwd": ["entity"], "scope_kwd": ["dataset"]},
             keywords,
             _scope_for_template,
             log_ctx=f"kb={dataset_id}",
         )
-        if resolved_kind == "mind_map":
+        if resolved_kind in {"knowledge_graph", "mind_map", "timeline"}:
             kw_entities = sgc.filter_entities_with_relations(kw_entities, kw_relations)
             if not kw_entities:
                 return True, empty
@@ -1914,11 +1957,11 @@ async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str, keyw
     templates_out: list[dict] = []
     for tid in kind_template_ids:
         try:
-            entities, relations = await sgc.build_bucket(index_nm, dataset_id, {"compilation_template_ids": [tid]})
+            entities, relations = await sgc.build_bucket(index_nm, dataset_id, {"compilation_template_ids": [tid], "scope_kwd": ["dataset"]})
         except Exception:
             logging.exception("get_dataset_structure: bucket build failed for kb=%s template=%s", dataset_id, tid)
             continue
-        if resolved_kind == "mind_map":
+        if resolved_kind in {"knowledge_graph", "mind_map", "timeline"}:
             entities = sgc.filter_entities_with_relations(entities, relations)
             if not entities:
                 continue
@@ -1960,10 +2003,10 @@ async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str, keyw
                 continue
             legacy_bucket["entities"].extend(graph.get("entities") or [])
             legacy_bucket["relations"].extend(graph.get("relations") or [])
-        if resolved_kind == "mind_map":
+        if resolved_kind in {"knowledge_graph", "mind_map", "timeline"}:
             legacy_bucket["entities"] = sgc.filter_entities_with_relations(legacy_bucket["entities"], legacy_bucket["relations"])
         if legacy_bucket["entities"] or legacy_bucket["relations"]:
-            if resolved_kind != "mind_map" or legacy_bucket["entities"]:
+            if resolved_kind not in {"knowledge_graph", "mind_map", "timeline"} or legacy_bucket["entities"]:
                 templates_out.append(legacy_bucket)
 
     return True, {"kind": kind, "templates": templates_out}

api/apps/services/structure_graph_common.py

@@ -133,11 +133,65 @@ def dedup_entities(entities: list[dict]) -> list[dict]:
     return out
 
 
+def _entity_response_id(entity: dict) -> str:
+    for field in ("id", "name", "slug"):
+        value = entity.get(field)
+        if isinstance(value, str) and value.strip():
+            return value.strip()
+    return ""
+
+
+def _endpoint_terms(value: str) -> list[str]:
+    value = value.strip()
+    if not value:
+        return []
+    return sorted({value, value.lower()})
+
+
+def normalize_relation_endpoints(entities: list[dict], relations: list[dict]) -> list[dict]:
+    """Align relation endpoints to the returned entity ids/names."""
+    if not entities or not relations:
+        return relations
+
+    lookup: dict[str, str] = {}
+    ambiguous: set[str] = set()
+    for entity in entities:
+        response_id = _entity_response_id(entity)
+        if not response_id:
+            continue
+        for field in ("id", "name", "slug"):
+            value = entity.get(field)
+            if not isinstance(value, str) or not value.strip():
+                continue
+            key = value.strip().lower()
+            if key in lookup and lookup[key] != response_id:
+                ambiguous.add(key)
+                continue
+            lookup[key] = response_id
+    for key in ambiguous:
+        lookup.pop(key, None)
+
+    normalized: list[dict] = []
+    for relation in relations:
+        if not isinstance(relation, dict):
+            continue
+        item = dict(relation)
+        for field in ("from", "to"):
+            value = item.get(field)
+            if isinstance(value, str):
+                item[field] = lookup.get(value.strip().lower(), value)
+        normalized.append(item)
+    return normalized
+
+
 def filter_entities_with_relations(entities: list[dict], relations: list[dict]) -> list[dict]:
     """Keep only entities that are referenced by at least one relation."""
     if not entities or not relations:
         return []
 
+    # Match case-insensitively: the dataset-scoped merge lowercases relation
+    # endpoints while entity names keep their original case, so exact matching
+    # would drop connected nodes from graph-like views.
     connected: set[str] = set()
     for relation in relations:
         if not isinstance(relation, dict):
@@ -145,7 +199,7 @@ def filter_entities_with_relations(entities: list[dict], relations: list[dict])
         for endpoint_key in ("from", "to"):
             endpoint = relation.get(endpoint_key)
             if isinstance(endpoint, str):
-                endpoint = endpoint.strip()
+                endpoint = endpoint.strip().lower()
                 if endpoint:
                     connected.add(endpoint)
 
@@ -164,7 +218,7 @@ def filter_entities_with_relations(entities: list[dict], relations: list[dict])
         for field in ("id", "name", "slug"):
             value = entity.get(field)
             if isinstance(value, str):
-                value = value.strip()
+                value = value.strip().lower()
                 if value:
                     keys.add(value)
         if keys & connected:
@@ -198,7 +252,8 @@ async def build_bucket(index_name, kb_id, scope: dict) -> tuple[list[dict], list
                 node = project_entity(row)
                 if node:
                     entities.append(node)
-        return dedup_entities(entities), relations
+        entities = dedup_entities(entities)
+        return entities, normalize_relation_endpoints(entities, relations)
 
     # Large bucket: sample. A = top entities by mention_count_int desc.
     order_by = OrderByExpr()
@@ -209,12 +264,13 @@ async def build_bucket(index_name, kb_id, scope: dict) -> tuple[list[dict], list
     ent_a_map, _ = await graph_search(index_name, kb_id, GRAPH_ENTITY_FIELDS, dict(scope, knowledge_graph_kwd=["entity"]), order_by, GRAPH_TOP_ENTITIES)
     set_a = [n for n in (project_entity(r) for r in ent_a_map.values()) if n]
     a_names = sorted({str(e.get("name") or "").strip() for e in set_a if str(e.get("name") or "").strip()})
+    a_name_terms = sorted({term for name in a_names for term in _endpoint_terms(name)})
 
     # relations whose source is one of A.
     relations = []
     target_names_lower: set[str] = set()
-    if a_names:
-        rel_map, _ = await graph_search(index_name, kb_id, GRAPH_RELATION_FIELDS, dict(scope, knowledge_graph_kwd=["relation"], from_entity_kwd=a_names), OrderByExpr(), GRAPH_EXPANSION_CAP)
+    if a_name_terms:
+        rel_map, _ = await graph_search(index_name, kb_id, GRAPH_RELATION_FIELDS, dict(scope, knowledge_graph_kwd=["relation"], from_entity_kwd=a_name_terms), OrderByExpr(), GRAPH_EXPANSION_CAP)
         for row in rel_map.values():
             edge = project_relation(row)
             if edge:
@@ -229,7 +285,8 @@ async def build_bucket(index_name, kb_id, scope: dict) -> tuple[list[dict], list
         tgt_map, _ = await graph_search(index_name, kb_id, GRAPH_ENTITY_FIELDS, dict(scope, knowledge_graph_kwd=["entity"], name_kwd=sorted(target_names_lower)), OrderByExpr(), GRAPH_EXPANSION_CAP)
         set_t = [n for n in (project_entity(r) for r in tgt_map.values()) if n]
 
-    return dedup_entities(set_a + set_t), relations
+    entities = dedup_entities(set_a + set_t)
+    return entities, normalize_relation_endpoints(entities, relations)
 
 
 async def keyword_subgraph(index_name, kb_id, embd_mdl, base_entity_condition, keywords, scope_for_template, log_ctx="") -> tuple[dict | None, list[dict], list[dict]]:
@@ -280,8 +337,9 @@ async def keyword_subgraph(index_name, kb_id, embd_mdl, base_entity_condition, k
     relations: list[dict] = []
     seen_rel: set[tuple[str, str, str]] = set()
     neighbor_names_lower: set[str] = set()
+    top_name_terms = _endpoint_terms(top_name)
     for field in ("from_entity_kwd", "to_entity_kwd"):
-        rel_map, _ = await graph_search(index_name, kb_id, GRAPH_RELATION_FIELDS, dict(scope, knowledge_graph_kwd=["relation"], **{field: [top_name]}), OrderByExpr(), GRAPH_EXPANSION_CAP)
+        rel_map, _ = await graph_search(index_name, kb_id, GRAPH_RELATION_FIELDS, dict(scope, knowledge_graph_kwd=["relation"], **{field: top_name_terms}), OrderByExpr(), GRAPH_EXPANSION_CAP)
         for row in rel_map.values():
             edge = project_relation(row)
             if not edge:
@@ -301,4 +359,5 @@ async def keyword_subgraph(index_name, kb_id, embd_mdl, base_entity_condition, k
         nb_map, _ = await graph_search(index_name, kb_id, GRAPH_ENTITY_FIELDS, dict(scope, knowledge_graph_kwd=["entity"], name_kwd=sorted(neighbor_names_lower)), OrderByExpr(), GRAPH_EXPANSION_CAP)
         entities.extend(n for n in (project_entity(r) for r in nb_map.values()) if n)
 
-    return bucket_meta, dedup_entities(entities), relations
+    entities = dedup_entities(entities)
+    return bucket_meta, entities, normalize_relation_endpoints(entities, relations)

api/db/services/document_service.py

@@ -512,6 +512,20 @@ class DocumentService(CommonService):
         except Exception as e:
             logging.error(f"Failed to delete chunks from doc store for document {doc.id}: {e}")
 
+        # Record doc deletion for incremental structure-merge ghost cleanup.
+        # Runs after the doc_id sweep so the marker (stored under
+        # deleted_doc_id to avoid matching the same sweep) survives.
+        try:
+            from rag.svr.task_executor_refactor.dataset_structure_merger import (
+                record_doc_deletion,
+            )
+
+            record_doc_deletion(tenant_id, doc.kb_id, doc.id)
+        except Exception as e:
+            logging.warning(
+                f"Failed to record doc deletion for structure merge: {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.