diff --git a/agent/component/llm.py b/agent/component/llm.py index dbafe0267e..cefe358709 100644 --- a/agent/component/llm.py +++ b/agent/component/llm.py @@ -88,10 +88,15 @@ class LLM(ComponentBase): def __init__(self, canvas, component_id, param: ComponentParamBase): super().__init__(canvas, component_id, param) - model_types = resolve_model_type(self._canvas.get_tenant_id(), self._param.llm_id) - model_type = "chat" if "chat" in model_types else model_types[0] - chat_model_config = resolve_model_config(self._canvas.get_tenant_id(), model_type, self._param.llm_id) - self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error) + try: + model_types = resolve_model_type(self._canvas.get_tenant_id(), self._param.llm_id) + model_type = "chat" if "chat" in model_types else model_types[0] + chat_model_config = resolve_model_config(self._canvas.get_tenant_id(), model_type, self._param.llm_id) + self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error) + except Exception as e: + logging.warning(f"Fail to load LLM configuration for component. {e}") + self.chat_mdl = None + self.imgs = [] def get_input_form(self) -> dict[str, dict]: diff --git a/api/apps/restful_apis/chunk_api.py b/api/apps/restful_apis/chunk_api.py index 3c63e170f7..d2aa634017 100644 --- a/api/apps/restful_apis/chunk_api.py +++ b/api/apps/restful_apis/chunk_api.py @@ -25,6 +25,7 @@ from pydantic import BaseModel, Field, validator from quart import request from api.apps import login_required +from api.apps.services import structure_graph_common as sgc from api.db.joint_services.tenant_model_service import ( split_model_name, resolve_model_config, @@ -650,57 +651,8 @@ async def get_document_structure_graph(tenant_id, dataset_id, document_id): template_meta_by_id[template_id] = meta template_meta_by_kind.setdefault(kind_norm, []).append(meta) - # Load every graph row for this doc in one shot. Each row corresponds - # to one (compile_kwd, template_id) tuple — written by - # ``_struct_upsert_graph_json``. index_name = search.index_name(dataset_tenant_id) - fields = [ - "content_with_weight", - "compile_kwd", - "compilation_template_ids", - "compilation_template_kind_kwd", - ] - try: - res = await thread_pool_exec( - settings.docStoreConn.search, - fields, - [], - {"doc_id": [document_id], "knowledge_graph_kwd": ["graph"]}, - [], - OrderByExpr(), - 0, - 1000, - index_name, - [dataset_id], - ) - rows = settings.docStoreConn.get_fields(res, fields) - - # The RAPTOR graph row is identified by ``compile_kwd`` - # alone — it intentionally doesn't carry ``knowledge_graph_kwd`` - # (which belongs to the KG feature). Query it separately and - # union into the same bucket map below. - res_raptor = await thread_pool_exec( - settings.docStoreConn.search, - fields, - [], - {"doc_id": [document_id], "compile_kwd": ["raptor_graph"]}, - [], - OrderByExpr(), - 0, - 16, - index_name, - [dataset_id], - ) - raptor_rows = settings.docStoreConn.get_fields(res_raptor, fields) - except Exception as e: - return server_error_response(e) - - # Merge the two field-maps so the grouping loop below treats them - # identically. Raptor rows clobber by id, which is fine — both - # sources produce stable per-row ids. - if raptor_rows: - rows = dict(rows or {}) - rows.update(raptor_rows) + keywords = (request.args.get("keywords") or "").strip() def _row_template_id(row: dict) -> str | None: raw = row.get("compilation_template_ids") @@ -712,30 +664,16 @@ async def get_document_structure_graph(tenant_id, dataset_id, document_id): return raw.strip() return None - # Group: template_id → {entities, relations, kind} - grouped: dict[str, dict] = {} - for row in (rows or {}).values(): - graph = {} - 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 + def _resolve_bucket(row: dict) -> tuple[dict, dict]: + """``(bucket_meta, scope_filter)`` for a graph/entity row. + ``bucket_meta`` = ``{template_id, template_name, kind}``; + ``scope_filter`` is the raw entity/relation-row filter (doc + + template, or legacy compile_kwd) WITHOUT ``knowledge_graph_kwd``. + """ tid = _row_template_id(row) compile_kwd_val = row.get("compile_kwd") or "" kind_val = row.get("compilation_template_kind_kwd") or compile_kwd_val - - # The RAPTOR graph row has no ``compilation_template_ids`` (it - # isn't derived from a user-authored template). Treat it as its - # own first-class bucket, not a legacy fallback. - is_raptor = compile_kwd_val == "raptor_graph" - if tid: bucket_id = tid row_kind_norm = _compilation_template_kind(kind_val) @@ -764,31 +702,127 @@ async def get_document_structure_graph(tenant_id, dataset_id, document_id): meta = kind_matches[0] bucket_name = (meta or {}).get("template_name") or bucket_id bucket_kind = (meta or {}).get("kind") or kind_val - elif is_raptor: - bucket_id = "raptor" - bucket_name = "RAPTOR Summary" - bucket_kind = "raptor" + scope = {"doc_id": [document_id], "compilation_template_ids": [bucket_id]} else: - # Legacy row: synthesize a stable id keyed by compile_kwd so - # multiple legacy kinds (e.g. ``list`` + ``hypergraph``) on - # the same doc surface as separate tabs. + # Legacy row (pre-dates the template stamp): keyed by compile_kwd + # so multiple legacy kinds surface as separate tabs. Scope excludes + # rows that DO carry a template id so it can't shadow a real bucket. bucket_id = f"legacy:{compile_kwd_val}" bucket_name = f"Legacy ({compile_kwd_val})" bucket_kind = kind_val + scope = {"doc_id": [document_id], "compile_kwd": [compile_kwd_val], "must_not": {"exists": "compilation_template_ids"}} + return {"template_id": bucket_id, "template_name": bucket_name, "kind": bucket_kind}, scope - if bucket_id not in grouped: - grouped[bucket_id] = { - "template_id": bucket_id, - "template_name": bucket_name, - "kind": bucket_kind, - "entities": [], - "relations": [], - } - grouped[bucket_id]["entities"].extend(entities) - grouped[bucket_id]["relations"].extend(relations) + # ── keywords mode: global KNN → the top-1 entity's focused subgraph ── + if keywords: + try: + embd_id = DocumentService.get_embd_id(document_id) + model_config = resolve_model_config(dataset_tenant_id, LLMType.EMBEDDING.value, embd_id) + embd_mdl = TenantLLMService.model_instance(model_config) + except Exception: + logging.exception("structure graph: embedding bind failed for doc=%s", document_id) + return get_result(data={"templates": []}) + try: + bucket_meta, kw_entities, kw_relations = await sgc.keyword_subgraph( + index_name, + dataset_id, + embd_mdl, + {"doc_id": [document_id], "knowledge_graph_kwd": ["entity"]}, + keywords, + _resolve_bucket, + log_ctx=f"doc={document_id}", + ) + except Exception as e: + return server_error_response(e) + if not bucket_meta or (not kw_entities and not kw_relations): + return get_result(data={"templates": []}) + bucket = dict(bucket_meta) + bucket["entities"] = kw_entities + bucket["relations"] = kw_relations + return get_result(data={"templates": [bucket]}) + + # ── normal mode: per-template subgraph sampling from the raw rows ── + # Metadata-only scan of the per-doc graph blob rows (one per + # (compile_kwd, template_id)) purely to discover buckets and resolve their + # display name/kind — WITHOUT loading the (potentially huge) + # content_with_weight. Each bucket's entities/relations are then fetched + # from the raw ``knowledge_graph_kwd`` rows with subgraph sampling. + meta_fields = ["compile_kwd", "compilation_template_ids", "compilation_template_kind_kwd"] + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + meta_fields, + [], + {"doc_id": [document_id], "knowledge_graph_kwd": ["graph"]}, + [], + OrderByExpr(), + 0, + 1000, + index_name, + [dataset_id], + ) + meta_rows = settings.docStoreConn.get_fields(res, meta_fields) or {} + except Exception as e: + return server_error_response(e) + + # Discover unique buckets. A template can own multiple compile_kwd blob + # rows; scoping by template_id folds them together, matching prior behavior. + bucket_metas: dict[str, dict] = {} + bucket_scopes: dict[str, dict] = {} + for row in meta_rows.values(): + meta, scope = _resolve_bucket(row) + bid = meta["template_id"] + if bid not in bucket_metas: + bucket_metas[bid] = meta + bucket_scopes[bid] = scope + + grouped: dict[str, dict] = {} + for bid, meta in bucket_metas.items(): + try: + entities, relations = await sgc.build_bucket(index_name, dataset_id, bucket_scopes[bid]) + except Exception as e: + return server_error_response(e) + if not entities and not relations: + continue + grouped[bid] = {**meta, "entities": entities, "relations": relations} + + # RAPTOR summary graph: a standalone blob (no ``knowledge_graph_kwd``, and + # no raw entity/relation rows), so read its content directly and don't + # sample it. + try: + res_raptor = await thread_pool_exec( + settings.docStoreConn.search, + ["content_with_weight", "compile_kwd"], + [], + {"doc_id": [document_id], "compile_kwd": ["raptor_graph"]}, + [], + OrderByExpr(), + 0, + 16, + index_name, + [dataset_id], + ) + raptor_rows = settings.docStoreConn.get_fields(res_raptor, ["content_with_weight", "compile_kwd"]) or {} + except Exception: + logging.exception("structure graph: RAPTOR blob load failed for doc=%s", document_id) + raptor_rows = {} + for row in raptor_rows.values(): + try: + graph = json.loads(row.get("content_with_weight") or "{}") + except Exception: + continue + if not isinstance(graph, dict): + continue + r_entities = graph.get("entities") or [] + r_relations = graph.get("relations") or [] + if not r_entities and not r_relations: + continue + rb = grouped.setdefault("raptor", {"template_id": "raptor", "template_name": "RAPTOR Summary", "kind": "raptor", "entities": [], "relations": []}) + rb["entities"].extend(r_entities) + rb["relations"].extend(r_relations) # Order: configured templates first (in the user's chosen order), - # then any legacy buckets after. + # then any discovered / legacy / raptor buckets after. ordered_ids: list[str] = [] for tid in configured_ids: if tid in grouped and tid not in ordered_ids: diff --git a/api/apps/restful_apis/dataset_api.py b/api/apps/restful_apis/dataset_api.py index f062fceafa..c931e9e607 100644 --- a/api/apps/restful_apis/dataset_api.py +++ b/api/apps/restful_apis/dataset_api.py @@ -647,22 +647,36 @@ async def list_wiki_topics(tenant_id, dataset_id): async def get_wiki_graph(tenant_id, dataset_id): """Return an incremental slice of the canvas graph for this dataset. - GET /api/v1/datasets//artifacts/graph[?node=] + GET /api/v1/datasets//artifacts/graph[?node=][&keywords=][&top_n=] - ``node`` omitted: overview centred on the heaviest-weighted - entities, expanded outward until ``MAX_LOADING_ENTITY`` is hit. + entities, expanded outward until the entity budget is hit. - ``node`` provided: subgraph centred on that entity, including all outgoing relations and their ``to`` targets (also capped). + - ``keywords`` (overview only; ignored with ``node``): seed the overview + from the best BM25 matches instead of the heaviest-weighted entities. + - ``top_n`` (a.k.a. ``topN``): override the entity budget (default 128). + Only entities referenced by at least one relation are returned. Success: ``{"code": 0, "data": {"entities":[…],"relations":[…]}}``. """ try: node = request.args.get("node", None) if isinstance(node, str): node = node.strip() or None + keywords = request.args.get("keywords", "") + top_n_arg = request.args.get("top_n") or request.args.get("topN") + top_n = None + if top_n_arg is not None: + try: + top_n = int(top_n_arg) + except (TypeError, ValueError): + top_n = None success, result = await dataset_api_service.get_wiki_graph( dataset_id, tenant_id, node=node, + keywords=keywords, + top_n=top_n, ) if success: return get_result(data=result) @@ -704,10 +718,12 @@ async def get_dataset_structure(tenant_id, dataset_id): message=f"Unsupported structure kind: {kind!r}. Expected one of: graph, mindmap, timeline, session_essence, session_graph.", code=RetCode.ARGUMENT_ERROR, ) + keywords = request.args.get("keywords", "") success, result = await dataset_api_service.get_dataset_structure( dataset_id, tenant_id, kind, + keywords=keywords, ) if success: return get_result(data=result) diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 4f780a3ccd..1687cf5729 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -53,9 +53,9 @@ _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", + "graph": "structure_graph", "raptor": "raptor", - "mindmap": "mindmap", + "mindmap": "structure_mindmap", "artifact": "artifact", "skill": "skill", # Structure merge types carry their own task_type (== index_type) so the @@ -1723,16 +1723,20 @@ def _resolve_dataset_structure_kind(kind) -> str | None: return _DATASET_STRUCTURE_KIND_ALIASES.get(kind.strip().lower().replace("-", "_")) -async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str): +async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str, keywords: 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. + ``session_essence`` / ``session_graph``. The ``knowledge_graph_kwd="dataset_graph"`` + blob rows (one per template, written by ``rebuild_dataset_structure_graph_json``) + are used only to DISCOVER the requested kind's template buckets; each bucket's + entities/relations are then fetched from the raw KB-wide ``entity`` / ``relation`` + rows with subgraph sampling — identical to the per-document endpoint but scoped + KB-wide (no ``doc_id`` filter), since dataset-merge templates dedup those rows + across documents. + + ``keywords`` (optional): return the single best-matching entity's 1-hop + subgraph (top-1 + neighbors + touching relations) across the kind, via KNN. Returns ``(True, {"kind": , "templates": [...]})`` or ``(False, message)`` on auth/validation failure. @@ -1754,30 +1758,10 @@ async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str): from common.doc_store.doc_store_base import OrderByExpr from api.db.services.compilation_template_service import CompilationTemplateService + from api.db.services.tenant_llm_service import TenantLLMService + from api.apps.services import structure_graph_common as sgc - 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 + keywords = (keywords or "").strip() def _row_template_id(row: dict) -> str | None: raw = row.get("compilation_template_ids") @@ -1811,40 +1795,148 @@ async def get_dataset_structure(dataset_id: str, tenant_id: str, kind: str): template_kind_cache[tid] = top_kind return top_kind - grouped: dict[str, dict] = {} - for row in rows.values(): + def _bucket_meta_for(tid: str, row_kind: str = "") -> dict: + if tid not in template_name_cache: + _template_meta(tid) + return { + "template_id": tid, + "template_name": template_name_cache.get(tid, tid), + "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 + + 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 + # ── keywords mode: global KNN across the kind → top-1's focused subgraph. ── + if keywords: + if not kind_template_ids: + return True, empty try: - graph = json.loads(row.get("content_with_weight") or "{}") + model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING.value, kb.embd_id) + embd_mdl = TenantLLMService.model_instance(model_config) except Exception: + logging.exception("get_dataset_structure: embedding bind failed for kb=%s", dataset_id) + return True, empty + + def _scope_for_template(row: dict): + 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]} + + 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"]}, + keywords, + _scope_for_template, + log_ctx=f"kb={dataset_id}", + ) + if resolved_kind == "mind_map": + kw_entities = sgc.filter_entities_with_relations(kw_entities, kw_relations) + if not kw_entities: + return True, empty + if not bucket_meta or (not kw_entities and not kw_relations): + return True, empty + bucket = dict(bucket_meta) + bucket["entities"] = kw_entities + bucket["relations"] = kw_relations + return True, {"kind": kind, "templates": [bucket]} + + # ── normal mode: per-template subgraph sampling from raw KB-wide rows. ── + 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]}) + except Exception: + logging.exception("get_dataset_structure: bucket build failed for kb=%s template=%s", dataset_id, tid) continue - if not isinstance(graph, dict): - continue - entities = graph.get("entities") or [] - relations = graph.get("relations") or [] + if resolved_kind == "mind_map": + entities = sgc.filter_entities_with_relations(entities, relations) + if not entities: + continue if not entities and not relations: continue + meta = _bucket_meta_for(tid) + templates_out.append({**meta, "entities": entities, "relations": relations}) - 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) + # Legacy template-less dataset_graph rows have no template id to scope raw + # rows by, so fall back to their blob content directly (no sampling). Rare — + # rebuild_dataset_structure_graph_json always stamps a template id. + if has_templateless: + try: + res_l = await thread_pool_exec( + settings.docStoreConn.search, + ["content_with_weight", "compilation_template_kind_kwd"], + [], + {"knowledge_graph_kwd": [_DATASET_STRUCTURE_ROW_KWD], "must_not": {"exists": "compilation_template_ids"}}, + [], + OrderByExpr(), + 0, + 1000, + index_nm, + [dataset_id], + ) + legacy_rows = settings.docStoreConn.get_fields(res_l, ["content_with_weight", "compilation_template_kind_kwd"]) or {} + except Exception: + logging.exception("get_dataset_structure: legacy blob fetch failed for kb=%s", dataset_id) + legacy_rows = {} + legacy_bucket = {"template_id": f"kind:{resolved_kind}", "template_name": f"kind:{resolved_kind}", "kind": resolved_kind, "entities": [], "relations": []} + for row in legacy_rows.values(): + if _resolve_dataset_structure_kind((row.get("compilation_template_kind_kwd") or "").strip()) != resolved_kind: + continue + try: + graph = json.loads(row.get("content_with_weight") or "{}") + except Exception: + continue + if not isinstance(graph, dict): + continue + legacy_bucket["entities"].extend(graph.get("entities") or []) + legacy_bucket["relations"].extend(graph.get("relations") or []) + if resolved_kind == "mind_map": + 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"]: + templates_out.append(legacy_bucket) - templates_out = [g for g in grouped.values() if g["entities"] or g["relations"]] return True, {"kind": kind, "templates": templates_out} @@ -2871,15 +2963,16 @@ async def _wiki_search_entity_page( dataset_id: str, offset: int, limit: int, + keywords: str = "", ): - """One page of artifact_entity rows, ordered by weight_int DESC.""" - from common.doc_store.doc_store_base import OrderByExpr + """One page of artifact_entity rows. - order_by = OrderByExpr() - try: - order_by.desc("weight_int") - except Exception: - order_by = OrderByExpr() + Without ``keywords``: ordered by ``weight_int DESC`` (heaviest nodes first). + With ``keywords``: BM25 full-text match over ``content_ltks`` (slug + + summary), ordered by relevance. ``artifact_entity`` rows are BM25-only (no + embedding vector), so this is a lexical search, not a dense KNN. + """ + from common.doc_store.doc_store_base import OrderByExpr select_fields = [ "id", @@ -2888,12 +2981,31 @@ async def _wiki_search_entity_page( "source_chunk_ids", "content_with_weight", ] + + keywords = (keywords or "").strip() + match_expressions: list = [] + order_by = OrderByExpr() + if keywords: + try: + match_text, _ = settings.retriever.qryr.question(keywords, min_match=0.1) + match_expressions = [match_text] + except Exception: + logging.exception("get_wiki_graph: failed to build keyword query for kb=%s", dataset_id) + match_expressions = [] + if not match_expressions: + # No keywords (or query build failed) → heaviest-weighted first. When a + # text match is present the store ranks by BM25 score instead. + try: + order_by.desc("weight_int") + except Exception: + order_by = OrderByExpr() + res = await thread_pool_exec( settings.docStoreConn.search, select_fields, [], {"compile_kwd": [_WIKI_GRAPH_ENTITY_KWD]}, - [], + match_expressions, order_by, offset, limit, @@ -2975,9 +3087,16 @@ async def get_wiki_graph( dataset_id: str, tenant_id: str, node: str | None = None, + keywords: str | None = None, + top_n: int | None = None, ): """Load the canvas graph payload incrementally from per-row data. + ``top_n`` overrides the entity budget (default ``_WIKI_GRAPH_MAX_LOADING_ENTITY``). + ``keywords`` (overview mode only; ignored when ``node`` is given) seeds the + graph from the best BM25 matches on ``artifact_entity`` rows instead of the + heaviest-weighted ones. Only entities referenced by a relation are returned. + Two modes: * **Overview** (``node`` is None) — paginate ``artifact_entity`` rows @@ -3010,7 +3129,18 @@ async def get_wiki_graph( return True, empty index_nm, _ = pack - cap = _WIKI_GRAPH_MAX_LOADING_ENTITY + from api.apps.services import structure_graph_common as sgc + + keywords = (keywords or "").strip() + # Entity budget: caller-overridable, clamped to a sane range so a bad param + # can neither disable the cap nor blow up the response. + if top_n is not None: + try: + cap = max(1, min(int(top_n), 1024)) + except (TypeError, ValueError): + cap = _WIKI_GRAPH_MAX_LOADING_ENTITY + else: + cap = _WIKI_GRAPH_MAX_LOADING_ENTITY page_size = _WIKI_GRAPH_ENTITY_PAGE_SIZE # ``entities`` preserves first-seen order so the canvas paints the @@ -3111,11 +3241,11 @@ async def get_wiki_graph( to_map = {} for row in (to_map or {}).values(): payload = _wiki_entity_payload(row) - if payload and len(entities) < cap: + if payload and len(entities) < cap * 2: _add_entity(payload) return True, { - "entities": list(entities.values()), + "entities": sgc.filter_entities_with_relations(list(entities.values()), relations), "relations": relations, } @@ -3130,6 +3260,7 @@ async def get_wiki_graph( dataset_id, offset, page_size, + keywords=keywords, ) except Exception: logging.exception( @@ -3223,7 +3354,7 @@ async def get_wiki_graph( page += 1 return True, { - "entities": list(entities.values()), + "entities": sgc.filter_entities_with_relations(list(entities.values()), relations), "relations": relations, } diff --git a/api/apps/services/structure_graph_common.py b/api/apps/services/structure_graph_common.py new file mode 100644 index 0000000000..32a4e59f4d --- /dev/null +++ b/api/apps/services/structure_graph_common.py @@ -0,0 +1,304 @@ +# +# 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. +# + +"""Shared structure-graph subgraph sampling. + +Both the per-document (``/datasets//documents//structure/graph``) and +the dataset-wide (``/datasets//artifacts_structure``) endpoints render +per-template structure graphs. For large graphs we don't return every +entity/relation — we fetch a representative subgraph from the raw +``knowledge_graph_kwd`` rows (which carry ``mention_count_int`` / ``name_kwd`` / +``from_entity_kwd`` / ``to_entity_kwd`` / ``q__vec``) so the response — and +the frontend render — stay bounded. + +The two endpoints differ only in *scope*: the document endpoint filters raw +rows by ``doc_id``; the dataset endpoint queries KB-wide (dataset-merge +templates dedup entity/relation rows across documents). That difference lives +entirely in the ``scope`` / ``base_entity_condition`` dicts the caller passes — +everything else is shared here. +""" + +import json +import logging + +from common import settings +from common.doc_store.doc_store_base import OrderByExpr +from common.misc_utils import thread_pool_exec + + +# Below this combined (entities + relations) count for a bucket, return all rows. +GRAPH_FULL_THRESHOLD = 1024 +# Size of the top-mention entity seed set (set A) for large buckets. +GRAPH_TOP_ENTITIES = 256 +# Upper bound on the relation / neighbor-entity expansion so a hub node can't +# blow up the response. +GRAPH_EXPANSION_CAP = 4096 + +GRAPH_ENTITY_FIELDS = ["id", "content_with_weight", "name_kwd", "mention_count_int", "source_chunk_ids"] +GRAPH_RELATION_FIELDS = ["id", "content_with_weight", "from_entity_kwd", "to_entity_kwd"] +GRAPH_ALL_FIELDS = ["id", "content_with_weight", "name_kwd", "mention_count_int", "source_chunk_ids", "from_entity_kwd", "to_entity_kwd", "knowledge_graph_kwd"] + + +async def graph_search(index_name, kb_id, select_fields, condition, order_by, limit, match_expressions=None): + """One raw-row search. Returns ``(field_map, total)`` where ``total`` is the + full match count (not the returned slice).""" + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + condition, + match_expressions or [], + order_by, + 0, + max(int(limit or 0), 1), + index_name, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + total = settings.docStoreConn.get_total(res) + return field_map, int(total or 0) + + +def project_entity(row: dict) -> dict | None: + """Project a raw ``knowledge_graph_kwd="entity"`` row to the graph-node shape + the frontend already consumes, surfacing ``mention_count_int`` as + ``mention_count``.""" + from rag.advanced_rag.knowlege_compile.structure import _struct_graph_entity + + try: + payload = json.loads(row.get("content_with_weight") or "{}") + except Exception: + return None + if not isinstance(payload, dict): + return None + node = _struct_graph_entity(payload, row.get("source_chunk_ids")) + if not node: + return None + mc = row.get("mention_count_int") + if isinstance(mc, list): # Infinity returns *_int scalars fine, but be defensive + mc = mc[0] if mc else None + try: + if mc is not None: + node["mention_count"] = int(mc) + except (TypeError, ValueError): + pass + return node + + +def project_relation(row: dict) -> dict | None: + """Project a raw ``knowledge_graph_kwd="relation"`` row to the edge shape. + Prefers the payload (matching the blob projection); falls back to the + authoritative ``*_entity_kwd`` columns.""" + from rag.advanced_rag.knowlege_compile.structure import _struct_graph_relation + + try: + payload = json.loads(row.get("content_with_weight") or "{}") + except Exception: + payload = {} + if isinstance(payload, dict): + node = _struct_graph_relation(payload) + if node: + return node + src = str(row.get("from_entity_kwd") or "").strip() + tgt = str(row.get("to_entity_kwd") or "").strip() + if not src or not tgt: + return None + typ = payload.get("type") if isinstance(payload, dict) else None + return {"from": src, "to": tgt, "type": str(typ).strip() if typ else "related"} + + +def dedup_entities(entities: list[dict]) -> list[dict]: + """Order-preserving dedup by (lowercased name, type).""" + out: list[dict] = [] + seen: set[tuple[str, str]] = set() + for e in entities: + key = (str(e.get("name") or "").strip().lower(), str(e.get("type") or "").strip().lower()) + if not key[0] or key in seen: + continue + seen.add(key) + out.append(e) + return out + + +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 [] + + connected: set[str] = set() + for relation in relations: + if not isinstance(relation, dict): + continue + for endpoint_key in ("from", "to"): + endpoint = relation.get(endpoint_key) + if isinstance(endpoint, str): + endpoint = endpoint.strip() + if endpoint: + connected.add(endpoint) + + if not connected: + return [] + + filtered: list[dict] = [] + for entity in entities: + if not isinstance(entity, dict): + continue + keys: set[str] = set() + # Structure-graph nodes are name-keyed and their relations reference + # names; artifact-graph nodes are slug-keyed and their relations + # reference slugs. Check all three identity fields so the same filter + # serves both callers. + for field in ("id", "name", "slug"): + value = entity.get(field) + if isinstance(value, str): + value = value.strip() + if value: + keys.add(value) + if keys & connected: + filtered.append(entity) + return filtered + + +async def build_bucket(index_name, kb_id, scope: dict) -> tuple[list[dict], list[dict]]: + """Build one bucket's ``(entities, relations)`` from raw rows. + + ``scope`` is the filter WITHOUT ``knowledge_graph_kwd`` — e.g. + ``{"doc_id":[id], "compilation_template_ids":[tid]}`` (document scope) or + ``{"compilation_template_ids":[tid]}`` (dataset scope). Small buckets are + returned whole; large ones are sampled: top-``GRAPH_TOP_ENTITIES`` entities + by ``mention_count_int``, the relations sourced from them, and those + relations' target entities. + """ + both_cond = dict(scope, knowledge_graph_kwd=["entity", "relation"]) + _, total = await graph_search(index_name, kb_id, ["id"], both_cond, OrderByExpr(), 1) + + if total < GRAPH_FULL_THRESHOLD: + field_map, _ = await graph_search(index_name, kb_id, GRAPH_ALL_FIELDS, both_cond, OrderByExpr(), total or 1) + entities: list[dict] = [] + relations: list[dict] = [] + for row in field_map.values(): + if row.get("knowledge_graph_kwd") == "relation": + edge = project_relation(row) + if edge: + relations.append(edge) + else: + node = project_entity(row) + if node: + entities.append(node) + return dedup_entities(entities), relations + + # Large bucket: sample. A = top entities by mention_count_int desc. + order_by = OrderByExpr() + try: + order_by.desc("mention_count_int") + except Exception: + order_by = OrderByExpr() + 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()}) + + # 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) + for row in rel_map.values(): + edge = project_relation(row) + if edge: + relations.append(edge) + tgt = str(edge.get("to") or "").strip().lower() + if tgt: + target_names_lower.add(tgt) + + # target entities of those relations (case-insensitive via name_kwd). + set_t = [] + if target_names_lower: + 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 + + +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]]: + """KNN the entity rows matching ``base_entity_condition`` for ``keywords``; + return ``(top1_bucket_meta, entities, relations)`` for the top-1 entity's + 1-hop subgraph (top-1 + neighbors + touching relations). ``(None, [], [])`` + when nothing matches or embedding is unavailable. + + ``base_entity_condition`` scopes the KNN (e.g. ``{"doc_id":[id], + "knowledge_graph_kwd":["entity"]}`` or ``{"compilation_template_ids":[...], + "knowledge_graph_kwd":["entity"]}``). ``scope_for_template(row)`` resolves + ``(bucket_meta, scope_filter)`` for the matched row (scope WITHOUT + ``knowledge_graph_kwd``). + """ + from common.doc_store.doc_store_base import MatchDenseExpr + + try: + qv, _ = await thread_pool_exec(embd_mdl.encode_queries, keywords) + vec = list(qv) + except Exception: + logging.exception("structure graph: keyword embedding failed (%s)", log_ctx) + return None, [], [] + if not vec: + return None, [], [] + + match_expr = MatchDenseExpr( + vector_column_name=f"q_{len(vec)}_vec", + embedding_data=vec, + embedding_data_type="float", + distance_type="cosine", + topn=1, + extra_options={"similarity": 0.0}, + ) + top_fields = GRAPH_ENTITY_FIELDS + ["compilation_template_ids", "compile_kwd", "compilation_template_kind_kwd"] + top_map, _ = await graph_search(index_name, kb_id, top_fields, base_entity_condition, OrderByExpr(), 1, match_expressions=[match_expr]) + if not top_map: + return None, [], [] + top_row = next(iter(top_map.values())) + top_node = project_entity(top_row) + if not top_node: + return None, [], [] + top_name = str(top_node.get("name") or "").strip() + if not top_name: + return None, [], [] + bucket_meta, scope = scope_for_template(top_row) + + # Relations where the top-1 entity is source OR target (two term queries). + relations: list[dict] = [] + seen_rel: set[tuple[str, str, str]] = set() + neighbor_names_lower: set[str] = set() + 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) + for row in rel_map.values(): + edge = project_relation(row) + if not edge: + continue + key = (edge.get("from", ""), edge.get("to", ""), edge.get("type", "")) + if key in seen_rel: + continue + seen_rel.add(key) + relations.append(edge) + for endpoint in (edge.get("from", ""), edge.get("to", "")): + endpoint = str(endpoint).strip() + if endpoint and endpoint.lower() != top_name.lower(): + neighbor_names_lower.add(endpoint.lower()) + + entities = [top_node] + if neighbor_names_lower: + 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 diff --git a/conf/infinity_mapping.json b/conf/infinity_mapping.json index d46af006ec..8dc4ed02f0 100644 --- a/conf/infinity_mapping.json +++ b/conf/infinity_mapping.json @@ -21,6 +21,7 @@ "position_int": {"type": "varchar", "default": ""}, "weight_int": {"type": "integer", "default": 0}, "weight_flt": {"type": "float", "default": 0.0}, + "mention_count_int": {"type": "integer", "default": 0}, "chunk_order_int": {"type": "integer", "default": 0}, "rank_int": {"type": "integer", "default": 0}, "rank_flt": {"type": "float", "default": 0}, diff --git a/rag/advanced_rag/harness/tools/navigation.py b/rag/advanced_rag/harness/tools/navigation.py index 989c0be680..898514e3bc 100644 --- a/rag/advanced_rag/harness/tools/navigation.py +++ b/rag/advanced_rag/harness/tools/navigation.py @@ -144,7 +144,7 @@ def _render_structure(entities: list[dict], relations: list[dict]) -> str: if not name: continue typ = (e.get("type") or "other").strip() - desc = " ".join((e.get("discription") or "").split()) + desc = " ".join((e.get("description") or "").split()) lines.append(f"- {name} ({typ})" + (f": {desc}" if desc else "")) if relations: lines.append("\nRelations:") @@ -427,7 +427,7 @@ async def dataset_navigate(tools, topic: str, keywords: str = "", doc_scope: lis { "name": node_name, "type": h.get("type") or "dataset_nav", - "discription": h.get("description") or "", + "description": h.get("description") or "", } ) for did in h.get("doc_ids") or []: @@ -537,7 +537,7 @@ def _kg_parse_entity(row: dict) -> dict | None: return { "name": name, "type": (payload.get("type") or "other"), - "discription": (payload.get("discription") or payload.get("description") or ""), + "description": (payload.get("description") or payload.get("description") or ""), "aliases": aliases, "source_chunk_ids": list(row.get("source_chunk_ids") or []), "doc_id": row.get("doc_id") or "", diff --git a/rag/advanced_rag/knowlege_compile/runner.py b/rag/advanced_rag/knowlege_compile/runner.py index 8747fd99d3..be6bec3e3e 100644 --- a/rag/advanced_rag/knowlege_compile/runner.py +++ b/rag/advanced_rag/knowlege_compile/runner.py @@ -167,7 +167,7 @@ def _page_index_graph_summary(graph: dict, limit: int = 80) -> str: if not isinstance(entity, dict): continue name = str(entity.get("name") or "").strip() - description = str(entity.get("discription") or entity.get("description") or "").strip() + description = str(entity.get("description") or entity.get("description") or "").strip() text = f"{name}: {description}".strip(": ").strip() if text: lines.append(text) diff --git a/rag/advanced_rag/knowlege_compile/structure.py b/rag/advanced_rag/knowlege_compile/structure.py index 87ef5e189f..ef59ef43b4 100644 --- a/rag/advanced_rag/knowlege_compile/structure.py +++ b/rag/advanced_rag/knowlege_compile/structure.py @@ -445,7 +445,7 @@ def _struct_graph_entity(payload: dict, source_chunk_ids: list | None = None) -> if not isinstance(aliases, list): aliases = [] aliases = [str(a).strip() for a in aliases if str(a).strip()] - description = payload.get("description") or payload.get("discription") or payload.get("definition_excerpt") or "" + description = payload.get("description") or payload.get("description") or payload.get("definition_excerpt") or "" if isinstance(source_chunk_ids, str): source_chunk_ids = [source_chunk_ids] source_chunk_ids = _struct_union_chunk_ids(source_chunk_ids) @@ -455,7 +455,7 @@ def _struct_graph_entity(payload: dict, source_chunk_ids: list | None = None) -> "name": name, "source_chunk_ids": source_chunk_ids, "type": typ or "other", - "discription": str(description).strip() if description is not None else "", + "description": str(description).strip() if description is not None else "", } @@ -489,8 +489,8 @@ def _struct_merge_graph_entities(entities: list[dict]) -> list[dict]: for alias in entity.get("aliases") or []: if alias not in aliases: aliases.append(alias) - if not target.get("discription") and entity.get("discription"): - target["discription"] = entity["discription"] + if not target.get("description") and entity.get("description"): + target["description"] = entity["description"] target["source_chunk_ids"] = _struct_union_chunk_ids( target.get("source_chunk_ids"), entity.get("source_chunk_ids"), @@ -586,6 +586,26 @@ def _struct_to_doc_storage_doc( f"q_{len(vec_list)}_vec": vec_list, "id": row_id, } + + # Surface two payload fields as queryable top-level columns so the store can + # filter/sort on them without parsing ``content_with_weight``. Both are + # copies: the originals stay in the payload, so payload consumers (merge, + # graph projection) are unaffected and ``row_id`` — which hashes the payload + # JSON — does not shift. Merges rebuild through this same function, so the + # columns re-derive from the merged payload automatically. + try: + mention_count = int(payload.get("mention_count") or 1) + except (TypeError, ValueError): + mention_count = 1 + doc["mention_count_int"] = mention_count + + # Lower-cased name for case-insensitive exact lookups. Relations carry no + # ``name`` (they use source/target), so the column is only stamped when the + # payload actually has one. + name_value = _struct_entity_name(payload) + if name_value: + doc["name_kwd"] = name_value.lower() + if template_id_str: doc["compilation_template_ids"] = [template_id_str] if compilation_template_kind: diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index 713da1bcb3..2e98c711d0 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -340,7 +340,7 @@ def _wiki_type_rules(fields: list) -> str: rule = rule.strip() if isinstance(rule, str) else "" lines.append(f"type: {typ}") if description: - lines.append(f" - discription: {description}") + lines.append(f" - description: {description}") if rule: lines.append(f" - rule: {rule}") return "\n".join(lines) diff --git a/rag/flow/compiler/compiler.py b/rag/flow/compiler/compiler.py index ccbea4e02f..4092348a70 100644 --- a/rag/flow/compiler/compiler.py +++ b/rag/flow/compiler/compiler.py @@ -51,7 +51,6 @@ class CompilerParam(ProcessParamBase, LLMParam): self.compilation_template_group_ids = [] 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] diff --git a/rag/svr/task_executor_refactor/dataset_structure_merger.py b/rag/svr/task_executor_refactor/dataset_structure_merger.py index e05977f708..47d93b64e0 100644 --- a/rag/svr/task_executor_refactor/dataset_structure_merger.py +++ b/rag/svr/task_executor_refactor/dataset_structure_merger.py @@ -163,8 +163,8 @@ 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 = bool(config.get("dataset_merge")) if isinstance(config, dict) else False + # 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 except Exception: