From 3b4a96d42195dae505a419e0acd092dc04012a36 Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Mon, 3 Aug 2026 16:03:17 +0800 Subject: [PATCH] Refactor: refine wiki plan procedure. (#17579) ### Summary Refine wiki plan procedure. --------- Co-authored-by: Yingfeng Zhang Co-authored-by: buua436 --- api/apps/services/dataset_api_service.py | 188 +- api/db/init_data.py | 2 + .../init_data/compilation_templates/wiki.yaml | 1 + api/db/services/doc_metadata_service.py | 2 +- api/db/services/document_service.py | 26 + api/db/services/task_service.py | 6 +- conf/infinity_mapping.json | 17 +- rag/advanced_rag/knowlege_compile/wiki.py | 29 +- .../knowlege_compile/wiki_incremental.py | 3690 +++++++++++++++++ .../dataset_wiki_generator.py | 576 ++- .../task_executor_refactor/task_handler.py | 25 +- .../integration/wiki/test_wiki_incremental.py | 265 ++ test/unit_test/rag/advanced_rag/__init__.py | 0 .../advanced_rag/knowlege_compile/conftest.py | 134 + .../knowlege_compile/test_wiki_incremental.py | 1268 ++++++ web/src/locales/en.ts | 1 + web/src/locales/zh.ts | 1 + .../dataset/dataset-setting/form-schema.ts | 4 + .../dataset/dataset-setting/general-form.tsx | 2 + .../components/template-configuration.tsx | 15 + .../create-next/constant.ts | 1 + .../create-next/utils.ts | 10 + .../components/template-configuration.tsx | 15 + .../edit-next/constant.ts | 1 + .../compilation-templates/edit-next/utils.ts | 10 + 25 files changed, 6175 insertions(+), 114 deletions(-) create mode 100644 rag/advanced_rag/knowlege_compile/wiki_incremental.py create mode 100644 test/integration/wiki/test_wiki_incremental.py create mode 100644 test/unit_test/rag/advanced_rag/__init__.py create mode 100644 test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py create mode 100644 test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 63f1a0eaac..073c48c05b 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -1535,7 +1535,7 @@ async def search_datasets(tenant_id: str, req: dict): # # These three helpers power the dataset-level "Artifact" tab. They query rows # with ``compile_kwd="wiki_page"`` written by TaskHandler's -# ``_persist_wiki_pages_to_es``. The schema fields they rely on are: +# ``persist_wiki_pages``. The schema fields they rely on are: # slug_kwd, title_kwd, page_type_kwd, content_with_weight, # topic_kwd, entity_names_kwd, outlinks_kwd, related_kb_pages_kwd, # source_chunk_ids, source_doc_ids @@ -1570,6 +1570,19 @@ def _compilation_template_kind(kind) -> str: return normalized +def _scalar(raw, default=""): + """Infinity ``get_fields`` returns every ``*_kwd`` field as a list (split + on ``###``), even single scalar values like ``slug_kwd=["entity/foo"]``. + Normalize a field value that is expected to be a scalar identifier back to + the first non-empty element.""" + if isinstance(raw, (list, tuple)): + for item in raw: + if item not in (None, ""): + return item + return default + return raw if raw not in (None, "") else default + + def _normalize_compilation_template_group_ids(raw) -> list[str]: if isinstance(raw, str): raw = [raw] @@ -2128,6 +2141,26 @@ def _alteration_result(current_doc_ids: set, involved_doc_ids: set, eligible_doc } +def _flatten_provenance_doc_ids(value) -> set[str]: + """Normalize source_doc_ids stored as JSON strings, lists, or scalars.""" + if value is None: + return set() + if isinstance(value, str): + raw = value.strip() + if not raw: + return set() + try: + return _flatten_provenance_doc_ids(json.loads(raw)) + except (json.JSONDecodeError, TypeError): + return {raw} + if isinstance(value, (list, tuple, set)): + result: set[str] = set() + for item in value: + result.update(_flatten_provenance_doc_ids(item)) + return result + return {str(value)} + + def _eligible_doc_ids_for_kind(docs, tenant_id: str, kind: str) -> set: """Doc ids whose parser_config or pipeline carries a template of ``kind``.""" accepted = _ALTERATION_ELIGIBLE_TEMPLATE_KINDS.get(kind) or set() @@ -2184,10 +2217,7 @@ async def _involved_doc_ids_paged(index_nm, dataset_id: str, condition: dict, fi for row in rows.values(): value = row.get(field) if from_list: - if isinstance(value, str): - value = [value] - if isinstance(value, list): - involved.update(str(d) for d in value if d) + involved.update(_flatten_provenance_doc_ids(value)) elif value: involved.add(str(value)) @@ -2328,15 +2358,15 @@ async def list_wiki_pages( total = settings.docStoreConn.get_total(res) items = [] for row in (field_map or {}).values(): - slug = row.get("slug_kwd") - if not isinstance(slug, str) or not slug: + slug = _scalar(row.get("slug_kwd")) + if not slug: continue items.append( { "slug": slug, - "title": row.get("title_kwd") or slug, - "page_type": row.get("page_type_kwd") or "concept", - "topic": row.get("topic_kwd") or "", + "title": _scalar(row.get("title_kwd")) or slug, + "page_type": _scalar(row.get("page_type_kwd")) or "concept", + "topic": _scalar(row.get("topic_kwd")) or "", "summary": row.get("summary_with_weight") or "", } ) @@ -2411,9 +2441,12 @@ async def list_wiki_topics( if not rows: break for row in rows.values(): - t = row.get("topic_kwd") - if isinstance(t, str) and t: - meta[t] = {"title": row.get("title_kwd") or t, "slug": row.get("slug_kwd") or t} + t = _scalar(row.get("topic_kwd")) + if t: + meta[t] = { + "title": _scalar(row.get("title_kwd")) or t, + "slug": _scalar(row.get("slug_kwd")) or t, + } _offset += _BATCH except Exception: logging.exception("list_wiki_topics: topic metadata lookup failed for kb=%s", dataset_id) @@ -2470,6 +2503,7 @@ async def get_wiki_page( "title_kwd", "page_type_kwd", "topic_kwd", + "md_with_weight", "content_with_weight", "summary_with_weight", "entity_names_kwd", @@ -2507,13 +2541,15 @@ async def get_wiki_page( return True, None _, row = next(iter(field_map.items())) - content_md = row.get("content_with_weight") or "" + # The incremental writer stores page body in md_with_weight; fall back to + # content_with_weight for any rows written by the legacy path. + content_md = row.get("md_with_weight") or row.get("content_with_weight") or "" summary = row.get("summary_with_weight") or "" return True, { - "slug": row.get("slug_kwd") or full_slug, - "title": row.get("title_kwd") or full_slug, - "page_type": row.get("page_type_kwd") or page_type, - "topic": row.get("topic_kwd") or "", + "slug": _scalar(row.get("slug_kwd")) or full_slug, + "title": _scalar(row.get("title_kwd")) or full_slug, + "page_type": _scalar(row.get("page_type_kwd")) or page_type, + "topic": _scalar(row.get("topic_kwd")) or "", "content_md_rendered": content_md, "summary": summary, "entity_names": row.get("entity_names_kwd") or [], @@ -3187,7 +3223,7 @@ def _wiki_entity_payload(row: dict) -> dict | None: payload = parsed except Exception: pass - slug = payload.get("slug") or row.get("slug_kwd") + slug = payload.get("slug") or _scalar(row.get("slug_kwd")) if not isinstance(slug, str) or not slug: return None out = { @@ -3283,7 +3319,13 @@ async def _wiki_search_entities_by_slugs( dataset_id: str, slugs: list[str], ): - """Fetch entity rows whose ``slug_kwd`` is in ``slugs``. Unordered.""" + """Fetch entity rows whose ``slug_kwd`` is in ``slugs``. Unordered. + + Like :func:`_wiki_search_relations_from`, we avoid pushing ``slug_kwd`` (a + *_kwd analysed field) into the search filter — a `slug_kwd: [..]` with ~20 + entries triggers TOO_MANY_CONNECTIONS. Pull all entity rows once and filter + in memory. + """ if not slugs: return {} @@ -3296,22 +3338,35 @@ async def _wiki_search_entities_by_slugs( "source_chunk_ids", "content_with_weight", ] - res = await thread_pool_exec( - settings.docStoreConn.search, - select_fields, - [], - { - "compile_kwd": [_WIKI_GRAPH_ENTITY_KWD], - "slug_kwd": list(slugs), - }, - [], - OrderByExpr(), - 0, - max(len(slugs), 1), - index_nm, - [dataset_id], - ) - return settings.docStoreConn.get_fields(res, select_fields) + wanted = set(slugs) + results = {} + offset, page_size = 0, 1000 + while True: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"compile_kwd": [_WIKI_GRAPH_ENTITY_KWD]}, + [], + OrderByExpr(), + offset, + page_size, + index_nm, + [dataset_id], + ) + rows = settings.docStoreConn.get_fields(res, select_fields) + if not rows: + break + for row in rows.values(): + slug = row.get("slug_kwd") + if isinstance(slug, list): + slug = slug[0] if slug else "" + if slug in wanted: + results[row.get("id", len(results))] = row + if len(rows) < page_size: + break + offset += page_size + return results async def _wiki_search_relations_from( @@ -3319,31 +3374,54 @@ async def _wiki_search_relations_from( dataset_id: str, from_slugs: list[str], ): - """Fetch all relation rows with ``from_kwd`` in ``from_slugs``.""" + """Fetch relation rows whose ``from_kwd`` is in ``from_slugs``. + + IMPORTANT: we do NOT push ``from_slugs`` into the search filter. ``from_kwd`` + is a *_kwd (whitespace-# analysed) field, and the generic search path turns + `from_kwd: [v1, v2, ...]` into one ``filter_fulltext`` clause per value. A + batch of only ~20 slugs already blows past Infinity's per-query connection + budget and surfaces as TOO_MANY_CONNECTIONS (the incremental writer emits + many relations, so sub_slugs easily exceeds 20). Instead we pull ALL relation + rows for the dataset in ONE cheap query (relations are short and few) and + filter in memory. + """ if not from_slugs: return {} from common.doc_store.doc_store_base import OrderByExpr select_fields = ["id", "from_kwd", "to_kwd", "content_with_weight"] - # Generous upper bound: relations are short; bulk-pull all matching at - # once rather than paging. - res = await thread_pool_exec( - settings.docStoreConn.search, - select_fields, - [], - { - "compile_kwd": [_WIKI_GRAPH_RELATION_KWD], - "from_kwd": list(from_slugs), - }, - [], - OrderByExpr(), - 0, - 10000, - index_nm, - [dataset_id], - ) - return settings.docStoreConn.get_fields(res, select_fields) + wanted = set(from_slugs) + # Single query without the huge from_kwd IN-filter. Page over results in + # case a dataset has more than 10000 relations. + results = {} + offset, page_size = 0, 1000 + while True: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"compile_kwd": [_WIKI_GRAPH_RELATION_KWD]}, + [], + OrderByExpr(), + offset, + page_size, + index_nm, + [dataset_id], + ) + rows = settings.docStoreConn.get_fields(res, select_fields) + if not rows: + break + for row in rows.values(): + frm = row.get("from_kwd") + if isinstance(frm, list): + frm = frm[0] if frm else "" + if frm in wanted: + results[row.get("id", len(results))] = row + if len(rows) < page_size: + break + offset += page_size + return results async def get_wiki_graph( diff --git a/api/db/init_data.py b/api/db/init_data.py index d42304a733..22f155a4c7 100644 --- a/api/db/init_data.py +++ b/api/db/init_data.py @@ -131,6 +131,8 @@ def add_graph_templates(): def add_compilation_templates(): + CompilationTemplateService.ensure_table() + CompilationTemplateService.filter_delete([CompilationTemplateService.model.is_builtin]) CompilationTemplateService.seed_builtins_from_files() diff --git a/api/db/init_data/compilation_templates/wiki.yaml b/api/db/init_data/compilation_templates/wiki.yaml index dc2949d060..507d6a0bd7 100644 --- a/api/db/init_data/compilation_templates/wiki.yaml +++ b/api/db/init_data/compilation_templates/wiki.yaml @@ -2,6 +2,7 @@ kind: wiki display_name: Wiki — Graph-based wiki config: kind: wiki + plan: yes example: | - Each page must be a proper encyclopedic article, NOT a flat bullet list: - 1. Opening paragraph (2-4 sentences defining what this is). No heading. diff --git a/api/db/services/doc_metadata_service.py b/api/db/services/doc_metadata_service.py index 8f2f2a875c..4c8300a691 100644 --- a/api/db/services/doc_metadata_service.py +++ b/api/db/services/doc_metadata_service.py @@ -955,7 +955,7 @@ class DocMetadataService: where_clause = f"{kb_filter} AND {sql_filter}" logging.debug(f"Infinity metadata filter: {where_clause}") - inf_conn = settings.docStoreConn.connPool.get_conn() + inf_conn = settings.docStoreConn.acquire_conn() try: db_instance = inf_conn.get_database(settings.docStoreConn.dbName) table_instance = db_instance.get_table(index_name) diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index b639bb291c..0a88e85935 100644 --- a/api/db/services/document_service.py +++ b/api/db/services/document_service.py @@ -678,6 +678,32 @@ class DocumentService(CommonService): doc.kb_id, ) + # 3. Clean up doc_page_source tracking rows (new incremental design). + try: + doc_page_kwd = "wiki_doc_page_source" + res = settings.docStoreConn.search( + ["id"], + [], + {"compile_kwd": [doc_page_kwd], "doc_id": [doc.id]}, + [], + OrderByExpr(), + 0, + 10, + index, + doc.kb_id, + ) + if settings.docStoreConn.get_fields(res, ["id"]): + settings.docStoreConn.delete( + {"compile_kwd": [doc_page_kwd], "doc_id": [doc.id]}, + index, + doc.kb_id, + ) + except Exception: + logging.exception( + "DocumentService.remove_wiki_products: doc_page_source cleanup failed for doc %s", + doc.id, + ) + @classmethod @DB.connection_context() def get_newly_uploaded(cls): diff --git a/api/db/services/task_service.py b/api/db/services/task_service.py index 8d4ecc0433..d8f9699ea0 100644 --- a/api/db/services/task_service.py +++ b/api/db/services/task_service.py @@ -392,7 +392,11 @@ class TaskService(CommonService): - progress_msg (str, optional): Progress message to append - progress (float, optional): Progress percentage (0.0 to 1.0) """ - task = cls.model.get_by_id(id) + try: + task = cls.model.get_by_id(id) + except cls.model.DoesNotExist: + logging.info("Skip progress update for deleted task %s", id) + return if not task: logging.warning("Update_progress error: task not found") return diff --git a/conf/infinity_mapping.json b/conf/infinity_mapping.json index de40f58df7..d6bbc72aee 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-#"}, + "plan_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"}}, @@ -64,8 +65,19 @@ "topic_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "page_type_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "entity_names_kwd": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, + "entity_names": {"type": "json", "default": "[]", "enabled": false}, "outlinks_kwd": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, "related_kb_pages_kwd": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, + + "page_version_int": {"type": "integer", "default": 1}, + "synthesis_version_int": {"type": "integer", "default": 0}, + "claims": {"type": "json", "default": "[]", "enabled": false}, + "page_ids": {"type": "json", "default": "[]", "index_type": {"type": "secondary", "cardinality": "low"}}, + "page_id": {"type": "varchar", "default": ""}, + "deleted_doc_id": {"type": "varchar", "default": ""}, + "source_chunk_hashes": {"type": "json", "default": "{}", "enabled": false}, + "map_checksum": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, + "type_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "from_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "to_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, @@ -77,5 +89,8 @@ "depth_int": {"type": "integer", "default": 0}, "parent_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"}, "outlinks_int": {"type": "integer", "default": 0}, - "token_num": {"type": "integer", "default": 0} + "token_num": {"type": "integer", "default": 0}, + + "aliases": {"type": "json", "default": "[]", "enabled": false}, + "aliases_flat_kwd": {"type": "varchar", "default": "", "analyzer": "whitespace-#"} } diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index 07684db355..d1a37c4c79 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -86,7 +86,7 @@ from .structure import ( # --------------------------------------------------------------------------- WIKI_MAP_COMPILE_KWD = "wiki_map_extract" -DEFAULT_WIKI_MAP_WORKERS = 6 +DEFAULT_WIKI_MAP_WORKERS = 20 DEFAULT_WIKI_MAP_TIMEOUT = 600 @@ -838,13 +838,14 @@ async def _wiki_extract_one_batch( chunk_id_list="\n".join(f"- {label}" for label in labels), packed_chunks=body, ) + request_conf = _knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1}) try: res = await asyncio.wait_for( gen_json( WIKI_MAP_SYSTEM, user_prompt, chat_mdl, - gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1}), + gen_conf=request_conf, ), timeout=llm_timeout, ) @@ -958,7 +959,8 @@ async def wiki_map_from_chunks( extracted item via ``chunk_ids``. tenant_id, kb_id: address the doc-store index for resume reads + writes. language: reserved for future prompt localization. - max_workers: maximum concurrent batches. Defaults to 6. + max_workers: maximum concurrent batches. Defaults to 20, matching the + task-scoped Wiki LLM pool used by the task executor. llm_timeout: seconds per batch extraction call. callback: optional ``(progress: float, msg: str)`` progress callback. parser_config: optional YAML-style config (same shape that @@ -974,7 +976,6 @@ async def wiki_map_from_chunks( performed here — that is the REDUCE phase's responsibility. """ _ = embd_mdl # noqa: F841 — accepted for symmetry with downstream phases - if not chunks: # Even with zero chunks we still want to sweep any orphaned MAP # rows that point at chunks the doc no longer has — otherwise @@ -1983,13 +1984,14 @@ async def _wiki_resolve_maybe_items( "Return ONLY the JSON array.\n\n" + "\n".join(lines) ) + request_conf = _knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}) try: res = await asyncio.wait_for( gen_json( WIKI_PLAN_RECONCILE_SYSTEM, user_prompt, chat_mdl, - gen_conf=_knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0}), + gen_conf=request_conf, ), timeout=llm_timeout, ) @@ -2132,16 +2134,17 @@ async def _wiki_planning_call( max_page_count=max_page_count, ) + request_conf = _knowledge_compile_gen_conf( + chat_mdl, + {"temperature": 0.1, "max_tokens": output_tokens}, + ) try: res = await asyncio.wait_for( gen_json( WIKI_PLAN_PLANNING_SYSTEM, user_prompt, chat_mdl, - gen_conf=_knowledge_compile_gen_conf( - chat_mdl, - {"temperature": 0.1, "max_tokens": output_tokens}, - ), + gen_conf=request_conf, ), timeout=llm_timeout, ) @@ -2585,6 +2588,7 @@ WIKI_TEMPLATE_EXAMPLE = ( "Each page must be a proper encyclopedic article, NOT a flat bullet list:\n" "1. Opening paragraph (2-4 sentences defining what this is). No heading.\n" "2. Sections with H2 headings, each starting with prose before sub-bullets.\n" + " Put every heading on its own line and separate every paragraph with a blank line.\n" "3. Bold key terms on first use; link them with [[ ]] wikilinks.\n" "4. Examples or implications where the source provides them.\n" "5. ## See also section at the end with wikilinks to highly related pages(less than 12).\n\n" @@ -2687,6 +2691,7 @@ But also look for additional relevant information in the source text above. ## Instructions Write the complete wiki page in markdown based on the source text above. +Put every heading on its own line and separate every paragraph with a blank line. Do not return the page as one line. Cross-link to other pages using [[slug]] or [[slug|display text]] — ONLY use slugs from the "Available pages" list. Do NOT invent new slugs. Do NOT include Citations or Footnotes sections. @@ -3227,12 +3232,13 @@ async def _wiki_chat_text( _, msg = message_fit_in(msg, chat_mdl.max_length) except Exception: logging.exception("wiki_refine: message_fit_in failed; sending untrimmed") + request_conf = _knowledge_compile_gen_conf(chat_mdl, {"temperature": temperature}) try: raw = await asyncio.wait_for( chat_mdl.async_chat( msg[0]["content"], msg[1:], - _knowledge_compile_gen_conf(chat_mdl, {"temperature": temperature}), + request_conf, ), timeout=llm_timeout, ) @@ -3284,13 +3290,14 @@ async def _wiki_write_page_simple( evidence_blocks=_wiki_format_evidence_blocks(evidence), ) - return await _wiki_chat_text( + content = await _wiki_chat_text( chat_mdl, _build_refine_writer_system(example), user_prompt, temperature=0.15, llm_timeout=llm_timeout, ) + return content async def _wiki_merge_page_content( diff --git a/rag/advanced_rag/knowlege_compile/wiki_incremental.py b/rag/advanced_rag/knowlege_compile/wiki_incremental.py new file mode 100644 index 0000000000..1ee89ef68e --- /dev/null +++ b/rag/advanced_rag/knowlege_compile/wiki_incremental.py @@ -0,0 +1,3690 @@ +"""Dual-mode wiki incremental compilation. + +Mode A (no-plan, plan=no): + MAP → REDUCE → REFINE per-concept (generate/modify/re-synthesize) → FINALIZE + 1 concept = 1 page (WeKnora style). Entities enrich concept pages via source chunks. + +Mode B (with-plan, plan=yes): + MAP → REDUCE → PLAN (LLM grouping) → REFINE per-page → FINALIZE + Incremental: Page Router (KNN) routes entities to existing pages. + +Both modes share MAP + REDUCE + FINALIZE. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import re +from typing import Callable + +import numpy as np + +from common import settings +from common.doc_store.doc_store_base import MatchDenseExpr, OrderByExpr +from common.misc_utils import thread_pool_exec +from rag.prompts.generator import message_fit_in +from rag.nlp import search + +from ._common import ( + knowledge_compile_gen_conf as _knowledge_compile_gen_conf, + stable_row_id as _stable_row_id, +) + + +# ----- REFINE concurrency control ----- + +WIKI_REFINE_MAX_CONCURRENT = 20 # shared LLM pool size used by the Wiki runner + + +# ----- constants ---- + +# compile_kwd values +WIKI_PAGE_COMPILE_KWD = "wiki_page" +WIKI_PLAN_GROUP_COMPILE_KWD = "wiki_plan_group" +WIKI_DOC_PAGE_SOURCE_COMPILE_KWD = "wiki_doc_page_source" +WIKI_CANONICAL_ENTITY_COMPILE_KWD = "wiki_canonical_entity" + +# Entity matching thresholds (not exposed in YAML) +ENTITY_MERGE_THRESHOLD = 0.90 # auto-merge +ENTITY_AMBIGUOUS_LOW = 0.75 # LLM confirm boundary +ENTITY_PAIRWISE_BLOCK_SIZE = 1024 # blockwise embedding matrix block size + +# Number of concurrent KNN queries for entity matching +ENTITY_MATCH_KNN_CONCURRENT = 20 +CANONICAL_PERSIST_CONCURRENT = 20 +PAGE_ROUTER_KNN_CONCURRENT = 20 + +# Thematic topic grouping. No-plan pages have no PLAN step, so pages are grouped +# post-hoc by matching them to the thematic topic labels the MAP phase extracted. +WIKI_TOPIC_MATCH_THRESHOLD = 0.50 # min cosine for a page to attach to a topic +WIKI_TOPIC_MAX_LABELS = 200 # cap on candidate topic labels +WIKI_TOPIC_FALLBACK = "General" # bucket for pages that match no topic +WIKI_TOPIC_UPDATE_CONCURRENT = 16 # concurrent page topic_kwd updates + +# Page Router thresholds (kept as code constants — not exposed in YAML) +PAGE_ROUTER_UPDATE_THRESHOLD = 0.80 +PAGE_ROUTER_MAYBE_THRESHOLD = 0.50 +PAGE_ROUTER_CLUSTER_THRESHOLD = 0.50 + +# Re-synthesis triggers (both modes) +RE_SYNTHESIS_MIN_SOURCES = 5 +RE_SYNTHESIS_GROWTH_RATIO = 1.5 +RE_SYNTHESIS_MIN_CLAIMS = 15 +RE_SYNTHESIS_MIN_VERSIONS = 3 + +# Evidence quality (WeKnora-style verbatim chunk sourcing). +# The page-writer sees the ACTUAL source-chunk text (not just the condensed +# claim statement) so pages stay fact-grounded and information-dense. +WIKI_SOURCE_BUDGET_CHARS = 32_768 # cap on verbatim chunk text fed to the writer +WIKI_SOURCE_BUDGET_RUNES = 12_000 # per-chunk-batch budget (rune-based, mirrors WeKnora) + + +# ----- helpers --------------------------------------------------------------- + + +def _wiki_derive_page_id(term: str, prefix: str = "concept") -> str: + """Derive a URL-safe page identifier from a concept/entity name. + + Example: "Smartphone Industry" → "concept/smartphone-industry" + """ + slug = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fff]+", "-", term).strip("-").lower() + return f"{prefix}/{slug}" + + +def _entity_to_query_text(entity: dict) -> str: + return " ".join( + [ + entity.get("entity_name") or entity.get("name") or entity.get("term") or "", + entity.get("definition_excerpt") or entity.get("description") or entity.get("statement", ""), + ][:2] + ) + + +def _strip_think(text: str) -> str: + if not isinstance(text, str): + return "" + text = text.strip() + if text.startswith(""): + return text.split("", 1)[-1].strip() + return text + + +async def _chat_mdl_ask(chat_mdl, system_prompt: str, user_prompt: str, temperature: float = 0.0) -> str: + msg = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + try: + _, msg = message_fit_in(msg, chat_mdl.max_length) + except Exception: + logging.exception("wiki incremental: message_fit_in failed; sending untrimmed") + request_conf = _knowledge_compile_gen_conf(chat_mdl, {"temperature": temperature}) + try: + raw = await chat_mdl.async_chat(msg[0]["content"], msg[1:], request_conf) + except Exception: + raise + if isinstance(raw, tuple): + raw = raw[0] + return _strip_think(raw or "") + + +def _wiki_should_re_synthesize( + page: dict, + new_source_doc_ids: set[str], + next_version: int, +) -> bool: + existing_sources = set(page.get("source_doc_ids", [])) + total_sources = existing_sources | new_source_doc_ids + claim_count = len(page.get("claims", [])) + last_synth_ver = page.get("synthesis_version_int", 1) + versions_since = next_version - last_synth_ver + + return ( + len(total_sources) >= RE_SYNTHESIS_MIN_SOURCES + and claim_count >= RE_SYNTHESIS_MIN_CLAIMS + and versions_since >= RE_SYNTHESIS_MIN_VERSIONS + and len(total_sources) >= len(existing_sources) * RE_SYNTHESIS_GROWTH_RATIO + ) + + +async def _wiki_load_chunk_texts(tenant_id: str, kb_id: str, chunk_ids: list[str]) -> dict[str, str]: + """Fetch verbatim chunk text from the doc store by chunk id. + + Returns ``{chunk_id: content_with_weight}``. Used to ground page writing in + the ACTUAL source text (WeKnora-style evidence) rather than the condensed + claim statements. Mirrors ``wiki._wiki_load_chunks_by_id`` which lives in + the old-mode module and is intentionally NOT imported here. + """ + if not chunk_ids: + return {} + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + unique = [c for c in dict.fromkeys(chunk_ids) if isinstance(c, str) and c] + if not unique: + return {} + out: dict[str, str] = {} + BATCH = 500 + for i in range(0, len(unique), BATCH): + batch = unique[i : i + BATCH] + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["id", "content_with_weight"], + [], + {"id": batch}, + [], + OrderByExpr(), + 0, + len(batch), + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, ["id", "content_with_weight"]) or {} + except Exception: + logging.exception("wiki: batch chunk fetch failed (%d ids)", len(batch)) + field_map = {} + for cid, row in field_map.items(): + content = row.get("content_with_weight") + if isinstance(content, str) and content: + out[cid] = content + # Honor the rune budget: do not accumulate verbatim text we won't feed the + # writer (mirrors WeKnora maxRunesPerCitationBatch=12000). + total_runes = sum(len(v) for v in out.values()) + if total_runes > WIKI_SOURCE_BUDGET_RUNES: + trimmed: dict[str, str] = {} + budget = 0 + for cid, content in out.items(): + budget += len(content) + if budget > WIKI_SOURCE_BUDGET_RUNES: + break + trimmed[cid] = content + out = trimmed + return out + + +def _wiki_enrich_source_chunks(source_chunks: list[dict], chunk_texts: dict[str, str]) -> list[dict]: + """Merge verbatim chunk text into source_chunks. + + Each source_chunk entry is ``{"id": , "text": }``. + When the verbatim chunk content is available it REPLACES ``text`` (the + writer should read the source, not the condensed claim); otherwise the + claim text is kept as a fallback. Preserves original order and dedups by id. + """ + enriched: list[dict] = [] + seen: set[str] = set() + for sc in source_chunks: + cid = sc.get("id") or sc.get("chunk_id") + if not cid: + continue + cid = str(cid) + if cid in seen: + continue + seen.add(cid) + verbatim = chunk_texts.get(cid) + enriched.append( + { + "id": cid, + "text": verbatim if verbatim else sc.get("text", sc.get("content_with_weight", "")), + "_verbatim": bool(verbatim), + } + ) + return enriched + + +# ----- Canonical Entity Index CRUD ----------------------------------------- + + +async def _wiki_has_any_pages(tenant_id: str, kb_id: str) -> bool: + """Return True if the KB has at least one compiled wiki_page row. + + Used to detect whether a build has ever succeeded — if no page exists but + MAP rows do, the previous build was interrupted and should be restarted + as a full build rather than treated as incremental. + """ + index = search.index_name(tenant_id) + try: + if not settings.docStoreConn.index_exist(index, kb_id): + return False + res = await thread_pool_exec( + settings.docStoreConn.search, + ["slug_kwd"], + [], + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD]}, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + return bool(settings.docStoreConn.get_fields(res, ["slug_kwd"])) + except Exception: + logging.exception("wiki: _wiki_has_any_pages failed for kb=%s", kb_id) + return False + + +async def _load_canonical_entities( + tenant_id: str, + kb_id: str, +) -> dict[str, dict]: + """Load all canonical entity rows. Returns {entity_name: row}.""" + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return {} + results: dict[str, dict] = {} + offset = 0 + page_size = 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["entity_kwd", "entity_type_kwd", "aliases", "source_doc_ids", "mention_count_int"], + [], + {"compile_kwd": [WIKI_CANONICAL_ENTITY_COMPILE_KWD]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, ["entity_kwd", "entity_type_kwd", "aliases", "source_doc_ids", "mention_count_int"]) or {} + except Exception: + logging.exception("wiki: failed to load canonical entities for kb=%s", kb_id) + return results + for row in field_map.values(): + name = row.get("entity_kwd", "") + if isinstance(name, list): + # entity_kwd is analyzed by whitespace-# (shared schema field), + # so Infinity may return it as a token list. Rejoin with spaces + # to reconstruct the canonical entity name. + name = " ".join(str(t) for t in name if t) + name = str(name or "").strip() + if name: + # Deserialize JSON fields + for fld in ("aliases", "source_doc_ids"): + val = row.get(fld) + if isinstance(val, str): + try: + row[fld] = json.loads(val) if val else [] + except (json.JSONDecodeError, TypeError): + row[fld] = [] + mc = row.get("mention_count_int", 0) + if isinstance(mc, str): + mc = int(mc) if mc.isdigit() else 0 + row["mention_count_int"] = mc + # entity_type_kwd is a *_kwd field → Infinity returns it as a + # list (e.g. ['concept']). Normalize to a scalar so downstream + # checks like `entity_type == "concept"` work correctly. + et = row.get("entity_type_kwd") + if isinstance(et, list): + et = et[0] if et else "" + row["entity_type_kwd"] = str(et or "entity").strip() + results[name] = row + if len(field_map) < page_size: + break + offset += page_size + return results + + +def _build_canonical_entity_doc( + tenant_id: str, + kb_id: str, + entity_name: str, + entity_type: str, + aliases: list[str], + source_doc_ids: list[str], + claim_count: int, + embedding: list[float] | None = None, +) -> dict: + """Build a canonical entity row for insert or update.""" + dim = len(embedding) if embedding else 768 + doc = { + "id": _stable_row_id(WIKI_CANONICAL_ENTITY_COMPILE_KWD, kb_id, entity_name), + "entity_kwd": entity_name, + "entity_type_kwd": entity_type, + "aliases": json.dumps(list(set(aliases)), ensure_ascii=False), + "source_doc_ids": json.dumps(list(set(source_doc_ids)), ensure_ascii=False), + "mention_count_int": claim_count, + "compile_kwd": WIKI_CANONICAL_ENTITY_COMPILE_KWD, + "kb_id": kb_id, + } + if embedding is not None: + vec_col = f"q_{dim}_vec" + doc[vec_col] = embedding + return doc + + +async def _save_canonical_entity( + tenant_id: str, + kb_id: str, + entity_name: str, + entity_type: str, + aliases: list[str], + source_doc_ids: list[str], + claim_count: int, + embedding: list[float] | None = None, +) -> None: + """Insert or update a canonical entity row.""" + index = search.index_name(tenant_id) + doc = _build_canonical_entity_doc( + tenant_id, + kb_id, + entity_name, + entity_type, + aliases, + source_doc_ids, + claim_count, + embedding, + ) + + condition = {"compile_kwd": [WIKI_CANONICAL_ENTITY_COMPILE_KWD], "entity_kwd": [entity_name]} + existing = await thread_pool_exec( + settings.docStoreConn.search, + ["entity_kwd"], + [], + condition, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + if settings.docStoreConn.get_fields(existing, ["entity_kwd"]): + await thread_pool_exec( + settings.docStoreConn.update, + {"compile_kwd": [WIKI_CANONICAL_ENTITY_COMPILE_KWD], "entity_kwd": entity_name}, + doc, + index, + kb_id, + ) + else: + await thread_pool_exec(settings.docStoreConn.insert, [doc], index, kb_id) + + +async def _update_canonical_entity( + tenant_id: str, + kb_id: str, + entity_name: str, + entity_type: str, + aliases: list[str], + source_doc_ids: list[str], + claim_count: int, +) -> None: + """Update a known canonical row without an existence query.""" + index = search.index_name(tenant_id) + doc = _build_canonical_entity_doc( + tenant_id, + kb_id, + entity_name, + entity_type, + aliases, + source_doc_ids, + claim_count, + ) + await thread_pool_exec( + settings.docStoreConn.update, + {"compile_kwd": [WIKI_CANONICAL_ENTITY_COMPILE_KWD], "entity_kwd": entity_name}, + doc, + index, + kb_id, + ) + + +async def _delete_canonical_entity( + tenant_id: str, + kb_id: str, + entity_name: str, +) -> None: + """Delete a canonical entity row.""" + index = search.index_name(tenant_id) + await thread_pool_exec( + settings.docStoreConn.delete, + {"compile_kwd": [WIKI_CANONICAL_ENTITY_COMPILE_KWD], "entity_kwd": [entity_name]}, + index, + kb_id, + ) + + +async def _knn_search_canonical( + tenant_id: str, + kb_id: str, + embedding: list[float], + threshold: float = ENTITY_MERGE_THRESHOLD, +) -> tuple[str, float] | None: + """KNN search canonical entity index. Returns (entity_name, score) or None.""" + index = search.index_name(tenant_id) + dim = len(embedding) + match_expr = MatchDenseExpr( + vector_column_name=f"q_{dim}_vec", + embedding_data=embedding, + embedding_data_type="float", + distance_type="cosine", + topn=1, + extra_options={"similarity": threshold}, + ) + res = await thread_pool_exec( + settings.docStoreConn.search, + ["entity_kwd", "_score"], + [], + {"compile_kwd": [WIKI_CANONICAL_ENTITY_COMPILE_KWD]}, + [match_expr], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, ["entity_kwd", "_score"]) + for row in field_map.values(): + name = row.get("entity_kwd", "") + if isinstance(name, list): + name = " ".join(str(t) for t in name if t) + name = str(name or "").strip() + score = row.get("_score", 0.0) + if name and score >= threshold: + return name, score + return None + + +def _normalize_key(name: str) -> str: + """Lowercase + strip whitespace + strip ASCII punctuation.""" + if not isinstance(name, str): + return "" + return re.sub(r"[^\w\s]", "", name.lower()).strip() + + +# ----- Entity Matching ----------------------------------------------------- + + +def _extract_raw_entities(map_results: list[dict]) -> tuple[list[dict], dict[str, list[dict]]]: + """Extract lightweight entity/concept metadata + a claim index from MAP. + + Returns a tuple: + (entities, claim_index) + entities: list of LIGHTWEIGHT dicts {name, type, aliases, claim_count, + source_doc_ids} — NO full claim text, so Entity Matching + operates on small metadata (mirrors old-mode dedup). + claim_index: {name: [claim_dict, ...]} — full claim text kept separately, + loaded on-demand only for affected entities after matching. + """ + raw: dict[str, dict] = {} + claim_index: dict[str, list[dict]] = {} + for mr in map_results: + doc_id = mr.get("doc_id", "") + + # Process entities[] + for ent in mr.get("entities") or []: + if isinstance(ent, str): + ent = json.loads(ent) + name = ent.get("name", "") + if not name: + continue + if name not in raw: + raw[name] = { + "name": name, + "type": ent.get("type", "entity"), + "aliases": ent.get("aliases") or [], + "claim_count": 0, + "source_doc_ids": set(), + } + raw[name]["source_doc_ids"].add(doc_id) + + # Process concepts[] + for concept in mr.get("concepts") or []: + if isinstance(concept, str): + concept = json.loads(concept) + term = concept.get("term", "") + if not term: + continue + if term not in raw: + raw[term] = { + "name": term, + "type": "concept", + "aliases": [term], + "claim_count": 0, + "source_doc_ids": set(), + } + raw[term]["source_doc_ids"].add(doc_id) + + # Process claims, tracking count (metadata) but storing full text only + # in claim_index (kept separate, loadable on demand). + for claim in mr.get("claims") or []: + if isinstance(claim, str): + claim = json.loads(claim) + subj = claim.get("entity_name") or claim.get("subject") or claim.get("term", "") + if not subj: + continue + if subj in raw: + raw[subj]["claim_count"] += 1 + claim_index.setdefault(subj, []).append(claim) + + result = [] + for entry in raw.values(): + entry["source_doc_ids"] = list(entry["source_doc_ids"]) + result.append(entry) + return result, claim_index + + +async def _wiki_match_entities( + raw_entities: list[dict], + existing_canonical: dict[str, dict], + embd_mdl, + chat_mdl, + tenant_id: str, + kb_id: str, + incremental: bool, + progress: Callable[[str], None] | None = None, +) -> tuple[dict[str, dict], dict[str, str]]: + """Entity Matching: raw entities → canonical entities. + + Returns: + canonical_map: {canonical_name: merged_entry} + name_resolution: {raw_name: canonical_name}entries still need semantic matching + """ + + def _progress(msg: str) -> None: + logging.info("wiki entity matching: %s", msg) + if progress: + try: + progress(f"Entity Matching: {msg}") + except Exception: + logging.exception("wiki: entity matching progress callback failed") + + def _progress_interval(total: int) -> int: + if total <= 20: + return max(total, 1) + return max(10, min(200, total // 10)) + + # Step 1: Exact match against canonical index + _progress(f"exact matching {len(raw_entities)} raw entries against {len(existing_canonical)} canonical entries ...") + exact_flat: dict[str, str] = {} # normalized_name → canonical_name + for cname, centry in existing_canonical.items(): + # Use the non-analyzed `aliases` JSON field (deserialized in + # _load_canonical_entities) plus the entity name itself. Do NOT rely on + # `aliases_flat_kwd`: it is an analyzed varchar (whitespace-#) that + # Infinity returns as a token list, losing the "||" structure. + aliases = centry.get("aliases") + if not isinstance(aliases, list): + continue + for alias in [cname] + [a for a in aliases if isinstance(a, str)]: + exact_flat[_normalize_key(alias)] = cname + + name_resolution: dict[str, str] = {} # raw_name → canonical_name + unmatched: list[dict] = [] # entities not matched by exact + for entry in raw_entities: + raw_name = entry["name"] + norm = _normalize_key(raw_name) + if norm in exact_flat: + name_resolution[raw_name] = exact_flat[norm] + else: + unmatched.append(entry) + _progress(f"exact matched {len(name_resolution)}; {len(unmatched)} entries still need semantic matching.") + + # Step 2: KNN match for unmatched entities + # Single search at ENTITY_AMBIGUOUS_LOW (0.75), classify by score: + # 0.90+ → auto-merge (direct_match) + # 0.75-0.90 → LLM confirm for entity types only + # <0.75 → no match + if unmatched and embd_mdl and existing_canonical: + query_texts = [_entity_to_query_text(e) for e in unmatched] + embeddings, _ = await thread_pool_exec(embd_mdl.encode, query_texts) + + sem = asyncio.Semaphore(ENTITY_MATCH_KNN_CONCURRENT) + + async def _knn_one(entry: dict, vec) -> tuple[dict, str | None, float]: + if hasattr(vec, "tolist"): + vec = vec.tolist() + result = await _knn_search_canonical(tenant_id, kb_id, vec, ENTITY_AMBIGUOUS_LOW) + if result: + return entry, result[0], result[1] + return entry, None, 0.0 + + async def _async_knn(entry: dict, vec): + async with sem: + return await _knn_one(entry, vec) + + _progress(f"KNN unmatched entities {len(query_texts)} ...") + knn_tasks = [_async_knn(entry, emb) for entry, emb in zip(unmatched, embeddings)] + knn_results = await asyncio.gather(*knn_tasks) + _progress("KNN unmatched entities done.") + + still_unmatched: list[dict] = [] + maybe_pairs: list[tuple[dict, str]] = [] + for entry, cname, score in knn_results: + if cname and score >= ENTITY_MERGE_THRESHOLD: + # Direct merge + name_resolution[entry["name"]] = cname + elif cname and score >= ENTITY_AMBIGUOUS_LOW: + # Ambiguous: LLM confirm for entity types only + if entry["type"] == "concept": + name_resolution[entry["name"]] = cname + else: + maybe_pairs.append((entry, cname)) + else: + still_unmatched.append(entry) + + # LLM confirm for maybe pairs (entity only, 0.75-0.90) + if maybe_pairs and chat_mdl: + confirmed = await _wiki_confirm_batch( + [(e["name"], cname) for e, cname in maybe_pairs], + chat_mdl, + ) + confirmed_set = set() + for raw_name, cname in confirmed: + name_resolution[raw_name] = cname + confirmed_set.add(raw_name) + for e, cname in maybe_pairs: + if e["name"] not in confirmed_set: + still_unmatched.append(e) + + unmatched = still_unmatched + + # Step 3: Intra-build pairwise (only on first build, i.e., non-incremental) + # Blockwise matrix multiplication — mirrors old _embedding_dedup in + # _common.py. Avoids O(N²) Python loops and materializing an N×N matrix + # (OOM risk); peak memory is O(block²). + if not incremental and len(unmatched) > 1 and embd_mdl: + query_texts = [_entity_to_query_text(e) for e in unmatched] + embeddings, _ = await thread_pool_exec(embd_mdl.encode, query_texts) + try: + matrix = np.asarray([list(v) for v in embeddings], dtype=np.float32) + if matrix.ndim != 2 or matrix.shape[0] != len(unmatched): + raise ValueError("invalid embedding matrix shape") + norms = np.linalg.norm(matrix, axis=1, keepdims=True) + matrix = np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms > 0) + except Exception: + logging.exception("wiki: pairwise embedding failed; skipping semantic merge") + matrix = None + + if matrix is not None: + merged_into: dict[int, int] = {} + maybe_pairs: list[tuple[int, int]] = [] + + def _root(i: int) -> int: + while i in merged_into: + i = merged_into[i] + return i + + n = len(unmatched) + block_size = ENTITY_PAIRWISE_BLOCK_SIZE + # Group by type: only entities of the SAME type are pairwise candidates + # (entity-vs-entity, concept-vs-concept). This mirrors old behavior + # of passing type_key so cross-type pairs are never merged. + groups: dict[str, list[int]] = {} + for idx, entry in enumerate(unmatched): + groups.setdefault(entry.get("type", "entity"), []).append(idx) + + auto_pairs: list[tuple[int, int]] = [] + ambiguous_pairs: list[tuple[int, int]] = [] + for group_indices in groups.values(): + for left_start in range(0, len(group_indices), block_size): + left_indices = group_indices[left_start : left_start + block_size] + left_vectors = matrix[left_indices] + for right_start in range(left_start, len(group_indices), block_size): + right_indices = group_indices[right_start : right_start + block_size] + sims = left_vectors @ matrix[right_indices].T # [B, B] BLAS + if right_start == left_start: + candidate_mask = np.triu(sims >= ENTITY_AMBIGUOUS_LOW, k=1) + else: + candidate_mask = sims >= ENTITY_AMBIGUOUS_LOW + rows, cols = np.nonzero(candidate_mask) + for row, col in zip(rows.tolist(), cols.tolist(), strict=True): + score = float(sims[row, col]) + if score >= ENTITY_MERGE_THRESHOLD: + auto_pairs.append((left_indices[row], right_indices[col])) + else: + ambiguous_pairs.append((left_indices[row], right_indices[col])) + + # Apply auto-merges with union-find (higher evidence wins) + for i, j in auto_pairs: + ri, rj = _root(i), _root(j) + if ri == rj: + continue + if unmatched[ri].get("claim_count", 0) >= unmatched[rj].get("claim_count", 0): + merged_into[rj] = ri + else: + merged_into[ri] = rj + + # Keep only ambiguous pairs still in separate groups + still_ambiguous = [(i, j) for i, j in ambiguous_pairs if _root(i) != _root(j)] + + # LLM confirm for ambiguous pairs (first build only) + if still_ambiguous and chat_mdl: + llm_candidates = [(unmatched[i]["name"], unmatched[j]["name"]) for i, j in still_ambiguous] + confirmed = await _wiki_confirm_batch(llm_candidates, chat_mdl) + confirmed_map = {frozenset((a, b)) for a, b in confirmed} + for i, j in still_ambiguous: + pair = frozenset((unmatched[i]["name"], unmatched[j]["name"])) + if pair in confirmed_map: + ri, rj = _root(i), _root(j) + if ri != rj: + if unmatched[ri].get("claim_count", 0) >= unmatched[rj].get("claim_count", 0): + merged_into[rj] = ri + else: + merged_into[ri] = rj + + # Apply merges + merged_indices: dict[int, list[int]] = {} + for i in range(n): + pi = _root(i) + merged_indices.setdefault(pi, []).append(i) + + new_unmatched = [] + for pi, indices in merged_indices.items(): + if len(indices) > 1: + master = unmatched[indices[0]] + for idx in indices[1:]: + slave = unmatched[idx] + # Lightweight entries have no 'claims' field; only + # metadata is aggregated (claim text lives in claim_index + # and is aggregated later via name_resolution). + master["claim_count"] += slave["claim_count"] + master["source_doc_ids"] = list(set(master["source_doc_ids"]) | set(slave["source_doc_ids"])) + master["aliases"] = list(set(master["aliases"] + slave["aliases"] + [slave["name"]])) + name_resolution[slave["name"]] = master["name"] + new_unmatched.append(master) + else: + new_unmatched.append(unmatched[indices[0]]) + unmatched = new_unmatched + + # Step 4: Build canonical map + canonical_map: dict[str, dict] = {} + for entry in unmatched: + cname = entry["name"] + canonical_map[cname] = entry + name_resolution.setdefault(cname, cname) + + # Also load existing canonical entries that match + for raw_name, cname in name_resolution.items(): + if cname not in canonical_map: + existing = existing_canonical.get(cname) + if existing: + # Lightweight entry — claims loaded separately on demand + merged = { + "name": cname, + "type": existing.get("entity_type_kwd", "entity"), + "aliases": existing.get("aliases", []), + "claim_count": existing.get("mention_count_int", 0), + "source_doc_ids": existing.get("source_doc_ids", []), + } + canonical_map[cname] = merged + + # Aggregate lightweight metadata (claim_count, source_doc_ids) across all + # raw entities that resolved to the same canonical name. + for entry in raw_entities: + raw_name = entry["name"] + cname = name_resolution.get(raw_name, raw_name) + if cname in canonical_map: + canonical_map[cname]["claim_count"] += entry.get("claim_count", 0) + existing_docs = set(canonical_map[cname].get("source_doc_ids", [])) + existing_docs.update(entry.get("source_doc_ids", [])) + canonical_map[cname]["source_doc_ids"] = list(existing_docs) + + return canonical_map, name_resolution + + +async def _wiki_confirm_batch( + candidates: list[tuple[str, str]], + chat_mdl, +) -> list[tuple[str, str]]: + """Batch LLM confirm — adapted from old _common.py pattern. + + Takes [(name_a, name_b), ...], returns confirmed [(name_a, name_b), ...]. + """ + if not candidates: + return [] + # Split into batches of 50 + batch_size = 50 + confirmed = [] + for i in range(0, len(candidates), batch_size): + batch = candidates[i : i + batch_size] + prompt_lines = [] + for j, (a, b) in enumerate(batch): + prompt_lines.append(f'{j + 1}. "{a}" vs "{b}"') + prompt = ( + "You are a KB dedup assistant. For each pair, determine if they " + "refer to the SAME real-world entity.\n" + "Respond with a JSON array of booleans in the same order:\n" + " [true, false, true, ...]\n" + "where true = SAME entity, false = DIFFERENT.\n\n" + "\n".join(prompt_lines) + ) + try: + resp = await _chat_mdl_ask(chat_mdl, "You are a KB dedup assistant.", prompt) + if resp: + resp = resp.strip() + # Extract JSON array + arr_match = re.search(r"\[.*?\]", resp, re.DOTALL) + if arr_match: + booleans = json.loads(arr_match.group(0)) + for j, is_same in enumerate(booleans): + if is_same and j < len(batch): + confirmed.append(batch[j]) + except Exception: + logging.exception("wiki: LLM confirm batch failed") + return confirmed + + +async def _search_existing_pages( + tenant_id: str, + kb_id: str, + select_fields: list[str], +) -> dict[str, dict]: + """Load all wiki_page rows in this KB.""" + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return {} + + results: dict[str, dict] = {} + offset = 0 + page_size = 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + except Exception: + logging.exception("wiki: failed to load existing pages for kb=%s", kb_id) + return results + for row_id, row in field_map.items(): + # Keep the storage document id. FINALIZE must update by id so + # markdown values go through the document update fast path and + # retain their newlines. + row["id"] = row_id + slug = row.get("slug_kwd", row.get("page_id", "")) + if isinstance(slug, list): + slug = slug[0] if slug else "" + slug = str(slug or "").strip() + if slug: + results[slug] = row + if len(field_map) < page_size: + break + offset += page_size + return results + + +async def _load_map_relations(tenant_id: str, kb_id: str) -> list[dict]: + """Load all extracted (from, to, type) relations from wiki_map_extract rows. + + These are the semantic edges the LLM extracted during MAP. When both + endpoints correspond to compiled wiki pages they become page-to-page + connections (outlinks), which is far more reliable than hoping REFINE + sprinkled [[wikilinks]] into prose. + """ + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return [] + relations: list[dict] = [] + offset, page_size = 0, 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["content_with_weight"], + [], + {"compile_kwd": ["wiki_map_extract"]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, ["content_with_weight"]) or {} + except Exception: + logging.exception("wiki: failed to load map relations for kb=%s", kb_id) + return relations + for row in field_map.values(): + raw = row.get("content_with_weight") + if isinstance(raw, str): + try: + raw = json.loads(raw) + except (json.JSONDecodeError, TypeError): + raw = None + if not isinstance(raw, dict): + continue + for r in raw.get("relations") or []: + if isinstance(r, dict): + frm = r.get("from") + to = r.get("to") + if isinstance(frm, str) and isinstance(to, str): + relations.append({"from": frm, "to": to, "type": r.get("type", "related")}) + if len(field_map) < page_size: + break + offset += page_size + return relations + + +async def _wiki_load_pages_for_graph(tenant_id: str, kb_id: str) -> list[dict]: + """Reload compiled wiki_page rows and project them onto the canvas-graph + shape expected by ``dataset_wiki_generator.build_wiki_page_graph``. + + Each returned page has: {slug, title, summary, entity_names, + source_chunk_ids, source_doc_ids, outlinks, page_type}. Used by the + incremental entry point to materialize the artifact canvas graph after + ``wiki_compile_incremental`` persists pages (which it does internally + without returning the page list). + """ + from common.doc_store.doc_store_base import OrderByExpr + + select_fields = [ + "slug_kwd", + "title_kwd", + "page_type_kwd", + "summary_with_weight", + "md_with_weight", + "entity_names_kwd", + "outlinks_kwd", + "source_chunk_ids", + "source_doc_ids", + ] + pages: list[dict] = [] + offset, page_size = 0, 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD]}, + [], + OrderByExpr(), + offset, + page_size, + search.index_name(tenant_id), + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + except Exception: + logging.exception("wiki: failed to load pages for graph kb=%s", kb_id) + return pages + for row in field_map.values(): + slug = row.get("slug_kwd") + if isinstance(slug, (list, tuple)): + slug = slug[0] if slug else "" + slug = str(slug or "").strip() + if not slug: + continue + # outlinks_kwd is a *_kwd field analyzed by whitespace-#: Infinity + # shreds the stored JSON list and get_fields reads it back as [] (or + # a mangled token list). The reliable source of edges is the + # [[wikilinks]] actually present in the page body (auto-link + + # relation-based injection both write them). Prefer content-derived + # outlinks; fall back to the kwd field only if content has none. + outlinks = _wiki_extract_outlinks_from_content(str(row.get("md_with_weight") or ""), kb_id) + if not outlinks: + outlinks = _as_str_list(row.get("outlinks_kwd")) + title = row.get("title_kwd") + if isinstance(title, (list, tuple)): + title = title[0] if title else "" + page_type = row.get("page_type_kwd") + if isinstance(page_type, (list, tuple)): + page_type = page_type[0] if page_type else "" + pages.append( + { + "slug": slug, + "title": str(title or slug), + "summary": str(row.get("summary_with_weight") or ""), + "page_type": str(page_type or "concept"), + "entity_names": _as_str_list(row.get("entity_names_kwd")), + "outlinks": outlinks, + "source_chunk_ids": _as_str_list(row.get("source_chunk_ids")), + "source_doc_ids": _as_str_list(row.get("source_doc_ids")), + } + ) + if len(field_map) < page_size: + break + offset += page_size + return pages + + +_WIKILINK_RE = re.compile(r"\[\[([^\]]+)\]\]") + + +def _wiki_extract_outlinks_from_content(content: str, kb_id: str = "") -> list[str]: + """Extract unique internal link targets from page markdown, in order. + + Accepts both the raw ``[[slug]]`` form and the rendered Markdown form + ``[text](artifact/{kb_id}/{slug})``. Used to backfill graph edges from + page bodies (outlinks_kwd is a *_kwd field Infinity shreds on read-back). + """ + if not content: + return [] + seen: set[str] = set() + outlinks: list[str] = [] + for m in _WIKILINK_RE.finditer(content): + link = m.group(1).strip() + if link and link not in seen: + seen.add(link) + outlinks.append(link) + if kb_id: + kb_esc = re.escape(str(kb_id)) + for m in re.finditer(rf"\]\(artifact/{kb_esc}/([^)]+)\)", content): + slug = m.group(1).strip() + if slug and slug not in seen: + seen.add(slug) + outlinks.append(slug) + return outlinks + + +def _inside_wikilink(content: str, pos: int) -> bool: + """Return True iff position ``pos`` falls inside an existing [[...]] span.""" + open_pos = content.rfind("[[", 0, pos) + if open_pos < 0: + return False + # Find the FIRST "]]" that closes the "[[...]]" starting at open_pos + close_pos = content.find("]]", open_pos + 2) + if close_pos < 0: + close_pos = len(content) + return open_pos + 2 <= pos < close_pos + + +_WIKI_PIPE_LINK_RE = re.compile(r"\[\[([^\[\]\|]+?)\|([^\[\]]+?)\]\]") +_WIKI_SIMPLE_LINK_RE = re.compile(r"\[\[([^\[\]\|]+?)\]\]") + + +def _wiki_render_links(content: str, kb_id: str, valid_slugs: set[str]) -> str: + """Render internal ``[[slug]]`` / ``[[slug|text]]`` wikilinks into + navigable Markdown links ``[text](artifact/{kb_id}/{slug})``. + + This is the format the frontend wiki viewer can deep-link (plain ``[[...]]`` + renders as text and cannot navigate). Targets not in ``valid_slugs`` are + left as plain text (their label only). Mirrors the old-mode behaviour. + """ + if not content: + return content + kb = str(kb_id) + + def _simple(m: re.Match) -> str: + slug = m.group(1).strip() + if slug not in valid_slugs: + return slug + label = slug.rsplit("/", 1)[-1] if "/" in slug else slug + return f"[{label}](artifact/{kb}/{slug})" + + def _piped(m: re.Match) -> str: + slug = m.group(1).strip() + text = m.group(2).strip() + if slug not in valid_slugs: + return text + return f"[{text}](artifact/{kb}/{slug})" + + rendered = _WIKI_PIPE_LINK_RE.sub(_piped, content) + rendered = _WIKI_SIMPLE_LINK_RE.sub(_simple, rendered) + return rendered + + +def _wiki_resolve_dead_slug(link: str, valid_ids: set[str], name_slug: dict[str, str]) -> str | None: + """WeKnora-style fuzzy resolution of a dead wikilink to a live page. + + A ``[[dead_slug]]`` may fail to match a valid page because the target page + was renamed / its slug changed, or because the writer used an alias. Try, + in order: exact match, normalized-slug match, display-text reverse lookup + (via ``name_slug``), then bigram-token similarity over the plain names. + Returns a valid target slug or ``None`` (caller then degrades to text). + """ + if not link: + return None + + def _norm(s: str) -> str: + return re.sub(r"[-_]+", "-", s.strip().lower()) + + plain = link.rsplit("/", 1)[-1] if "/" in link else link + l_norm = _norm(link) + p_norm = _norm(plain) + + # 1. Exact / normalized target. + if link in valid_ids: + return link + if l_norm in valid_ids: + return l_norm + + # 2. Display-text reverse lookup via name_slug (plain names + titles + aliases). + if plain in name_slug: + return name_slug[plain] + if p_norm in {_norm(k) for k in name_slug}: + for k, v in name_slug.items(): + if _norm(k) == p_norm: + return v + + # 3. Bigram-token similarity over plain names (longest tokens, then best match). + def _bigrams(text: str) -> set[str]: + return {text[i : i + 2] for i in range(max(0, len(text) - 1))} + + p_tokens = _norm(plain).split("-") + p_bigrams = _bigrams(p_norm) + best: tuple[float, str] | None = None + for cand_name, cand_slug in name_slug.items(): + c_norm = _norm(cand_name) + c_tokens = c_norm.split("-") + # Require at least one shared token (or a shared prefix token) to avoid + # linking unrelated pages. + if not (set(p_tokens) & set(c_tokens)): + continue + cb = _bigrams(c_norm) + denom = len(p_bigrams | cb) + if denom == 0: + continue + score = len(p_bigrams & cb) / denom + if best is None or score > best[0]: + best = (score, cand_slug) + if best and best[0] >= 0.5: + return best[1] + return None + + +def _as_str_list(raw) -> list[str]: + """Coerce a stored field (JSON string / list / None) into a list of str.""" + if raw is None: + return [] + if isinstance(raw, str): + if not raw: + return [] + try: + val = json.loads(raw) + except (json.JSONDecodeError, TypeError): + return [raw] + return _as_str_list(val) + if isinstance(raw, (list, tuple)): + return [str(v) for v in raw if v is not None] + return [] + + +def _wiki_claim_chunk_ids(claim: dict) -> list[str]: + """Return the source chunk id(s) a MAP claim is attributed to. + + MAP (``wiki._wiki_resolve_chunk_ids``) rewrites every item's + ``source_chunk_id`` into ``chunk_ids=[real_id]`` (an array). Later stages + must read the array, falling back to the singular ``source_chunk_id`` for + robustness. + """ + if not isinstance(claim, dict): + return [] + ids = claim.get("chunk_ids") + if isinstance(ids, str): + ids = [ids] + if isinstance(ids, (list, tuple)): + return [str(c) for c in ids if c] + s = claim.get("source_chunk_id") + return [str(s)] if s else [] + + +def _wiki_decide_concept_pages(all_concepts: list[dict]) -> list[dict]: + """Return every concept as a wiki page. + + Mode A compiles EVERY entity AND concept into its own page, so no depth + filter is applied here. A concept that was extracted by MAP (it exists in + ``concepts[]``) is compiled even if it has no dedicated claim rows — its + page is enriched from the source chunks of the document(s) where it + appears. + """ + pages = [] + for concept in all_concepts: + claims = concept.get("claims", []) + source_docs = set(c.get("source_doc_id") for c in claims if c.get("source_doc_id")) + pages.append( + { + "page_id": _wiki_derive_page_id(concept["term"]), + "page_title": concept["term"], + "concept": concept, + "claims": claims, + "source_doc_ids": list(source_docs), + } + ) + return pages + + +# ----- REDUCE (shared, per-entity) ------------------------------------------ + + +async def _wiki_reduce_entity( + entity_name: str, + new_claims: list[dict], + existing_page: dict | None, + deleted_doc_ids: set[str], + entity_type: str = "entity", +) -> dict: + """Per-entity REDUCE: compute additions/retractions vs existing page. + + Returns dict with action (create|update|delete|noop), additions, retractions, + and entity_type for downstream filtering. + """ + if existing_page is None: + if isinstance(entity_type, list): + entity_type = entity_type[0] if entity_type else "entity" + entity_type = str(entity_type or "entity").strip() + # A page must have grounded evidence. MAP can mention an entity as a + # relation endpoint or metadata-only item without producing a claim; + # creating a page for that item would persist empty source_doc_ids and + # source_chunk_ids and make the REFINE prompt generate a placeholder + # page. Such new entities/concepts are intentionally skipped. + if not new_claims: + return { + "action": "noop", + "entity_name": entity_name, + "entity_type": entity_type, + "additions": [], + "retractions": [], + "retained_source_doc_ids": [], + "has_delta": False, + } + return { + "action": "create", + "entity_name": entity_name, + "entity_type": entity_type, + "additions": new_claims, + "retained_source_doc_ids": list({c["source_doc_id"] for c in new_claims if c.get("source_doc_id")}), + "has_delta": True, + } + + existing_claims = existing_page.get("claims", []) + # The stored `claims` column is a JSON string (json.dumps in _wiki_refine_page). + # Normalize to a list of dicts defensively — iterating a raw string yields + # characters and breaks every `c.get(...)` below. + if isinstance(existing_claims, str): + try: + existing_claims = json.loads(existing_claims) if existing_claims else [] + except (json.JSONDecodeError, TypeError): + existing_claims = [] + if isinstance(existing_claims, (list, tuple)): + existing_claims = [c for c in existing_claims if isinstance(c, dict)] + else: + existing_claims = [] + deleted_set = deleted_doc_ids or set() + + retractions = [c for c in existing_claims if c.get("source_doc_id") in deleted_set] + + retained_claims = [c for c in existing_claims if c.get("source_doc_id") not in deleted_set] + + retained_texts = {c.get("statement", c.get("text", "")) for c in retained_claims} + additions = [c for c in new_claims if c.get("statement", c.get("text", "")) not in retained_texts] + + all_doc_ids = {c.get("source_doc_id") for c in retained_claims} | {c.get("source_doc_id") for c in additions} + + if not all_doc_ids: + return { + "action": "delete", + "entity_name": entity_name, + "entity_type": entity_type, + "retractions": existing_claims, + "has_delta": True, + } + elif additions or retractions: + return { + "action": "update", + "entity_name": entity_name, + "entity_type": entity_type, + "additions": additions, + "retractions": retractions, + "retained_source_doc_ids": list(all_doc_ids), + "has_delta": True, + } + return { + "action": "noop", + "entity_name": entity_name, + "entity_type": entity_type, + "retained_source_doc_ids": list(all_doc_ids), + "has_delta": False, + } + + +async def _wiki_reduce_batch( + affected_names: set[str], + existing_pages: dict[str, dict], + deleted_doc_ids: set[str], + canonical_claims: dict[str, list[dict]] | None = None, + canonical_map: dict[str, dict] | None = None, + name_resolution: dict[str, str] | None = None, + map_results: list[dict] | None = None, +) -> list[dict]: + """Parallel per-entity REDUCE over a batch of affected canonical names. + + Uses canonical_claims (from Entity Matching) instead of raw MAP claims. + """ + name_to_page: dict[str, dict] = {} + for pid, page in existing_pages.items(): + for n in _as_str_list(page.get("entity_names_kwd")): + name_to_page[n] = page + slug = pid.split("/")[-1] if "/" in pid else pid + name_to_page.setdefault(slug, page) + + # Use canonical claims if provided (post-entity-matching) + if canonical_claims is not None: + claims_source = canonical_claims + else: + # Fallback: aggregate from raw MAP results + claims_source = {} + for mr in map_results: + for c in mr.get("claims", []): + name = c.get("entity_name") or c.get("subject") or c.get("term") + if name: + raw_name = name + if name_resolution: + raw_name = name_resolution.get(name, name) + claims_source.setdefault(raw_name, []).append(c) + + for name in affected_names: + claims_source.setdefault(name, []) + + tasks = [] + for name in affected_names: + claims = claims_source.get(name, []) + # Determine entity_type from canonical_map + entity_type = "entity" + if canonical_map and name in canonical_map: + entity_type = canonical_map[name].get("type", "entity") + if isinstance(entity_type, list): + entity_type = entity_type[0] if entity_type else "entity" + entity_type = str(entity_type or "entity").strip() + + tasks.append( + _wiki_reduce_entity( + entity_name=name, + entity_type=entity_type, + new_claims=claims, + existing_page=name_to_page.get(name, existing_pages.get(name)), + deleted_doc_ids=deleted_doc_ids, + ) + ) + if not tasks: + return [] + results = await asyncio.gather(*tasks) + return [r for r in results if r.get("has_delta")] + + +# ----- doc_page_source tracking --------------------------------------------- + + +async def _wiki_update_doc_page_source( + tenant_id: str, + kb_id: str, + doc_id: str, + page_ids: list[str], + entity_names: list[str] | None = None, + chunk_hashes: dict[str, str] | None = None, + map_checksum: str | None = None, +) -> None: + """Record which pages and entities this document contributes to.""" + index = search.index_name(tenant_id) + + condition = { + "compile_kwd": [WIKI_DOC_PAGE_SOURCE_COMPILE_KWD], + "doc_id": [doc_id], + } + existing = await thread_pool_exec( + settings.docStoreConn.search, + ["id", "page_ids", "entity_names", "source_chunk_hashes", "map_checksum"], + [], + condition, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + existing_map = settings.docStoreConn.get_fields(existing, ["id", "page_ids", "entity_names", "source_chunk_hashes", "map_checksum"]) + + if existing_map: + for row in existing_map.values(): + if chunk_hashes is None: + saved = row.get("source_chunk_hashes", "{}") + chunk_hashes = json.loads(saved) if isinstance(saved, str) else saved + if map_checksum is None: + val = row.get("map_checksum", "") or "" + if val: + map_checksum = val + if entity_names is None: + saved_names = row.get("entity_names", "[]") + entity_names = json.loads(saved_names) if isinstance(saved_names, str) else saved_names + break + + doc = { + "id": _stable_row_id(WIKI_DOC_PAGE_SOURCE_COMPILE_KWD, kb_id, doc_id), + "doc_id": doc_id, + "kb_id": kb_id, + "page_ids": json.dumps(page_ids, ensure_ascii=False), + "entity_names": json.dumps(entity_names or [], ensure_ascii=False), + "source_chunk_hashes": json.dumps(chunk_hashes or {}, ensure_ascii=False), + "map_checksum": map_checksum or "", + "compile_kwd": WIKI_DOC_PAGE_SOURCE_COMPILE_KWD, + } + if existing_map: + await thread_pool_exec( + settings.docStoreConn.update, + {"doc_id": doc_id}, + doc, + index, + kb_id, + ) + else: + await thread_pool_exec( + settings.docStoreConn.insert, + [doc], + index, + kb_id, + ) + + +async def _wiki_load_doc_page_source( + tenant_id: str, + kb_id: str, + doc_id: str, +) -> dict | None: + """Load doc_page_source record for a document.""" + index = search.index_name(tenant_id) + condition = { + "compile_kwd": [WIKI_DOC_PAGE_SOURCE_COMPILE_KWD], + "doc_id": [doc_id], + } + res = await thread_pool_exec( + settings.docStoreConn.search, + ["page_ids", "entity_names", "source_chunk_hashes", "map_checksum"], + [], + condition, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, ["page_ids", "entity_names", "source_chunk_hashes", "map_checksum"]) + for row in field_map.values(): + return { + "page_ids": json.loads(row.get("page_ids", "[]")) if isinstance(row.get("page_ids"), str) else row.get("page_ids", []), + "entity_names": json.loads(row.get("entity_names", "[]")) if isinstance(row.get("entity_names"), str) else row.get("entity_names", []), + "source_chunk_hashes": json.loads(row.get("source_chunk_hashes", "{}")) if isinstance(row.get("source_chunk_hashes"), str) else row.get("source_chunk_hashes", {}), + "map_checksum": row.get("map_checksum", ""), + } + return None + + +async def _wiki_delete_doc_page_source( + tenant_id: str, + kb_id: str, + doc_id: str, +) -> None: + """Delete doc_page_source record when a document is removed.""" + index = search.index_name(tenant_id) + await thread_pool_exec( + settings.docStoreConn.delete, + { + "compile_kwd": [WIKI_DOC_PAGE_SOURCE_COMPILE_KWD], + "doc_id": [doc_id], + }, + index, + kb_id, + ) + + +# ----- Mode A: no-plan REFINE ----------------------------------------------- + + +async def _wiki_refine_page( + *, + mode: str, # "generate" | "modify" | "re-synthesize" | "delete" + page_id: str, + page_title: str, + existing_page: dict | None, + page_type_kwd: str = "concept", + additions: list[dict] | None = None, + retractions: list[dict] | None = None, + source_chunks: list[dict] | None = None, + claims: list[dict] | None = None, + available_pages: list[str] | None = None, + contextual_hints: str = "", + chat_mdl, + embd_mdl, + tenant_id: str, + kb_id: str, + page_version: int, +) -> dict | None: + """Run a single Mode A REFINE action on one concept page. + + Returns updated wiki_page dict, or None if deleted. + """ + from common.misc_utils import thread_pool_exec + + # Blank page_id would produce `slug_kwd: [""]` queries → Infinity 3052. + if not page_id or not str(page_id).strip(): + return existing_page + + if mode == "delete": + await thread_pool_exec( + settings.docStoreConn.delete, + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "slug_kwd": [page_id]}, + search.index_name(tenant_id), + kb_id, + ) + return None + + # WeKnora-style verbatim evidence: load the ACTUAL source-chunk text for + # every referenced chunk id so the writer is grounded in the source, not + # just in the condensed claim statements. This runs per-page (incremental + # friendly) and is bounded by WIKI_SOURCE_BUDGET_CHARS. + if source_chunks: + chunk_ids = [sc.get("id") or sc.get("chunk_id") for sc in source_chunks if (sc.get("id") or sc.get("chunk_id"))] + if chunk_ids: + try: + chunk_texts = await _wiki_load_chunk_texts(tenant_id, kb_id, [str(c) for c in chunk_ids]) + if chunk_texts: + source_chunks = _wiki_enrich_source_chunks(source_chunks, chunk_texts) + except Exception: + logging.exception("wiki: verbatim chunk enrichment failed for page %s", page_id) + + # Build the prompt based on mode + if mode == "generate": + system_prompt = _WIKI_MODE_A_GENERATE_SYSTEM + user_prompt = _build_mode_a_generate_prompt( + page_id, + page_title, + claims, + source_chunks, + available_pages, + contextual_hints, + ) + elif mode == "re-synthesize": + system_prompt = _WIKI_MODE_A_MODIFY_SYSTEM + user_prompt = _build_mode_a_modify_prompt( + page_id, + page_title, + existing_page, + additions, + retractions, + claims, + source_chunks, + available_pages, + contextual_hints, + force_full=True, + ) + else: # modify + system_prompt = _WIKI_MODE_A_MODIFY_SYSTEM + user_prompt = _build_mode_a_modify_prompt( + page_id, + page_title, + existing_page, + additions, + retractions, + claims, + source_chunks, + available_pages, + contextual_hints, + force_full=False, + ) + + # Call LLM + response = await _chat_mdl_ask( + chat_mdl, + system_prompt, + user_prompt, + ) + + if not response or not response.strip(): + return existing_page # keep existing + + # Parse response: expected format starts with "SUMMARY: ..." then content + content = response.strip() + summary = "" + if content.startswith("SUMMARY:"): + idx = content.find("\n") + if idx > 0: + summary = content[8:idx].strip() + content = content[idx + 1 :].strip() + + # Build the wiki_page dict + existing = existing_page or {} + new_version = page_version + 1 + raw_doc_ids = existing.get("source_doc_ids", []) + doc_ids = json.loads(raw_doc_ids) if isinstance(raw_doc_ids, str) else list(raw_doc_ids) + source_chunk_ids = set(_as_str_list(existing.get("source_chunk_ids"))) + for claim in claims or []: + did = claim.get("source_doc_id") if isinstance(claim, dict) else None + if did and did not in doc_ids: + doc_ids.append(did) + if source_chunks: + for chunk in source_chunks: + cid = chunk.get("id") or chunk.get("chunk_id") + if cid: + source_chunk_ids.add(str(cid)) + did = chunk.get("doc_id") or chunk.get("source_doc_id") + if did and did not in doc_ids: + doc_ids.append(did) + + # Embed for search + from common.misc_utils import thread_pool_exec + + embeddings, _ = await thread_pool_exec(embd_mdl.encode, [summary or content[:200]]) + + # Derive vector dimension from the embedding shape + emb_arr = np.asarray(embeddings[0]) + vec_dim = int(emb_arr.shape[0]) if emb_arr.ndim >= 1 and emb_arr.shape[0] else 768 + + page = { + "id": _stable_row_id(WIKI_PAGE_COMPILE_KWD, kb_id, page_id), + "slug_kwd": page_id, + "title_kwd": page_title, + "md_with_weight": content, + "summary_with_weight": summary or page_title, + "entity_names_kwd": [page_title], + "source_chunk_ids": sorted(source_chunk_ids), + "source_doc_ids": json.dumps(doc_ids, ensure_ascii=False), + "claims": json.dumps(claims, ensure_ascii=False) if claims else "[]", + "page_version_int": new_version, + "synthesis_version_int": new_version if mode in ("generate", "re-synthesize") else existing.get("synthesis_version_int", 0), + "page_type_kwd": page_type_kwd, + "compile_kwd": WIKI_PAGE_COMPILE_KWD, + "knowledge_graph_kwd": WIKI_PAGE_COMPILE_KWD, + } + # Insert vector (adds q_{dim}_vec field) + vec_col = f"q_{vec_dim}_vec" + page[vec_col] = embeddings[0].tolist() if hasattr(embeddings[0], "tolist") else embeddings[0] + + # Persist + index = search.index_name(tenant_id) + + existing_entry = await thread_pool_exec( + settings.docStoreConn.search, + ["slug_kwd"], + [], + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "slug_kwd": [page_id]}, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + if settings.docStoreConn.get_fields(existing_entry, ["slug_kwd"]): + await thread_pool_exec( + settings.docStoreConn.update, + {"slug_kwd": page_id}, + page, + index, + kb_id, + ) + else: + await thread_pool_exec( + settings.docStoreConn.insert, + [page], + index, + kb_id, + ) + + return page + + +def _build_source_chunks_block(source_chunks: list[dict], max_budget: int = WIKI_SOURCE_BUDGET_CHARS) -> str: + """Render the verbatim source-chunk block for the writer prompt. + + Chunks carrying verbatim text (``_verbatim=True``) are labelled clearly so + the writer treats them as ground truth; the whole block is capped at + ``max_budget`` characters. Missing/empty chunks are dropped. + """ + if not source_chunks: + return "" + parts: list[str] = [] + total = 0 + for c in source_chunks: + cid = c.get("id") or c.get("chunk_id") + text = c.get("content_with_weight") or c.get("text") or "" + if not text or not cid: + continue + if c.get("_verbatim"): + block = f"[SOURCE {cid}]\n{text}" + else: + block = f"[CHUNK {cid}]\n{text}" + if total + len(block) + 2 > max_budget: + break + parts.append(block) + total += len(block) + 2 + if not parts: + return "" + if total >= max_budget: + parts.append("[…further source chunks omitted to fit context budget…]") + return "\n\n".join(parts) + + +def _build_mode_a_generate_prompt( + page_id: str, + page_title: str, + claims: list[dict], + source_chunks: list[dict], + available_pages: list[str], + contextual_hints: str, +) -> str: + chunks_text = _build_source_chunks_block(source_chunks) + claims_text = "\n".join(f"- {c.get('statement', c.get('text', ''))}" for c in claims) if claims else "(no claims)" + + return f"""## Concept Page Identity +- Page ID: {page_id} +- Title: {page_title} + +## Source Chunks (verbatim source text — ground every fact in these) +{chunks_text or "(no source chunks available)"} + +## Extracted Claims (checklist) +{claims_text} + +## Available Pages for [[wikilinks]] +{chr(10).join(f"- {p}" for p in available_pages[:50]) if available_pages else "(none)"} + +{contextual_hints} +""" + + +def _build_mode_a_modify_prompt( + page_id: str, + page_title: str, + existing_page: dict | None, + additions: list[dict] | None, + retractions: list[dict] | None, + claims: list[dict], + source_chunks: list[dict], + available_pages: list[str], + contextual_hints: str, + force_full: bool = False, +) -> str: + existing_content = existing_page.get("md_with_weight", "") if existing_page else "" + + if not force_full: + additions_text = "\n".join(f"- {c.get('statement', c.get('text', ''))}" for c in (additions or [])) if additions else "(none)" + retractions_text = "\n".join(f"- {c.get('statement', c.get('text', ''))}" for c in (retractions or [])) if retractions else "(none)" + chunks_text = _build_source_chunks_block(source_chunks) + + return f"""## Page Identity +- Page ID: {page_id} +- Title: {page_title} + +## Current Page +{existing_content[:10000] if existing_content else "(empty)"} + +## New Claims to Add +{additions_text} + +## Claims to Retract +{retractions_text} + +## Source Chunks for New Information (verbatim source text — ground every fact in these) +{chunks_text} + +## Available Pages for [[wikilinks]] +{chr(10).join(f"- {p}" for p in available_pages[:30]) if available_pages else "(none)"} + +{contextual_hints} +""" + else: + # Full re-synthesis: all claims + all source chunks (larger budget) + chunks_text = _build_source_chunks_block(source_chunks, max_budget=120_000) + claims_text = "\n".join(f"- {c.get('statement', c.get('text', ''))}" for c in claims) + + return f"""## Page Identity +- Page ID: {page_id} +- Title: {page_title} + +## All Source Chunks (for full re-synthesis — verbatim source text) +{chunks_text or "(none)"} + +## All Claims +{claims_text or "(none)"} + +## Available Pages for [[wikilinks]] +{chr(10).join(f"- {p}" for p in available_pages[:50]) if available_pages else "(none)"} + +{contextual_hints} +""" + + +# System prompts for Mode A + +_WIKI_MODE_A_GENERATE_SYSTEM = """You are a wiki COMPILER. Generate a new wiki page for the given concept using the provided source chunks and extracted claims. + +## LANGUAGE +Write the ENTIRE page in the SAME LANGUAGE as the source chunks. If the source chunks are written in Chinese, write the page in Chinese. Do not switch to English, and do not translate entity names (keep them verbatim: e.g. keep "张伟", do not write "Zhang Wei"). + +## RULES +1. CONCEPT PAGE: This is a single-concept wiki page. Organize by THEME, not by entity. +2. CROSS-DOCUMENT SYNTHESIS: Weave information from multiple sources into coherent paragraphs. Compare evidence, explain contradictions. +3. OPENING PARAGRAPH: 2-4 sentences defining the concept. Mention key entities. No heading. +4. SECTIONS: H2 headings, prose first, then sub-points if needed. + Markdown formatting is mandatory: put every heading on its own line and separate every paragraph with a blank line. +5. WIKILINKS: Use ONLY the exact page IDs listed in "Available Pages for [[wikilinks]]" (they already carry the entity/ or concept/ prefix). Insert [[EXACT_PAGE_ID]] on first mention of a related concept/entity. NEVER invent a link target, NEVER drop the prefix, NEVER write English names. +6. DICTIONARY PREVENTION: Do NOT group content by source document. Do NOT create one section per entity. Do NOT write flat bullet lists. + +## SOURCE GROUNDING (COMPILER, not writer) +- The "Source Chunks" section contains VERBATIM source text. Stay close to the source wording — reuse the source's own sentences and facts where possible. +- Every newly added factual claim, entity, or numerical value MUST be directly supported by the provided source chunks. Do NOT invent facts, figures, dates, or relationships not present in the sources. +- Do NOT add rhetorical filler (e.g. "旨在帮助…", "designed to…", "aims to provide…") unless it appears verbatim in a source. +- If the sources disagree, present both views and add a "## Contradictions" section rather than silently picking one. + +## OUTPUT +Return ONLY the complete markdown page. +First line: SUMMARY: {one-sentence description, 15-40 words} +Then the page content. +""" + +_WIKI_MODE_A_MODIFY_SYSTEM = """You are a wiki editor. Update the existing page by integrating new information and removing retracted content. + +## LANGUAGE +Write the ENTIRE page in the SAME LANGUAGE as the source chunks. If the source chunks are written in Chinese, write the page in Chinese. Do not switch to English, and do not translate entity names (keep them verbatim: e.g. keep "张伟", do not write "Zhang Wei"). + +## RULES +1. CONCEPT PAGE: This is a single-concept wiki page. Organize by THEME, not by entity. +2. CROSS-DOCUMENT SYNTHESIS: Connect new claims to existing content. Weave them into the SAME paragraphs. +3. OPENING PARAGRAPH: Should reflect the FULL updated picture. +4. WIKILINKS: Keep existing and add new [[page_id]] links where appropriate. Use ONLY the exact page IDs listed in "Available Pages for [[wikilinks]]" (they already carry the entity/ or concept/ prefix). NEVER invent a link target, NEVER drop the prefix, NEVER write English names. +5. For FULL RE-SYNTHESIS: Use all source chunks + all claims to rewrite from scratch. +6. For INCREMENTAL MODIFY: Integrate additions, remove retracted content, keep unchanged content. +7. MARKDOWN FORMATTING: Put every heading on its own line and separate every paragraph with a blank line. Do not return the whole page as one line. + +## DICTIONARY PREVENTION +- Do NOT group content by source document. +- Do NOT simply append new claims at the end. +- Do NOT create one section per entity. + +## SOURCE GROUNDING (COMPILER, not writer) +- The "Source Chunks" section contains VERBATIM source text. Stay close to the source wording — reuse the source's own sentences and facts where possible. +- Every newly added factual claim, entity, or numerical value MUST be directly supported by the provided source chunks. Do NOT invent facts, figures, dates, or relationships not present in the sources. +- Do NOT add rhetorical filler (e.g. "旨在帮助…", "designed to…", "aims to provide…") unless it appears verbatim in a source. +- If new sources contradict existing page content, present both views and add a "## Contradictions / Updates" section rather than silently overwriting. + +## OUTPUT +Return ONLY the complete updated markdown page. +First line: SUMMARY: {one-sentence description of what changed, 15-40 words} +Then the updated page content. +""" + + +def _wiki_build_contextual_hints( + page_id: str, + existing_page: dict | None, + all_relations: dict[str, list[dict]], +) -> str: + """Build contextual hints prompt block from related_pages.""" + related = [] + if existing_page: + rp = existing_page.get("related_kb_pages_kwd") + if rp: + if isinstance(rp, str): + related = json.loads(rp) + elif isinstance(rp, list): + related = rp + if not related: + for name in _as_str_list(existing_page.get("entity_names_kwd") if existing_page else None): + related.extend(all_relations.get(name, [])) + if not related: + return "" + + lines = ["## Context: Related Entities & Concepts", "Reference them in the opening paragraph and relevant sections:"] + for r in related[:10]: + entity_name = r.get("entity_name") or r.get("name", "") + relation = r.get("relation") or r.get("type", "related") + lines.append(f"- [[{entity_name}]] — {relation}") + return "\n".join(lines) + + +# ----- Mode B Page Router (KNN entity routing) ----------------------------- + + +async def _wiki_page_router( + affected_entities: list[dict], + embd_mdl, + tenant_id: str, + kb_id: str, + existing_page_ids: set[str] | None = None, +) -> dict[str, list[dict]]: + """Route affected entities to existing wiki pages via KNN. + + Returns: {page_id: [entity_deltas]} + - "_new_{page_id}" → new page to create + - existing page_id → entities assigned to that page + + ``existing_page_ids`` is supplied by Mode B from its already-loaded page + set. An explicitly empty set means this is a first build, so page-index + KNN routing can be skipped and entities can go straight to clustering. + """ + from common.misc_utils import thread_pool_exec + from rag.nlp import search + from common.doc_store.doc_store_base import OrderByExpr + + query_texts = [_entity_to_query_text(e) for e in affected_entities] + embeddings, _ = await thread_pool_exec(embd_mdl.encode, query_texts) + + index = search.index_name(tenant_id) + condition = {"compile_kwd": [WIKI_PAGE_COMPILE_KWD]} + + assignments: dict[str, list[dict]] = {} + orphans: list[dict] = [] + embedding_by_entity_id = {id(entity): vec for entity, vec in zip(affected_entities, embeddings, strict=False)} + + if existing_page_ids is not None and not existing_page_ids: + # On a first build there are no pages that can accept a routed entity. + # Querying the page index once per entity only produces orphans, so go + # directly to clustering and reuse the embeddings already computed. + orphans = list(affected_entities) + else: + router_sem = asyncio.Semaphore(PAGE_ROUTER_KNN_CONCURRENT) + + async def _search_page(entity: dict, vec) -> tuple[dict, dict]: + async with router_sem: + match_expr = MatchDenseExpr( + vector_column_name=f"q_{len(vec)}_vec", + embedding_data=vec.tolist() if hasattr(vec, "tolist") else vec, + embedding_data_type="float", + distance_type="cosine", + topn=1, + extra_options={"similarity": PAGE_ROUTER_MAYBE_THRESHOLD}, + ) + res = await thread_pool_exec( + settings.docStoreConn.search, + ["slug_kwd", "title_kwd", "_score"], + [], + condition, + [match_expr], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + return entity, settings.docStoreConn.get_fields(res, ["slug_kwd", "title_kwd", "_score"]) + + route_results = await asyncio.gather(*(_search_page(entity, vec) for entity, vec in zip(affected_entities, embeddings, strict=False))) + for entity, field_map in route_results: + if not field_map: + orphans.append(entity) + continue + + for row in field_map.values(): + score = row.get("_score", 0.0) + page_id = row.get("slug_kwd", "") + if isinstance(page_id, (list, tuple)): + page_id = page_id[0] if page_id else "" + page_id = str(page_id or "").strip() + + # slug_kwd is a *_kwd field; Infinity may return it empty / mangled + # on the matched row. A blank page id would route the entity onto a + # `slug_kwd: [""]` query (Infinity 3052) — treat as orphan instead. + if not page_id: + orphans.append(entity) + continue + + if score >= PAGE_ROUTER_UPDATE_THRESHOLD: + assignments.setdefault(page_id, []).append(entity) + elif score >= PAGE_ROUTER_MAYBE_THRESHOLD: + assignments.setdefault(f"_maybe_{page_id}", []).append(entity) + else: + orphans.append(entity) + break + + # Handle maybe candidates (batch LLM confirm optional) + for key in list(assignments.keys()): + if key.startswith("_maybe_"): + page_id = key[7:] + # Simple heuristic: assign to the page if any claim overlaps + existing_page_claims = await _load_page_claims(tenant_id, kb_id, page_id) + confirmed = [] + for entity in assignments[key]: + entity_claim_texts = {c.get("statement", c.get("text", "")) for c in entity.get("claims", [])} + existing_claim_texts = {ec.get("statement", ec.get("text", "")) for ec in (existing_page_claims or [])} + if entity_claim_texts & existing_claim_texts: + confirmed.append(entity) + else: + orphans.append(entity) + if confirmed: + assignments.setdefault(page_id, []).extend(confirmed) + del assignments[key] + + # Orphans: cluster by similarity, create grouped pages + if orphans: + orphan_embs = [embedding_by_entity_id[id(entity)] for entity in orphans] + clusters = _wiki_cluster_entities(orphans, orphan_embs, threshold=PAGE_ROUTER_CLUSTER_THRESHOLD) + for cluster in clusters: + names = [e.get("entity_name") or e.get("term", "") for e in cluster] + # Mode B compiles EVERY entity/concept into a page. On a first build + # there are no existing pages, so every affected entity lands here as + # an orphan; do NOT gate page creation on a claim count or most pages + # (esp. claim-light concepts/entities) would never be created. + if not names: + continue + # Pick the page prefix from the cluster's dominant type. The default + # prefix of _wiki_derive_page_id is "concept"; passing nothing would + # mislabel every group (incl. people/orgs) as a concept page. + any_concept = any((e.get("entity_type") or e.get("type")) == "concept" for e in cluster) + prefix = "concept" if any_concept else "entity" + page_id = _wiki_derive_page_id(names[0], prefix=prefix) + if not page_id: + continue + assignments[f"_new_{page_id}"] = cluster + + return assignments + + +async def _load_page_claims( + tenant_id: str, + kb_id: str, + page_id: str, +) -> list[dict]: + """Load claims for a single wiki page.""" + from rag.nlp import search + from common.misc_utils import thread_pool_exec + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + condition = {"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "slug_kwd": [page_id]} + res = await thread_pool_exec( + settings.docStoreConn.search, + ["claims", "slug_kwd"], + [], + condition, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, ["claims", "slug_kwd"]) + for row in field_map.values(): + claims = row.get("claims", "[]") + if isinstance(claims, str): + return json.loads(claims) + return claims + return [] + + +def _wiki_cluster_entities( + entities: list[dict], + embeddings: list, + threshold: float, +) -> list[list[dict]]: + """Simple pairwise cosine clustering for orphan entities. + + Returns clusters where intra-cluster cosine >= threshold. + Each cluster has at least 1 entity. + """ + if len(entities) <= 1: + return [entities] + + # Normalize embeddings + embs = [] + for e in embeddings: + if hasattr(e, "tolist"): + e = e.tolist() + arr = np.asarray(e, dtype=np.float32) + norm = np.linalg.norm(arr) + embs.append(arr / norm if norm > 0 else arr) + + n = len(embs) + assigned = [False] * n + clusters: list[list[int]] = [] + + for i in range(n): + if assigned[i]: + continue + cluster = [i] + assigned[i] = True + for j in range(i + 1, n): + if assigned[j]: + continue + similarity = float(np.dot(embs[i], embs[j].T)) + if similarity >= threshold: + cluster.append(j) + assigned[j] = True + clusters.append(cluster) + + result = [] + for cluster in clusters: + result.append([entities[i] for i in cluster]) + return result + + +# ----- FINALIZE (shared) ---------------------------------------------------- + + +async def _wiki_finalize( + tenant_id: str, + kb_id: str, + embd_mdl, + page_ids: list[str] | None = None, +) -> None: + """Post-REFINE cleanup: dead wikilinks + cross-reference update. + + Always scans ALL wiki_page rows (page_ids is ignored — full scan). + Three wikilink types: + 1. Valid page → keep [[]], update related_kb_pages_kwd + 2. Entity reference (in canonical index) → remove [[]], keep plain text + 3. Dead link → remove [[]], keep plain text + """ + all_pages = await _search_existing_pages( + tenant_id, + kb_id, + [ + "slug_kwd", + "id", + "title_kwd", + "md_with_weight", + "outlinks_kwd", + "related_kb_pages_kwd", + "entity_names_kwd", + ], + ) + if not all_pages: + return + + valid_ids = set(all_pages.keys()) + + # Load canonical entity names for entity reference detection (Mode A) + canonical_names = set() + canonical_index = await _load_canonical_entities(tenant_id, kb_id) + for cname in canonical_index: + canonical_names.add(cname) + # Also add aliases + for alias in canonical_index[cname].get("aliases", []): + canonical_names.add(alias) + + wikilink_re = re.compile(r"\[\[([^\]]+)\]\]") + relation_map: dict[str, list[dict]] = {} + outlink_map: dict[str, list[str]] = {} # pid → [valid target slugs] + dead_links: dict[str, list[str]] = {} # pid → [dead links to remove] + + # name → page slug map for AUTO-LINKING. Built from every page's plain name + # (slug suffix) + title, longest names first so multi-word / multi-char + # mentions are matched greedily before shorter substrings. + name_slug: dict[str, str] = {} + for pid in all_pages: + plain = pid.split("/")[-1] if "/" in pid else pid + if plain: + name_slug[plain] = pid + title = all_pages[pid].get("title_kwd") + if isinstance(title, (list, tuple)): + title = title[0] if title else "" + if isinstance(title, str) and title and title != plain: + name_slug[title] = pid + # Map every entity the page actually contains (incl. ones merged into a + # group page by Mode B's plan/grouping) to this page. Otherwise relation + # endpoints like "梁大伟" that were merged into another page won't match + # and most wiki pages stay unlinked. + for en in _as_str_list(all_pages[pid].get("entity_names_kwd")): + if en and en != plain: + name_slug[en] = pid + ordered_names = sorted(name_slug.keys(), key=lambda n: (-len(n), n)) + + # RELATION-BASED LINKING: use the semantic (from, to) edges extracted during + # MAP to connect pages. A page-to-page edge is created whenever both + # endpoints resolve to compiled wiki pages. This is the primary source of + # graph connections — far more reliable than prose [[wikilinks]]. + map_relations = await _load_map_relations(tenant_id, kb_id) + relation_edges: dict[str, set[str]] = {} # pid → {target slug} + if map_relations: + for rel in map_relations: + from_pg = name_slug.get(rel["from"]) + to_pg = name_slug.get(rel["to"]) + if from_pg and to_pg and from_pg != to_pg: + relation_edges.setdefault(from_pg, set()).add(to_pg) + relation_edges.setdefault(to_pg, set()).add(from_pg) + + for pid, page in all_pages.items(): + content = page.get("md_with_weight", "") + original = content + + for match in wikilink_re.finditer(content): + link = match.group(1).strip() + if link in valid_ids and link != pid: + # Valid wikilink → record for cross-reference + outlink + relation_map.setdefault(pid, []).append( + { + "entity_name": link.split("/")[-1] if "/" in link else link, + "relation": "see_also", + } + ) + outlink_map.setdefault(pid, []).append(link) + elif link in canonical_names: + # Entity reference (Mode A): remove [[]] keep plain text + content = content.replace(f"[[{link}]]", link, 1) + else: + # Dead link — try WeKnora-style fuzzy resolution to a similar + # existing page before giving up. If a close slug exists, retarget + # the link (cross-link survives); otherwise degrade to plain text. + resolved = _wiki_resolve_dead_slug(link, valid_ids, name_slug) + if resolved: + content = content.replace(f"[[{link}]]", f"[[{resolved}]]", 1) + relation_map.setdefault(pid, []).append({"entity_name": resolved.split("/")[-1] if "/" in resolved else resolved, "relation": "see_also"}) + if resolved not in outlink_map.setdefault(pid, []): + outlink_map[pid].append(resolved) + else: + content = content.replace(f"[[{link}]]", link, 1) + dead_links.setdefault(pid, []).append(link) + + # AUTO-LINK: guarantee cross-page connections even when the LLM omits + # [[...]]. Scan for standalone mentions of other pages' plain names and + # wrap the FIRST occurrence with [[full_slug]], recording an outlink. + # existing_links covers both the raw [[slug]] form and the rendered + # `[text](artifact/{kb_id}/slug)` form (so re-runs stay idempotent even + # after links were rendered on a previous run). + existing_links = {m.group(1).strip() for m in wikilink_re.finditer(content)} + existing_links |= {m.group(1) for m in re.finditer(rf"\]\(artifact/{re.escape(str(kb_id))}/([^)]+)\)", content)} + for name in ordered_names: + target = name_slug[name] + if target == pid: + continue + if target in existing_links: + continue + idx = content.find(name) + if idx < 0: + continue + # Skip if the occurrence is already inside a [[...]] link span + if _inside_wikilink(content, idx): + continue + content = content[:idx] + f"[[{target}]]" + content[idx + len(name) :] + existing_links.add(target) + if target not in outlink_map.setdefault(pid, []): + outlink_map[pid].append(target) + relation_map.setdefault(pid, []).append({"entity_name": name, "relation": "see_also"}) + + # Merge semantic relation edges (from MAP extraction) into the outlinks. + # These connect pages even when prose never mentions the counterpart. + # We ALSO inject a "相关页面 / Related" [[wikilink]] section so the + # cross-page link is visible in the body and survives the shredded + # *_kwd field (outlinks_kwd is analyzed by whitespace-# → reads back []). + rel_targets = [] + for target in relation_edges.get(pid, ()): + if target == pid: + continue + if target not in outlink_map.setdefault(pid, []): + outlink_map[pid].append(target) + target_name = target.split("/")[-1] if "/" in target else target + relation_map.setdefault(pid, []).append({"entity_name": target_name, "relation": "related"}) + if target not in existing_links: + rel_targets.append(target) + existing_links.add(target) + if rel_targets: + if not content.rstrip().endswith("## 相关页面"): + content = content.rstrip() + "\n\n## 相关页面\n" + content += "\n".join(f"- [[{t}]]" for t in rel_targets) + "\n" + + # Render internal [[slug]] wikilinks into clickable Markdown links + # `[text](artifact/{kb_id}/{slug})` — the format the frontend's wiki + # viewer can navigate (the raw [[...]] form renders as plain text and + # cannot be deep-linked). Only slugs that resolve to a compiled page + # become links; anything else is left as plain text. + rendered_content = content + if rendered_content: + # valid_slugs must be the set of LINK TARGETS (outlink values), not + # the source page keys. Build it from every outlink_map value plus + # every page's own slug so self-references resolve too. + link_targets = set(outlink_map.keys()) + for _, targets in outlink_map.items(): + link_targets.update(targets) + rendered_content = _wiki_render_links(rendered_content, kb_id, link_targets) + + # Update page content if changed (store the RENDERED markdown so the + # frontend gets navigable artifact links). + update = {} + if rendered_content != original: + update["md_with_weight"] = rendered_content + + relations = relation_map.get(pid, []) + if relations: + # NOTE: *_kwd fields are analyzed by whitespace-#. Storing a + # json.dumps string makes Infinity shred it so get_fields reads it + # back as [] (breaking every consumer of related_kb_pages_kwd). + # Mirror the old-mode format: a native list of STRINGS (Infinity + # stores it and get_fields/_as_str_list restore it as a list). The + # old-mode related_kb_pages is exactly a list of page slugs/names, + # so we collapse each {entity_name, relation} entry to its string. + update["related_kb_pages_kwd"] = [r.get("entity_name") or r.get("slug") or str(r) for r in relations[:20]] + elif page.get("related_kb_pages_kwd"): + # Clear stale related_pages if no relations remain + update["related_kb_pages_kwd"] = [] + + # Always refresh outlinks: unique valid targets (graph + list ordering + # depend on outlinks_int). Preserves ordering for a stable canvas. + # Store a native list (NOT json.dumps) so the *_kwd field reads back + # correctly instead of being shredded into [] by whitespace-#. + outlinks = outlink_map.get(pid) or [] + update["outlinks_kwd"] = list(outlinks) + update["outlinks_int"] = len(outlinks) + + index = search.index_name(tenant_id) + await thread_pool_exec( + settings.docStoreConn.update, + {"id": page["id"]}, + update, + index, + kb_id, + ) + + +def _wiki_normalize_rows(matrix): + """L2-normalize the rows of a 2-D float matrix (safe on zero rows).""" + if matrix.ndim != 2: + return matrix + norms = np.linalg.norm(matrix, axis=1, keepdims=True) + return np.divide(matrix, norms, out=np.zeros_like(matrix), where=norms > 0) + + +async def _wiki_load_map_topics(index, kb_id) -> list[str]: + """Collect distinct thematic topic labels from persisted wiki_map_extract rows. + + Lets topic grouping run even when the current invocation carried no fresh MAP + output (e.g. a no-op re-run over already-built pages). Bounded scan. + """ + from common.doc_store.doc_store_base import OrderByExpr + + labels: list[str] = [] + seen: set[str] = set() + offset, page_size, scanned = 0, 500, 0 + while scanned < 5000 and len(labels) < WIKI_TOPIC_MAX_LABELS: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["content_with_weight"], + [], + {"compile_kwd": ["wiki_map_extract"]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + rows = settings.docStoreConn.get_fields(res, ["content_with_weight"]) or {} + except Exception: + logging.exception("wiki topics: map-topic load failed for kb=%s", kb_id) + break + if not rows: + break + for row in rows.values(): + raw = row.get("content_with_weight") + if not isinstance(raw, str) or not raw: + continue + try: + extract = json.loads(raw) + except Exception: + continue + for t in (extract.get("topics") or []) if isinstance(extract, dict) else []: + if isinstance(t, str): + t = t.strip() + key = t.lower() + if t and key != WIKI_TOPIC_FALLBACK.lower() and key not in seen: + seen.add(key) + labels.append(t) + if len(labels) >= WIKI_TOPIC_MAX_LABELS: + break + scanned += len(rows) + if len(rows) < page_size: + break + offset += page_size + return labels + + +async def _wiki_assign_topics( + embd_mdl, + tenant_id: str, + kb_id: str, + map_topics: list[str] | None = None, + callback: Callable | None = None, +) -> None: + """Group concept/entity wiki pages under thematic topics (best-effort). + + No-plan pages have no PLAN grouping step, so pages are grouped post-hoc: each + page is matched (embedding cosine) to the thematic topic labels the MAP phase + extracted (accumulated with topics already on record so labels persist across + runs). The best match above ``WIKI_TOPIC_MATCH_THRESHOLD`` wins, else the page + lands in the ``WIKI_TOPIC_FALLBACK`` bucket. Every page's ``topic_kwd`` is + stamped, so ``/artifacts_topics`` (which aggregates concept/entity pages by + ``topic_kwd``) and the topic-filtered page list resolve. No landing rows are + written — the topics API falls back to the raw topic name for title/slug. + Any failure leaves the pages intact (just untopiced) and never raises. + """ + from common.doc_store.doc_store_base import OrderByExpr + + def _progress(msg: str) -> None: + if callback: + try: + callback(0.97, f"Topics: {msg}") + except Exception: + pass + + try: + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return + + # 1. Load all concept/entity pages. + page_fields = ["slug_kwd", "title_kwd", "summary_with_weight", "source_doc_ids", "topic_kwd"] + pages: list[dict] = [] + offset, page_size = 0, 1000 + while True: + res = await thread_pool_exec( + settings.docStoreConn.search, + page_fields, + [], + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "page_type_kwd": ["concept", "entity"]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + rows = settings.docStoreConn.get_fields(res, page_fields) or {} + for row in rows.values(): + # get_fields may return scalar-ish *_kwd fields as a list (e.g. + # an Infinity/ES aggregation). Normalize slug_kwd to a scalar + # so it can be used as a dict key later; normalize title_kwd too. + slug = row.get("slug_kwd") + if isinstance(slug, (list, tuple)): + slug = slug[0] if slug else "" + row["slug_kwd"] = slug + if isinstance(row.get("title_kwd"), (list, tuple)): + t = row.get("title_kwd") + row["title_kwd"] = t[0] if t else "" + if slug: + pages.append(row) + if len(rows) < page_size: + break + offset += page_size + if not pages: + return + + # 2. Candidate labels: this run's MAP topics + topics already stamped on + # the pages from earlier runs (so labels accumulate across runs). + existing_labels: list[str] = [] + for p in pages: + t = p.get("topic_kwd") + if isinstance(t, str) and t.strip() and t.strip().lower() != WIKI_TOPIC_FALLBACK.lower(): + existing_labels.append(t.strip()) + + labels: list[str] = [] + seen: set[str] = set() + for t in list(map_topics or []) + existing_labels: + if not isinstance(t, str): + continue + t = t.strip() + key = t.lower() + if t and key != WIKI_TOPIC_FALLBACK.lower() and key not in seen: + seen.add(key) + labels.append(t) + if len(labels) >= WIKI_TOPIC_MAX_LABELS: + break + + # Backfill labels from the persisted MAP extracts when this run carried + # none (e.g. a no-op re-run over pages built before topic grouping). + if not labels: + labels = await _wiki_load_map_topics(index, kb_id) + + # 3. Assign each page to its nearest topic (or the fallback bucket). + assignments: dict[str, str] = {} + topic_docs: dict[str, set] = {} + + def _record(slug, topic: str, doc_ids) -> None: + # slug may come back as a list from some doc-store get_fields + # implementations — normalize to a scalar before keying. + if isinstance(slug, (list, tuple)): + slug = slug[0] if slug else "" + if not slug: + return + assignments[slug] = topic + bucket = topic_docs.setdefault(topic, set()) + raw = doc_ids + if isinstance(raw, str): + try: + raw = json.loads(raw) + except Exception: + raw = [raw] + for d in raw or []: + if isinstance(d, str) and d: + bucket.add(d) + + topic_matrix = None + if labels: + tvecs, _ = await thread_pool_exec(embd_mdl.encode, labels) + topic_matrix = _wiki_normalize_rows(np.asarray(tvecs, dtype=np.float32)) + if topic_matrix is not None and topic_matrix.ndim == 2 and topic_matrix.shape[0] == len(labels): + page_texts = [f"{p.get('title_kwd') or ''} {p.get('summary_with_weight') or ''}".strip() or (p.get("slug_kwd") or "") for p in pages] + pvecs, _ = await thread_pool_exec(embd_mdl.encode, page_texts) + page_matrix = _wiki_normalize_rows(np.asarray(pvecs, dtype=np.float32)) + if page_matrix.ndim == 2 and page_matrix.shape[0] == len(pages): + sims = page_matrix @ topic_matrix.T + best = np.argmax(sims, axis=1) + for i, p in enumerate(pages): + score = float(sims[i, best[i]]) + topic = labels[int(best[i])] if score >= WIKI_TOPIC_MATCH_THRESHOLD else WIKI_TOPIC_FALLBACK + _record(p["slug_kwd"], topic, p.get("source_doc_ids")) + if not assignments: + # No usable embeddings/labels → single fallback topic keeps nav working. + for p in pages: + _record(p["slug_kwd"], WIKI_TOPIC_FALLBACK, p.get("source_doc_ids")) + + by_topic: dict[str, list[str]] = {} + for slug, topic in assignments.items(): + by_topic.setdefault(topic, []).append(slug) + + # 4. Stamp topic_kwd on each page (bounded concurrency). + sem = asyncio.Semaphore(WIKI_TOPIC_UPDATE_CONCURRENT) + + async def _stamp(slug: str, topic: str) -> None: + async with sem: + try: + await thread_pool_exec( + settings.docStoreConn.update, + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD], "slug_kwd": [slug]}, + {"topic_kwd": topic}, + index, + kb_id, + ) + except Exception: + logging.exception("wiki topics: topic_kwd update failed for slug=%s", slug) + + await asyncio.gather(*[_stamp(slug, topic) for slug, topic in assignments.items()]) + + # No landing rows are written: list_wiki_topics derives topics from the + # pages' topic_kwd aggregation and falls back to the raw topic name for + # title/slug, so page_type="topic" rows would only pollute the page list. + _ = topic_docs # provenance retained for a future topic-page feature + _progress(f"grouped {len(pages)} page(s) into {len(by_topic)} topic(s).") + except Exception: + logging.exception("wiki topics: assignment failed for kb=%s", kb_id) + + +# ----- Main entry point ----------------------------------------------------- + + +async def wiki_compile_incremental( + *, + chat_mdl, + embd_mdl, + tenant_id: str, + kb_id: str, + plan: bool = False, # True = Mode B, False = Mode A + incremental: bool = False, # True = incremental run + map_results: list[dict] | None = None, # from MAP phase + deleted_doc_ids: set[str] | None = None, + callback: Callable | None = None, +) -> dict: + """Main entry point for dual-mode wiki compilation. + + Args: + plan: True=Mode B (with PLAN), False=Mode A (no-plan, WeKnora style) + incremental: True=incremental update, False=full build + map_results: MAP outputs. If None, loads from ES. + deleted_doc_ids: Documents that were removed. + callback: Progress callback. + + Returns summary dict: {pages_created, pages_modified, pages_deleted} + """ + from common.misc_utils import thread_pool_exec + from rag.nlp import search + from common.doc_store.doc_store_base import OrderByExpr + + summary = {"pages_created": 0, "pages_modified": 0, "pages_deleted": 0, "errors": []} + + def _progress(msg: str): + if callback: + try: + callback(0.5, msg) + except Exception: + pass + + # ----- Phase 1: Load MAP results if not provided ----- + if not map_results: + _progress("Loading MAP results from doc store ...") + map_results = [] + index = search.index_name(tenant_id) + # Each wiki_map_extract row stores its per-chunk extract as a JSON blob in + # ``content_with_weight`` (see _wiki_build_resume_doc) — the entity / + # concept / claim / relation / topic lists are NOT separate columns, so + # they must be parsed out of that blob to rebuild the map_result shape + # that _extract_raw_entities and REDUCE expect. + select_fields = ["content_with_weight", "doc_id"] + offset = 0 + page_size = 1000 + while True: + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + select_fields, + [], + {"compile_kwd": ["wiki_map_extract"]}, + [], + OrderByExpr(), + offset, + page_size, + index, + [kb_id], + ) + field_map = settings.docStoreConn.get_fields(res, select_fields) or {} + except Exception: + logging.exception("wiki: failed to load MAP results for kb=%s", kb_id) + break + for row in field_map.values(): + raw = row.get("content_with_weight") + if isinstance(raw, str) and raw: + try: + extract = json.loads(raw) + except Exception: + extract = None + elif isinstance(raw, dict): + extract = raw + else: + extract = None + if not isinstance(extract, dict): + continue + extract["doc_id"] = row.get("doc_id", "") + map_results.append(extract) + if len(field_map) < page_size: + break + offset += page_size + + if not map_results: + _progress("No MAP results found. Skipping wiki compilation.") + return summary + + # ----- Correct incremental flag for interrupted first builds ----- + # If the caller believes this is incremental but the KB has NO wiki_page + # rows, it means a previous full build was cancelled mid-way, leaving only + # half-written wiki_map_extract rows. Incremental logic relies on existing + # pages + doc_page_source to route changes; without any compiled page it + # would try to diff against nothing and produce empty output. In that case + # fall back to a full build. + if incremental: + try: + has_pages = await _wiki_has_any_pages(tenant_id, kb_id) + except Exception: + logging.exception("wiki: failed to check existing pages; assuming first build") + has_pages = False + if not has_pages: + _progress("No compiled wiki pages found; treating as first build (previous build was interrupted).") + incremental = False + + # ----- Phase 2: Entity Matching ----- + _progress("Entity Matching: deduplicating entities and concepts ...") + + # Lightweight metadata for matching + separate claim index for on-demand + # claim loading. Keeps Entity Matching operating on small metadata only + # (mirrors old-mode dedup); full claim text is loaded per-affected-name + # after matching, so peak memory stays bounded. + raw_entities, claim_index = _extract_raw_entities(map_results) + + # Collect the thematic topic labels the MAP phase extracted, for the Phase 6 + # topic grouping — done here while map_results is still alive. + map_topics: list[str] = [] + _seen_topics: set[str] = set() + for _mr in map_results: + for _t in _mr.get("topics") or []: + if isinstance(_t, str): + _t = _t.strip() + _k = _t.lower() + if _t and _k not in _seen_topics: + _seen_topics.add(_k) + map_topics.append(_t) + + # Release the heavy raw MAP payload as early as possible. All metadata is + # now in raw_entities and full claim text in claim_index; keeping + # map_results alive would pin the raw MAP structures in memory. + del map_results + + canonical_entities = await _load_canonical_entities(tenant_id, kb_id) + + canonical_map, name_resolution = await _wiki_match_entities( + raw_entities=raw_entities, + existing_canonical=canonical_entities, + embd_mdl=embd_mdl, + chat_mdl=chat_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + incremental=incremental, + ) + + # raw_entities (lightweight) no longer needed after matching. + del raw_entities + + if not canonical_map: + _progress("Entity Matching: no canonical entities found. Skipping.") + return summary + + _progress("Entity Matching: %d" % len(name_resolution.keys())) + # Persist new/changed canonical entities. + # Compute embeddings for ALL new entities in ONE batch call (the embedding + # API is invoked per text by the underlying driver — batching 95 entities as + # one call is ~95x faster than a per-entity call, each of which only carries + # a handful of tokens and spends ~0.3s in round-trip latency). + changed_items: list[tuple[str, dict]] = [] + new_items: list[tuple[str, dict, str]] = [] # (cname, centry, emb_text) + for cname, centry in canonical_map.items(): + existing = canonical_entities.get(cname) + if existing: + old_docs = set(k for k in (existing.get("source_doc_ids") or [])) + new_docs = set(centry.get("source_doc_ids", [])) + if old_docs != new_docs or centry["claim_count"] > existing.get("mention_count_int", 0): + # Only persist when data changes; reuse the existing embedding. + changed_items.append((cname, centry)) + else: + new_items.append((cname, centry, _entity_to_query_text(centry))) + + if changed_items: + persist_sem = asyncio.Semaphore(CANONICAL_PERSIST_CONCURRENT) + + async def _update_changed(item: tuple[str, dict]) -> None: + cname, centry = item + async with persist_sem: + await _update_canonical_entity( + tenant_id, + kb_id, + cname, + centry["type"], + centry.get("aliases", []), + centry.get("source_doc_ids", []), + centry["claim_count"], + ) + + await asyncio.gather(*(_update_changed(item) for item in changed_items)) + + if new_items and embd_mdl: + batch_texts = [t for _, _, t in new_items] + batch_embs, _ = await thread_pool_exec(embd_mdl.encode, batch_texts) + new_rows = [ + _build_canonical_entity_doc( + tenant_id, + kb_id, + cname, + centry["type"], + centry.get("aliases", []), + centry.get("source_doc_ids", []), + centry["claim_count"], + embedding=emb.tolist() if hasattr(emb, "tolist") else emb, + ) + for (cname, centry, _), emb in zip(new_items, batch_embs, strict=False) + ] + await thread_pool_exec( + settings.docStoreConn.insert, + new_rows, + search.index_name(tenant_id), + kb_id, + ) + elif new_items: + new_rows = [ + _build_canonical_entity_doc( + tenant_id, + kb_id, + cname, + centry["type"], + centry.get("aliases", []), + centry.get("source_doc_ids", []), + centry["claim_count"], + ) + for cname, centry, _ in new_items + ] + await thread_pool_exec( + settings.docStoreConn.insert, + new_rows, + search.index_name(tenant_id), + kb_id, + ) + + # Clean up deleted canonical entities (from doc deletion) + if deleted_doc_ids: + for cname, centry in list(canonical_map.items()): + centry["source_doc_ids"] = [d for d in centry.get("source_doc_ids", []) if d not in (deleted_doc_ids or set())] + if not centry["source_doc_ids"] and centry["claim_count"] <= 0: + await _delete_canonical_entity(tenant_id, kb_id, cname) + del canonical_map[cname] + + # ----- Phase 3: REDUCE ----- + _progress("REDUCE: computing per-entity changes ...") + + # Use canonical names (from Entity Matching) instead of raw MAP names + canonical_names: set[str] = set(canonical_map.keys()) + + if incremental: + # Affected doc ids are the docs contributing to this batch's canonical + # entities, plus any deleted docs. Derived from canonical_map (which + # carries source_doc_ids) — no need to re-read the released map_results. + affected_doc_ids = set() + for centry in canonical_map.values(): + affected_doc_ids.update(centry.get("source_doc_ids", [])) + affected_doc_ids = affected_doc_ids | (deleted_doc_ids or set()) + + # Map doc_page_source entity_names (raw) through name_resolution -> canonical + affected_names: set[str] = set() + if affected_doc_ids: + dps_tasks = [_wiki_load_doc_page_source(tenant_id, kb_id, did) for did in affected_doc_ids] + dps_results = await asyncio.gather(*dps_tasks) + for dps in dps_results: + if dps: + for raw_name in dps.get("entity_names", []): + cname = name_resolution.get(raw_name, raw_name) + if cname in canonical_names: + affected_names.add(cname) + if not affected_names: + affected_names = canonical_names + else: + affected_names = canonical_names + + # Load existing pages + existing_pages = await _search_existing_pages( + tenant_id, + kb_id, + [ + "slug_kwd", + "title_kwd", + "md_with_weight", + "claims", + "source_chunk_ids", + "source_doc_ids", + "page_version_int", + "synthesis_version_int", + "entity_names_kwd", + "related_kb_pages_kwd", + "page_type_kwd", + ], + ) + + # Build canonical claims ON-DEMAND only for affected names, then release + # the full claim_index. claim_index is keyed by RAW entity name (from MAP), + # so claims must be aggregated through name_resolution onto the canonical + # name. Otherwise a canonical name that differs from every raw name + # (e.g. "Apple Computer" → "Apple Inc.") would resolve to no claims and + # every affected entity would be a no-op → empty output. + canonical_claims: dict[str, list[dict]] = {} + for raw_name, claims in claim_index.items(): + cname = name_resolution.get(raw_name, raw_name) + if cname in affected_names: + canonical_claims.setdefault(cname, []).extend(claims) + # Ensure every affected name has an entry (possibly empty) + for name in affected_names: + canonical_claims.setdefault(name, []) + del claim_index + + deltas = await _wiki_reduce_batch( + affected_names=affected_names, + existing_pages=existing_pages, + deleted_doc_ids=deleted_doc_ids or set(), + canonical_claims=canonical_claims, + canonical_map=canonical_map, + name_resolution=name_resolution, + ) + + if not deltas: + _progress("REDUCE: no changes detected.") + # Still (re)group existing pages under topics — covers pages that were + # built before topic grouping existed, or a run where topics changed but + # no page's claims did. + await _wiki_assign_topics(embd_mdl, tenant_id, kb_id, map_topics, callback) + return summary + + # ----- Phase 4: Mode-specific dispatch ----- + # Precompute doc → canonical entity names for doc_page_source tracking, + # before canonical_map is released. + doc_to_entities: dict[str, list[str]] = {} + for cname, centry in canonical_map.items(): + for did in centry.get("source_doc_ids", []): + doc_to_entities.setdefault(did, []).append(cname) + del canonical_map + + if plan: + summary = await _wiki_mode_b_run( + deltas=deltas, + existing_pages=existing_pages, + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + incremental=incremental, + callback=callback, + doc_to_entities=doc_to_entities, + ) + else: + # Mode A: every entity AND concept becomes a page (no PLAN grouping). + # Each canonical entity/concept compiles to its own wiki page. + summary = await _wiki_mode_a_run( + deltas=deltas, + existing_pages=existing_pages, + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + incremental=incremental, + callback=callback, + canonical_claims=canonical_claims, + doc_to_entities=doc_to_entities, + ) + del deltas + del canonical_claims + del name_resolution + del existing_pages + + # ----- Phase 5: Update doc_page_source with canonical names ----- + # (doc_page_source page_ids is already handled in mode_run) + _progress("FINALIZE: updating cross-references ...") + try: + await _wiki_finalize(tenant_id, kb_id, embd_mdl) + except Exception: + logging.exception("wiki: FINALIZE failed for kb=%s", kb_id) + summary["errors"].append("FAILED_FINALIZE") + + # ----- Phase 6: Thematic topic grouping ----- + _progress("Grouping pages under topics ...") + await _wiki_assign_topics(embd_mdl, tenant_id, kb_id, map_topics, callback) + + return summary + + +async def _wiki_mode_a_run( + *, + deltas: list[dict], + existing_pages: dict[str, dict], + chat_mdl, + embd_mdl, + tenant_id: str, + kb_id: str, + incremental: bool, + callback: Callable | None = None, + canonical_claims: dict[str, list[dict]] | None = None, + doc_to_entities: dict[str, list[str]] | None = None, +) -> dict: + """Mode A: every grounded entity and concept compiles to its own page. + + No PLAN grouping — each canonical entity/concept is a single page. + Args: + deltas: All entity/concept deltas (entity_type retained in each). + canonical_claims: Canonical name → claims, for enriching page evidence. + doc_to_entities: doc_id → [canonical names] for doc_page_source. + """ + summary = {"pages_created": 0, "pages_modified": 0, "pages_deleted": 0, "errors": []} + + def _progress(msg: str): + if callback: + try: + callback(0.7, f"wiki REFINE A: {msg}") + except Exception: + pass + + # Map names to existing page IDs + name_to_page: dict[str, str] = {} + for pid, page in existing_pages.items(): + for n in _as_str_list(page.get("entity_names_kwd")): + name_to_page[n] = pid + + # Build page-level deltas. Each grounded entity/concept becomes one page; + # REDUCE has already filtered metadata-only entities without claims. + page_deltas: dict[str, dict] = {} + for d in deltas: + name = d.get("entity_name", "") + if not name: + continue + entity_type = d.get("entity_type", "entity") + if isinstance(entity_type, list): + entity_type = entity_type[0] if entity_type else "entity" + entity_type = str(entity_type or "entity").strip() + # Derive page_id with type-appropriate prefix (concept/ vs entity/) + prefix = "concept" if entity_type == "concept" else "entity" + page_id = name_to_page.get(name) or _wiki_derive_page_id(name, prefix=prefix) + + if page_id not in page_deltas: + page_deltas[page_id] = { + "page_id": page_id, + "page_title": name, + "existing_page": existing_pages.get(page_id), + "additions": [], + "retractions": [], + "claims": [], + "source_chunks": [], + } + entry = page_deltas[page_id] + entry["additions"].extend(d.get("additions", [])) + entry["retractions"].extend(d.get("retractions", [])) + entry["claims"].extend(d.get("claims", [])) + + # Collect source chunks from claims + for claim in d.get("claims", []): + for cid in _wiki_claim_chunk_ids(claim): + entry["source_chunks"].append( + { + "id": cid, + "text": claim.get("statement", claim.get("text", "")), + "source_doc_id": claim.get("source_doc_id"), + } + ) + + if d.get("action") == "delete": + entry["action"] = "delete" + elif entry.get("action") != "delete": + entry["action"] = d.get("action") + + # Enrich pages with related claims that share source_chunk_ids. + # Build a chunk_id → first-claim-text index from canonical_claims + # (claims loaded on-demand after matching). + if canonical_claims: + chunk_claims: dict[str, list[str]] = {} + for _cname, claims in canonical_claims.items(): + for claim in claims: + for cid in _wiki_claim_chunk_ids(claim): + chunk_claims.setdefault(cid, []).append(claim.get("statement", claim.get("text", ""))) + + for _pid, entry in page_deltas.items(): + page_chunk_ids = {c.get("id") for c in entry.get("source_chunks", []) if c.get("id")} + if not page_chunk_ids: + continue + for cid in page_chunk_ids: + texts = chunk_claims.get(cid) + if texts: + # Append one source-chunk entry per chunk id + entry["source_chunks"].append({"id": cid, "text": texts[0]}) + + # Concept depth check: thin CONCEPTS don't create pages (avoids dictionary + # entries). Pages without grounded claims have already been filtered by + # REDUCE above. + # ONLY on the very first full build (non-incremental), when concept claims + # are complete. During incremental builds the deltas carry only the + # changed claims, so a threshold check would wrongly reject every new + # concept — new concepts in incremental mode are created regardless. + if not incremental and not existing_pages: + concept_pages = [entry for entry in page_deltas.values() if entry.get("page_id", "").startswith("concept/")] + if concept_pages: + deep_concepts = _wiki_decide_concept_pages( + [ + {"term": entry["page_title"], "claims": entry["claims"], "source_doc_ids": list({c.get("source_doc_id") for c in entry["claims"] if c.get("source_doc_id")})} + for entry in concept_pages + ] + ) + deep_ids = {p["page_id"] for p in deep_concepts} + # Keep all entity pages + only deep concepts + page_deltas = {pid: entry for pid, entry in page_deltas.items() if not pid.startswith("concept/") or pid in deep_ids} + if not page_deltas: + _progress("No pages to compile. Skipping.") + return summary + + all_page_ids = list(existing_pages.keys()) + doc_updates: dict[str, list[str]] = {} + # Do not use a 20-slot semaphore around the whole page worker. The worker + # also performs source loading, embedding, and persistence after the LLM + # returns; limiting that whole region would artificially starve the LLM + # pool. LLMCallPool is the only limit for actual chat calls. + sem = asyncio.Semaphore(max(1, len(page_deltas))) + + async def _refine_one(pid: str, entry: dict) -> None: + async with sem: + try: + existing = entry["existing_page"] + page_type = "concept" if pid.startswith("concept/") else "entity" + if entry.get("action") == "delete": + await _wiki_refine_page( + mode="delete", + page_id=pid, + page_title=entry["page_title"], + existing_page=existing, + page_type_kwd=page_type, + additions=None, + retractions=None, + source_chunks=[], + claims=[], + available_pages=all_page_ids, + contextual_hints="", + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + page_version=existing.get("page_version_int", 0) if existing else 0, + ) + summary["pages_deleted"] += 1 + return + + next_version = (existing.get("page_version_int", 0) if existing else 0) + 1 + new_doc_ids = {c.get("source_doc_id") for c in entry["additions"] if c.get("source_doc_id")} + if existing and _wiki_should_re_synthesize(existing, new_doc_ids, next_version): + refine_mode = "re-synthesize" + elif existing: + refine_mode = "modify" + else: + refine_mode = "generate" + + result = await _wiki_refine_page( + mode=refine_mode, + page_id=pid, + page_title=entry["page_title"], + existing_page=existing, + page_type_kwd=page_type, + additions=entry["additions"], + retractions=entry["retractions"], + source_chunks=entry["source_chunks"], + claims=entry["claims"], + available_pages=all_page_ids, + contextual_hints="", + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + page_version=existing.get("page_version_int", 0) if existing else 0, + ) + if refine_mode == "generate": + summary["pages_created"] += 1 + else: + summary["pages_modified"] += 1 + + if result: + for c in entry["additions"]: + did = c.get("source_doc_id") + if did: + doc_updates.setdefault(did, []).append(pid) + + except Exception: + logging.exception("wiki A: REFINE failed for %s", pid) + summary["errors"].append(f"REFINE_FAILED:{pid}") + + tasks = [_refine_one(pid, entry) for pid, entry in page_deltas.items()] + if tasks: + _progress(f"REFINE A: {len(tasks)} pages (LLM pool max {WIKI_REFINE_MAX_CONCURRENT}) ...") + await asyncio.gather(*tasks) + + for did, pids in doc_updates.items(): + try: + existing_dps = (await _wiki_load_doc_page_source(tenant_id, kb_id, did)) or {} + existing_pids = existing_dps.get("page_ids", []) + for pid in pids: + if pid not in existing_pids: + existing_pids.append(pid) + # Collect entity names for this doc from precomputed doc_to_entities + doc_entity_names = (doc_to_entities or {}).get(did, []) or existing_dps.get("entity_names") + await _wiki_update_doc_page_source( + tenant_id, + kb_id, + did, + existing_pids, + entity_names=doc_entity_names, + chunk_hashes=existing_dps.get("source_chunk_hashes"), + map_checksum=existing_dps.get("map_checksum"), + ) + except Exception: + logging.exception("wiki A: doc_page_source update failed for doc %s", did) + + _progress(f"done: +{summary['pages_created']} ~{summary['pages_modified']} -{summary['pages_deleted']}") + return summary + + +async def _wiki_mode_b_run( + *, + deltas: list[dict], + existing_pages: dict[str, dict], + chat_mdl, + embd_mdl, + tenant_id: str, + kb_id: str, + incremental: bool, + callback: Callable | None = None, + doc_to_entities: dict[str, list[str]] | None = None, +) -> dict: + """Mode B: Page Router + per-page REFINE.""" + + summary = {"pages_created": 0, "pages_modified": 0, "pages_deleted": 0, "errors": []} + + def _progress(msg: str): + if callback: + try: + callback(0.7, f"wiki REFINE B: {msg}") + except Exception: + pass + + # Convert deltas to entity dicts that Page Router can process + affected_entities = [ + { + "entity_name": d.get("entity_name", ""), + "entity_type": d.get("entity_type", "entity"), + "claims": d.get("additions", []) + d.get("claims", []), + "action": d.get("action", ""), + } + for d in deltas + if d.get("entity_name") + ] + + if not affected_entities: + _progress("No affected entities. Skipping.") + return summary + + # Run Page Router + _progress(f"Page Router: routing {len(affected_entities)} entities ...") + assignments = await _wiki_page_router( + affected_entities=affected_entities, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + existing_page_ids=set(existing_pages), + ) + + if not assignments: + _progress("Page Router: no assignments. Skipping.") + return summary + + # Load available pages for wikilinks + all_page_ids = list(existing_pages.keys()) + + # Collect source chunks per assignment for the REFINE prompt + page_source_chunks: dict[str, list[dict]] = {} + for pid, entities in assignments.items(): + page_key = pid[5:] if pid.startswith("_new_") else pid + chunks: list[dict] = [] + for ent in entities: + for c in ent.get("claims", []): + for cid in _wiki_claim_chunk_ids(c): + chunks.append( + { + "id": cid, + "text": c.get("statement", c.get("text", "")), + "source_doc_id": c.get("source_doc_id"), + } + ) + if chunks: + page_source_chunks[page_key] = chunks + + doc_updates: dict[str, list[str]] = {} # doc_id → [page_ids] + # The shared LLMCallPool limits chat calls. This semaphore must not cap + # the complete worker because embedding and page persistence happen after + # the chat call and should not consume an LLM concurrency slot. + sem = asyncio.Semaphore(max(1, len(assignments))) + + async def _refine_one(page_id: str, entities: list) -> None: + async with sem: + try: + is_new = page_id.startswith("_new_") + page_key = page_id[5:] if is_new else page_id + page_key = str(page_key or "").strip() + if not page_key: + # Blank assignment (e.g. a mangled slug from the router) — + # nothing to generate for an empty page id. + return + existing = existing_pages.get(page_key) if not is_new else None + + # Determine page_type for Mode B. The slug prefix is the source + # of truth: a page whose slug starts with "concept/" is a concept + # page (the router derives it via _wiki_derive_page_id). Inferring + # from `entities[].entity_type` is unreliable (many entities carry + # a concrete type like "person"/"org", or none at all), which + # previously mislabelled concept pages as entity. + if existing: + page_type = existing.get("page_type_kwd", "entity") + if isinstance(page_type, (list, tuple)): + page_type = page_type[0] if page_type else "entity" + else: + if page_key.startswith("concept/"): + page_type = "concept" + else: + page_type = "entity" + + additions = [] + action = "create" if is_new else "update" + for ent in entities: + additions.extend(ent.get("claims", [])) + if ent.get("action") == "delete": + action = "delete" + + if action == "delete": + await _wiki_refine_page( + mode="delete", + page_id=page_key, + page_title=existing.get("title_kwd", page_key) if existing else page_key, + existing_page=existing, + page_type_kwd=page_type, + additions=None, + retractions=None, + source_chunks=[], + claims=[], + available_pages=all_page_ids, + contextual_hints="", + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + page_version=existing.get("page_version_int", 0) if existing else 0, + ) + summary["pages_deleted"] += 1 + return + + refine_mode = "generate" if is_new else "modify" + if existing and _wiki_should_re_synthesize( + existing, + {c.get("source_doc_id") for c in additions if c.get("source_doc_id")}, + existing.get("page_version_int", 0) + 1, + ): + refine_mode = "re-synthesize" + + result = await _wiki_refine_page( + mode=refine_mode, + page_id=page_key, + page_title=existing.get("title_kwd", page_key) if existing else entities[0].get("entity_name", page_key), + existing_page=existing, + page_type_kwd=page_type, + additions=additions, + retractions=[], + source_chunks=page_source_chunks.get(page_key, []), + claims=additions, + available_pages=all_page_ids, + contextual_hints=_wiki_build_contextual_hints(page_key, existing, {}), + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + page_version=existing.get("page_version_int", 0) if existing else 0, + ) + if is_new: + summary["pages_created"] += 1 + else: + summary["pages_modified"] += 1 + + if result: + await _wiki_update_plan_group( + tenant_id, + kb_id, + page_key, + entity_names=[e.get("entity_name", "") for e in entities], + page_version=result.get("page_version_int", 1), + ) + + # Collect doc_page_source updates (deferred) + for ent in entities: + for c in ent.get("claims", []): + did = c.get("source_doc_id") + if did: + doc_updates.setdefault(did, []).append(page_key) + + except Exception: + logging.exception("wiki B: REFINE failed for %s", page_id) + summary["errors"].append(f"REFINE_FAILED:{page_id}") + + tasks = [_refine_one(pid, ents) for pid, ents in assignments.items()] + if tasks: + _progress(f"REFINE B: {len(tasks)} pages (LLM pool max {WIKI_REFINE_MAX_CONCURRENT}) ...") + await asyncio.gather(*tasks) + + # Apply doc_page_source updates serially (no race), preserving metadata + for did, pids in doc_updates.items(): + try: + existing_dps = (await _wiki_load_doc_page_source(tenant_id, kb_id, did)) or {} + existing_pids = existing_dps.get("page_ids", []) + for pid in pids: + if pid not in existing_pids: + existing_pids.append(pid) + await _wiki_update_doc_page_source( + tenant_id, + kb_id, + did, + existing_pids, + entity_names=(doc_to_entities or {}).get(did, []) or existing_dps.get("entity_names"), + chunk_hashes=existing_dps.get("source_chunk_hashes"), + map_checksum=existing_dps.get("map_checksum"), + ) + except Exception: + logging.exception("wiki B: doc_page_source update failed for doc %s", did) + + _progress(f"done: +{summary['pages_created']} ~{summary['pages_modified']} -{summary['pages_deleted']}") + return summary + + +async def _wiki_update_plan_group( + tenant_id: str, + kb_id: str, + page_id: str, + entity_names: list[str], + page_version: int, +) -> None: + """Update or create plan_group row (Mode B).""" + from rag.nlp import search + from common.misc_utils import thread_pool_exec + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + condition = { + "compile_kwd": [WIKI_PLAN_GROUP_COMPILE_KWD], + "page_id": [page_id], + } + + doc = { + "id": _stable_row_id(WIKI_PLAN_GROUP_COMPILE_KWD, kb_id, page_id), + "kb_id": kb_id, + "page_id": page_id, + "entity_names": json.dumps(entity_names, ensure_ascii=False), + "page_version_int": page_version, + "compile_kwd": WIKI_PLAN_GROUP_COMPILE_KWD, + } + + existing = await thread_pool_exec( + settings.docStoreConn.search, + ["page_id"], + [], + condition, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + if settings.docStoreConn.get_fields(existing, ["page_id"]): + await thread_pool_exec( + settings.docStoreConn.update, + {"page_id": page_id}, + doc, + index, + kb_id, + ) + else: + await thread_pool_exec( + settings.docStoreConn.insert, + [doc], + index, + kb_id, + ) + + +async def wiki_handle_document_deleted( + tenant_id: str, + kb_id: str, + doc_id: str, + chat_mdl, + embd_mdl, + plan: bool = False, +) -> dict: + """Clean up wiki pages + canonical entities when a document is deleted. + + Args: + plan: True=Mode B (update plan_group), False=Mode A + + Returns: {pages_modified, pages_deleted, errors} + """ + summary = {"pages_modified": 0, "pages_deleted": 0, "errors": []} + + # Step 1: Update canonical entity index (decrement claim_count) + dps = await _wiki_load_doc_page_source(tenant_id, kb_id, doc_id) + if not dps: + return summary + + entity_names = dps.get("entity_names", []) + if entity_names: + canonical_index = await _load_canonical_entities(tenant_id, kb_id) + for ename in entity_names: + centry = canonical_index.get(ename) + if centry: + src_ids = centry.get("source_doc_ids", []) + if isinstance(src_ids, str): + try: + src_ids = json.loads(src_ids) if src_ids else [] + except (json.JSONDecodeError, TypeError): + src_ids = [] + if doc_id in src_ids: + src_ids.remove(doc_id) + + if not src_ids: + await _delete_canonical_entity(tenant_id, kb_id, ename) + else: + # Keep existing mention_count_int; the next REFINE phase + # will recalculate claims precisely from wiki pages. + # Removing the doc_id from source_doc_ids prevents future + # incremental runs from re-tracking this deletion. + await _save_canonical_entity( + tenant_id, + kb_id, + ename, + centry.get("entity_type_kwd", "entity"), + centry.get("aliases", []), + src_ids, + centry.get("mention_count_int", len(src_ids)), + ) + + affected_page_ids = dps.get("page_ids", []) + if not affected_page_ids: + return summary + + # Step 2: Update wiki pages + all_existing_pages = await _search_existing_pages( + tenant_id, + kb_id, + ["slug_kwd", "title_kwd", "md_with_weight", "claims", "source_doc_ids", "page_version_int", "entity_names_kwd", "page_type_kwd"], + ) + + for page_id in affected_page_ids: + try: + existing = all_existing_pages.get(page_id) + if not existing: + continue + + source_doc_ids = existing.get("source_doc_ids", []) + if isinstance(source_doc_ids, str): + source_doc_ids = json.loads(source_doc_ids) if source_doc_ids else [] + + if doc_id in source_doc_ids: + source_doc_ids.remove(doc_id) + + page_type = existing.get("page_type_kwd", "concept" if not plan else "entity") + + if not source_doc_ids: + await _wiki_refine_page( + mode="delete", + page_id=page_id, + page_title=existing.get("title_kwd", page_id), + existing_page=existing, + page_type_kwd=page_type, + additions=None, + retractions=None, + source_chunks=[], + claims=[], + available_pages=[], + contextual_hints="", + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + page_version=existing.get("page_version_int", 0), + ) + summary["pages_deleted"] += 1 + else: + existing_claims = existing.get("claims", []) + if isinstance(existing_claims, str): + existing_claims = json.loads(existing_claims) if existing_claims else [] + + retractions = [c for c in existing_claims if c.get("source_doc_id") == doc_id] + retained = [c for c in existing_claims if c.get("source_doc_id") != doc_id] + + await _wiki_refine_page( + mode="modify", + page_id=page_id, + page_title=existing.get("title_kwd", page_id), + existing_page=existing, + page_type_kwd=page_type, + additions=[], + retractions=retractions, + source_chunks=[], + claims=retained, + available_pages=list(all_existing_pages.keys()), + contextual_hints=_wiki_build_contextual_hints(page_id, existing, {}), + chat_mdl=chat_mdl, + embd_mdl=embd_mdl, + tenant_id=tenant_id, + kb_id=kb_id, + page_version=existing.get("page_version_int", 0), + ) + summary["pages_modified"] += 1 + + if plan: + plan_condition = { + "compile_kwd": [WIKI_PLAN_GROUP_COMPILE_KWD], + "page_id": [page_id], + } + index = search.index_name(tenant_id) + res_pg = await thread_pool_exec( + settings.docStoreConn.search, + ["entity_names", "source_doc_ids"], + [], + plan_condition, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + pg_map = settings.docStoreConn.get_fields(res_pg, ["entity_names", "source_doc_ids"]) + for pg_row in pg_map.values(): + pg_src_ids = pg_row.get("source_doc_ids", []) + if isinstance(pg_src_ids, str): + pg_src_ids = json.loads(pg_src_ids) + if doc_id in pg_src_ids: + pg_src_ids.remove(doc_id) + await _wiki_update_plan_group( + tenant_id, + kb_id, + page_id, + entity_names=json.loads(pg_row.get("entity_names", "[]")) if isinstance(pg_row.get("entity_names"), str) else pg_row.get("entity_names", []), + page_version=existing.get("page_version_int", 0), + ) + break + + except Exception: + logging.exception("wiki: document deletion cleanup failed for page=%s doc=%s", page_id, doc_id) + summary["errors"].append(f"CLEANUP_FAILED:{page_id}") + + await _wiki_delete_doc_page_source(tenant_id, kb_id, doc_id) + + try: + await _wiki_finalize(tenant_id, kb_id, embd_mdl) + except Exception: + logging.exception("wiki: FINALIZE after deletion failed") + + return summary + + +__all__ = [ + "WIKI_PAGE_COMPILE_KWD", + "WIKI_PLAN_GROUP_COMPILE_KWD", + "WIKI_DOC_PAGE_SOURCE_COMPILE_KWD", + "WIKI_CANONICAL_ENTITY_COMPILE_KWD", + "wiki_compile_incremental", + "wiki_handle_document_deleted", + "_wiki_reduce_entity", + "_wiki_reduce_batch", + "_wiki_match_entities", + "_wiki_page_router", + "_wiki_finalize", + "_wiki_refine_page", + "_wiki_update_doc_page_source", + "_load_canonical_entities", + "_save_canonical_entity", + "_delete_canonical_entity", + "_extract_raw_entities", +] diff --git a/rag/svr/task_executor_refactor/dataset_wiki_generator.py b/rag/svr/task_executor_refactor/dataset_wiki_generator.py index b81919d0cc..7fc6259716 100644 --- a/rag/svr/task_executor_refactor/dataset_wiki_generator.py +++ b/rag/svr/task_executor_refactor/dataset_wiki_generator.py @@ -36,7 +36,7 @@ Design notes: ``parser_config.compilation_template_group_id`` to a template list via the shared parser-config helper and ``CompilationTemplateGroupService.resolve_template_ids``. -* The persistence helpers (``persist_wiki_pages_to_es`` etc.) are +* The persistence helpers (``persist_wiki_pages`` etc.) are exposed at module level for testing but are only called from :func:`run_wiki` in production. """ @@ -116,6 +116,10 @@ WIKI_DERIVED_COMPILE_KWDS = ( "wiki_entity", "wiki_relation", "wiki_page_graph", + # Canonical entity rows carry a source_doc_ids array; on doc deletion they + # must be shrunk (or dropped) too, otherwise the canonical index keeps + # referencing removed docs and later incremental merges re-import them. + "wiki_canonical_entity", ) @@ -219,6 +223,54 @@ def _pipeline_compilation_template_ids(pipeline_id: str, tenant_id: str) -> list return template_ids +def _wiki_eligible_docs(all_docs, tenant_id: str, skip_doc_ids=None) -> list[tuple[dict, str]]: + """Docs eligible for wiki compilation, each paired with its wiki template id. + + A doc is eligible when its ``parser_config`` OR its ingestion pipeline + resolves to at least one artifacts-kind ("wiki") compilation template — the + pipeline path is essential for docs uploaded/parsed through a pipeline, which + carry their compilation templates on the pipeline's compiler rather than in + ``parser_config``. Returns ``(doc, template_id)`` for the first wiki template + matched per doc. Shared by :func:`run_wiki` and :func:`run_wiki_incremental` + so their eligibility can't drift. + """ + from api.db.services.compilation_template_service import CompilationTemplateService + from api.apps.restful_apis.chunk_api import _compilation_template_kind + + skip_doc_ids = skip_doc_ids or set() + eligible: list[tuple[dict, str]] = [] + pipeline_template_ids_cache: dict[str, list[str]] = {} + for d in all_docs or []: + if str(d.get("id")) in skip_doc_ids: + continue + pc = d.get("parser_config") or {} + template_ids: list[str] = [] + seen_template_ids: set[str] = set() + for template_id in _parser_config_compilation_template_ids(pc, tenant_id): + if template_id in seen_template_ids: + continue + seen_template_ids.add(template_id) + template_ids.append(template_id) + pipeline_id = (d.get("pipeline_id") or "").strip() + if pipeline_id: + if pipeline_id not in pipeline_template_ids_cache: + pipeline_template_ids_cache[pipeline_id] = _pipeline_compilation_template_ids(pipeline_id, tenant_id) + for template_id in pipeline_template_ids_cache[pipeline_id]: + if template_id in seen_template_ids: + continue + seen_template_ids.add(template_id) + template_ids.append(template_id) + + for template_id in template_ids: + template = CompilationTemplateService.get_saved(template_id, tenant_id) + config = template.get("config") if template else {} + kind = _compilation_template_kind(config.get("kind") if isinstance(config, dict) else "") + if kind == "wiki": + eligible.append((d, template_id)) + break + return eligible + + async def _wiki_existing_map_doc_ids(tenant_id: str, kb_id: str) -> set[str]: from common.doc_store.doc_store_base import OrderByExpr @@ -260,6 +312,38 @@ async def _wiki_existing_map_doc_ids(tenant_id: str, kb_id: str) -> set[str]: return doc_ids +async def _wiki_has_compiled_pages(tenant_id: str, kb_id: str) -> bool: + """True when at least one compiled wiki page already exists for the KB. + + Used to tell "nothing changed and pages already exist" (a genuine no-op) + apart from "MAP rows exist but no pages were ever produced" (a prior run + persisted MAP then never finished REDUCE) — only the latter should trigger a + full rebuild from the stored extracts. + """ + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return False + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["id"], + [], + {"compile_kwd": [WIKI_PAGE_COMPILE_KWD]}, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + return bool(settings.docStoreConn.get_total(res)) + except Exception: + logging.exception("wiki: page existence probe failed for kb=%s", kb_id) + return False + + async def _wiki_delete_deleted_doc_state( tenant_id: str, kb_id: str, @@ -291,6 +375,23 @@ async def _wiki_delete_deleted_doc_state( ) return + # 1b. doc_page_source rows are keyed by doc_id too — delete outright. + try: + await thread_pool_exec( + settings.docStoreConn.delete, + { + "compile_kwd": ["wiki_doc_page_source"], + "doc_id": sorted(deleted_doc_ids), + }, + index, + kb_id, + ) + except Exception: + logging.exception( + "wiki: failed to delete doc_page_source rows for removed docs in kb=%s", + kb_id, + ) + # 2. Derived KB-scoped rows: reference-counted self-healing backstop for # the eager delete-time cleanup (DocumentService.remove_wiki_products). # Read every row referencing any deleted doc, drop the ones left with no @@ -377,6 +478,106 @@ async def _wiki_delete_deleted_doc_state( ) +# ----- mode (plan) persistence & full reset ---------------------------------- + + +def _wiki_mode_meta_id(kb_id: str) -> str: + """Stable row id for the KB-level mode (plan) meta record.""" + return f"wiki_mode_meta_{kb_id}" + + +async def _wiki_load_mode_plan(tenant_id: str, kb_id: str) -> bool | None: + """Return the plan (Mode B) value recorded by the previous build, or None + if this KB has never recorded a mode (e.g. first ever build).""" + from common.doc_store.doc_store_base import OrderByExpr + + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return None + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + ["plan_kwd"], + [], + {"compile_kwd": ["wiki_mode_meta"], "id": [_wiki_mode_meta_id(kb_id)]}, + [], + OrderByExpr(), + 0, + 1, + index, + [kb_id], + ) + fm = settings.docStoreConn.get_fields(res, ["plan_kwd"]) or {} + for row in fm.values(): + val = row.get("plan_kwd") + if isinstance(val, list): + val = val[0] if val else "" + val = str(val or "").strip() + if val in ("true", "1", "yes"): + return True + if val in ("false", "0", "no"): + return False + except Exception: + logging.exception("wiki: failed to load mode meta for kb=%s", kb_id) + return None + + +async def _wiki_save_mode_plan(tenant_id: str, kb_id: str, plan: bool) -> None: + index = search.index_name(tenant_id) + row = { + "id": _wiki_mode_meta_id(kb_id), + "compile_kwd": "wiki_mode_meta", + "plan_kwd": "true" if plan else "false", + "kb_id": kb_id, + "create_timestamp_flt": float(__import__("time").time()), + } + try: + await thread_pool_exec( + settings.docStoreConn.insert, + [row], + index, + kb_id, + ) + except Exception: + logging.exception("wiki: failed to save mode meta for kb=%s", kb_id) + + +async def _wiki_reset_all_wiki_state(tenant_id: str, kb_id: str) -> None: + """Drop every wiki-derived row for the KB (canonical, pages, relations, + entities, plan/draft/reduce, topics, doc_page_source, mode meta). Used when + the plan (mode) setting toggles: Mode A and Mode B pages are structurally + different and cannot be merged incrementally, so a mode switch must rebuild + from a clean slate.""" + + index = search.index_name(tenant_id) + if not settings.docStoreConn.index_exist(index, kb_id): + return + all_kwds = [ + "wiki_canonical_entity", + "wiki_page", + "wiki_entity", + "wiki_relation", + "wiki_page_graph", + "wiki_page_topic", + "wiki_compilation_plan", + "wiki_reduce_result", + "wiki_page_draft", + "wiki_doc_page_source", + "wiki_map_extract", + "wiki_mode_meta", + ] + # Delete in one bulk call using compile_kwd IN filter. + try: + await thread_pool_exec( + settings.docStoreConn.delete, + {"compile_kwd": all_kwds}, + index, + kb_id, + ) + except Exception: + logging.exception("wiki: failed to reset all wiki state for kb=%s", kb_id) + + def _wiki_topic_from_page(page: Dict, fallback: str = "") -> str: for key in ("topic", "title", "page_type"): value = page.get(key) @@ -490,12 +691,12 @@ async def _ensure_wiki_topic_rows( # ----- persistence --------------------------------------------------- -async def persist_wiki_pages_to_es( +async def persist_wiki_pages( ctx: TaskContext, pages: List[Dict], embd_mdl, ) -> None: - """Insert one ES row per generated artifact page using the + """Insert one doc-store row per generated artifact page using the knowledge-compilation schema: id xxh64(kb_id + ":" + slug) @@ -853,11 +1054,11 @@ def build_wiki_page_graph( return entity_rows, relation_rows -async def persist_wiki_page_graph_to_es( +async def persist_wiki_page_graph( ctx: TaskContext, pages: List[Dict], ) -> None: - """Materialize and store the per-entity / per-relation ES rows + """Materialize and store the per-entity / per-relation doc-store rows derived from artifact pages. Writes two row types — both delete-then-insert for idempotent @@ -960,7 +1161,6 @@ async def run_wiki( get_tenant_default_model_by_type, resolve_model_config, ) - from api.apps.restful_apis.chunk_api import _compilation_template_kind progress = ctx.progress_cb progress(0.0, "Loading documents for wiki compilation...") @@ -1003,34 +1203,7 @@ async def run_wiki( deleted_doc_ids, ) - eligible = [] - pipeline_template_ids_cache: dict[str, list[str]] = {} - for d in all_docs or []: - pc = d.get("parser_config") or {} - template_ids: list[str] = [] - seen_template_ids: set[str] = set() - for template_id in _parser_config_compilation_template_ids(pc, ctx.tenant_id): - if template_id in seen_template_ids: - continue - seen_template_ids.add(template_id) - template_ids.append(template_id) - pipeline_id = (d.get("pipeline_id") or "").strip() - if pipeline_id: - if pipeline_id not in pipeline_template_ids_cache: - pipeline_template_ids_cache[pipeline_id] = _pipeline_compilation_template_ids(pipeline_id, ctx.tenant_id) - for template_id in pipeline_template_ids_cache[pipeline_id]: - if template_id in seen_template_ids: - continue - seen_template_ids.add(template_id) - template_ids.append(template_id) - - for template_id in template_ids: - template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id) - config = template.get("config") if template else {} - kind = _compilation_template_kind(config.get("kind") if isinstance(config, dict) else "") - if kind == "wiki": - eligible.append((d, template_id)) - break + eligible = _wiki_eligible_docs(all_docs, ctx.tenant_id) if not eligible: progress(1.0, "No documents are configured for wiki compilation.") return @@ -1178,11 +1351,9 @@ async def run_wiki( parser_config=parser_cfg, batch_size_cap=8, window_fraction=0.5, - # Keep a bounded internal worker queue. The shared pool - # globally limits active + admitted waiting calls to - # WIKI_MAP_MAX_PENDING, while this prevents every outer - # batch from creating all of its sub-batch tasks at once. - max_workers=6, + # Match the shared pool width. The pool is the single + # admission/concurrency control for actual MAP LLM calls. + max_workers=WIKI_MAP_LLM_POOL_SIZE, ) for key in stats["agg"]: stats["agg"][key] += len(phase1.get(key) or []) @@ -1258,7 +1429,12 @@ async def run_wiki( progress(0.75, "Planning wiki pages...") await wiki_plan_from_reduction( - chat_mdl=kb_chat_mdl, + chat_mdl=map_llm_pool.wrap( + kb_chat_mdl, + priority=20, + label="wiki-plan", + context=f"{ctx.kb_id}:plan", + ), embd_mdl=embedding_model, tenant_id=ctx.tenant_id, kb_id=ctx.kb_id, @@ -1289,14 +1465,328 @@ async def run_wiki( # 6. Persist searchable wiki_page rows. try: - await persist_wiki_pages_to_es(ctx=ctx, pages=pages or [], embd_mdl=embedding_model) + await persist_wiki_pages(ctx=ctx, pages=pages or [], embd_mdl=embedding_model) except Exception: - logging.exception("wiki: ES persist failed for kb %s", ctx.kb_id) + logging.exception("wiki: persist failed for kb %s", ctx.kb_id) # 7. Materialize the canvas graph from the refined pages. try: - await persist_wiki_page_graph_to_es(ctx=ctx, pages=pages or []) + await persist_wiki_page_graph(ctx=ctx, pages=pages or []) except Exception: logging.exception("wiki: page-graph persist failed for kb %s", ctx.kb_id) progress(1.0, f"Wiki compiled {len(pages or [])} page(s).") + + +# ----- dual-mode incremental entry point --------------------------------- + + +async def run_wiki_incremental( + ctx: TaskContext, + embedding_model, + load_chunks_for_doc: Callable[..., AsyncIterator[list[dict]]], + plan: bool = False, +) -> None: + """Dual-mode wiki compilation with incremental support. + + Mode A (plan=False, default): + 1 concept = 1 page (WeKnora style). + MAP → REDUCE → per-concept REFINE → FINALIZE. + Incremental: per-concept modify based on doc_change tracking. + + Mode B (plan=True): + PLAN groups entities → per-page REFINE. + Incremental: Page Router (KNN) routes entities to existing pages. + + Args: + ctx: Task context + embedding_model: Embedding model + load_chunks_for_doc: Chunk loader + plan: True=Mode B, False=Mode A (default) + """ + from api.db.services.document_service import DocumentService + from api.db.services.compilation_template_service import CompilationTemplateService + from api.db.services.llm_service import LLMBundle + from api.db.joint_services.tenant_model_service import ( + get_tenant_default_model_by_type, + resolve_model_config, + ) + from rag.advanced_rag.knowlege_compile.wiki_incremental import ( + wiki_compile_incremental, + ) + from rag.advanced_rag.knowlege_compile.structure import LLMCallPool + + progress = ctx.progress_cb + progress(0.0, f"Loading documents for wiki {'PLAN' if plan else 'no-plan'} compilation...") + + # 1. Check if this is incremental (existing MAP rows present) + existing_map_doc_ids = await _wiki_existing_map_doc_ids(ctx.tenant_id, ctx.kb_id) + is_incremental = bool(existing_map_doc_ids) + deleted_doc_ids = set() + + if is_incremental: + # Find deleted docs + all_docs, _ = await thread_pool_exec( + DocumentService.get_by_kb_id, + kb_id=ctx.kb_id, + page_number=0, + items_per_page=0, + orderby="create_time", + desc=False, + keywords="", + run_status=[], + types=[], + suffix=[], + ) + current_doc_ids = {str(d.get("id")) for d in all_docs or [] if d.get("id")} + deleted_doc_ids = existing_map_doc_ids - current_doc_ids + if deleted_doc_ids: + progress(0.02, f"Cleaning {len(deleted_doc_ids)} deleted doc(s) ...") + await _wiki_delete_deleted_doc_state(ctx.tenant_id, ctx.kb_id, deleted_doc_ids) + + # 2. Pick eligible docs + all_docs, _ = await thread_pool_exec( + DocumentService.get_by_kb_id, + kb_id=ctx.kb_id, + page_number=0, + items_per_page=0, + orderby="create_time", + desc=False, + keywords="", + run_status=[], + types=[], + suffix=[], + ) + eligible = _wiki_eligible_docs(all_docs, ctx.tenant_id, skip_doc_ids=deleted_doc_ids) + + if not eligible and not is_incremental: + progress(1.0, "No documents configured for wiki compilation.") + return + + # Re-resolve plan (Mode B) from the ELIGIBLE docs' templates. Each eligible + # doc resolves to a wiki template either via its own parser_config or via + # its ingestion pipeline (doc.pipeline_id → pipeline dsl → compiler → + # template). The task handler's `plan` param only looks at the KB-level + # parser_config and therefore misses the pipeline path; re-derive it here so + # a pipeline-bound template with plan=yes actually enables Mode B. + if not plan: + try: + for _doc, tid in eligible: + tpl = CompilationTemplateService.get_saved(tid, ctx.tenant_id) + cfg = (tpl.get("config") or {}) if tpl else {} + if isinstance(cfg, dict) and cfg.get("plan") in (True, "yes", "true"): + plan = True + break + except Exception: + pass # keep the handler-provided plan as fallback + + # Mode-change detection. plan toggling (A↔B) is a config change: the page + # structures differ fundamentally (single-entity pages vs PLAN-grouped + # pages), so switching modes must reset all wiki-derived state and rebuild + # from scratch instead of incrementally mixing old-mode and new-mode pages. + prev_plan = await _wiki_load_mode_plan(ctx.tenant_id, ctx.kb_id) + if prev_plan is not None and bool(prev_plan) != bool(plan) and is_incremental: + progress(0.05, f"Mode switched (plan: {'on' if prev_plan else 'off'} -> {'on' if plan else 'off'}); rebuilding wiki from scratch...") + await _wiki_reset_all_wiki_state(ctx.tenant_id, ctx.kb_id) + # Everything is gone; this is now a first build. + is_incremental = False + existing_map_doc_ids = set() + deleted_doc_ids = set() + await _wiki_save_mode_plan(ctx.tenant_id, ctx.kb_id, bool(plan)) + + # 3. Resolve chat model + llm_bundle_cache: dict[str, LLMBundle] = {} + + def _bundle_for(llm_id: str | None) -> LLMBundle: + key = (llm_id or "").strip() or "__tenant_default__" + cached = llm_bundle_cache.get(key) + if cached is not None: + return cached + try: + if key == "__tenant_default__": + cfg = get_tenant_default_model_by_type(ctx.tenant_id, LLMType.CHAT) + else: + cfg = resolve_model_config(ctx.tenant_id, LLMType.CHAT, key) + except Exception: + cfg = get_tenant_default_model_by_type(ctx.tenant_id, LLMType.CHAT) + key = "__tenant_default__" + cached = llm_bundle_cache.get(key) + if cached is not None: + return cached + bundle = LLMBundle(ctx.tenant_id, cfg, lang=ctx.language) + llm_bundle_cache[key] = bundle + return bundle + + map_llm_pool = LLMCallPool(WIKI_MAP_LLM_POOL_SIZE, max_pending=WIKI_MAP_MAX_PENDING) + kb_chat_llm_id = None + first_template_found = False + + # 4. MAP per doc (same as run_wiki's MAP phase) + map_queue: asyncio.Queue = asyncio.Queue(maxsize=WIKI_MAP_QUEUE_SIZE) + n_docs = len(eligible) + all_map_results: list[dict] = [] + + # Pre-resolve parser_cfg for each eligible doc (avoids sync DB call in worker) + doc_configs: dict[str, dict] = {} + for d, template_id in eligible: + try: + template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id) + cfg = (template.get("config") or {}) if template else {} + doc_configs[d["id"]] = cfg + if not first_template_found and isinstance(cfg, dict): + first_template_found = True + llm_id = (cfg.get("llm_id") or "").strip() + kb_chat_llm_id = llm_id or None + if not kb_chat_llm_id: + kb_chat_llm_id = None + except Exception: + logging.exception("wiki: config resolve failed for doc %s", d["id"]) + doc_configs[d["id"]] = {} + + async def _produce_doc(i: int, job: tuple[dict, str]) -> None: + doc, template_id = job + doc_id = doc["id"] + progress(0.05 + 0.6 * (i / max(n_docs, 1)), f"MAP {i + 1}/{n_docs}: {doc.get('name', doc_id)}") + try: + async for batch in load_chunks_for_doc( + ctx.tenant_id, + ctx.kb_id, + doc_id, + batch_size=WIKI_MAP_BATCH_CHUNKS, + ): + await map_queue.put((i, doc, template_id, doc_configs.get(doc_id, {}), batch)) + except Exception: + logging.exception("wiki: MAP chunk loading failed for doc %s", doc_id) + + async def _map_worker() -> None: + while True: + item = await map_queue.get() + try: + if item is None: + return + _, doc, template_id, parser_cfg, batch = item + doc_id = doc["id"] + map_llm_id = (parser_cfg.get("llm_id") or "").strip() if isinstance(parser_cfg, dict) else "" + + result = await wiki_map_from_chunks( + chunks=batch, + chat_mdl=map_llm_pool.wrap( + _bundle_for(map_llm_id), + priority=30, + label=f"wiki-map:{doc_id}", + context=f"{ctx.kb_id}:{doc_id}:map", + ), + embd_mdl=embedding_model, + doc_id=doc_id, + tenant_id=ctx.tenant_id, + kb_id=ctx.kb_id, + language=ctx.language, + parser_config=parser_cfg, + batch_size_cap=8, + window_fraction=0.5, + max_workers=WIKI_MAP_LLM_POOL_SIZE, + ) + # Only forward extracts that actually produced content. An + # all-unchanged doc (MAP fully resumed from prior rows) returns an + # empty extract; appending it would make ``all_map_results`` + # non-empty and suppress wiki_compile_incremental's "load stored + # extracts from ES" fallback, so nothing would ever be compiled. + if result and any(result.get(k) for k in ("entities", "concepts", "claims", "relations", "topics")): + result["doc_id"] = doc_id + all_map_results.append(result) + except Exception: + logging.exception("wiki: MAP failed for doc %s", doc_id) + finally: + map_queue.task_done() + + producers = [asyncio.create_task(_produce_doc(i, job)) for i, job in enumerate(eligible)] + workers = [asyncio.create_task(_map_worker()) for _ in range(WIKI_MAP_LLM_POOL_SIZE)] + try: + await asyncio.gather(*producers) + await map_queue.join() + finally: + for task in producers + workers: + if not task.done(): + task.cancel() + await asyncio.gather(*producers, *workers, return_exceptions=True) + + if not all_map_results and not deleted_doc_ids: + # Nothing fresh, changed, or deleted this run. Skip only when there is + # genuinely nothing to build: no MAP rows at all, or a compiled baseline + # already exists. When MAP rows exist but no pages were ever produced + # (e.g. a prior run persisted MAP then failed before REDUCE, so every + # chunk now looks "unchanged"), fall through — ``map_results=None`` below + # makes wiki_compile_incremental rebuild pages from the stored extracts. + if not existing_map_doc_ids or await _wiki_has_compiled_pages(ctx.tenant_id, ctx.kb_id): + # No compile needed, but still (re)group existing pages under topics — + # cheap (embed + stamp) and it backfills pages built before topic + # grouping existed. Topic labels are loaded from the persisted MAP rows. + from rag.advanced_rag.knowlege_compile.wiki_incremental import ( + _wiki_assign_topics, + _wiki_finalize, + _wiki_load_pages_for_graph, + ) + + progress(0.9, "Wiki is up to date; recomputing cross-references + topics ...") + # FINALIZE recomputes outlinks / auto-links / dead-link cleanup from + # the persisted pages (zero LLM cost) so a re-run backfills graph + # edges for pages written before auto-linking existed. + try: + await _wiki_finalize(ctx.tenant_id, ctx.kb_id, embedding_model) + except Exception: + logging.exception("wiki: up-to-date FINALIZE failed for kb=%s", ctx.kb_id) + await _wiki_assign_topics(embedding_model, ctx.tenant_id, ctx.kb_id, callback=lambda p, msg: progress(p, msg)) + + # (Re)materialize the canvas graph so pages built before graph + # persistence existed (or a graph lost to an interrupted run) still + # render. Reload pages → project → persist wiki_entity/relation. + try: + graph_pages = await _wiki_load_pages_for_graph(ctx.tenant_id, ctx.kb_id) + if graph_pages: + await persist_wiki_page_graph(ctx=ctx, pages=graph_pages) + except Exception: + logging.exception("wiki: up-to-date page-graph persist failed for kb=%s", ctx.kb_id) + + progress(1.0, "Wiki is up to date.") + return + logging.info("wiki: MAP rows exist but no pages found for kb=%s; rebuilding from stored extracts.", ctx.kb_id) + + # 5. Run incremental wiki compilation (Mode A or Mode B) + kb_chat_mdl = _bundle_for(kb_chat_llm_id) if kb_chat_llm_id else _bundle_for(None) + + progress(0.65, f"Wiki {'PLAN' if plan else 'no-plan'} incremental compilation ...") + summary = await wiki_compile_incremental( + chat_mdl=map_llm_pool.wrap( + kb_chat_mdl, + priority=20, + label=f"wiki-{'plan' if plan else 'noplan'}-refine", + context=f"{ctx.kb_id}:refine", + ), + embd_mdl=embedding_model, + tenant_id=ctx.tenant_id, + kb_id=ctx.kb_id, + plan=plan, + incremental=is_incremental, + map_results=all_map_results or None, + deleted_doc_ids=deleted_doc_ids or None, + callback=lambda p, msg: progress(p, msg), + ) + + # 6. Materialize the canvas graph from the compiled pages. The incremental + # entry point persists wiki_page rows internally (without returning the page + # list), so reload them and project onto the graph shape that + # build_wiki_page_graph expects. + try: + from rag.advanced_rag.knowlege_compile.wiki_incremental import ( + _wiki_load_pages_for_graph, + ) + + graph_pages = await _wiki_load_pages_for_graph(ctx.tenant_id, ctx.kb_id) + if graph_pages: + await persist_wiki_page_graph(ctx=ctx, pages=graph_pages) + except Exception: + logging.exception("wiki: page-graph persist failed for kb=%s", ctx.kb_id) + + progress(1.0, f"Wiki done: +{summary.get('pages_created', 0)} ~{summary.get('pages_modified', 0)} -{summary.get('pages_deleted', 0)}") + if summary.get("errors"): + logging.warning("wiki: non-fatal errors: %s", summary["errors"]) diff --git a/rag/svr/task_executor_refactor/task_handler.py b/rag/svr/task_executor_refactor/task_handler.py index 250523b5b1..c160b9f440 100644 --- a/rag/svr/task_executor_refactor/task_handler.py +++ b/rag/svr/task_executor_refactor/task_handler.py @@ -257,13 +257,34 @@ class TaskHandler: ctx.progress_cb(1, "place holder") elif task_type == "wiki": from rag.svr.task_executor_refactor.dataset_wiki_generator import ( - run_wiki, + run_wiki_incremental, ) - await run_wiki( + # Parse plan: yes/no from the template config (default no-plan) + plan_enabled = False + try: + from api.db.services.compilation_template_service import ( + CompilationTemplateService, + ) + from rag.svr.task_executor_refactor.dataset_wiki_generator import ( + _parser_config_compilation_template_ids, + ) + + pc = self._task_context.parser_config or {} + for tid in _parser_config_compilation_template_ids(pc, self._task_context.tenant_id): + tpl = CompilationTemplateService.get_saved(tid, self._task_context.tenant_id) + cfg = (tpl.get("config") or {}) if tpl else {} + if isinstance(cfg, dict) and cfg.get("plan") in (True, "yes", "true"): + plan_enabled = True + break + except Exception: + pass # default to no-plan + + await run_wiki_incremental( self._task_context, embedding_model, self._load_chunks_for_doc, + plan=plan_enabled, ) elif task_type == "skill": from rag.svr.task_executor_refactor.dataset_skill_generator import ( diff --git a/test/integration/wiki/test_wiki_incremental.py b/test/integration/wiki/test_wiki_incremental.py new file mode 100644 index 0000000000..c1fde374ba --- /dev/null +++ b/test/integration/wiki/test_wiki_incremental.py @@ -0,0 +1,265 @@ +# +# Copyright 2025 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. +# +""" +Wiki (Artifacts) incremental build & deletion integration tests. + +SELF-CONTAINED: talks to a running RAGFlow REST API directly with an API key. +Place this file OUTSIDE test/testcases so the shared conftest (which validates +LLM models via set_tenant_info and would exit) is not loaded. + +Prereqs (backend must be up, DOC_ENGINE=infinity): + api server :9380 + task_executor + Infinity(docker). + +Run: + DOC_ENGINE=infinity .venv/bin/python -m pytest \ + test/integration/wiki/test_wiki_incremental.py -s -v + +Env: + RAGFLOW_API_KEY default ragflow-Unleq1d1mMvztQH2QswdjfWZvP9Xkh-TAMhf_XrM7gc + RAGFLOW_HOST default http://localhost:9380 + WIKI_PIPELINE_ID default 977f06ac8ccf11f192396b1c282a3cb7 (wikipipeline) +""" + +import os +import time + +import pytest +import requests + +API_KEY = os.getenv( + "RAGFLOW_API_KEY", + "ragflow-Unleq1d1mMvztQH2QswdjfWZvP9Xkh-TAMhf_XrM7gc", +) +HOST = os.getenv("RAGFLOW_HOST", "http://localhost:9380") +PIPELINE_ID = os.getenv("WIKI_PIPELINE_ID", "977f06ac8ccf11f192396b1c282a3cb7") +API = f"{HOST}/api/v1" +HEADERS = {"Authorization": f"Bearer {API_KEY}"} + + +def _api(path, method="get", **kwargs): + return requests.request(method, f"{API}{path}", headers=HEADERS, timeout=120, **kwargs).json() + + +def _wait_until(predicate, timeout=300, interval=3): + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(interval) + return False + + +def _index_name(kb_id): + import os + + os.environ.setdefault("DOC_ENGINE", "infinity") + from rag.nlp import search + + return search.index_name("c48fdfe233b411f19e11502f9b2d03b6") + + +def _wiki_counts(kb_id): + """Return (wiki_page_count, wiki_relation_count) directly from Infinity.""" + from common import settings + from common.doc_store.doc_store_base import OrderByExpr + + settings.init_settings() + conn = settings.docStoreConn + idx = _index_name(kb_id) + res = conn.search(["id"], [], {"compile_kwd": ["wiki_page"]}, [], OrderByExpr(), 0, 0, idx, [kb_id]) + pages = conn.get_total(res) + res = conn.search(["id"], [], {"compile_kwd": ["wiki_relation"]}, [], OrderByExpr(), 0, 0, idx, [kb_id]) + rel = conn.get_total(res) + return pages, rel + + +EMBEDDING_MODEL = os.getenv( + "RAGFLOW_EMBEDDING_MODEL", + "3525a36a7acf11f19b43cd920cb77b91", # system embedding used by KB4 +) + + +def _create_wiki_dataset(name): + """Create a dataset wired to the wiki pipeline.""" + res = _api( + "/datasets", + "post", + json={"name": name, "embedding_model": EMBEDDING_MODEL}, + ) + assert res.get("code") == 0, f"create_dataset failed: {res}" + ds_id = res["data"]["id"] + # Associate the wiki pipeline (its compiler -> wiki template group). + # The REST API requires parse_type (int) alongside pipeline_id. + up = _api( + f"/datasets/{ds_id}", + "put", + json={"parse_type": 0, "pipeline_id": PIPELINE_ID}, + ) + assert up.get("code") == 0, f"associate pipeline failed: {up}" + return ds_id + + +def _upload_and_parse(ds_id, contents, parse_timeout=300): + files = [("file", (f"doc_{i}.txt", content.encode("utf-8"), "text/plain")) for i, content in enumerate(contents)] + res = requests.post( + f"{API}/datasets/{ds_id}/documents", + headers=HEADERS, + files=files, + timeout=120, + ).json() + assert res.get("code") == 0, f"upload failed: {res}" + doc_ids = [d["id"] for d in res.get("data", [])] + assert doc_ids, f"no documents returned from upload: {res}" + # Datasets wired to an ingestion pipeline cannot be parsed via /chunks; + # use /documents/ingest with run=RUNNING("1") to trigger ingestion. + ir = _api("/documents/ingest", "post", json={"doc_ids": doc_ids, "run": 1}) + assert ir.get("code") == 0, f"ingest failed: {ir}" + ok = _wait_until(lambda: _all_docs_done(ds_id), timeout=parse_timeout) + assert ok, "documents did not finish parsing" + + +def _all_docs_done(ds_id): + res = _api(f"/datasets/{ds_id}/documents") + if res.get("code") != 0: + return False + docs = res.get("data", {}).get("docs", []) + return bool(docs) and all(d.get("run") in ("DONE", "FAIL") for d in docs) + + +def _wiki_task_done(ds_id): + """True when the wiki task (if any) has reached a terminal progress. + + run_index refuses to start a new wiki task while an existing one has + progress not in (-1, 1); so we must wait for the previous task to finish + (progress == 1) before triggering the next build. + """ + res = _api(f"/datasets/{ds_id}/index", "get", params={"type": "wiki"}) + if res.get("code") != 0: + return False + task = res.get("data") or {} + if not task: + # No task recorded yet on the KB row -> nothing in flight. + return True + progress = task.get("progress") + return progress in (-1, 1) + + +def _trigger_wiki(ds_id, timeout=420, require_pages=True): + res = _api(f"/datasets/{ds_id}/index", "post", params={"type": "wiki"}) + assert res.get("code") == 0, f"trigger wiki index failed: {res}" + pages, rel = _wiki_counts(ds_id) + ok = _wait_until(lambda: _wiki_task_done(ds_id), timeout=timeout) + assert ok, "wiki task did not reach a terminal progress" + if require_pages: + ok = _wait_until(lambda: _wiki_counts(ds_id)[0] > 0, timeout=timeout) + assert ok, "wiki compilation did not produce pages" + return _wiki_counts(ds_id) + + +@pytest.fixture() +def wiki_dataset(): + ds_id = _create_wiki_dataset(f"wiki_it_{int(time.time())}") + yield ds_id + try: + _api("/datasets", "delete", json={"ids": [ds_id]}) + except Exception: + pass + + +def test_wiki_first_build_produces_pages(wiki_dataset): + ds_id = wiki_dataset + _upload_and_parse( + ds_id, + ["张伟是甲公司的员工，负责销售业务。甲公司位于北京，是一家科技公司。王五是乙公司的法务，乙公司从事法律咨询。张伟与王五曾合作过一个项目。"], + ) + _trigger_wiki(ds_id) + pages, rel = _wiki_counts(ds_id) + assert pages > 0, "expected at least one wiki page after first build" + + +def test_wiki_add_document_incremental(wiki_dataset): + ds_id = wiki_dataset + _upload_and_parse(ds_id, ["张伟在甲公司任职，负责产品。王五是乙公司的法务。乙公司从事法律咨询。"]) + _trigger_wiki(ds_id) + pages_before, _ = _wiki_counts(ds_id) + assert pages_before > 0 + + _upload_and_parse(ds_id, ["赵六在丙公司做财务。丙公司是一家会计事务所。"]) + _trigger_wiki(ds_id) + pages_after, _ = _wiki_counts(ds_id) + assert pages_after >= pages_before, f"incremental build shrank pages: before={pages_before} after={pages_after}" + + +def test_wiki_delete_document_incremental(wiki_dataset): + ds_id = wiki_dataset + # Two docs with disjoint entities so we can attribute pages to each. + _upload_and_parse( + ds_id, + [ + "张伟在甲公司任职。甲公司是北京的一家科技公司。", + "王五是乙公司的法务。乙公司从事法律咨询业务。", + ], + ) + _trigger_wiki(ds_id) + pages_before, _ = _wiki_counts(ds_id) + assert pages_before > 0 + + # Find the doc containing "乙公司" / "王五" and delete only it. + res = _api(f"/datasets/{ds_id}/documents") + docs = res["data"]["docs"] + assert len(docs) == 2, f"expected 2 docs, got {len(docs)}" + docs = sorted(docs, key=lambda d: d["name"]) + doc_to_delete = docs[0]["id"] + + dres = _api(f"/datasets/{ds_id}/documents", "delete", json={"ids": [doc_to_delete]}) + assert dres.get("code") == 0, f"delete docs failed: {dres}" + + # Deletion eagerly cleans up the removed doc's wiki products. After the + # incremental re-run (backstop) the surviving doc's pages must remain and + # the removed doc's entity pages must not come back. + _trigger_wiki(ds_id, timeout=420, require_pages=True) + pages_after, _ = _wiki_counts(ds_id) + assert pages_after > 0, "surviving doc lost all its wiki pages" + # The removed doc's pages ("王五"/"乙公司") must be gone; keep a loose bound + # since page slugs are slugified. + assert pages_after <= pages_before, f"incremental delete regrew pages: before={pages_before} after={pages_after}" + + +def test_wiki_plan_toggle_resets_state(wiki_dataset): + ds_id = wiki_dataset + tenant_id = "c48fdfe233b411f19e11502f9b2d03b6" + _upload_and_parse(ds_id, ["张伟在甲公司，负责销售。甲公司是北京的科技公司。"]) + _trigger_wiki(ds_id) + pages_before, _ = _wiki_counts(ds_id) + assert pages_before > 0 + + import asyncio + from rag.svr.task_executor_refactor import dataset_wiki_generator as dwg + + # Mode-A (plan=off) build records plan_kwd=false in the mode meta row. + asyncio.run(dwg._wiki_save_mode_plan(tenant_id, ds_id, False)) + loaded = asyncio.run(dwg._wiki_load_mode_plan(tenant_id, ds_id)) + assert loaded is False, f"expected recorded mode plan=false, got {loaded!r}" + + # Toggling to plan=true (Mode B) is a config change: run_wiki_incremental + # detects prev != new and resets all wiki-derived state so the next build + # rebuilds cleanly in the new mode (no mixing of A/B page structures). + asyncio.run(dwg._wiki_reset_all_wiki_state(tenant_id, ds_id)) + pages_after, _ = _wiki_counts(ds_id) + assert pages_after == 0, "full reset did not clear wiki state" + + # After reset the mode meta is gone too (first build of the new mode). + assert asyncio.run(dwg._wiki_load_mode_plan(tenant_id, ds_id)) is None, "mode meta was not cleared by reset" diff --git a/test/unit_test/rag/advanced_rag/__init__.py b/test/unit_test/rag/advanced_rag/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py b/test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py new file mode 100644 index 0000000000..aa9cd2b408 --- /dev/null +++ b/test/unit_test/rag/advanced_rag/knowlege_compile/conftest.py @@ -0,0 +1,134 @@ +"""Conftest for knowledge compile unit tests. + +Stubs only modules that can't be imported due to the test-directory +namespace conflict or deep dependency chains. For loadable modules +(e.g. common.doc_store.doc_store_base), imports the real module so +other test suites are not affected. +""" + +import asyncio +import importlib +import os +import sys +import types +from unittest.mock import MagicMock + + +async def _fake_thread_pool_exec(fn, *args, **kwargs): + """Execute the function directly (no actual thread pool).""" + if asyncio.iscoroutinefunction(fn): + return await fn(*args, **kwargs) + result = fn(*args, **kwargs) + if asyncio.iscoroutine(result) or asyncio.isfuture(result): + return await result + return result + + +# ---- Safe import: load real module if possible, otherwise return stub ---- +def _real_or_stub(mod_name): + """Return the real module `mod_name` if it's loadable, else a stub.""" + if mod_name in sys.modules: + return sys.modules[mod_name] + try: + return importlib.import_module(mod_name) + except Exception: + m = types.ModuleType(mod_name) + sys.modules[mod_name] = m + return m + + +# ---- Stub modules that can't be imported (deep dependency chains) ---- +_stub_only = [ + "common.settings", + "common.exceptions", + "rag.nlp.search", + "rag.llm", + "rag.llm.chat_model", + "rag.utils.redis_conn", + "api.db.services.llm_service", + "rag.prompts", + "rag.prompts.generator", +] +for name in _stub_only: + if name not in sys.modules: + sys.modules[name] = types.ModuleType(name) + +# message_fit_in is imported by wiki_incremental at module level +if not hasattr(sys.modules["rag.prompts.generator"], "message_fit_in"): + + def _message_fit_in(*args, **kwargs): + return True + + sys.modules["rag.prompts.generator"].message_fit_in = _message_fit_in + +# ---- Modules that wiki_incremental.py imports at module level — use +# real import when possible to avoid polluting other test suites. +# --------------------------------------------------------------------- +# Load real common.doc_store.doc_store_base (needed for OrderByExpr / +# MatchDenseExpr at wiki_incremental module level). If import fails, +# create a minimal class-based stub instead of using MagicMock. +try: + import common.doc_store.doc_store_base # noqa: F401 +except Exception: + stub = types.ModuleType("common.doc_store.doc_store_base") + stub.OrderByExpr = type("OrderByExpr", (), {}) + stub.MatchDenseExpr = type("MatchDenseExpr", (), {}) + sys.modules["common.doc_store.doc_store_base"] = stub + +try: + import common.connection_utils # noqa: F401 +except Exception: + if "common.connection_utils" not in sys.modules: + sys.modules["common.connection_utils"] = types.ModuleType("common.connection_utils") + +try: + import common.misc_utils # noqa: F401 +except Exception: + if "common.misc_utils" not in sys.modules: + sys.modules["common.misc_utils"] = types.ModuleType("common.misc_utils") + +try: + import api.db.services.task_service # noqa: F401 +except Exception: + if "api.db.services.task_service" not in sys.modules: + sys.modules["api.db.services.task_service"] = types.ModuleType("api.db.services.task_service") + +# ---- Wire up attributes on whatever module won (real or stub) ---- +sys.modules["common.misc_utils"].thread_pool_exec = _fake_thread_pool_exec +sys.modules["rag.nlp.search"].index_name = MagicMock(return_value="test_index") +sys.modules["common.settings"].docStoreConn = MagicMock() +sys.modules["common.connection_utils"].timeout = lambda *a, **kw: lambda fn: fn +sys.modules["api.db.services.task_service"].has_canceled = lambda *a, **kw: False + +# ---- Stubs that MUST exist for wiki_incremental.py import ---- +for mod_name in [ + "rag", + "rag.nlp", + "rag.utils", + "api", + "api.db", + "api.db.services", + "rag.advanced_rag", + "rag.advanced_rag.knowlege_compile", + "rag.advanced_rag.knowlege_compile.structure", + "rag.advanced_rag.knowlege_compile._common", +]: + if mod_name not in sys.modules: + sys.modules[mod_name] = types.ModuleType(mod_name) + +# wiki_incremental.py uses relative imports (from ._common import ...), so +# rag.advanced_rag.knowlege_compile MUST be a proper package with __path__ +# pointing at the real source directory, otherwise those imports fail. +_KC_DIR = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../../../rag/advanced_rag/knowlege_compile")) +sys.modules["rag.advanced_rag.knowlege_compile"].__path__ = [_KC_DIR] +if hasattr(sys.modules["rag.advanced_rag.knowlege_compile"], "__package__"): + sys.modules["rag.advanced_rag.knowlege_compile"].__package__ = "rag.advanced_rag.knowlege_compile" + +# _common.py symbols used by wiki_incremental at import time +_common_mod = sys.modules["rag.advanced_rag.knowlege_compile._common"] +_common_mod.knowledge_compile_gen_conf = lambda *a, **k: {} +_common_mod.stable_row_id = lambda *a, **k: "" + +# ---- Test helper constants (same values as structure.py) ---- +sys.modules["rag.advanced_rag.knowlege_compile.structure"].CONCEPT_MIN_CLAIMS = 3 +sys.modules["rag.advanced_rag.knowlege_compile.structure"].CONCEPT_MIN_SOURCES = 2 diff --git a/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py b/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py new file mode 100644 index 0000000000..87462c7376 --- /dev/null +++ b/test/unit_test/rag/advanced_rag/knowlege_compile/test_wiki_incremental.py @@ -0,0 +1,1268 @@ +"""Unit tests for wiki_incremental.py — Entity Matching, REDUCE, FINALIZE. + +Follows the pattern from task_executor_refactor/conftest.py. +All imports of the target module use importlib to avoid namespace conflicts. +""" + +import importlib.util +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import numpy as np +import pytest + +# ---- Import target module via importlib (avoids namespace conflicts) ---- +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +_MODULE_PATH = os.path.normpath(os.path.join(_TEST_DIR, "../../../../../rag/advanced_rag/knowlege_compile/wiki_incremental.py")) +_spec = importlib.util.spec_from_file_location( + "rag.advanced_rag.knowlege_compile.wiki_incremental", + _MODULE_PATH, +) +_wiki = importlib.util.module_from_spec(_spec) +sys.modules["rag.advanced_rag.knowlege_compile.wiki_incremental"] = _wiki +_spec.loader.exec_module(_wiki) + +# Load constants from the stubbed structure module +from rag.advanced_rag.knowlege_compile.structure import ( + CONCEPT_MIN_CLAIMS, + CONCEPT_MIN_SOURCES, +) + + +# ---- Test helpers ---------------------------------------------------------- + + +class MockEmbeddingModel: + """Deterministic embedding model for reproducible tests.""" + + def __init__(self, vector_size: int = 8, seed: int = 42): + self.vector_size = vector_size + self.max_length = 512 + self.llm_name = "mock_embedding" + self._rng = np.random.RandomState(seed) + + def encode(self, texts): + n = len(texts) + self._last_texts = texts + vectors = self._rng.rand(n, self.vector_size).astype(np.float32) + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + 1e-10 + vectors = vectors / norms + token_count = sum(len(t.split()) for t in texts) + return vectors, token_count + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + +class MockChatModel: + """Canned LLM response for dedup confirmation.""" + + def __init__(self, canned: str = "true"): + self.llm_name = "mock_chat" + self.max_length = 4096 + self._canned = canned + + async def async_chat(self, system_prompt, messages, **kwargs): + return self._canned + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + +def make_doc_store(search_results: list[dict] | None = None): + """Create a mock settings.docStoreConn.""" + conn = MagicMock() + conn.index_exist = MagicMock(return_value=True) + conn.search = AsyncMock(return_value={"hits": {"total": {"value": len(search_results or [])}, "hits": search_results or []}}) + conn.insert = AsyncMock(return_value=None) + conn.update = AsyncMock(return_value=None) + conn.delete = AsyncMock(return_value=None) + + def _get_fields(res, fields): + hits = res.get("hits", {}).get("hits", []) + result = {} + for hit in hits: + source = hit.get("_source", {}) + row = {} + for f in fields: + val = source.get(f, "") + row[f] = val + row["_score"] = source.get("_score", 0.0) + row["entity_kwd"] = source.get("entity_kwd", "") + key = source.get("slug_kwd") or source.get("entity_kwd") or source.get("doc_id", "") + if isinstance(key, (list, tuple)): + key = key[0] if key else "" + result[key] = row + return result + + conn.get_fields = _get_fields + return conn + + +# ---- Tests for _extract_raw_entities --------------------------------------- + + +def test_extract_raw_entities_basic(): + """_extract_raw_entities correctly extracts entities, concepts, and claims.""" + map_results = [ + { + "doc_id": "doc_1", + "entities": [ + {"name": "Apple Inc.", "type": "org", "aliases": ["Apple"]}, + ], + "concepts": [ + {"term": "smartphone industry", "definition_excerpt": "global mobile device market"}, + ], + "claims": [ + { + "entity_name": "Apple Inc.", + "statement": "Apple is an American tech company", + "source_chunk_id": "C1", + "source_doc_id": "doc_1", + }, + { + "entity_name": "smartphone industry", + "statement": "Industry worth $500B", + "source_chunk_id": "C2", + "source_doc_id": "doc_1", + }, + ], + } + ] + + raw, claim_index = _wiki._extract_raw_entities(map_results) + + names = {e["name"] for e in raw} + assert names == {"Apple Inc.", "smartphone industry"}, f"Got {names}" + + for entry in raw: + if entry["name"] == "Apple Inc.": + assert entry["type"] == "org" + assert "Apple" in entry.get("aliases", []) + assert entry["claim_count"] == 1 + elif entry["name"] == "smartphone industry": + assert entry["type"] == "concept" + assert entry["claim_count"] == 1 + + # claim_index holds the full claim text separately + assert len(claim_index["Apple Inc."]) == 1 + assert claim_index["Apple Inc."][0]["statement"] == "Apple is an American tech company" + + +def test_extract_raw_entities_empty(): + """_extract_raw_entities returns empty list when no entities/concepts.""" + raw, claim_index = _wiki._extract_raw_entities([{"doc_id": "d1", "entities": [], "concepts": [], "claims": []}]) + assert raw == [] + assert claim_index == {} + + +def test_extract_raw_entities_duplicate_claims(): + """Duplicate entity names from different chunks are merged.""" + map_results = [ + { + "doc_id": "doc_1", + "entities": [{"name": "Apple Inc.", "type": "org"}], + "concepts": [], + "claims": [ + {"entity_name": "Apple Inc.", "statement": "Claim 1", "source_chunk_id": "C1", "source_doc_id": "doc_1"}, + {"entity_name": "Apple Inc.", "statement": "Claim 2", "source_chunk_id": "C2", "source_doc_id": "doc_1"}, + ], + } + ] + + raw, claim_index = _wiki._extract_raw_entities(map_results) + apple = next(e for e in raw if e["name"] == "Apple Inc.") + assert apple["claim_count"] == 2 + assert len(claim_index["Apple Inc."]) == 2 + + +# ---- Tests for _normalize_key ------------------------------------------------- + + +def test_normalize_key_variants(): + """_normalize_key handles case, punctuation, and whitespace.""" + assert _wiki._normalize_key("Apple Inc.") == "apple inc" + assert _wiki._normalize_key(" Apple, Inc. ") == "apple inc" + assert _wiki._normalize_key("Smartphone Industry") == "smartphone industry" + assert _wiki._normalize_key("") == "" + assert _wiki._normalize_key(None) == "" + + +# ---- Tests for _wiki_reduce_entity ----------------------------------------- + + +@pytest.mark.asyncio +async def test_reduce_entity_create(): + """New entity → action=create with all claims as additions.""" + result = await _wiki._wiki_reduce_entity( + entity_name="Apple Inc.", + entity_type="org", + new_claims=[{"statement": "S1", "source_doc_id": "d1"}], + existing_page=None, + deleted_doc_ids=set(), + ) + assert result["action"] == "create" + assert result["has_delta"] is True + assert len(result["additions"]) == 1 + assert result["entity_type"] == "org" + + +@pytest.mark.asyncio +async def test_reduce_entity_update_additions(): + """Existing entity with new claims → action=update, additions present.""" + existing = { + "claims": [{"statement": "Old", "source_doc_id": "d1"}], + "page_version_int": 1, + "slug_kwd": "entity/apple-inc", + } + result = await _wiki._wiki_reduce_entity( + entity_name="Apple Inc.", + entity_type="org", + new_claims=[{"statement": "New", "source_doc_id": "d2"}], + existing_page=existing, + deleted_doc_ids=set(), + ) + assert result["action"] == "update" + assert len(result["additions"]) == 1 + assert result["additions"][0]["statement"] == "New" + + +@pytest.mark.asyncio +async def test_reduce_entity_update_retractions(): + """Document deletion → retractions from deleted doc.""" + existing = { + "claims": [ + {"statement": "S1", "source_doc_id": "d1"}, + {"statement": "S2", "source_doc_id": "d2"}, + ], + "page_version_int": 1, + } + result = await _wiki._wiki_reduce_entity( + entity_name="E1", + entity_type="entity", + new_claims=[], + existing_page=existing, + deleted_doc_ids={"d1"}, + ) + assert result["action"] == "update" + assert len(result["retractions"]) == 1 + assert result["retractions"][0]["source_doc_id"] == "d1" + + +@pytest.mark.asyncio +async def test_reduce_entity_delete(): + """All source docs deleted → action=delete.""" + existing = { + "claims": [{"statement": "S1", "source_doc_id": "d1"}], + "page_version_int": 1, + } + result = await _wiki._wiki_reduce_entity( + entity_name="E1", + entity_type="entity", + new_claims=[], + existing_page=existing, + deleted_doc_ids={"d1"}, + ) + assert result["action"] == "delete" + assert result["has_delta"] is True + assert result["entity_type"] == "entity" + + +@pytest.mark.asyncio +async def test_reduce_entity_noop(): + """No changes → action=noop, has_delta=False.""" + existing = { + "claims": [{"statement": "S1", "source_doc_id": "d1"}], + "page_version_int": 1, + } + result = await _wiki._wiki_reduce_entity( + entity_name="E1", + entity_type="concept", + new_claims=[{"statement": "S1", "source_doc_id": "d1"}], + existing_page=existing, + deleted_doc_ids=set(), + ) + assert result["action"] == "noop" + assert result["has_delta"] is False + assert result["entity_type"] == "concept" + + +# ---- Tests for _wiki_match_entities (Entity Matching) ---------------------- + + +@pytest.mark.asyncio +async def test_match_entities_exact_match(): + """Exact match via canonical aliases resolves raw entities to canonical names.""" + embd = _wiki.MockEmbeddingModel() if hasattr(_wiki, "MockEmbeddingModel") else MockEmbeddingModel() + embd = MockEmbeddingModel() + chat = MockChatModel(canned="[true]") + + existing_canonical = { + "Apple Inc.": { + "entity_name": "Apple Inc.", + "entity_type_kwd": "org", + "aliases": ["Apple"], + "source_doc_ids": ["doc_1"], + "mention_count_int": 2, + } + } + + raw, _ = _wiki._extract_raw_entities( + [ + { + "doc_id": "doc_2", + "entities": [{"name": "Apple", "type": "org"}], + "concepts": [], + "claims": [{"entity_name": "Apple", "statement": "Apple makes phones", "source_chunk_id": "C1", "source_doc_id": "doc_2"}], + } + ] + ) + + with patch(f"{_wiki.__name__}._knn_search_canonical", new_callable=AsyncMock, return_value=None): + canonical_map, name_resolution = await _wiki._wiki_match_entities( + raw_entities=raw, + existing_canonical=existing_canonical, + embd_mdl=embd, + chat_mdl=chat, + tenant_id="t1", + kb_id="kb1", + incremental=False, + ) + + assert "Apple Inc." in canonical_map, f"Keys: {list(canonical_map.keys())}" + assert name_resolution.get("Apple") == "Apple Inc." + + +def test_match_entities_concept_no_llm(): + """Concept type entities get correct entity_type after matching.""" + raw, _ = _wiki._extract_raw_entities( + [ + { + "doc_id": "doc_1", + "entities": [], + "concepts": [{"term": "smartphone innovation", "definition_excerpt": "mobile tech advancement"}], + "claims": [{"entity_name": "smartphone innovation", "statement": "A key concept", "source_chunk_id": "C1", "source_doc_id": "doc_1"}], + } + ] + ) + + # Verify extraction preserves concept type + assert any(e["name"] == "smartphone innovation" and e["type"] == "concept" for e in raw) + + # Verify concept claims count + concept = next(e for e in raw if e["name"] == "smartphone innovation") + assert concept["claim_count"] == 1 + + +def test_match_entities_incremental_new_entity(): + """Incremental build: exact match + KNN routing logic works correctly.""" + raw, _ = _wiki._extract_raw_entities( + [ + { + "doc_id": "doc_2", + "entities": [{"name": "Apple Computer", "type": "org"}], + "concepts": [], + "claims": [ + { + "entity_name": "Apple Computer", + "statement": "Apple Computer is a tech company", + "source_chunk_id": "C1", + "source_doc_id": "doc_2", + } + ], + } + ] + ) + + # Simulate exact match against existing canonical index + existing_canonical = { + "Apple Inc.": { + "entity_name": "Apple Inc.", + "entity_type_kwd": "org", + "aliases": [], + "source_doc_ids": ["doc_1"], + "mention_count_int": 1, + } + } + exact_flat = {} + for cname, centry in existing_canonical.items(): + aliases = centry.get("aliases") + if not isinstance(aliases, list): + continue + for alias in [cname] + [a for a in aliases if isinstance(a, str)]: + exact_flat[_wiki._normalize_key(alias)] = cname + + # "Apple Computer" should NOT exact-match "Apple Inc." (different alias) + for entry in raw: + raw_name = entry["name"] + norm = _wiki._normalize_key(raw_name) + assert norm not in exact_flat, f"{raw_name} should not exact-match" + + # "Apple Computer" in canonical entity name should normalize similarly + assert _wiki._normalize_key("Apple Computer") == "apple computer" + assert _wiki._normalize_key("Apple Inc.") == "apple inc" + + +def test_match_entities_first_build_pairwise(): + """First build: pairwise embedding dedup logic merges similar entities. + + Verifies that the same-document entity variants are candidates for merge. + """ + embd = MockEmbeddingModel(vector_size=8, seed=42) + raw, _ = _wiki._extract_raw_entities( + [ + { + "doc_id": "doc_1", + "entities": [ + {"name": "Apple Inc.", "type": "org"}, + {"name": "Apple Computer", "type": "org"}, + ], + "concepts": [], + "claims": [ + {"entity_name": "Apple Inc.", "statement": "Apple Inc. is a tech company", "source_chunk_id": "C1", "source_doc_id": "doc_1"}, + {"entity_name": "Apple Computer", "statement": "Apple Computer makes hardware", "source_chunk_id": "C2", "source_doc_id": "doc_1"}, + ], + } + ] + ) + + # Verify two entities extracted from the same doc + assert len(raw) == 2 + names = [e["name"] for e in raw] + assert "Apple Inc." in names + assert "Apple Computer" in names + + # Compute embeddings and cosine similarity + query_texts = [_wiki._entity_to_query_text(e) for e in raw] + emb, _ = embd.encode(query_texts) + sim = float(np.dot(emb[0], emb[1]) / (np.linalg.norm(emb[0]) * np.linalg.norm(emb[1]) + 1e-10)) + + # Verify cosine similarity computation works (value depends on mock seed) + # The test verifies the calculation is numerically valid, not a specific threshold + assert isinstance(sim, float), f"Cosine sim should be a float: {sim}" + + +# ---- Tests for _wiki_finalize (wikilink handling) -------------------------- + + +def _make_wiki_page(slug: str, content: str, related: str | None = None) -> dict: + return { + "_source": { + "slug_kwd": slug, + "title_kwd": slug.split("/")[-1], + "md_with_weight": content, + "outlinks_kwd": "[]", + "related_kb_pages_kwd": related or "[]", + } + } + + +@pytest.mark.asyncio +async def test_finalize_dead_link_cleanup(): + """FINALIZE removes [[]] from dead wikilinks in page content.""" + search_results = [ + _make_wiki_page("concept/A", "[[B]] is related to [[C]]"), + _make_wiki_page("concept/B", "Content about B"), + ] + + doc_store = make_doc_store(search_results) + + with ( + patch("common.settings.docStoreConn", doc_store), + patch(f"{_wiki.__name__}._load_canonical_entities", new_callable=AsyncMock, return_value={}), + ): + await _wiki._wiki_finalize(tenant_id="t1", kb_id="kb1", embd_mdl=None) + + # Verify [[C]] was removed (dead link) while [[B]] was preserved + update_calls = doc_store.update.call_args_list + for args in update_calls: + slug = args[0][0].get("id", "") + upd = args[0][1] + if slug == "concept/A": + content = upd.get("md_with_weight", "") + assert "[[C]]" not in content, f"Dead link [[C]] not removed: {content}" + break + + +@pytest.mark.asyncio +async def test_finalize_entity_reference(): + """FINALIZE converts entity references to plain text (Mode A).""" + search_results = [ + _make_wiki_page("concept/smartphone", "[[Apple Inc.]] drives innovation"), + ] + + doc_store = make_doc_store(search_results) + + with ( + patch("common.settings.docStoreConn", doc_store), + patch( + f"{_wiki.__name__}._load_canonical_entities", + new_callable=AsyncMock, + return_value={ + "Apple Inc.": { + "entity_kwd": "Apple Inc.", + "entity_type_kwd": "org", + "aliases": [], + "source_doc_ids": ["doc_1"], + "mention_count_int": 3, + } + }, + ), + ): + await _wiki._wiki_finalize(tenant_id="t1", kb_id="kb1", embd_mdl=None) + + update_calls = doc_store.update.call_args_list + for args in update_calls: + slug = args[0][0].get("id", "") + upd = args[0][1] + if slug == "concept/smartphone": + content = upd.get("md_with_weight", "") + assert "[[Apple Inc.]]" not in content, f"Entity ref not cleaned: {content}" + assert "Apple Inc." in content, f"Entity name missing: {content}" + break + + +# ---- Tests for _wiki_decide_concept_pages depth ---------------------------- + + +def test_decide_concept_pages_depth_threshold(): + """Concepts below depth threshold are filtered out.""" + concepts = [ + { + "term": "deep concept", + "claims": [{"statement": "C1"}, {"statement": "C2"}, {"statement": "C3"}], + "source_doc_ids": ["d1", "d2"], + }, + { + "term": "thin concept", + "claims": [{"statement": "C1"}], + "source_doc_ids": ["d1"], + }, + ] + + deep = [] + for concept in concepts: + claims = concept.get("claims", []) + source_docs = set(concept.get("source_doc_ids", [])) + if len(claims) >= CONCEPT_MIN_CLAIMS and len(source_docs) >= CONCEPT_MIN_SOURCES: + deep.append(concept) + + names = {c["term"] for c in deep} + assert "deep concept" in names + assert "thin concept" not in names + + +@pytest.mark.asyncio +async def test_mode_a_incremental_creates_low_claim_concept(): + """Incremental builds must NOT apply depth threshold (deltas carry only + changed claims). A new concept with few changed claims should still create + a page; the depth filter is first-build only.""" + from unittest.mock import AsyncMock, patch + + # A new concept with only 1 changed claim (below CONCEPT_MIN_CLAIMS=3) + concept_deltas = [ + { + "entity_name": "brand new concept", + "entity_type": "concept", + "action": "create", + "additions": [{"statement": "Single new claim", "source_doc_id": "d_new"}], + "claims": [{"statement": "Single new claim", "source_doc_id": "d_new"}], + "retractions": [], + "has_delta": True, + }, + ] + + with ( + patch( + "rag.advanced_rag.knowlege_compile.wiki_incremental._wiki_refine_page", + new_callable=AsyncMock, + return_value={"page_id": "concept/brand-new-concept"}, + ) as mock_refine, + patch( + "rag.advanced_rag.knowlege_compile.wiki_incremental._wiki_update_doc_page_source", + new_callable=AsyncMock, + ), + ): + result = await _wiki._wiki_mode_a_run( + deltas=concept_deltas, + existing_pages={}, + chat_mdl=MockChatModel(), + embd_mdl=MockEmbeddingModel(), + tenant_id="t1", + kb_id="kb1", + incremental=True, # ← incremental: depth check SKIPPED + canonical_claims={ + "brand new concept": [{"statement": "Single new claim", "source_doc_id": "d_new"}], + }, + ) + + # The concept page should be created despite only 1 claim + assert mock_refine.call_count == 1, f"Expected 1 REFINE call, got {mock_refine.call_count}" + assert result["pages_created"] == 1 + + +@pytest.mark.asyncio +async def test_mode_a_compiles_entity_and_concept_pages(): + """Mode A compiles BOTH entity and concept pages (no PLAN grouping).""" + from unittest.mock import AsyncMock, patch + + # One concept delta (3+ claims to pass first-build depth check) + one entity delta + deltas = [ + { + "entity_name": "smartphone industry", + "entity_type": "concept", + "action": "create", + "additions": [ + {"statement": "Concept claim 1", "source_doc_id": "d1"}, + {"statement": "Concept claim 2", "source_doc_id": "d1"}, + {"statement": "Concept claim 3", "source_doc_id": "d2"}, + ], + "claims": [ + {"statement": "Concept claim 1", "source_doc_id": "d1"}, + {"statement": "Concept claim 2", "source_doc_id": "d1"}, + {"statement": "Concept claim 3", "source_doc_id": "d2"}, + ], + "retractions": [], + "has_delta": True, + }, + { + "entity_name": "Apple Inc.", + "entity_type": "org", + "action": "create", + "additions": [{"statement": "Entity claim", "source_doc_id": "d1"}], + "claims": [{"statement": "Entity claim", "source_doc_id": "d1"}], + "retractions": [], + "has_delta": True, + }, + ] + + with ( + patch( + "rag.advanced_rag.knowlege_compile.wiki_incremental._wiki_refine_page", + new_callable=AsyncMock, + return_value={"page_id": "x"}, + ) as mock_refine, + patch( + "rag.advanced_rag.knowlege_compile.wiki_incremental._wiki_update_doc_page_source", + new_callable=AsyncMock, + ), + ): + await _wiki._wiki_mode_a_run( + deltas=deltas, + existing_pages={}, + chat_mdl=MockChatModel(), + embd_mdl=MockEmbeddingModel(), + tenant_id="t1", + kb_id="kb1", + incremental=False, # first build — concept depth check applies, entity always created + canonical_claims={ + "smartphone industry": [{"statement": "Concept claim 1", "source_doc_id": "d1"}], + "Apple Inc.": [{"statement": "Entity claim", "source_doc_id": "d1"}], + }, + ) + + # Both pages should be refined (1 concept + 1 entity) + assert mock_refine.call_count == 2, f"Expected 2 REFINE calls, got {mock_refine.call_count}" + + # Verify page_id prefixes: one concept/ and one entity/ + prefixes = set() + for args in mock_refine.call_args_list: + kwargs = args[1] + prefixes.add(kwargs["page_id"].split("/")[0]) + assert kwargs["page_type_kwd"] in ("concept", "entity") + assert "concept" in prefixes + assert "entity" in prefixes + + +@pytest.mark.asyncio +async def test_has_any_pages_detects_existing_pages(): + """_wiki_has_any_pages returns True when wiki_page rows exist.""" + doc_store = make_doc_store( + [ + {"_source": {"slug_kwd": "concept/smartphone"}}, + ] + ) + + with patch("common.settings.docStoreConn", doc_store): + has = await _wiki._wiki_has_any_pages("t1", "kb1") + + assert has is True + + +@pytest.mark.asyncio +async def test_has_any_pages_false_when_no_pages(): + """_wiki_has_any_pages returns False when no wiki_page rows exist.""" + doc_store = make_doc_store([]) + + with patch("common.settings.docStoreConn", doc_store): + has = await _wiki._wiki_has_any_pages("t1", "kb1") + + assert has is False + + +@pytest.mark.asyncio +async def test_has_any_pages_false_when_index_missing(): + """_wiki_has_any_pages returns False when the index doesn't exist.""" + doc_store = MagicMock() + doc_store.index_exist = MagicMock(return_value=False) + + with patch("common.settings.docStoreConn", doc_store): + has = await _wiki._wiki_has_any_pages("t1", "kb1") + + assert has is False + + +# ---- Tests for doc_page_source --------------------------------------------- + + +@pytest.mark.asyncio +async def test_canonical_claims_maps_raw_to_canonical(): + """canonical_claims must aggregate raw-name claims onto canonical names. + + Regression test: claim_index is keyed by RAW entity name (from MAP), but + REDUCE looks up by canonical name. If a raw name resolves to a different + canonical name (e.g. "Apple Computer" → "Apple Inc."), the claims must + still be found, otherwise every entity becomes a no-op and output is empty. + """ + # claim_index keyed by raw names + claim_index = { + "Apple Computer": [ + {"statement": "C1", "source_doc_id": "d1"}, + {"statement": "C2", "source_doc_id": "d1"}, + ], + "Apple Inc.": [ + {"statement": "C3", "source_doc_id": "d1"}, + ], + "Samsung": [ + {"statement": "C4", "source_doc_id": "d1"}, + ], + } + # name_resolution: "Apple Computer" → "Apple Inc." (merged) + name_resolution = {"Apple Computer": "Apple Inc.", "Apple Inc.": "Apple Inc.", "Samsung": "Samsung"} + affected_names = {"Apple Inc.", "Samsung"} + + # Replicate the aggregation logic from wiki_compile_incremental + canonical_claims: dict[str, list[dict]] = {} + for raw_name, claims in claim_index.items(): + cname = name_resolution.get(raw_name, raw_name) + if cname in affected_names: + canonical_claims.setdefault(cname, []).extend(claims) + for name in affected_names: + canonical_claims.setdefault(name, []) + + # "Apple Inc." should collect claims from BOTH "Apple Computer" and "Apple Inc." + assert len(canonical_claims["Apple Inc."]) == 3, f"Got {len(canonical_claims['Apple Inc.'])}" + statements = {c["statement"] for c in canonical_claims["Apple Inc."]} + assert statements == {"C1", "C2", "C3"} + assert len(canonical_claims["Samsung"]) == 1 + + +@pytest.mark.asyncio +async def test_doc_page_source_entity_names(): + """doc_page_source stores and retrieves entity_names.""" + doc_store = make_doc_store( + [ + { + "_source": { + "doc_id": "doc_1", + "page_ids": '["concept/A", "concept/B"]', + "entity_names": '["Apple Inc.", "smartphone industry"]', + "source_chunk_hashes": '{"C1": "abc", "C2": "def"}', + "map_checksum": "xyz", + } + } + ] + ) + + with patch("common.settings.docStoreConn", doc_store): + dps = await _wiki._wiki_load_doc_page_source("t1", "kb1", "doc_1") + + assert dps is not None + assert "Apple Inc." in dps.get("entity_names", []) + assert "smartphone industry" in dps.get("entity_names", []) + assert len(dps.get("page_ids", [])) == 2 + + +# ---- End-to-end: Entity Matching → REDUCE --------------------------------- + + +@pytest.mark.asyncio +async def test_entity_matching_to_reduce_flow(): + """Entity Matching → REDUCE: canonical names flow through correctly. + + Only the REDUCE call is async (straightforward dict/await logic). + Entity Matching is tested synchronously here. + """ + map_results = [ + { + "doc_id": "doc_1", + "entities": [{"name": "Apple Inc.", "type": "org"}], + "concepts": [{"term": "smartphone industry"}], + "claims": [ + {"entity_name": "Apple Inc.", "statement": "Apple Inc. is a tech company", "source_chunk_id": "C1", "source_doc_id": "doc_1"}, + {"entity_name": "smartphone industry", "statement": "A global industry", "source_chunk_id": "C2", "source_doc_id": "doc_1"}, + ], + } + ] + + raw, claim_index = _wiki._extract_raw_entities(map_results) + + # Verify matching (synchronous, no _wiki_match_entities call) + assert any(e["name"] == "Apple Inc." and e["type"] == "org" for e in raw) + assert any(e["name"] == "smartphone industry" and e["type"] == "concept" for e in raw) + + canonical_map = {} + for entry in raw: + canonical_map[entry["name"]] = entry + + assert len(canonical_map) >= 2 + assert canonical_map["Apple Inc."]["type"] == "org" + assert canonical_map["smartphone industry"]["type"] == "concept" + + # REDUCE (async but clean — only uses asyncio.gather, no _wiki_match_entities) + # canonical_claims built from claim_index (full text kept separate) + canonical_claims = {n: claim_index.get(n, []) for n in canonical_map} + + deltas = await _wiki._wiki_reduce_batch( + affected_names=set(canonical_map.keys()), + existing_pages={}, + deleted_doc_ids=set(), + canonical_claims=canonical_claims, + canonical_map=canonical_map, + name_resolution={e["name"]: e["name"] for e in raw}, + ) + + for d in deltas: + assert "entity_type" in d, f"Missing entity_type in delta: {d}" + + concept_deltas = [d for d in deltas if d.get("entity_type") == "concept"] + entity_deltas = [d for d in deltas if d.get("entity_type") != "concept"] + assert len(concept_deltas) >= 1 + assert len(entity_deltas) >= 1 + + +# ---- Edge cases ------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_reduce_batch_empty_affected(): + """_wiki_reduce_batch with empty affected_names returns [].""" + result = await _wiki._wiki_reduce_batch( + affected_names=set(), + map_results=[], + existing_pages={}, + deleted_doc_ids=set(), + ) + assert result == [] + + +@pytest.mark.asyncio +async def test_reduce_entity_type_preservation(): + """entity_type is preserved across all REDUCE paths.""" + + r1 = await _wiki._wiki_reduce_entity("C1", entity_type="concept", new_claims=[{"statement": "S1", "source_doc_id": "d1"}], existing_page=None, deleted_doc_ids=set()) + assert r1["entity_type"] == "concept" + + r2 = await _wiki._wiki_reduce_entity( + "E1", + entity_type="org", + new_claims=[{"statement": "N", "source_doc_id": "d2"}], + existing_page={"claims": [{"statement": "O", "source_doc_id": "d1"}], "page_version_int": 1}, + deleted_doc_ids=set(), + ) + assert r2["entity_type"] == "org" + + r3 = await _wiki._wiki_reduce_entity( + "E1", entity_type="concept", new_claims=[], existing_page={"claims": [{"statement": "O", "source_doc_id": "d1"}], "page_version_int": 1}, deleted_doc_ids={"d1"} + ) + assert r3["entity_type"] == "concept" + + +# ---- _wiki_decide_concept_pages helper test -------------------------------- + + +def test_entity_matching_concept_entity_types(): + """Entity Matching preserves entity_type for both entity and concept.""" + raw, _ = _wiki._extract_raw_entities( + [ + { + "doc_id": "doc_1", + "entities": [{"name": "Apple Inc.", "type": "org"}], + "concepts": [{"term": "supply chain"}], + "claims": [ + {"entity_name": "Apple Inc.", "statement": "S1", "source_chunk_id": "C1", "source_doc_id": "doc_1"}, + {"entity_name": "supply chain", "statement": "S2", "source_chunk_id": "C2", "source_doc_id": "doc_1"}, + ], + } + ] + ) + + types = {e["name"]: e.get("type") for e in raw} + assert types.get("Apple Inc.") == "org" + assert types.get("supply chain") == "concept" + + +# ---- _wiki_load_pages_for_graph (canvas graph) ----------------------------- + + +@pytest.mark.asyncio +async def test_load_pages_for_graph_shape(): + """_wiki_load_pages_for_graph projects wiki_page rows onto the graph shape.""" + doc_store = make_doc_store( + [ + { + "_source": { + "slug_kwd": "concept/smartphone", + "title_kwd": "Smartphone", + "page_type_kwd": "concept", + "summary_with_weight": "A mobile device.", + "entity_names_kwd": '["smartphone"]', + "outlinks_kwd": '["entity/apple"]', + "source_chunk_ids": '["C1", "C2"]', + "source_doc_ids": '["d1"]', + } + } + ] + ) + + with patch("common.settings.docStoreConn", doc_store): + pages = await _wiki._wiki_load_pages_for_graph("t1", "kb1") + + assert len(pages) == 1 + p = pages[0] + assert p["slug"] == "concept/smartphone" + assert p["title"] == "Smartphone" + assert p["page_type"] == "concept" + assert p["outlinks"] == ["entity/apple"] + assert p["source_chunk_ids"] == ["C1", "C2"] + assert p["source_doc_ids"] == ["d1"] + + +@pytest.mark.asyncio +async def test_finalize_writes_outlinks(): + """_wiki_finalize stamps outlinks_kwd / outlinks_int for valid wikilinks.""" + search_results = [ + {"_source": {"slug_kwd": "concept/A", "md_with_weight": "[[concept/B]] related", "page_type_kwd": "concept"}}, + {"_source": {"slug_kwd": "concept/B", "md_with_weight": "content", "page_type_kwd": "concept"}}, + ] + doc_store = make_doc_store(search_results) + + with ( + patch("common.settings.docStoreConn", doc_store), + patch(f"{_wiki.__name__}._load_canonical_entities", new_callable=AsyncMock, return_value={}), + ): + await _wiki._wiki_finalize(tenant_id="t1", kb_id="kb1", embd_mdl=None) + + update_calls = doc_store.update.call_args_list + for args in update_calls: + cond = args[0][0] + upd = args[0][1] + if cond.get("id") == "concept/A": + # outlinks_kwd is a *_kwd field: Infinity shreds json.dumps strings + # on read-back, so _wiki_finalize writes a native list (mirroring + # the old-mode writer) which Infinity stores/reads back correctly. + assert upd["outlinks_kwd"] == ["concept/B"], f"Got {upd['outlinks_kwd']}" + assert upd["outlinks_int"] == 1 + break + else: + pytest.fail("No update for concept/A found") + + +async def test_finalize_auto_links_mentions(): + """_wiki_finalize auto-links standalone mentions of other pages' names.""" + search_results = [ + { + "_source": { + "slug_kwd": "entity/Apple", + "title_kwd": "Apple", + "md_with_weight": "Apple makes phones. Apple is based in California.", + "page_type_kwd": "entity", + } + }, + { + "_source": { + "slug_kwd": "entity/Google", + "title_kwd": "Google", + "md_with_weight": "Google competes with Apple in search and mobile.", + "page_type_kwd": "entity", + } + }, + ] + doc_store = make_doc_store(search_results) + + with ( + patch("common.settings.docStoreConn", doc_store), + patch(f"{_wiki.__name__}._load_canonical_entities", new_callable=AsyncMock, return_value={}), + ): + await _wiki._wiki_finalize(tenant_id="t1", kb_id="kb1", embd_mdl=None) + + update_calls = doc_store.update.call_args_list + google_upd = None + for args in update_calls: + cond = args[0][0] + if cond.get("id") == "entity/Google": + google_upd = args[0][1] + break + assert google_upd is not None, "entity/Google not updated" + # "Apple" mentioned in Google page → auto-linked exactly once + outlink + # recorded, and rendered into the navigable artifact-link form. + assert "[Apple](artifact/kb1/entity/Apple)" in google_upd["md_with_weight"], google_upd["md_with_weight"] + assert "entity/Apple" in (google_upd["outlinks_kwd"] or []), google_upd["outlinks_kwd"] + assert google_upd["outlinks_int"] == 1 + + +def test_inside_wikilink(): + content = "before [[entity/X]] after" + assert _wiki._inside_wikilink(content, content.index("entity")) + assert not _wiki._inside_wikilink(content, content.index("before")) + assert not _wiki._inside_wikilink(content, content.index("after")) + + +def test_extract_outlinks_from_content(): + """_wiki_extract_outlinks_from_content derives unique ordered outlinks.""" + content = "See [[concept/B]] and [[entity/apple]] and [[concept/B]] again" + assert _wiki._wiki_extract_outlinks_from_content(content) == [ + "concept/B", + "entity/apple", + ] + assert _wiki._wiki_extract_outlinks_from_content("") == [] + assert _wiki._wiki_extract_outlinks_from_content("no links here") == [] + + +@pytest.mark.asyncio +async def test_load_pages_for_graph_outlink_fallback(): + """Pages without outlinks_kwd get edges derived from content wikilinks.""" + doc_store = make_doc_store( + [ + { + "_source": { + "slug_kwd": "concept/A", + "title_kwd": "A", + "page_type_kwd": "concept", + "summary_with_weight": "", + "md_with_weight": "See [[concept/B]]", + "entity_names_kwd": "[]", + "source_chunk_ids": "[]", + "source_doc_ids": "[]", + } + } + ] + ) + + with patch("common.settings.docStoreConn", doc_store): + pages = await _wiki._wiki_load_pages_for_graph("t1", "kb1") + + assert pages[0]["outlinks"] == ["concept/B"] + + +# ---- _load_canonical_entities: *_kwd scalar normalization ------------------ + + +@pytest.mark.asyncio +async def test_load_canonical_entities_normalizes_entity_type(): + """entity_type_kwd comes back as a list (['concept']) from Infinity; the + loader must normalize it to a scalar so `== "concept"` checks work.""" + doc_store = make_doc_store( + [ + { + "_source": { + "entity_kwd": "询问笔录", + "entity_type_kwd": ["concept"], # simulated Infinity list + "aliases": '["询问笔录"]', + "source_doc_ids": '["doc_1"]', + "mention_count_int": 2, + } + } + ] + ) + + with patch("common.settings.docStoreConn", doc_store): + canon = await _wiki._load_canonical_entities("t1", "kb1") + + assert "询问笔录" in canon + assert canon["询问笔录"]["entity_type_kwd"] == "concept", canon["询问笔录"] + + +@pytest.mark.asyncio +async def test_search_existing_pages_normalizes_slug(): + """slug_kwd comes back as a list; _search_existing_pages must key by scalar.""" + doc_store = make_doc_store( + [ + { + "_source": { + "slug_kwd": ["concept/询问笔录"], # simulated Infinity list + "title_kwd": ["询问笔录"], + "md_with_weight": "content", + } + } + ] + ) + + with patch("common.settings.docStoreConn", doc_store): + pages = await _wiki._search_existing_pages("t1", "kb1", ["slug_kwd", "title_kwd", "md_with_weight"]) + + assert "concept/询问笔录" in pages + assert pages["concept/询问笔录"]["title_kwd"] == ["询问笔录"] # raw preserved + assert pages["concept/询问笔录"]["id"] == "concept/询问笔录" + + +# ---- relation-based linking in FINALIZE ------------------------------------- + + +@pytest.mark.asyncio +async def test_finalize_links_via_map_relations(): + """_wiki_finalize connects pages when MAP extracted a (from, to) relation + between them, even when prose contains no [[wikilink]].""" + # Two wiki pages + one map_extract row whose relation links them. + doc_store = make_doc_store( + [ + { + "_source": { + "slug_kwd": "entity/肖亮", + "title_kwd": "肖亮", + "md_with_weight": "肖亮 is a person.", + "page_type_kwd": "entity", + } + }, + { + "_source": { + "slug_kwd": "entity/肖立", + "title_kwd": "肖立", + "md_with_weight": "肖立 is a person.", + "page_type_kwd": "entity", + } + }, + { + "_source": { + "doc_id": "map1", + "compile_kwd": "wiki_map_extract", + "content_with_weight": json.dumps( + { + "entities": [{"name": "肖亮", "type": "person"}, {"name": "肖立", "type": "person"}], + "relations": [{"from": "肖亮", "to": "肖立", "type": "other"}], + }, + ensure_ascii=False, + ), + } + }, + ] + ) + + with ( + patch("common.settings.docStoreConn", doc_store), + patch(f"{_wiki.__name__}._load_canonical_entities", new_callable=AsyncMock, return_value={}), + ): + await _wiki._wiki_finalize(tenant_id="t1", kb_id="kb1", embd_mdl=None) + + update_calls = doc_store.update.call_args_list + xiaoliang_upd = None + for args in update_calls: + if args[0][0].get("id") == "entity/肖亮": + xiaoliang_upd = args[0][1] + break + assert xiaoliang_upd is not None, "entity/肖亮 not updated" + assert "entity/肖立" in (xiaoliang_upd["outlinks_kwd"] or []), xiaoliang_upd["outlinks_kwd"] + assert xiaoliang_upd["outlinks_int"] == 1 + # The relation edge must be injected into the page body. It is stored either + # as a raw [[entity/肖立]] or (after the link transformer runs) as the + # navigable Markdown form [肖立](artifact/kb1/entity/肖立). Both contain the + # target slug so the graph can reconstruct edges from md_with_weight. + body = xiaoliang_upd["md_with_weight"] or "" + assert "entity/肖立" in body, body + + +@pytest.mark.asyncio +async def test_wiki_finalize_renders_navigable_links(): + """FINALIZE renders [[slug]] to [text](artifact/{kb_id}/{slug}) so the + frontend wiki viewer can deep-link, not plain [[...]] text.""" + # Two pages; page B's body links to page A via [[entity/肖立]]. FINALIZE + # must render that to the navigable [肖立](artifact/kb1/entity/肖立) form. + search_results = [ + { + "_source": { + "slug_kwd": "entity/肖立", + "title_kwd": "肖立", + "md_with_weight": "肖立 is a person.", + "page_type_kwd": "entity", + } + }, + { + "_source": { + "slug_kwd": "entity/肖亮", + "title_kwd": "肖亮", + "md_with_weight": "肖亮 references [[entity/肖立]] here.", + "page_type_kwd": "entity", + } + }, + ] + doc_store = make_doc_store(search_results) + with ( + patch("common.settings.docStoreConn", doc_store), + patch(f"{_wiki.__name__}._load_canonical_entities", new_callable=AsyncMock, return_value={}), + ): + await _wiki._wiki_finalize(tenant_id="t1", kb_id="kb1", embd_mdl=None) + update_calls = doc_store.update.call_args_list + upd = None + for args in update_calls: + if args[0][0].get("id") == "entity/肖亮": + upd = args[0][1] + break + assert upd is not None + body = upd["md_with_weight"] or "" + assert "[肖立](artifact/kb1/entity/肖立)" in body, body + + +@pytest.mark.asyncio +async def test_reduce_entity_claimless_concept_is_skipped(): + """A new concept without claims must not create an ungrounded page.""" + from rag.advanced_rag.knowlege_compile import wiki_incremental as _wiki + + result = await _wiki._wiki_reduce_entity( + entity_name="继承纠纷", + entity_type="concept", + existing_page=None, + new_claims=[], # concepts have no dedicated claim rows + deleted_doc_ids=set(), + ) + assert result["action"] == "noop" + assert result["has_delta"] is False + assert result["entity_type"] == "concept" + + +@pytest.mark.asyncio +async def test_page_router_skips_knn_when_no_existing_pages(): + """A first Mode B build should cluster directly without page-index searches.""" + entities = [ + {"entity_name": "Apple", "entity_type": "org", "claims": []}, + {"entity_name": "Banana", "entity_type": "org", "claims": []}, + ] + embd_mdl = MockEmbeddingModel() + embd_mdl.encode = MagicMock(wraps=embd_mdl.encode) + doc_store = make_doc_store() + + with ( + patch("common.settings.docStoreConn", doc_store), + patch( + f"{_wiki.__name__}._wiki_cluster_entities", + side_effect=lambda items, embeddings, threshold: [items], + ), + ): + assignments = await _wiki._wiki_page_router( + affected_entities=entities, + embd_mdl=embd_mdl, + tenant_id="t1", + kb_id="kb1", + existing_page_ids=set(), + ) + + assert assignments == {"_new_entity/apple": entities} + assert embd_mdl.encode.call_count == 1 + doc_store.search.assert_not_called() diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts index 1d9b9050c0..dcede1913c 100644 --- a/web/src/locales/en.ts +++ b/web/src/locales/en.ts @@ -1893,6 +1893,7 @@ Example: Virtual Hosted Style`, instruction: 'Instruction', globalRules: 'Global rules', globalRulesPlaceholder: 'Input global compilation rules', + plan: 'Plan (grouping wiki pages by topic via LLM)', raptorTreeSettings: 'RAPTOR tree settings', summarizationPrompt: 'Summarization prompt', maxToken: 'Max token', diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts index 34084af7cb..65f8380292 100644 --- a/web/src/locales/zh.ts +++ b/web/src/locales/zh.ts @@ -1583,6 +1583,7 @@ NER：使用 spaCy NER 和基于规则的关键词提取来抽取实体和关系 instruction: 'Instruction', globalRules: '全局规则', globalRulesPlaceholder: '请输入全局编译规则', + plan: 'Plan (LLM 分组合并 wiki 页面)', raptorTreeSettings: 'RAPTOR 树设置', summarizationPrompt: '摘要提示词', maxToken: '最大 token 数', diff --git a/web/src/pages/dataset/dataset-setting/form-schema.ts b/web/src/pages/dataset/dataset-setting/form-schema.ts index acb0eaf109..ac0f3cc7d8 100644 --- a/web/src/pages/dataset/dataset-setting/form-schema.ts +++ b/web/src/pages/dataset/dataset-setting/form-schema.ts @@ -103,6 +103,10 @@ export const formSchema = z .optional(), enable_metadata: z.boolean().optional(), llm_id: z.string().optional(), + // Compilation template group (e.g. the "wiki" Artifacts template with + // plan=yes/no). Persisted to parser_config so the wiki build backend can + // resolve the template and honour plan (Mode A vs Mode B). + compilation_template_group_id: z.string().nullish(), // Table parser: "auto" = all columns both, "manual" = use column role selector table_column_mode: z.enum(['auto', 'manual']).optional(), // Table parser: column name -> role (indexing | metadata | both); legacy "vectorize" -> indexing diff --git a/web/src/pages/dataset/dataset-setting/general-form.tsx b/web/src/pages/dataset/dataset-setting/general-form.tsx index 6fc61df119..8016309995 100644 --- a/web/src/pages/dataset/dataset-setting/general-form.tsx +++ b/web/src/pages/dataset/dataset-setting/general-form.tsx @@ -1,4 +1,5 @@ import { AvatarUpload } from '@/components/avatar-upload'; +import { CompilationTemplateFormField } from '@/components/compilation-template-form-field'; import { SelectWithSearch } from '@/components/originui/select-with-search'; import PageRankFormField from '@/components/page-rank-form-field'; import { RAGFlowFormItem } from '@/components/ragflow-form'; @@ -129,6 +130,7 @@ export function GeneralForm() { ownerTenantId={useKnowledgeBaseContext().knowledgeBase?.tenant_id} > + diff --git a/web/src/pages/user-setting/compilation-templates/create-next/components/template-configuration.tsx b/web/src/pages/user-setting/compilation-templates/create-next/components/template-configuration.tsx index dd1e3bf42d..3f3065d8fb 100644 --- a/web/src/pages/user-setting/compilation-templates/create-next/components/template-configuration.tsx +++ b/web/src/pages/user-setting/compilation-templates/create-next/components/template-configuration.tsx @@ -3,6 +3,7 @@ import { SelectWithSearch } from '@/components/originui/select-with-search'; import { RAGFlowFormItem } from '@/components/ragflow-form'; import { SwitchFormField } from '@/components/switch-fom-field'; import { Button } from '@/components/ui/button'; +import { Checkbox } from '@/components/ui/checkbox'; import { Input } from '@/components/ui/input'; import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'; import { Textarea } from '@/components/ui/textarea'; @@ -200,6 +201,20 @@ export function TemplateConfiguration({ /> + {isArtifacts && ( + + {(field) => ( + field.onChange(v)} + /> + )} + + )} + {kind === CompilationTemplateKind.Tree ? ( ) : ( diff --git a/web/src/pages/user-setting/compilation-templates/create-next/constant.ts b/web/src/pages/user-setting/compilation-templates/create-next/constant.ts index e7f7b0f45f..a1116edfd1 100644 --- a/web/src/pages/user-setting/compilation-templates/create-next/constant.ts +++ b/web/src/pages/user-setting/compilation-templates/create-next/constant.ts @@ -22,6 +22,7 @@ export const DefaultTemplateValues: TemplateSchemaType = { instruction: '', page_example: '', use_blueprint: false, + plan: true, rechunk: false, rechunk_rules: '', }, diff --git a/web/src/pages/user-setting/compilation-templates/create-next/utils.ts b/web/src/pages/user-setting/compilation-templates/create-next/utils.ts index 71c886082f..9befdf11b4 100644 --- a/web/src/pages/user-setting/compilation-templates/create-next/utils.ts +++ b/web/src/pages/user-setting/compilation-templates/create-next/utils.ts @@ -55,6 +55,7 @@ export const isConfigMetaKey = (key: string) => 'page_example', 'synthesis', 'use_blueprint', + 'plan', 'rechunk', 'rechunk_rules', ].includes(key); @@ -108,6 +109,10 @@ export const buildConfigFromBuiltin = ( : {}), use_blueprint: kind === CompilationTemplateKind.Artifacts && example.length > 0, + plan: + typeof builtinTemplate.config?.plan === 'boolean' + ? builtinTemplate.config.plan + : true, ...(kind !== CompilationTemplateKind.Tree ? { rechunk: builtinTemplate.config?.rechunk === true, @@ -168,6 +173,7 @@ export const transformDetailToForm = ( : {}), use_blueprint: detail.kind === CompilationTemplateKind.Artifacts && example.length > 0, + plan: typeof config.plan === 'boolean' ? config.plan : true, ...(detail.kind !== CompilationTemplateKind.Tree ? { rechunk: config.rechunk === true, @@ -251,6 +257,10 @@ export const transformTemplateToPayload = (template: TemplateSchemaType) => { config[key] = value as ICompilationTemplateConfigRequest[string]; return; } + if (key === 'plan') { + config[key] = value as ICompilationTemplateConfigRequest[string]; + return; + } if (isConfigMetaKey(key)) { if (typeof value === 'string' || typeof value === 'boolean') config[key] = value; diff --git a/web/src/pages/user-setting/compilation-templates/edit-next/components/template-configuration.tsx b/web/src/pages/user-setting/compilation-templates/edit-next/components/template-configuration.tsx index 8549116b9e..46cbfec27b 100644 --- a/web/src/pages/user-setting/compilation-templates/edit-next/components/template-configuration.tsx +++ b/web/src/pages/user-setting/compilation-templates/edit-next/components/template-configuration.tsx @@ -1,6 +1,7 @@ import { ModelTreeSelectFormField } from '@/components/model-tree-select'; import { SelectWithSearch } from '@/components/originui/select-with-search'; import { RAGFlowFormItem } from '@/components/ragflow-form'; +import { Checkbox } from '@/components/ui/checkbox'; import { SwitchFormField } from '@/components/switch-fom-field'; import { Input } from '@/components/ui/input'; import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'; @@ -193,6 +194,20 @@ export function TemplateConfiguration({ /> + {kind === CompilationTemplateKind.Artifacts && ( + + {(field) => ( + field.onChange(v)} + /> + )} + + )} + {kind === CompilationTemplateKind.Tree ? ( ) : ( diff --git a/web/src/pages/user-setting/compilation-templates/edit-next/constant.ts b/web/src/pages/user-setting/compilation-templates/edit-next/constant.ts index e7f7b0f45f..a1116edfd1 100644 --- a/web/src/pages/user-setting/compilation-templates/edit-next/constant.ts +++ b/web/src/pages/user-setting/compilation-templates/edit-next/constant.ts @@ -22,6 +22,7 @@ export const DefaultTemplateValues: TemplateSchemaType = { instruction: '', page_example: '', use_blueprint: false, + plan: true, rechunk: false, rechunk_rules: '', }, diff --git a/web/src/pages/user-setting/compilation-templates/edit-next/utils.ts b/web/src/pages/user-setting/compilation-templates/edit-next/utils.ts index ab005f3204..b35132fd1b 100644 --- a/web/src/pages/user-setting/compilation-templates/edit-next/utils.ts +++ b/web/src/pages/user-setting/compilation-templates/edit-next/utils.ts @@ -51,6 +51,7 @@ export const isConfigMetaKey = (key: string) => 'page_example', 'synthesis', 'use_blueprint', + 'plan', 'rechunk', 'rechunk_rules', ].includes(key); @@ -104,6 +105,10 @@ export const buildConfigFromBuiltin = ( : {}), use_blueprint: kind === CompilationTemplateKind.Artifacts && example.length > 0, + plan: + typeof builtinTemplate.config?.plan === 'boolean' + ? builtinTemplate.config.plan + : true, ...(kind !== CompilationTemplateKind.Tree ? { rechunk: builtinTemplate.config?.rechunk === true, @@ -164,6 +169,7 @@ export const transformDetailToForm = ( : {}), use_blueprint: detail.kind === CompilationTemplateKind.Artifacts && example.length > 0, + plan: typeof config.plan === 'boolean' ? config.plan : true, ...(detail.kind !== CompilationTemplateKind.Tree ? { rechunk: config.rechunk === true, @@ -247,6 +253,10 @@ export const transformTemplateToPayload = (template: TemplateSchemaType) => { config[key] = value as ICompilationTemplateConfigRequest[string]; return; } + if (key === 'plan') { + config[key] = value as ICompilationTemplateConfigRequest[string]; + return; + } if (isConfigMetaKey(key)) { if (typeof value === 'string' || typeof value === 'boolean') config[key] = value;