diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index cb9f05af5e..b92a6180ab 100644 --- a/api/apps/services/dataset_api_service.py +++ b/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} diff --git a/api/apps/services/structure_graph_common.py b/api/apps/services/structure_graph_common.py index 32a4e59f4d..1cedb546f5 100644 --- a/api/apps/services/structure_graph_common.py +++ b/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) diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index 3f75f94b7d..660e90c9c0 100644 --- a/api/db/services/document_service.py +++ b/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. diff --git a/conf/infinity_mapping.json b/conf/infinity_mapping.json index 8dc4ed02f0..c8a0b1bfa3 100644 --- a/conf/infinity_mapping.json +++ b/conf/infinity_mapping.json @@ -45,6 +45,7 @@ "extra": {"type": "varchar", "default": ""}, "compile_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, + "scope_kwd": {"type": "varchar", "default": "doc", "analyzer": "whitespace-#"}, "source_chunk_ids": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, "source_doc_ids": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, "compilation_template_ids": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, diff --git a/rag/advanced_rag/knowlege_compile/structure.py b/rag/advanced_rag/knowlege_compile/structure.py index d95a98b0a1..bc45452bf7 100644 --- a/rag/advanced_rag/knowlege_compile/structure.py +++ b/rag/advanced_rag/knowlege_compile/structure.py @@ -572,6 +572,8 @@ def _struct_to_doc_storage_doc( target_field: str | None = None, compilation_template_id: str | None = None, compilation_template_kind: str | None = None, + scope: str = "doc", + doc_ids: list[str] | None = None, ) -> dict: """Build one ES doc for an extracted entity or relation. @@ -582,11 +584,11 @@ def _struct_to_doc_storage_doc( ``from_entity_kwd`` / ``to_entity_kwd``. compilation_template_id / compilation_template_kind: stamped onto every row so the document-structure endpoint can group by - template id and the UI can render one tab per template. The - id is stored as a single-element list under - ``compilation_template_ids`` because the same logical entity - *could* later be claimed by multiple templates during a - cross-template merge (rare, but the schema is forward-compat). + template id and the UI can render one tab per template. + scope: ``"doc"`` for document-level rows, ``"dataset"`` for the KB-wide + merged entity rows written by the Build button. + doc_ids: only set for ``scope="dataset"`` rows — lists the documents + that contributed to this merged entity. """ content_with_weight = json.dumps(payload, ensure_ascii=False) if hasattr(vec, "tolist"): @@ -602,13 +604,18 @@ def _struct_to_doc_storage_doc( # Mix the template id into the stable row id so two templates with the # same compile_kwd don't collide on identical payloads (e.g. two # different list-kind templates that each extract "headline X"). + # Dataset-scope rows get the scope in the row seed to avoid colliding + # with doc-scope rows for the same entity. row_seed_extras = [template_id_str] if template_id_str else [] + if scope == "dataset": + row_seed_extras.append("dataset") row_id = _stable_row_id(content_with_weight, doc_id_str, *row_seed_extras) doc = { "content_with_weight": content_with_weight, "compile_kwd": compile_kwd, "knowledge_graph_kwd": kind, + "scope_kwd": scope, "doc_id": doc_id_str, "source_chunk_ids": list(chunk_ids or []), "content_ltks": content_ltks, @@ -616,6 +623,8 @@ def _struct_to_doc_storage_doc( f"q_{len(vec_list)}_vec": vec_list, "id": row_id, } + if scope == "dataset" and doc_ids: + doc["doc_ids_kwd"] = list(doc_ids) # Surface two payload fields as queryable top-level columns so the store can # filter/sort on them without parsing ``content_with_weight``. Both are @@ -2226,44 +2235,116 @@ async def _struct_upsert_dataset_graph_json( compile_kwd: str, compilation_template_id: str | None = None, structure_kind: str | None = None, + embd_mdl=None, ) -> None: + """Write dataset-level entity/relation rows from a merged graph. + + Replaces the old ``dataset_graph`` JSON blob with individual searchable + entity/relation rows (``scope_kwd="dataset"``). Each row carries its own + embedding and tokenized text, so both KNN and full-text search can hit it. + """ from common import settings from rag.nlp import search as _rag_search + from ._common import encode as _encode, tokenize_for_search as _tokenize_for_search, stable_row_id as _stable_row_id 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), + kb_id_str = str(kb_id) + + # Write a single metadata row so the artifacts_structure discovery endpoint + # can find this template without scanning entity rows. + meta_id = _dataset_struct_graph_row_id(kb_id, compile_kwd, compilation_template_id) + meta_row = { + "id": meta_id, "compile_kwd": compile_kwd, "knowledge_graph_kwd": "dataset_graph", - "doc_id": kb_id, - "kb_id": kb_id, + "scope_kwd": "dataset", + "doc_id": kb_id_str, + "kb_id": kb_id_str, "available_int": 0, + "compilation_template_ids": [compilation_template_id] if compilation_template_id else [], } - 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]) + meta_row["compilation_template_kind_kwd"] = str(structure_kind) + old = await thread_pool_exec(settings.docStoreConn.get, meta_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, - ) + await thread_pool_exec(settings.docStoreConn.update, {"id": meta_id}, {k: v for k, v in meta_row.items() if k != "id"}, index, kb_id) else: - await thread_pool_exec(settings.docStoreConn.insert, [row], index, kb_id) + await thread_pool_exec(settings.docStoreConn.insert, [meta_row], index, kb_id) + + # Write individual entity/relation rows (scope_kwd="dataset", searchable). + rows = [] + for ent in graph.get("entities") or []: + payload = {"name": ent.get("name", ""), "type": ent.get("type", "other"), "description": ent.get("description", "")} + ent_name = (ent.get("name") or "").strip() + desc = ent.get("description") or ent_name + ltks, sm_ltks = _tokenize_for_search(desc) + mention_count = ent.get("mention_count", 1) + source_chunk_ids = ent.get("source_chunk_ids") or [] + doc_ids = ent.get("doc_ids_kwd") or [] + row_id = _stable_row_id(ent_name.lower(), kb_id_str, compile_kwd, compilation_template_id or "", "dataset") + row = { + "id": row_id, + "content_with_weight": json.dumps(payload, ensure_ascii=False), + "compile_kwd": compile_kwd, + "knowledge_graph_kwd": "entity", + "scope_kwd": "dataset", + "doc_id": kb_id_str, + "kb_id": kb_id_str, + "source_chunk_ids": source_chunk_ids, + "content_ltks": ltks, + "content_sm_ltks": sm_ltks, + "mention_count_int": mention_count, + "name_kwd": ent_name.lower(), + "available_int": 1, + } + if compilation_template_id: + row["compilation_template_ids"] = [compilation_template_id] + if structure_kind: + row["compilation_template_kind_kwd"] = str(structure_kind) + if doc_ids: + row["doc_ids_kwd"] = doc_ids + # Re-embed: use embd_mdl if available, otherwise skip the vector. + if embd_mdl: + vecs = await _encode(embd_mdl, [desc]) + if vecs and len(vecs[0]) > 0: + dim = len(vecs[0]) + row[f"q_{dim}_vec"] = list(vecs[0]) + rows.append(row) + + for rel in graph.get("relations") or []: + src = str(rel.get("from", "")).strip() + tgt = str(rel.get("to", "")).strip() + if not src or not tgt: + continue + rel_type = str(rel.get("type", "related")).strip() + payload = {"from": src, "to": tgt, "type": rel_type} + desc = f"{src} {rel_type} {tgt}" + ltks, sm_ltks = _tokenize_for_search(desc) + rel_key = f"{src.lower()} -> {rel_type.lower()} -> {tgt.lower()}" + row_id = _stable_row_id(rel_key, kb_id_str, compile_kwd, compilation_template_id or "", "dataset") + doc_ids = rel.get("doc_ids_kwd") or [] + row = { + "id": row_id, + "content_with_weight": json.dumps(payload, ensure_ascii=False), + "compile_kwd": compile_kwd, + "knowledge_graph_kwd": "relation", + "scope_kwd": "dataset", + "doc_id": kb_id_str, + "kb_id": kb_id_str, + "from_entity_kwd": src.lower(), + "to_entity_kwd": tgt.lower(), + "content_ltks": ltks, + "content_sm_ltks": sm_ltks, + "available_int": 1, + } + if compilation_template_id: + row["compilation_template_ids"] = [compilation_template_id] + if doc_ids: + row["doc_ids_kwd"] = doc_ids + rows.append(row) + + if rows: + await thread_pool_exec(settings.docStoreConn.insert, rows, index, kb_id) async def rebuild_dataset_structure_graph_json( @@ -2272,19 +2353,16 @@ async def rebuild_dataset_structure_graph_json( compile_kwd: str, compilation_template_id: str | None = None, structure_kind: str | None = None, + embd_mdl=None, ) -> dict: - """Rebuild and persist the KB-wide (dataset) structure graph for one - (compile_kwd, template_id) pair. + """Rebuild and persist the KB-wide dataset structure graph. - 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. + Reads every document's entity/relation rows in the KB (no ``doc_id`` + filter) and writes individual ``scope_kwd="dataset"`` entity/relation + rows with embeddings and tokenized text — making them searchable. ``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.""" + ``session_graph``); it is stamped on the meta row so the API can filter.""" graph = await _struct_rebuild_graph_json( tenant_id, kb_id, @@ -2299,6 +2377,7 @@ async def rebuild_dataset_structure_graph_json( compile_kwd, compilation_template_id, structure_kind=structure_kind, + embd_mdl=embd_mdl, ) return graph diff --git a/rag/svr/task_executor_refactor/dataset_structure_merger.py b/rag/svr/task_executor_refactor/dataset_structure_merger.py index 47d93b64e0..edc06c075a 100644 --- a/rag/svr/task_executor_refactor/dataset_structure_merger.py +++ b/rag/svr/task_executor_refactor/dataset_structure_merger.py @@ -14,143 +14,821 @@ # limitations under the License. # -"""KB-wide structure-graph merge task. +"""KB-wide structure-graph merge task (incremental, bucketed). -Runs when the user POSTs to ``/datasets//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. +Triggered when the user POSTs to ``/datasets//index`` with a structure +index type. Scans ``scope_kwd="doc"`` entity rows grouped by ``(name, type)``, +merges them into ``scope_kwd="dataset"`` rows with bucketing for bounded memory. -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. +Key design: + - 256 hash buckets — all rows accumulated in memory, merged once after scan + - Incremental: only processes doc_graph rows changed since ``last_build_time`` + - Document deletions tracked via ``record_doc_deletion()`` – ghost entities + are cleaned up incrementally at the start of the next build + - Template changes still require a full rebuild + - Resulting dataset_graph rows are searchable (``available_int=1``) """ from __future__ import annotations +import json import logging from typing import Optional +import xxhash + 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.advanced_rag.knowlege_compile._common import ( + encode as _encode, + tokenize_for_search as _tokenize_for_search, + stable_row_id as _stable_row_id, ) from rag.svr.task_executor_refactor.task_context import TaskContext +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- -# 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 -} +_BUCKET_COUNT = 256 +_BUCKET_FLUSH_THRESHOLD = 500 +_PAGE_SIZE = 1000 +_SCOPE_KWD_DOC = "doc" +_SCOPE_KWD_DATASET = "dataset" +_DELETION_META_KWD = "doc_deleted" -STRUCTURE_MERGE_TASK_TYPES = frozenset(_STRUCTURE_TASK_TYPE_TO_KIND) +# Row id seeds for metadata +_META_ROW_KWD = "kg_build_meta" + +STRUCTURE_MERGE_TASK_TYPES = frozenset( + { + "structure_graph", + "structure_mindmap", + "timeline", + "session_graph", + "session_essence", + "structure", + } +) 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``).""" +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _index_name(tenant_id: str) -> str: + return search.index_name(tenant_id) + + +def _meta_row_id(kb_id: str, compile_kwd: str, template_id: str | None) -> str: + return xxhash.xxh64( + f"{_META_ROW_KWD}:{kb_id}:{compile_kwd}:{template_id or ''}".encode("utf-8", "surrogatepass"), + ).hexdigest() + + +def _dataset_entity_row_id(kb_id: str, compile_kwd: str, template_id: str | None, name: str) -> str: + return _stable_row_id(name, kb_id, compile_kwd, template_id or "", _SCOPE_KWD_DATASET) + + +def _dataset_relation_row_id(kb_id: str, compile_kwd: str, template_id: str | None, src: str, tgt: str, rel_type: str) -> str: + key = f"{src.lower()} -> {rel_type.lower()} -> {tgt.lower()}" + return _stable_row_id(key, kb_id, compile_kwd, template_id or "", _SCOPE_KWD_DATASET) + + +# --------------------------------------------------------------------------- +# Document-store I/O +# --------------------------------------------------------------------------- + + +async def _index_search( + tenant_id: str, + kb_id: str, + condition: dict, + fields: list[str], + limit: int = 10000, + offset: int = 0, +) -> list[dict]: from common.doc_store.doc_store_base import OrderByExpr - index = search.index_name(tenant_id) + index = _index_name(tenant_id) if not settings.docStoreConn.index_exist(index, kb_id): - return set() + return [] + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + fields, + [], + condition, + [], + OrderByExpr(), + offset, + limit, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, fields) or {} + return list(field_map.values()) + except Exception: + logging.exception("structure_merge: search failed for kb=%s", kb_id) + return [] - select_fields = ["id", "compile_kwd", "compilation_template_ids"] - pairs: set[tuple[str, str]] = set() + +async def _index_delete(tenant_id: str, kb_id: str, condition: dict) -> None: + index = _index_name(tenant_id) + try: + await thread_pool_exec(settings.docStoreConn.delete, condition, index, kb_id) + except Exception: + logging.exception("structure_merge: delete failed for kb=%s cond=%s", kb_id, condition) + + +async def _index_insert(tenant_id: str, kb_id: str, rows: list[dict]) -> None: + if not rows: + return + index = _index_name(tenant_id) + try: + await thread_pool_exec(settings.docStoreConn.insert, rows, index, kb_id) + except Exception: + logging.exception("structure_merge: insert failed for kb=%s", kb_id) + + +# --------------------------------------------------------------------------- +# Document deletion tracking (sync, called from DocumentService.remove_document) +# --------------------------------------------------------------------------- + + +def record_doc_deletion(tenant_id: str, kb_id: str, doc_id: str) -> None: + """Record a document deletion for incremental ghost cleanup. + + Creates a lightweight ES meta row so that the next incremental merge build + can look up which docs were deleted and clean up orphaned dataset-level + entity rows (those whose *only* source document was the deleted doc). + + The deleted doc ID is stored in **deleted_doc_id** (not ``doc_id``) so the + row survives the ``doc_id``-based chunk sweep that runs in + ``DocumentService.remove_document``. + + Must be callable from a synchronous context (``DocumentService.remove_document``). + """ + import datetime + + try: + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return + row = { + "id": f"{_DELETION_META_KWD}:{kb_id}:{doc_id}", + "kb_id": kb_id, + "deleted_doc_id": doc_id, + "knowledge_graph_kwd": _DELETION_META_KWD, + "compile_kwd": "*", + "create_timestamp_flt": datetime.datetime.now().timestamp(), + } + settings.docStoreConn.insert([row], index, kb_id) + except Exception: + logging.exception("structure_merge: failed to record doc deletion kb=%s doc=%s", kb_id, doc_id) + + +# --------------------------------------------------------------------------- +# Merge logic +# --------------------------------------------------------------------------- + + +async def _load_last_build_time(tenant_id: str, kb_id: str, compile_kwd: str, template_id: str | None) -> float | None: + """Read the last build timestamp from a metadata row.""" + meta_id = _meta_row_id(kb_id, compile_kwd, template_id) + index = _index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return None + try: + row = await thread_pool_exec(settings.docStoreConn.get, meta_id, index, [kb_id]) + if row: + ts = row.get("build_timestamp_flt") + if ts is not None and isinstance(ts, (int, float)) and ts > 0: + return float(ts) + except Exception: + logging.exception("structure_merge: failed to load build time for kb=%s compile=%s", kb_id, compile_kwd) + return None + + +async def _save_build_time(tenant_id: str, kb_id: str, compile_kwd: str, template_id: str | None, timestamp: float) -> None: + """Persist the build timestamp.""" + import datetime + + meta_id = _meta_row_id(kb_id, compile_kwd, template_id) + index = _index_name(tenant_id) + payload = { + "id": meta_id, + "kb_id": kb_id, + "doc_id": kb_id, + "compile_kwd": compile_kwd, + "knowledge_graph_kwd": _META_ROW_KWD, + "build_timestamp_flt": timestamp, + "create_time": datetime.datetime.now().isoformat(), + } + existing = await thread_pool_exec(settings.docStoreConn.get, meta_id, index, [kb_id]) + if existing: + await thread_pool_exec(settings.docStoreConn.update, {"id": meta_id}, payload, index, kb_id) + else: + await thread_pool_exec(settings.docStoreConn.insert, [payload], index, kb_id) + + +def _bucket_id(name: str, bucket_count: int = _BUCKET_COUNT) -> int: + """Determine the hash bucket for an entity name.""" + return xxhash.xxh64(name.encode("utf-8", "surrogatepass")).intdigest() % bucket_count + + +async def _merge_bucket( + tenant_id: str, + kb_id: str, + compile_kwd: str, + template_id: str | None, + bucket_rows: list[dict], + embd_mdl, + structure_kind: str | None = None, +) -> list[dict]: + """Flush one bucket: group by (name, type), merge, build dataset rows. + + Returns the list of entity/relation rows to insert. + """ + from rag.advanced_rag.knowlege_compile.structure import _struct_entity_name + + # Group by (name, type) + groups: dict[tuple[str, str], list[dict]] = {} + for row in bucket_rows: + payload = json.loads(row.get("content_with_weight") or "{}") + name = _struct_entity_name(payload) or "" + typ = payload.get("type", "other") + key = (name.lower(), str(typ).strip()) + groups.setdefault(key, []).append(row) + + rows_out: list[dict] = [] + for (name_lower, typ), rows in groups.items(): + # Collect all source_chunk_ids, doc_ids, and original-cased names + all_chunks: list[str] = [] + all_docs: list[str] = [] + all_descriptions: list[str] = [] + all_original_names: list[str] = [] + mention_count = 0 + for r in rows: + payload = json.loads(r.get("content_with_weight") or "{}") + chunks = r.get("source_chunk_ids") or [] + for c in chunks: + if c not in all_chunks: + all_chunks.append(c) + doc = r.get("doc_id") or "" + if doc and doc not in all_docs: + all_docs.append(doc) + mention_count += int(r.get("mention_count_int") or payload.get("mention_count", 1)) + desc = payload.get("description", "") + if desc and desc not in all_descriptions: + all_descriptions.append(desc) + orig = _struct_entity_name(payload) or "" + if orig and orig not in all_original_names: + all_original_names.append(orig) + + # Use the longest description (most complete) and best original-cased name + best_desc = max(all_descriptions, key=len) if all_descriptions else name_lower + best_orig_name = max(all_original_names, key=len) if all_original_names else name_lower + payload_out = {"name": best_orig_name, "type": typ, "description": best_desc, "mention_count": mention_count} + ltks, sm_ltks = _tokenize_for_search(best_desc) + + row_id = _dataset_entity_row_id(kb_id, compile_kwd, template_id, name_lower) + row = { + "id": row_id, + "content_with_weight": json.dumps(payload_out, ensure_ascii=False), + "compile_kwd": compile_kwd, + "knowledge_graph_kwd": "entity", + "scope_kwd": _SCOPE_KWD_DATASET, + "doc_id": kb_id, + "kb_id": kb_id, + "name_kwd": name_lower, + "source_chunk_ids": all_chunks, + "doc_ids_kwd": all_docs, + "content_ltks": ltks, + "content_sm_ltks": sm_ltks, + "mention_count_int": mention_count, + "available_int": 1, + } + if template_id: + row["compilation_template_ids"] = [template_id] + if structure_kind: + row["compilation_template_kind_kwd"] = str(structure_kind) + # Re-embed + if embd_mdl: + vecs = await _encode(embd_mdl, [best_desc]) + if vecs and len(vecs[0]) > 0: + dim = len(vecs[0]) + row[f"q_{dim}_vec"] = list(vecs[0]) + rows_out.append(row) + + return rows_out + + +async def _merge_relations( + tenant_id: str, + kb_id: str, + compile_kwd: str, + template_id: str | None, + rel_bucket_rows: list[dict], + embd_mdl, + structure_kind: str | None = None, +) -> list[dict]: + """Group doc-graph relation rows by (src, rel_type, tgt) and merge. + + Each group produces one dataset-scoped relation row. Duplicate + ``source_chunk_ids`` and ``doc_ids`` are deduplicated per group. + When *embd_mdl* is provided the best description is re-embedded and + stored as ``q__vec``, matching the entity path. + """ + groups: dict[tuple[str, str, str], list[dict]] = {} + for row in rel_bucket_rows: + payload = json.loads(row.get("content_with_weight") or "{}") + src = (payload.get("from") or row.get("from_entity_kwd") or "").strip().lower() + tgt = (payload.get("to") or row.get("to_entity_kwd") or "").strip().lower() + rel_type = (payload.get("type") or "related").strip().lower() + key = (src, rel_type, tgt) + groups.setdefault(key, []).append(row) + + rows_out: list[dict] = [] + for (src, rel_type, tgt), rows in groups.items(): + all_chunks: list[str] = [] + all_docs: list[str] = [] + # Use the longest description for tokenization + all_desc: list[str] = [] + for r in rows: + payload = json.loads(r.get("content_with_weight") or "{}") + chunks = r.get("source_chunk_ids") or [] + for c in chunks: + if c not in all_chunks: + all_chunks.append(c) + doc = r.get("doc_id") or "" + if doc and doc not in all_docs: + all_docs.append(doc) + desc = payload.get("description") or f"{src} {rel_type} {tgt}" + if desc not in all_desc: + all_desc.append(desc) + best_desc = max(all_desc, key=len) + ltks, sm_ltks = _tokenize_for_search(best_desc) + + row_id = _dataset_relation_row_id(kb_id, compile_kwd, template_id, src, tgt, rel_type) + row = { + "id": row_id, + "content_with_weight": json.dumps({"from": src, "to": tgt, "type": rel_type}, ensure_ascii=False), + "compile_kwd": compile_kwd, + "knowledge_graph_kwd": "relation", + "scope_kwd": _SCOPE_KWD_DATASET, + "doc_id": kb_id, + "kb_id": kb_id, + "from_entity_kwd": src, + "to_entity_kwd": tgt, + "source_chunk_ids": all_chunks, + "doc_ids_kwd": all_docs, + "content_ltks": ltks, + "content_sm_ltks": sm_ltks, + "available_int": 1, + } + if template_id: + row["compilation_template_ids"] = [template_id] + if structure_kind: + row["compilation_template_kind_kwd"] = str(structure_kind) + # Re-embed (matching _merge_bucket) + if embd_mdl: + vecs = await _encode(embd_mdl, [best_desc]) + if vecs and len(vecs[0]) > 0: + dim = len(vecs[0]) + row[f"q_{dim}_vec"] = list(vecs[0]) + rows_out.append(row) + + return rows_out + + +# --------------------------------------------------------------------------- +# Deletion-driven ghost cleanup +# --------------------------------------------------------------------------- + + +async def _cleanup_deleted_docs( + tenant_id: str, + kb_id: str, + compile_kwd: str, + template_id: str | None, +) -> bool: + """Clean dataset-level entity rows whose *only* source doc was deleted. + + Called once per (compile_kwd, template_id) at the beginning of an + incremental build. Works in two passes: + + 1. Read all pending ``doc_deleted`` meta rows for this kb (stored under + ``deleted_doc_id``). + 2. Find dataset entity/relation rows whose ``doc_ids_kwd`` references a + deleted doc. If after removing the deleted doc ID the ``doc_ids_kwd`` + list is empty the row is a **ghost** and gets deleted. + + Shared entities (those still referenced by other live docs) keep their + existing content — the next incremental re-merge will correct any stale + descriptions or mention counts. + + Returns ``True`` when Pass 2 completed without store errors, ``False`` + otherwise. The caller must delete the processed meta rows only after + **all** (compile_kwd, template_id) pairs have been handled successfully. + """ + from common.doc_store.doc_store_base import OrderByExpr + + index = _index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return False + + # ── Pass 1: collect pending deletions ──────────────────────────── + cursor: dict[str, str] = {} # meta_row_id → deleted_doc_id offset = 0 - page_size = 1000 while True: try: res = await thread_pool_exec( settings.docStoreConn.search, - select_fields, + ["id", "deleted_doc_id"], [], - {"knowledge_graph_kwd": ["entity", "relation"]}, + {"kb_id": [kb_id], "knowledge_graph_kwd": [_DELETION_META_KWD]}, [], OrderByExpr(), offset, - page_size, + _PAGE_SIZE, index, [kb_id], ) - field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + fm = settings.docStoreConn.get_fields(res, ["id", "deleted_doc_id"]) or {} except Exception: - logging.exception("structure_merge: failed to scan entity/relation rows for kb=%s", kb_id) break - if not field_map: + if not fm: 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] + for row in fm.values(): + did = row.get("deleted_doc_id") + if did: + cursor[row["id"]] = did + if len(fm) < _PAGE_SIZE: + break + offset += _PAGE_SIZE + + if not cursor: + return True # no markers — nothing to do + deleted_doc_ids = list(set(cursor.values())) + + # ── Pass 2: find affected dataset rows and remove ghosts ───────── + dataset_del_cond: dict = { + "scope_kwd": [_SCOPE_KWD_DATASET], + "compile_kwd": [compile_kwd], + "kb_id": [kb_id], + "doc_ids_kwd": deleted_doc_ids, + } + if template_id: + dataset_del_cond["compilation_template_ids"] = [template_id] + + ids_to_delete: list[str] = [] + ids_to_update: list[tuple[str, list[str]]] = [] # (row_id, new_doc_ids) + pass2_ok = True + offset = 0 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["id", "doc_ids_kwd"], + [], + dataset_del_cond, + [], + OrderByExpr(), + offset, + _PAGE_SIZE, + index, + [kb_id], + ) + fm = settings.docStoreConn.get_fields(res, ["id", "doc_ids_kwd"]) or {} + except Exception: + pass2_ok = False + break + if not fm: + break + for row in fm.values(): + current: list = row.get("doc_ids_kwd") or [] + if not isinstance(current, list): + current = [] + remaining = [d for d in current if d not in deleted_doc_ids] + if not remaining: + ids_to_delete.append(row["id"]) + elif len(remaining) < len(current): + ids_to_update.append((row["id"], remaining)) + if len(fm) < _PAGE_SIZE: + break + offset += _PAGE_SIZE + + if not pass2_ok: + return False # marker rows preserved for retry + + # Execute deletions + if ids_to_delete: + logging.info( + "structure_merge: deleting %d ghost dataset rows for kb=%s compile=%s", + len(ids_to_delete), + kb_id, + compile_kwd, + ) + await _index_delete(tenant_id, kb_id, {"id": ids_to_delete}) + + # Execute updates (remove stale doc_id from doc_ids_kwd) + if ids_to_update: + logging.info( + "structure_merge: updating doc_ids_kwd on %d shared dataset rows for kb=%s compile=%s", + len(ids_to_update), + kb_id, + compile_kwd, + ) + for rid, remaining in ids_to_update: + try: + await thread_pool_exec( + settings.docStoreConn.update, + {"id": [rid]}, + {"doc_ids_kwd": remaining}, + index, + kb_id, + ) + except Exception: + logging.exception( + "structure_merge: failed to update doc_ids_kwd on row %s", + rid, + ) + + # Markers are NOT deleted here — caller (run_structure_merge) deletes + # them once after all (compile_kwd, template_id) pairs succeed. + return True + + +async def _consume_deletion_markers(tenant_id: str, kb_id: str) -> None: + """Remove all pending ``doc_deleted`` meta rows for *kb_id*. + + Called from :func:`run_structure_merge` *after* every eligible + ``(compile_kwd, template_id)`` pair has been processed successfully. + """ + from common.doc_store.doc_store_base import OrderByExpr + + index = _index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return + offset = 0 + all_marker_ids: list[str] = [] + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["id"], + [], + {"kb_id": [kb_id], "knowledge_graph_kwd": [_DELETION_META_KWD]}, + [], + OrderByExpr(), + offset, + _PAGE_SIZE, + index, + [kb_id], + ) + fm = settings.docStoreConn.get_fields(res, ["id"]) or {} + except Exception: + break + if not fm: + break + all_marker_ids.extend(row["id"] for row in fm.values() if row.get("id")) + if len(fm) < _PAGE_SIZE: + break + offset += _PAGE_SIZE + if all_marker_ids: + await _index_delete(tenant_id, kb_id, {"id": all_marker_ids}) + + +async def _do_build( + tenant_id: str, + kb_id: str, + compile_kwd: str, + template_id: str | None, + structure_kind: str | None, + embd_mdl, + incremental: bool = True, +) -> bool: + """Core build logic — read doc_graph rows, bucket, merge, write dataset rows. + + When *incremental* is True, only processes rows changed since the last build + timestamp. When False (full rebuild), reads ALL doc_graph rows for the + given (compile_kwd, template) and deletes existing dataset rows first. + + Document deletions are tracked via :func:`record_doc_deletion` and cleaned + up incrementally (:func:`_cleanup_deleted_docs`) so that deleting a + document no longer forces a full rebuild. + + Returns ``True`` if deletion markers were processed successfully (or no + markers existed), ``False`` if a store error prevented cleanup. + """ + cleanup_ok = True + + index = _index_name(tenant_id) + + # ── Determine the time window ───────────────────────────────────── + last_build_time = None + if incremental: + last_build_time = await _load_last_build_time(tenant_id, kb_id, compile_kwd, template_id) + + # ── Full rebuild: delete existing dataset rows ──────────────────── + if not incremental or last_build_time is None: + del_cond: dict = { + "scope_kwd": [_SCOPE_KWD_DATASET], + "compile_kwd": [compile_kwd], + "kb_id": [kb_id], + } + if template_id: + del_cond["compilation_template_ids"] = [template_id] + offset = 0 + while True: + existing = await _index_search( + tenant_id, + kb_id, + del_cond, + ["id"], + limit=1000, + offset=offset, + ) + if not existing: + break + ids = [r["id"] for r in existing if r.get("id")] + if ids: + await _index_delete(tenant_id, kb_id, {"id": ids}) + if len(existing) < 1000: + break + offset += 1000 + # Also delete the meta rows + meta_id = _meta_row_id(kb_id, compile_kwd, template_id) + try: + await thread_pool_exec(settings.docStoreConn.delete, {"id": [meta_id]}, index, kb_id) + except Exception: + pass + + # ── Incremental: clean up ghosts from deleted docs ─────────────── + if incremental and last_build_time is not None: + cleanup_ok = await _cleanup_deleted_docs(tenant_id, kb_id, compile_kwd, template_id) + + # ── Scan doc_graph rows ────────────────────────────────────────── + # Backward compat: existing doc_graph rows don't have scope_kwd. + # Search for (scope_kwd="doc" OR (not exists scope_kwd)) AND compile_kwd. + base_cond = { + "compile_kwd": [compile_kwd], + "knowledge_graph_kwd": ["entity", "relation"], + } + if template_id: + base_cond["compilation_template_ids"] = [template_id] + scan_all = not incremental or last_build_time is None + if not scan_all: + base_cond["scope_kwd"] = [_SCOPE_KWD_DOC] + + # Build timestamp filter for incremental mode + ts_cond = None + if not scan_all and last_build_time is not None: + ts_cond = {"create_timestamp_flt": {"gte": last_build_time}} + + limit = _PAGE_SIZE + offset = 0 + buckets: list[list] = [[] for _ in range(_BUCKET_COUNT)] + rel_buckets: list[list] = [[] for _ in range(_BUCKET_COUNT)] + + last_page = False + while not last_page: + rows = await _index_search( + tenant_id, + kb_id, + base_cond, + [ + "id", + "content_with_weight", + "source_chunk_ids", + "doc_id", + "name_kwd", + "mention_count_int", + "compilation_template_ids", + "create_timestamp_flt", + "from_entity_kwd", + "to_entity_kwd", + ], + limit=limit, + offset=offset, + ) + if not rows: + break + for row in rows: + # Time filter (incremental mode) + if ts_cond: + ts = row.get("create_timestamp_flt", 0) + if isinstance(ts, (int, float)) and ts < ts_cond["create_timestamp_flt"]["gte"]: + continue + is_rel = row.get("knowledge_graph_kwd") == "relation" or "from_entity_kwd" in row + if is_rel: + # Bucketize relations by stable hash of src/type/tgt + payload = json.loads(row.get("content_with_weight") or "{}") + rsrc = (payload.get("from") or row.get("from_entity_kwd") or "").strip().lower() + rtgt = (payload.get("to") or row.get("to_entity_kwd") or "").strip().lower() + rtype = (payload.get("type") or "related").strip().lower() + rbid = _bucket_id(f"{rsrc}:{rtype}:{rtgt}") + if rbid < _BUCKET_COUNT: + rel_buckets[rbid].append(row) else: - tids = [] - for tid in tids: - pairs.add((compile_kwd, tid)) - if len(field_map) < page_size: - break - offset += page_size - return pairs + name = row.get("name_kwd", "") + if not name: + continue + bid = _bucket_id(name) + if bid < _BUCKET_COUNT: + buckets[bid].append(row) + if len(rows) < limit: + last_page = True + else: + offset += limit + + # Merge entity buckets into dataset rows + for bi in range(_BUCKET_COUNT): + if buckets[bi]: + out = await _merge_bucket( + tenant_id, + kb_id, + compile_kwd, + template_id, + buckets[bi], + embd_mdl, + structure_kind, + ) + if out: + for start in range(0, len(out), _BUCKET_FLUSH_THRESHOLD): + await _index_insert(tenant_id, kb_id, out[start : start + _BUCKET_FLUSH_THRESHOLD]) + + # Merge relation buckets into dataset rows + for bi in range(_BUCKET_COUNT): + if rel_buckets[bi]: + rel_out = await _merge_relations( + tenant_id, + kb_id, + compile_kwd, + template_id, + rel_buckets[bi], + embd_mdl, + structure_kind, + ) + if rel_out: + for start in range(0, len(rel_out), _BUCKET_FLUSH_THRESHOLD): + await _index_insert(tenant_id, kb_id, rel_out[start : start + _BUCKET_FLUSH_THRESHOLD]) + + # ── Save build timestamp ───────────────────────────────────────── + import datetime + + await _save_build_time(tenant_id, kb_id, compile_kwd, template_id, datetime.datetime.now().timestamp()) + return cleanup_ok 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. - """ + """Entry point — called from task_handler.py when ``is_structure_merge_task`` matches.""" 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) + target_kind = { + "structure_graph": "knowledge_graph", + "structure_mindmap": "mind_map", + "timeline": "timeline", + "session_graph": "session_graph", + "session_essence": "session_essence", + }.get(task_type) merge_all = task_type == "structure" - def _canceled() -> bool: - try: - return bool(ctx.has_canceled_func(ctx.id)) - except Exception: - return False + # Resolve embedding model (required for re-embedding dataset rows) + embd_mdl = None + try: + from api.db.services.llm_service import LLMBundle + from api.apps.services.dataset_api_service import resolve_model_config + from common.constants import LLMType - 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.") + from api.db.services.knowledgebase_service import KnowledgebaseService + + ok, kb = KnowledgebaseService.get_by_id(ctx.kb_id) + if ok and kb: + model_config = resolve_model_config(ctx.tenant_id, LLMType.EMBEDDING.value, kb.embd_id) + embd_mdl = LLMBundle(ctx.tenant_id, model_config, lang=ctx.language) + else: + raise RuntimeError(f"Knowledge base {ctx.kb_id} not found") + except Exception: + logging.exception("structure_merge: failed to bind embedding model for kb=%s", ctx.kb_id) + progress(1.0, "Failed to bind embedding model — aborting.") + return + + progress(0.0, "Scanning doc_graph rows...") + pairs = await _collect_structure_pairs(ctx.tenant_id, ctx.kb_id) + if not pairs: + progress(1.0, "No doc_graph rows found.") 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]]: @@ -163,12 +841,10 @@ async def run_structure_merge(ctx: TaskContext) -> None: 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 = True # 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 + keep = kind_ok except Exception: - logging.exception("structure_merge: failed to resolve template %s for kb=%s", template_id, ctx.kb_id) + logging.exception("structure_merge: failed to resolve template %s", template_id) result = (keep, structure_kind) template_meta[template_id] = result return result @@ -181,34 +857,81 @@ async def run_structure_merge(ctx: TaskContext) -> None: 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}.") + progress(1.0, f"No eligible templates for {kind_label}.") return total = len(eligible) rebuilt = 0 + all_cleanup_ok = True for i, (compile_kwd, template_id, structure_kind) in enumerate(eligible): - if _canceled(): - progress(-1, "Task has been canceled.") + if ctx.has_canceled_func(ctx.id): + progress(1.0, f"Cancelled after {rebuilt}/{total} dataset graph(s).") return - progress( - 0.05 + 0.9 * (i / total), - f"Merging structure graph {i + 1}/{total} (compile_kwd={compile_kwd}) ...", - ) + progress(0.05 + 0.9 * (i / total), f"Building dataset graph {i + 1}/{total} ...") try: - await rebuild_dataset_structure_graph_json( - ctx.tenant_id, - ctx.kb_id, - compile_kwd, - compilation_template_id=template_id, - structure_kind=structure_kind, - ) + cleanup_ok = await _do_build(ctx.tenant_id, ctx.kb_id, compile_kwd, template_id, structure_kind, embd_mdl) + if not cleanup_ok: + all_cleanup_ok = False 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, - ) + all_cleanup_ok = False + logging.exception("structure_merge: build failed for kb=%s compile_kwd=%s", ctx.kb_id, compile_kwd) - progress(1.0, f"Merged {rebuilt}/{total} structure graph(s).") + # Delete pending deletion markers only after every pair has been + # processed successfully. If any pair had a store error the markers + # survive so the next run retries the cleanup. + if all_cleanup_ok: + await _consume_deletion_markers(ctx.tenant_id, ctx.kb_id) + + progress(1.0, f"Built {rebuilt}/{total} dataset graph(s).") + + +async def _collect_structure_pairs(tenant_id: str, kb_id: str) -> set[tuple[str, str]]: + """Collect distinct (compile_kwd, template_id) pairs from doc_graph entity rows.""" + from common.doc_store.doc_store_base import OrderByExpr + + index = _index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return set() + + # NOTE: No scope_kwd filter here because old doc_graph rows that predate + # the scope_kwd field won't have it — applying ["doc"] would miss those + # rows and prevent their pairs from ever being rebuilt. The scan in + # _do_build adds its own scope_kwd constraint. + pairs: set[tuple[str, str]] = set() + offset = 0 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["id", "compile_kwd", "compilation_template_ids"], + [], + {"knowledge_graph_kwd": ["entity"]}, + [], + OrderByExpr(), + offset, + _PAGE_SIZE, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, ["id", "compile_kwd", "compilation_template_ids"]) or {} + except Exception: + 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") + tids = [] + if isinstance(raw_tids, list): + tids = [t for t in raw_tids if isinstance(t, str) and t] + elif isinstance(raw_tids, str) and raw_tids: + tids = [raw_tids] + for tid in tids: + pairs.add((compile_kwd, tid)) + if len(field_map) < _PAGE_SIZE: + break + offset += _PAGE_SIZE + return pairs