Feat: Add knowledge compilation workflows (#16515)

## Summary
- Add knowledge compilation template APIs, services, and builtin
template seed data
- Add advanced knowledge compile structure/artifact/RAPTOR workflow
support
- Update parsing, dataset/document APIs, and supporting services for
compilation workflows

rag/advanced_rag/knowlege_compile/__init__.py

@@ -0,0 +1,43 @@
+#
+#  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.
+#
+
+from .structure import compile_structure_from_text, merge_compiled_structures
+from .wiki import (
+    WIKI_DRAFT_COMPILE_KWD,
+    WIKI_MAP_COMPILE_KWD,
+    WIKI_PAGE_COMPILE_KWD,
+    WIKI_PLAN_COMPILE_KWD,
+    WIKI_REDUCE_COMPILE_KWD,
+    wiki_map_from_chunks,
+    wiki_plan_from_reduction,
+    wiki_reduce_from_extracts,
+    wiki_refine_from_plan,
+)
+
+
+__all__ = [
+    "compile_structure_from_text",
+    "merge_compiled_structures",
+    "wiki_map_from_chunks",
+    "wiki_reduce_from_extracts",
+    "wiki_plan_from_reduction",
+    "wiki_refine_from_plan",
+    "WIKI_MAP_COMPILE_KWD",
+    "WIKI_REDUCE_COMPILE_KWD",
+    "WIKI_PLAN_COMPILE_KWD",
+    "WIKI_PAGE_COMPILE_KWD",
+    "WIKI_DRAFT_COMPILE_KWD",
+]

rag/advanced_rag/knowlege_compile/_common.py

@@ -0,0 +1,913 @@
+#
+#  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.
+#
+"""Shared helpers for the knowlege_compile pipelines (structure + wiki).
+
+Both ``structure.py`` (compile_structure_from_text / merge_compiled_structures)
+and ``wiki.py`` (the MAP→REDUCE→PLAN→REFINE artifact pipeline) need the same set
+of plumbing: encode-through-LLMBundle, stable id minting, search-tokenizer
+pairs, order-preserving chunk-id unions, defensive LLMBundle validation, the
+``chat_mdl.max_length * INPUT_UTILIZATION - prompt_overhead`` token-budget
+calculation, and thin ES I/O wrappers.
+
+Anything in this module is meant to be:
+  - LLMBundle-aware but provider-agnostic;
+  - Safe to import from either pipeline without circular references;
+  - Synchronous unless an awaitable behaviour is required.
+
+Heavier shared logic that is conceptually identical but happens to differ in
+shape between the two pipelines (e.g. pairwise-cosine dedup, LLM "are these
+the same?" batching) intentionally stays in each pipeline file for now —
+extract those only when their shapes converge.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import string
+from typing import Any, Awaitable, Callable, Iterable, Optional
+
+import xxhash
+
+from common.misc_utils import thread_pool_exec
+from common.token_utils import num_tokens_from_string
+from rag.nlp import rag_tokenizer
+from rag.prompts.generator import INPUT_UTILIZATION, gen_json, split_chunks
+
+
+# ---------------------------------------------------------------------------
+# ID minting
+# ---------------------------------------------------------------------------
+
+
+def stable_row_id(*parts) -> str:
+    """xxh64 hexdigest of ``":".join(parts)`` — stable per part tuple, used
+    as the ES row id when we want idempotent upserts.
+
+    ``None`` parts become empty strings, everything else is ``str()``-ified.
+    """
+    key = ":".join("" if p is None else str(p) for p in parts)
+    return xxhash.xxh64(key.encode("utf-8", "surrogatepass")).hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# Embedding
+# ---------------------------------------------------------------------------
+
+
+async def encode(embd_mdl, texts: list[str]) -> list:
+    """``LLMBundle.encode`` wrapped in ``thread_pool_exec``.
+
+    Returns the embeddings list (drops the ``used_tokens`` count); empty
+    input returns ``[]``. Caller is responsible for ensuring ``embd_mdl``
+    is a real bundle — use :func:`ensure_llm_bundle` to validate at entry.
+    """
+    if not texts:
+        return []
+    embeddings, _ = await thread_pool_exec(embd_mdl.encode, texts)
+    return list(embeddings)
+
+
+# ---------------------------------------------------------------------------
+# Tokenization for keyword search
+# ---------------------------------------------------------------------------
+
+
+def tokenize_for_search(text: str) -> tuple[str, str]:
+    """Returns ``(content_ltks, content_sm_ltks)`` for a piece of text.
+
+    Empty / non-string input returns ``("", "")``. Used wherever we write a
+    searchable ES row that needs both tokenizations.
+    """
+    if not isinstance(text, str) or not text:
+        return "", ""
+    ltks = rag_tokenizer.tokenize(text)
+    if not ltks:
+        return "", ""
+    sm = rag_tokenizer.fine_grained_tokenize(ltks)
+    return ltks, sm
+
+
+# ---------------------------------------------------------------------------
+# Order-preserving union of string lists
+# ---------------------------------------------------------------------------
+
+
+def union_ordered(*lists: Optional[Iterable]) -> list[str]:
+    """Concatenate iterables and dedupe, preserving first-seen order.
+    Falsy values and non-strings are silently dropped.
+    """
+    seen_set: set[str] = set()
+    seen: list[str] = []
+    for lst in lists:
+        if not lst:
+            continue
+        for v in lst:
+            if not v or not isinstance(v, str):
+                continue
+            if v in seen_set:
+                continue
+            seen_set.add(v)
+            seen.append(v)
+    return seen
+
+
+# ---------------------------------------------------------------------------
+# Token-budget calculation for split_chunks
+# ---------------------------------------------------------------------------
+
+
+def make_input_budget(
+    chat_mdl,
+    *prompts: str,
+    floor: int = 1024,
+    utilization: float = INPUT_UTILIZATION,
+) -> int:
+    """``chat_mdl.max_length * utilization - num_tokens(sum of prompts)``,
+    floored at ``floor``.
+
+    Mirrors the budget idiom used by ``compile_structure_from_text`` and
+    ``wiki_map_from_chunks``: caller passes the constant prompt scaffolding
+    (system prompt + user template) — ``split_chunks`` then sizes batches
+    to leave that much room.
+    """
+    overhead = num_tokens_from_string("".join(p or "" for p in prompts))
+    budget = int(chat_mdl.max_length * utilization) - overhead
+    return max(budget, floor)
+
+
+# ---------------------------------------------------------------------------
+# Defensive LLMBundle validation
+# ---------------------------------------------------------------------------
+
+
+def ensure_llm_bundle(mdl, method: str, *, label: str = "model"):
+    """Return ``mdl`` if it exposes ``method``; otherwise try to unwrap a
+    tuple, otherwise return ``None`` and log an error.
+
+    Common cause for tuple inputs at call sites: ``LLMBundle.encode()`` and
+    similar methods return ``(embeddings, used_tokens)``. If a caller stores
+    the *result* of ``encode()`` into a variable named like
+    ``embedding_model`` and passes that in, we end up with a tuple here.
+    We unwrap with a warning so the pipeline keeps working while the caller
+    is fixed.
+    """
+    if hasattr(mdl, method):
+        return mdl
+    if isinstance(mdl, tuple) and mdl and hasattr(mdl[0], method):
+        logging.warning(
+            "%s arrived as a %s; unwrapping to first element (check the call site — was %s()'s return value passed instead of the LLMBundle?)",
+            label,
+            type(mdl).__name__,
+            method,
+        )
+        return mdl[0]
+    logging.error(
+        "%s has no .%s method (type=%s); aborting",
+        label,
+        method,
+        type(mdl).__name__,
+    )
+    return None
+
+
+# ---------------------------------------------------------------------------
+# ES I/O wrappers
+# ---------------------------------------------------------------------------
+
+
+async def es_search(
+    select_fields: list[str],
+    condition: dict,
+    *,
+    tenant_id: str,
+    kb_ids: list[str],
+    match_expressions: list | None = None,
+    offset: int = 0,
+    limit: int = 1000,
+    label: str = "es_search",
+) -> dict:
+    """Thin wrapper around ``docStoreConn.search`` + ``get_fields``.
+
+    Returns ``{row_id: row_dict}``. Returns ``{}`` on failure (with a
+    logged exception). ``label`` is included in the failure log so each
+    call site is identifiable.
+    """
+    from common import settings
+    from common.doc_store.doc_store_base import OrderByExpr
+    from rag.nlp import search as _rag_search
+
+    index = _rag_search.index_name(tenant_id)
+    try:
+        res = await thread_pool_exec(
+            settings.docStoreConn.search,
+            select_fields,
+            [],
+            condition,
+            match_expressions or [],
+            OrderByExpr(),
+            offset,
+            limit,
+            index,
+            kb_ids,
+        )
+        return settings.docStoreConn.get_fields(res, select_fields) or {}
+    except Exception:
+        logging.exception("%s failed (condition=%r)", label, condition)
+        return {}
+
+
+async def es_insert(
+    rows: list[dict],
+    tenant_id: str,
+    kb_id: str,
+    *,
+    label: str = "es_insert",
+) -> None:
+    """Bulk insert wrapped in ``thread_pool_exec``. Logs on failure."""
+    if not rows:
+        return
+    from common import settings
+    from rag.nlp import search as _rag_search
+
+    index = _rag_search.index_name(tenant_id)
+    try:
+        await thread_pool_exec(settings.docStoreConn.insert, rows, index, kb_id)
+    except Exception:
+        logging.exception("%s failed (%d row(s))", label, len(rows))
+
+
+async def es_delete(
+    condition: dict,
+    tenant_id: str,
+    kb_id: str,
+    *,
+    label: str = "es_delete",
+) -> None:
+    """Bulk delete wrapped in ``thread_pool_exec``. Best-effort; logs on
+    failure (some callers rely on id-based upsert as a fallback)."""
+    from common import settings
+    from rag.nlp import search as _rag_search
+
+    index = _rag_search.index_name(tenant_id)
+    try:
+        await thread_pool_exec(settings.docStoreConn.delete, condition, index, kb_id)
+    except Exception:
+        logging.debug("%s failed (condition=%r); caller may rely on id-upsert", label, condition)
+
+
+async def es_upsert_one(
+    filter_condition: dict,
+    row: dict,
+    tenant_id: str,
+    kb_id: str,
+    *,
+    label: str = "es_upsert_one",
+) -> None:
+    """Delete-by-filter then insert. Used when an in-place update would
+    require knowing the existing row's id and we'd rather drop+re-create.
+
+    Best-effort delete (failures are debug-logged) followed by the insert.
+    Set ``row["id"]`` to a stable value derived from the filter
+    (:func:`stable_row_id`) so id-based dedup at the connector catches any
+    race that bypasses the delete.
+    """
+    await es_delete(filter_condition, tenant_id, kb_id, label=f"{label}.delete")
+    await es_insert([row], tenant_id, kb_id, label=f"{label}.insert")
+
+
+# ---------------------------------------------------------------------------
+# Doc-vector field discovery
+# ---------------------------------------------------------------------------
+
+
+def find_vec_field(doc: dict) -> tuple[Optional[str], Optional[list]]:
+    """Locate the ``q_<dim>_vec`` field on an ES doc dict. Returns
+    ``(field_name, vec)`` or ``(None, None)`` if the doc carries no
+    embedding."""
+    for k, v in doc.items():
+        if isinstance(k, str) and k.startswith("q_") and k.endswith("_vec"):
+            return k, v
+    return None, None
+
+
+# ---------------------------------------------------------------------------
+# Chunked-LLM pipeline engine
+# ---------------------------------------------------------------------------
+#
+# Both artifact MAP and compile_structure_from_text follow the same outer shape:
+#
+#   1. Filter chunks (drop empty text, optionally skip a "resume" set);
+#   2. Pack remaining chunks into batches via ``split_chunks`` sized to leave
+#      room for the prompt scaffolding;
+#   3. Run an LLM-driven ``process_batch`` over each batch in parallel under
+#      an ``asyncio.Semaphore(max_workers)``;
+#   4. Aggregate the per-batch results into a single value.
+#
+# The inner LLM call shape diverges between the pipelines — artifact uses a
+# single ``gen_json`` per batch with ``[CHUNK_ID Cn]``-labelled bodies,
+# structure uses two ``gen_json`` calls (nodes then edges) with ``---``
+# separators and no per-chunk attribution. That divergence lives in each
+# pipeline's ``process_batch`` closure; this engine only owns the scaffold.
+
+
+def _default_chunk_text(chunk: dict) -> str:
+    if not isinstance(chunk, dict):
+        return ""
+    text = chunk.get("text") or chunk.get("content_with_weight") or chunk.get("content") or ""
+    return text if isinstance(text, str) else ""
+
+
+def _default_label(position_in_batch: int) -> str:
+    return f"C{position_in_batch + 1}"
+
+
+def build_chunk_batches(
+    chunks: list[dict],
+    chat_mdl,
+    *,
+    prompt_overhead_tokens: int,
+    resume_chunk_ids: Optional[set[str]] = None,
+    scrub_text: Optional[Callable[[str], str]] = None,
+    label_fn: Callable[[int], str] = _default_label,
+    chunk_text_picker: Optional[Callable[[dict], str]] = None,
+    budget_floor: int = 1024,
+    batch_size_cap: Optional[int] = None,
+    window_fraction: Optional[float] = None,
+) -> tuple[list[list[dict]], dict]:
+    """Filter chunks, pack into batches, return per-batch entries.
+
+    Each batch entry is ``{"label": str, "chunk_id": str, "text": str}``
+    where ``label`` is per-batch positional (default ``C1``, ``C2``, …) and
+    ``text`` is the post-scrub chunk body. Empty or resume-skipped chunks
+    are dropped.
+
+    Two packing modes:
+      - **Default (split_chunks)**: ``input_budget`` derived from
+        ``chat_mdl.max_length * INPUT_UTILIZATION - prompt_overhead_tokens``.
+        Used by ``structure.py`` and the legacy artifact MAP path.
+      - **Cap+fraction (greedy)**: when ``batch_size_cap`` is provided,
+        chunks are packed greedily with two cutoffs — chunk-count exceeds
+        ``batch_size_cap`` OR accumulated tokens exceed
+        ``chat_mdl.max_length * window_fraction``. This is the artifact
+        compilation rule (BS=8, window=0.5).
+
+    Returns ``(batches, info)`` where ``info`` is a small stats dict.
+    """
+    if not chunks:
+        return [], {"total": 0, "kept": 0, "skipped_resume": 0, "skipped_empty": 0, "input_budget": 0, "n_batches": 0}
+
+    picker = chunk_text_picker or _default_chunk_text
+    resume_set = resume_chunk_ids or set()
+
+    chunk_ids: list[str] = []
+    chunk_texts: list[str] = []
+    skipped_resume = 0
+    skipped_empty = 0
+
+    for chunk in chunks:
+        cid = chunk.get("id") or chunk.get("chunk_id")
+        if not cid:
+            skipped_empty += 1
+            continue
+        if cid in resume_set:
+            skipped_resume += 1
+            continue
+        text = picker(chunk)
+        if not text or not text.strip():
+            skipped_empty += 1
+            continue
+        if scrub_text is not None:
+            text = scrub_text(text)
+            if not text or not text.strip():
+                skipped_empty += 1
+                continue
+        chunk_ids.append(cid)
+        chunk_texts.append(text)
+
+    if not chunk_texts:
+        return [], {
+            "total": len(chunks),
+            "kept": 0,
+            "skipped_resume": skipped_resume,
+            "skipped_empty": skipped_empty,
+            "input_budget": 0,
+            "n_batches": 0,
+        }
+
+    batches: list[list[dict]] = []
+    input_budget: int
+
+    if batch_size_cap is not None:
+        # Artifact mode — greedy bin-packing with chunk-count + token caps.
+        fraction = window_fraction if window_fraction is not None else 0.5
+        token_cap = max(int(chat_mdl.max_length * fraction), budget_floor)
+        input_budget = token_cap
+
+        current: list[dict] = []
+        current_tks = 0
+        for idx, text in enumerate(chunk_texts):
+            tks = num_tokens_from_string(text)
+            would_overflow_count = len(current) >= batch_size_cap
+            would_overflow_tokens = current and (current_tks + tks > token_cap)
+            if would_overflow_count or would_overflow_tokens:
+                batches.append(current)
+                current = []
+                current_tks = 0
+            current.append(
+                {
+                    "label": label_fn(len(current)),
+                    "chunk_id": chunk_ids[idx],
+                    "text": text,
+                }
+            )
+            current_tks += tks
+        if current:
+            batches.append(current)
+    else:
+        input_budget = max(
+            int(chat_mdl.max_length * INPUT_UTILIZATION) - prompt_overhead_tokens,
+            budget_floor,
+        )
+
+        raw_batches = split_chunks(chunk_texts, input_budget) or []
+        for batch in raw_batches:
+            packed: list[dict] = []
+            for position, item in enumerate(batch):
+                for idx, text in item.items():
+                    packed.append(
+                        {
+                            "label": label_fn(position),
+                            "chunk_id": chunk_ids[idx],
+                            "text": text,
+                        }
+                    )
+            if packed:
+                batches.append(packed)
+
+    info = {
+        "total": len(chunks),
+        "kept": len(chunk_texts),
+        "skipped_resume": skipped_resume,
+        "skipped_empty": skipped_empty,
+        "input_budget": input_budget,
+        "n_batches": len(batches),
+    }
+    return batches, info
+
+
+async def run_chunked_pipeline(
+    batches: list[list[dict]],
+    *,
+    process_batch: Callable[..., Awaitable[Any]],
+    aggregate: Optional[Callable[[list[Any]], Any]] = None,
+    max_workers: int = 6,
+    callback: Optional[Callable] = None,
+    log_prefix: str = "chunked_pipeline",
+) -> Any:
+    """Run ``process_batch`` over each batch in parallel.
+
+    ``process_batch`` is called as
+    ``await process_batch(entries: list[dict], batch_idx: int, total: int)``
+    and may return anything; ``aggregate`` (if given) is called with the
+    list of per-batch results and its return value is the engine's return.
+    Without ``aggregate`` the raw per-batch results list is returned.
+
+    Cancel-on-error semantics: if any task raises, all sibling tasks are
+    cancelled and the exception propagates.
+    """
+    if not batches:
+        return aggregate([]) if aggregate else []
+
+    total = len(batches)
+    semaphore = asyncio.Semaphore(max_workers) if max_workers and max_workers > 0 else None
+
+    async def _one(idx: int, entries: list[dict]) -> Any:
+        async def _do() -> Any:
+            return await process_batch(entries, idx, total)
+
+        if semaphore is not None:
+            async with semaphore:
+                return await _do()
+        return await _do()
+
+    tasks = [asyncio.create_task(_one(i, b)) for i, b in enumerate(batches) if b]
+    if not tasks:
+        return aggregate([]) if aggregate else []
+
+    try:
+        results = await asyncio.gather(*tasks, return_exceptions=False)
+    except Exception:
+        for t in tasks:
+            t.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        raise
+
+    if callback:
+        try:
+            callback(1.0, f"{log_prefix}: {total} batch(es) complete")
+        except Exception:
+            logging.debug("%s: completion callback failed", log_prefix, exc_info=True)
+
+    return aggregate(results) if aggregate else results
+
+
+# ---------------------------------------------------------------------------
+# Bulk dedup engine — exact + embedding + LLM disambiguation
+# ---------------------------------------------------------------------------
+#
+# Replaces wiki's _wiki_exact_dedup_entities / _wiki_exact_dedup_concepts /
+# _wiki_embedding_dedup_entities / _wiki_resolve_ambiguous_entities /
+# _wiki_apply_merges with one parameterised engine. structure.py's
+# merge_compiled_structures uses a different algorithm (incremental
+# kept-set + per-pair LLM judgement) and stays as-is.
+
+
+_PUNCT_TABLE = str.maketrans("", "", string.punctuation)
+
+
+DEFAULT_DISAMBIGUATE_SYSTEM = "You are a named-entity resolution assistant. Return only JSON."
+
+
+def normalize_key(name) -> str:
+    """Lowercase + strip whitespace + strip ASCII punctuation. Used as the
+    bucket key for exact dedup."""
+    if not isinstance(name, str):
+        return ""
+    return name.lower().strip().translate(_PUNCT_TABLE)
+
+
+def _exact_dedup_by_key(
+    items: list[dict],
+    *,
+    name_key: str,
+    type_key: Optional[str] = None,
+    aggregate_extra: Optional[Callable[[list[dict]], dict]] = None,
+) -> list[dict]:
+    """Group items by ``(normalize(item[name_key]), item.get(type_key))``.
+
+    Canonical record per group:
+      - ``<name_key>``: the most-common spelling across the group
+      - ``<type_key>`` (if given): the group's shared value
+      - ``aliases``: sorted union of every name + every input alias, minus
+        the canonical name
+      - ``mention_count``: sum of input ``mention_count`` values (defaults
+        to ``1`` per missing)
+      - ``chunk_ids``: order-preserving union
+      - ``_norm``: the normalized key (stripped by ``bulk_dedup_items``)
+      - any extras from ``aggregate_extra(group)``
+    """
+    groups: dict[tuple, list[dict]] = {}
+    for it in items:
+        if not isinstance(it, dict):
+            continue
+        norm = normalize_key(it.get(name_key, ""))
+        if not norm:
+            continue
+        key = (norm, it.get(type_key) if type_key else None)
+        groups.setdefault(key, []).append(it)
+
+    canonical: list[dict] = []
+    for (norm, type_val), group in groups.items():
+        name_counts: dict[str, int] = {}
+        for it in group:
+            n = it.get(name_key, "")
+            if isinstance(n, str) and n:
+                name_counts[n] = name_counts.get(n, 0) + 1
+        best = max(name_counts, key=lambda k: name_counts[k]) if name_counts else ""
+
+        aliases: set[str] = set()
+        chunk_id_lists: list[list] = []
+        mention_count = 0
+        for it in group:
+            n = it.get(name_key, "")
+            if isinstance(n, str) and n:
+                aliases.add(n)
+            for a in it.get("aliases") or []:
+                if isinstance(a, str) and a:
+                    aliases.add(a)
+            chunk_id_lists.append(it.get("chunk_ids") or [])
+            mention_count += int(it.get("mention_count") or 1)
+        aliases.discard(best)
+
+        record: dict = {
+            name_key: best,
+            "aliases": sorted(aliases),
+            "mention_count": mention_count,
+            "chunk_ids": union_ordered(*chunk_id_lists),
+            "_norm": norm,
+        }
+        if type_key:
+            record[type_key] = type_val
+        if aggregate_extra is not None:
+            try:
+                extras = aggregate_extra(group) or {}
+                if isinstance(extras, dict):
+                    record.update(extras)
+            except Exception:
+                logging.exception("bulk_dedup: aggregate_extra failed for group %r", norm)
+        canonical.append(record)
+
+    return canonical
+
+
+async def _embedding_dedup(
+    canonical: list[dict],
+    embd_mdl,
+    *,
+    name_key: str,
+    type_key: Optional[str] = None,
+    merge_threshold: float = 0.90,
+    ambiguous_low: float = 0.75,
+) -> tuple[dict[int, int], list[tuple[int, int]], Optional[list]]:
+    """Vectorised pairwise cosine; same-type-only when ``type_key`` given.
+
+    Returns ``(merged_into, ambiguous_pairs, vectors)``. ``merged_into``
+    is a union-find map ``index → parent_index``. ``ambiguous_pairs`` is the
+    [ambiguous_low, merge_threshold) bucket (after removing pairs already
+    linked by auto-merges). ``vectors`` is ``None`` on embedding failure
+    (caller should skip dedup).
+    """
+    n = len(canonical)
+    if n <= 1:
+        return {}, [], []
+
+    names = [it.get(name_key, "") for it in canonical]
+    try:
+        vectors = await encode(embd_mdl, names)
+    except Exception:
+        logging.exception("bulk_dedup: embedding batch failed")
+        return {}, [], None
+    if vectors is None or len(vectors) != n:
+        return {}, [], None
+
+    try:
+        from sklearn.metrics.pairwise import cosine_similarity
+        import numpy as np
+
+        matrix = np.asarray([list(v) for v in vectors], dtype=float)
+        sims = cosine_similarity(matrix)
+    except Exception:
+        logging.exception("bulk_dedup: pairwise cosine failed; skipping")
+        return {}, [], vectors
+
+    merged_into: dict[int, int] = {}
+
+    def _root(i: int) -> int:
+        while i in merged_into:
+            i = merged_into[i]
+        return i
+
+    auto_pairs: list[tuple[int, int]] = []
+    ambiguous_pairs: list[tuple[int, int]] = []
+
+    for i in range(n):
+        for j in range(i + 1, n):
+            if type_key and canonical[i].get(type_key) != canonical[j].get(type_key):
+                continue
+            s = float(sims[i, j])
+            if s >= merge_threshold:
+                auto_pairs.append((i, j))
+            elif s >= ambiguous_low:
+                ambiguous_pairs.append((i, j))
+
+    for i, j in auto_pairs:
+        ri, rj = _root(i), _root(j)
+        if ri == rj:
+            continue
+        if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
+            merged_into[rj] = ri
+        else:
+            merged_into[ri] = rj
+
+    still_ambiguous = [(i, j) for i, j in ambiguous_pairs if _root(i) != _root(j)]
+    return merged_into, still_ambiguous, vectors
+
+
+async def _resolve_ambiguous_pairs(
+    canonical: list[dict],
+    ambiguous_pairs: list[tuple[int, int]],
+    merged_into: dict[int, int],
+    chat_mdl,
+    *,
+    name_key: str,
+    type_key: Optional[str] = None,
+    batch_size: int = 50,
+    llm_timeout: int = 60,
+    system_prompt: str = DEFAULT_DISAMBIGUATE_SYSTEM,
+) -> dict[int, int]:
+    """LLM-judged disambiguation in batches; returns updated ``merged_into``."""
+    if not ambiguous_pairs:
+        return merged_into
+
+    def _root(i: int) -> int:
+        while i in merged_into:
+            i = merged_into[i]
+        return i
+
+    for start in range(0, len(ambiguous_pairs), batch_size):
+        batch = ambiguous_pairs[start : start + batch_size]
+        batch = [(i, j) for i, j in batch if _root(i) != _root(j)]
+        if not batch:
+            continue
+
+        lines: list[str] = []
+        for k, (i, j) in enumerate(batch):
+            a_type = f" ({canonical[i].get(type_key, '')})" if type_key else ""
+            b_type = f" ({canonical[j].get(type_key, '')})" if type_key else ""
+            lines.append(f'{k + 1}. "{canonical[i].get(name_key, "")}"{a_type} vs "{canonical[j].get(name_key, "")}"{b_type}')
+
+        user_prompt = (
+            "For each pair below, determine if they refer to the same real-world entity.\n"
+            f"Return a JSON array of exactly {len(batch)} booleans "
+            "(true = same entity, false = different).\n"
+            "Return ONLY the JSON array.\n\n" + "\n".join(lines)
+        )
+
+        try:
+            res = await asyncio.wait_for(
+                gen_json(system_prompt, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}),
+                timeout=llm_timeout,
+            )
+        except asyncio.TimeoutError:
+            logging.warning("bulk_dedup: disambiguation timed out (%d pairs)", len(batch))
+            continue
+        except Exception:
+            logging.exception("bulk_dedup: disambiguation call failed (%d pairs)", len(batch))
+            continue
+
+        decisions = None
+        if isinstance(res, list):
+            decisions = res
+        elif isinstance(res, dict):
+            for v in res.values():
+                if isinstance(v, list):
+                    decisions = v
+                    break
+        if not isinstance(decisions, list):
+            logging.warning("bulk_dedup: disambiguation returned unexpected shape: %r", type(res))
+            continue
+
+        for k, (i, j) in enumerate(batch):
+            verdict = decisions[k] if k < len(decisions) else False
+            if not verdict:
+                continue
+            ri, rj = _root(i), _root(j)
+            if ri == rj:
+                continue
+            if canonical[ri].get("mention_count", 0) >= canonical[rj].get("mention_count", 0):
+                merged_into[rj] = ri
+            else:
+                merged_into[ri] = rj
+
+    return merged_into
+
+
+def _apply_dedup_merges(
+    canonical: list[dict],
+    merged_into: dict[int, int],
+    *,
+    name_key: str,
+) -> list[dict]:
+    """Union-find collapse: sum ``mention_count``, union ``aliases`` and
+    ``chunk_ids`` per canonical."""
+
+    def _root(i: int) -> int:
+        while i in merged_into:
+            i = merged_into[i]
+        return i
+
+    roots: set[int] = {_root(i) for i in range(len(canonical))}
+    out: list[dict] = []
+    for ri in roots:
+        base = dict(canonical[ri])
+        aliases: set[str] = set(base.get("aliases") or [])
+        chunk_id_lists: list[list] = [base.get("chunk_ids") or []]
+        mention_count = int(base.get("mention_count") or 0)
+        for i, it in enumerate(canonical):
+            if i == ri or _root(i) != ri:
+                continue
+            mention_count += int(it.get("mention_count") or 0)
+            aliases.update(it.get("aliases") or [])
+            n = it.get(name_key)
+            if isinstance(n, str) and n:
+                aliases.add(n)
+            chunk_id_lists.append(it.get("chunk_ids") or [])
+        aliases.discard(base.get(name_key) or "")
+        base["aliases"] = sorted(aliases)
+        base["mention_count"] = mention_count
+        base["chunk_ids"] = union_ordered(*chunk_id_lists)
+        out.append(base)
+    return out
+
+
+async def bulk_dedup_items(
+    items: list[dict],
+    *,
+    name_key: str,
+    type_key: Optional[str] = None,
+    chat_mdl=None,
+    embd_mdl=None,
+    merge_threshold: float = 0.90,
+    ambiguous_low: float = 0.75,
+    ambiguous_batch_size: int = 50,
+    disambiguate_system_prompt: str = DEFAULT_DISAMBIGUATE_SYSTEM,
+    llm_timeout: int = 60,
+    aggregate_extra: Optional[Callable[[list[dict]], dict]] = None,
+    strip_norm_key: bool = True,
+) -> list[dict]:
+    """Three-phase dedup → canonical items.
+
+    Phase 1 (always): exact dedup by ``(normalize(item[name_key]),
+    item.get(type_key))`` — groups by normalized key, sums mention_count,
+    unions aliases and chunk_ids, optionally adds extras via
+    ``aggregate_extra(group)``.
+
+    Phase 2 (when ``embd_mdl`` is provided AND ``len(canonical) > 1``):
+    vectorised pairwise cosine over the canonical ``name_key`` values.
+    Pairs at similarity ≥ ``merge_threshold`` auto-merge; pairs in
+    ``[ambiguous_low, merge_threshold)`` move to phase 3. When ``type_key``
+    is given, pairs are only considered when both endpoints share the same
+    type. Embedding failures cause this phase (and 3) to be skipped.
+
+    Phase 3 (when ``chat_mdl`` is provided AND ambiguous pairs remain):
+    batched LLM disambiguation via ``gen_json`` — each batch asks for a
+    JSON array of booleans. True verdicts join the union-find.
+
+    Apply: union-find collapse — sum mention_count, union aliases /
+    chunk_ids per canonical.
+
+    Setting both ``chat_mdl`` and ``embd_mdl`` to ``None`` makes this an
+    exact-dedup-only call (which is what artifact uses for concepts).
+    """
+    canonical = _exact_dedup_by_key(
+        items,
+        name_key=name_key,
+        type_key=type_key,
+        aggregate_extra=aggregate_extra,
+    )
+
+    if len(canonical) > 1 and embd_mdl is not None:
+        merged_into, ambig, vectors = await _embedding_dedup(
+            canonical,
+            embd_mdl,
+            name_key=name_key,
+            type_key=type_key,
+            merge_threshold=merge_threshold,
+            ambiguous_low=ambiguous_low,
+        )
+        if vectors is None:
+            logging.warning("bulk_dedup: embedding phase skipped — keeping exact-dedup result")
+        else:
+            if chat_mdl is not None and ambig:
+                merged_into = await _resolve_ambiguous_pairs(
+                    canonical,
+                    ambig,
+                    merged_into,
+                    chat_mdl,
+                    name_key=name_key,
+                    type_key=type_key,
+                    batch_size=ambiguous_batch_size,
+                    llm_timeout=llm_timeout,
+                    system_prompt=disambiguate_system_prompt,
+                )
+            canonical = _apply_dedup_merges(canonical, merged_into, name_key=name_key)
+
+    if strip_norm_key:
+        for it in canonical:
+            it.pop("_norm", None)
+    return canonical
+
+
+__all__ = [
+    "stable_row_id",
+    "encode",
+    "tokenize_for_search",
+    "union_ordered",
+    "make_input_budget",
+    "ensure_llm_bundle",
+    "es_search",
+    "es_insert",
+    "es_delete",
+    "es_upsert_one",
+    "find_vec_field",
+    # New engines
+    "normalize_key",
+    "build_chunk_batches",
+    "run_chunked_pipeline",
+    "bulk_dedup_items",
+    "DEFAULT_DISAMBIGUATE_SYSTEM",
+]

rag/advanced_rag/knowlege_compile/dataset_nav.py

@@ -0,0 +1,437 @@
+#
+#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+"""Dataset-level navigation markdown for tree-kind compilations.
+
+After a doc finishes a ``tree``-kind compilation template the helper
+``upsert_dataset_nav_doc`` here appends (or refreshes) one line in the
+KB's nav markdown — one line per doc, each line carrying the doc id +
+a short summary lifted from the per-doc tree's root.
+
+Storage: a single ES row per KB under ``compile_kwd="dataset_nav"``,
+``available_int=0`` (so retrievers never surface it). The markdown
+body lives in ``md_with_weight``; ``doc_count_int`` and ``doc_ids_kwd``
+mirror the markdown's order for fast cap-check and dedup.
+
+Concurrency: every write is wrapped in a ``RedisDistributedLock``
+keyed by ``f"dataset_nav:{kb_id}"`` — multiple task executors
+finishing tree templates for the same KB in parallel must not
+interleave their read-modify-writes.
+
+The router/retrieval side that *consumes* this markdown is
+intentionally out of scope here.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+from typing import Any, Iterable
+
+import xxhash
+
+from common.token_utils import num_tokens_from_string
+from rag.utils.redis_conn import RedisDistributedLock
+
+
+# Hard cap on the number of docs we record in the nav markdown.
+# Beyond this we no-op on adds; the next doc to drop out of the KB
+# frees a slot via ``remove_dataset_nav_doc``.
+MAX_DATASET_NAV_DOCS = 128
+
+# Hard cap on the per-doc summary length, in tokens. Long summaries
+# bloat the markdown and slow downstream LLM passes that ingest the
+# whole nav blob; 128 tokens is enough for 1-2 sentences in either
+# Chinese or English text.
+MAX_DOC_SUMMARY_TOKENS = 128
+
+_COMPILE_KWD = "dataset_nav"
+
+# Lock TTL — long enough that an ES round-trip can't expire it mid-write
+# but short enough that a crashed executor doesn't pin the KB.
+_LOCK_TIMEOUT_S = 30
+_LOCK_BLOCKING_TIMEOUT_S = 5
+
+
+def _nav_row_id(kb_id: str) -> str:
+    """Stable per-KB row id. Mirrors the pattern used by ``skill_all``."""
+    return xxhash.xxh64(
+        f"dataset_nav:{kb_id}".encode("utf-8", "surrogatepass"),
+    ).hexdigest()
+
+
+def _nav_lock_key(kb_id: str) -> str:
+    return f"dataset_nav:{kb_id}"
+
+
+# Each line of the markdown looks like ``- **<doc_id>**: <summary>``.
+# The ``doc_id`` part is anchored at the start of a bullet so a simple
+# regex can locate the line on remove without touching adjacent lines.
+_LINE_RE = re.compile(r"^- \*\*([^*]+)\*\*:.*$")
+
+
+def _format_line(doc_id: str, summary: str) -> str:
+    # Strip newlines from the summary so each doc stays on a single
+    # markdown line. Multi-line summaries break the dedup regex and
+    # confuse downstream consumers that split on ``\n``.
+    one_line = summary.replace("\n", " ").replace("\r", " ").strip()
+    return f"- **{doc_id}**: {one_line}"
+
+
+def _truncate_summary(text: str) -> str:
+    """Trim ``text`` to ``MAX_DOC_SUMMARY_TOKENS`` tokens.
+
+    Uses the project's tokenizer so the cap matches what the LLM will
+    see. Falls back to a generous character cap if tokenization is
+    unavailable.
+    """
+    if not text:
+        return ""
+    text = text.strip()
+    try:
+        n = num_tokens_from_string(text)
+    except Exception:
+        # Best-effort character cap — 4 chars per token is a safe lower
+        # bound for English; Chinese is closer to 1 char per token but
+        # 4x still keeps the row size sane.
+        return text[: MAX_DOC_SUMMARY_TOKENS * 4]
+    if n <= MAX_DOC_SUMMARY_TOKENS:
+        return text
+    # Binary-search the right character cut so we land at exactly the
+    # token budget. ``num_tokens_from_string`` is cheap enough that a
+    # handful of probes per call is fine.
+    lo, hi = 0, len(text)
+    while lo < hi:
+        mid = (lo + hi + 1) // 2
+        try:
+            tn = num_tokens_from_string(text[:mid])
+        except Exception:
+            tn = mid // 4
+        if tn <= MAX_DOC_SUMMARY_TOKENS:
+            lo = mid
+        else:
+            hi = mid - 1
+    return text[:lo].rstrip()
+
+
+def _extract_root_summary_from_tree(tree: dict | None) -> str:
+    """Pull the doc-level abstract out of a RAPTOR-built tree.
+
+    Convention used by ``_raptor_tree_to_graph``: the root node's
+    ``title`` field carries the LLM summary at the highest layer.
+    Internal nodes lower in the tree carry their own per-cluster
+    summaries. We just take the root.
+    """
+    if not isinstance(tree, dict):
+        return ""
+    title = tree.get("title") or ""
+    if isinstance(title, str) and title.strip():
+        return title.strip()
+    # Some RAPTOR shapes use ``summary`` or ``content`` instead.
+    for alt in ("summary", "content_with_weight", "content"):
+        v = tree.get(alt)
+        if isinstance(v, str) and v.strip():
+            return v.strip()
+    return ""
+
+
+def _parse_existing_lines(md: str) -> list[tuple[str, str]]:
+    """Return ``(doc_id, raw_line)`` tuples in markdown order.
+
+    We keep the *raw* line so callers that just want to update one
+    doc's line don't have to re-derive the formatting. Lines that
+    don't match the per-doc shape (e.g. headers, blank lines) are
+    skipped — they're never written by this module, but a future
+    schema bump might add them and we shouldn't crash on it.
+    """
+    out: list[tuple[str, str]] = []
+    if not md:
+        return out
+    for line in md.splitlines():
+        m = _LINE_RE.match(line)
+        if not m:
+            continue
+        out.append((m.group(1), line))
+    return out
+
+
+def _render_md(entries: Iterable[tuple[str, str]]) -> str:
+    return "\n".join(line for _doc_id, line in entries)
+
+
+def _row_id_field(row: dict | None) -> dict:
+    if row and isinstance(row, dict):
+        return row
+    return {}
+
+
+async def _get_existing(tenant_id: str, kb_id: str) -> dict | None:
+    """Read the existing nav row, or ``None`` if it doesn't exist yet."""
+    from common import settings
+    from common.misc_utils import thread_pool_exec
+    from rag.nlp import search as _rag_search
+
+    index = _rag_search.index_name(tenant_id)
+    if not settings.docStoreConn.index_exist(index, kb_id):
+        return None
+    try:
+        existing = await thread_pool_exec(
+            settings.docStoreConn.get,
+            _nav_row_id(kb_id),
+            index,
+            [kb_id],
+        )
+    except Exception:
+        logging.exception(
+            "dataset_nav: read failed for kb=%s",
+            kb_id,
+        )
+        return None
+    return _row_id_field(existing) or None
+
+
+async def _write_row(tenant_id: str, kb_id: str, payload: dict) -> None:
+    """Upsert the nav row in the doc store."""
+    from common import settings
+    from common.misc_utils import thread_pool_exec
+    from rag.nlp import search as _rag_search
+
+    index = _rag_search.index_name(tenant_id)
+    row_id = _nav_row_id(kb_id)
+    payload = {
+        "id": row_id,
+        "kb_id": kb_id,
+        "doc_id": kb_id,
+        "compile_kwd": _COMPILE_KWD,
+        "knowledge_graph_kwd": "graph",
+        "available_int": 0,
+        **payload,
+    }
+    existing = await thread_pool_exec(
+        settings.docStoreConn.get,
+        row_id,
+        index,
+        [kb_id],
+    )
+    if existing:
+        await thread_pool_exec(
+            settings.docStoreConn.update,
+            {"id": row_id},
+            {k: v for k, v in payload.items() if k != "id"},
+            index,
+            kb_id,
+        )
+    else:
+        await thread_pool_exec(
+            settings.docStoreConn.insert,
+            [payload],
+            index,
+            kb_id,
+        )
+
+
+# --------------------------------------------------------------------
+# Public surface
+# --------------------------------------------------------------------
+
+
+async def upsert_dataset_nav_doc(
+    tenant_id: str,
+    kb_id: str,
+    doc_id: str,
+    summary_or_tree: Any,
+) -> None:
+    """Add or refresh a doc's line in the KB's nav markdown.
+
+    ``summary_or_tree`` can be:
+      - a plain string (taken as-is and truncated to ``MAX_DOC_SUMMARY_TOKENS``)
+      - a tree dict (the root summary is extracted via
+        ``_extract_root_summary_from_tree``)
+
+    The 128-doc cap is enforced here: if the doc isn't already in the
+    markdown and the row is full, the call is a no-op. Existing docs
+    always get their summary updated regardless of count.
+
+    Called from ``_run_tree_templates`` after each successful
+    ``_struct_upsert_graph_json``.
+    """
+    if not doc_id or not kb_id:
+        return
+
+    if isinstance(summary_or_tree, dict):
+        summary = _extract_root_summary_from_tree(summary_or_tree)
+    elif isinstance(summary_or_tree, str):
+        summary = summary_or_tree
+    else:
+        summary = ""
+    summary = _truncate_summary(summary)
+    if not summary:
+        # Nothing to record — a tree with no root summary means the
+        # RAPTOR pass produced a degenerate result; safer to skip than
+        # to write an empty line.
+        logging.info(
+            "dataset_nav: skipping doc=%s (kb=%s) — no usable summary",
+            doc_id,
+            kb_id,
+        )
+        return
+
+    new_line = _format_line(doc_id, summary)
+    lock = RedisDistributedLock(
+        _nav_lock_key(kb_id),
+        timeout=_LOCK_TIMEOUT_S,
+        blocking_timeout=_LOCK_BLOCKING_TIMEOUT_S,
+    )
+    try:
+        await lock.spin_acquire()
+    except Exception:
+        logging.exception(
+            "dataset_nav: lock acquire failed for kb=%s; proceeding lock-free",
+            kb_id,
+        )
+
+    try:
+        existing = await _get_existing(tenant_id, kb_id)
+        md = (existing or {}).get("md_with_weight") or ""
+        entries = _parse_existing_lines(md)
+
+        replaced = False
+        for i, (existing_doc_id, _) in enumerate(entries):
+            if existing_doc_id == doc_id:
+                entries[i] = (doc_id, new_line)
+                replaced = True
+                break
+        if not replaced:
+            if len(entries) >= MAX_DATASET_NAV_DOCS:
+                logging.info(
+                    "dataset_nav: kb=%s already at cap (%d); skipping doc=%s",
+                    kb_id,
+                    MAX_DATASET_NAV_DOCS,
+                    doc_id,
+                )
+                return
+            entries.append((doc_id, new_line))
+
+        payload = {
+            "md_with_weight": _render_md(entries),
+            "doc_count_int": len(entries),
+            "doc_ids_kwd": [doc_id for doc_id, _ in entries],
+        }
+        try:
+            await _write_row(tenant_id, kb_id, payload)
+        except Exception:
+            logging.exception(
+                "dataset_nav: write failed for kb=%s doc=%s",
+                kb_id,
+                doc_id,
+            )
+    finally:
+        try:
+            lock.release()
+        except Exception:
+            logging.exception("dataset_nav: lock release failed for kb=%s", kb_id)
+
+
+async def remove_dataset_nav_doc(
+    tenant_id: str,
+    kb_id: str,
+    doc_id: str,
+) -> None:
+    """Remove ``doc_id``'s line from the KB's nav markdown.
+
+    Called from ``DocumentService.remove_document`` so the markdown
+    stays in sync with the KB's doc set. No-op if the row doesn't
+    exist or the doc isn't represented in it.
+    """
+    if not doc_id or not kb_id:
+        return
+
+    lock = RedisDistributedLock(
+        _nav_lock_key(kb_id),
+        timeout=_LOCK_TIMEOUT_S,
+        blocking_timeout=_LOCK_BLOCKING_TIMEOUT_S,
+    )
+    try:
+        await lock.spin_acquire()
+    except Exception:
+        logging.exception(
+            "dataset_nav: lock acquire failed for kb=%s; proceeding lock-free",
+            kb_id,
+        )
+
+    try:
+        existing = await _get_existing(tenant_id, kb_id)
+        if not existing:
+            return
+        md = existing.get("md_with_weight") or ""
+        entries = _parse_existing_lines(md)
+        before = len(entries)
+        entries = [(d, line) for (d, line) in entries if d != doc_id]
+        if len(entries) == before:
+            return
+
+        payload = {
+            "md_with_weight": _render_md(entries),
+            "doc_count_int": len(entries),
+            "doc_ids_kwd": [d for d, _ in entries],
+        }
+        try:
+            await _write_row(tenant_id, kb_id, payload)
+        except Exception:
+            logging.exception(
+                "dataset_nav: remove-write failed for kb=%s doc=%s",
+                kb_id,
+                doc_id,
+            )
+    finally:
+        try:
+            lock.release()
+        except Exception:
+            logging.exception("dataset_nav: lock release failed for kb=%s", kb_id)
+
+
+def remove_dataset_nav_doc_sync(
+    tenant_id: str,
+    kb_id: str,
+    doc_id: str,
+) -> None:
+    """Sync wrapper around ``remove_dataset_nav_doc``.
+
+    ``DocumentService.remove_document`` is synchronous (Peewee-driven)
+    and the doc-store helpers it calls are sync too. We need a sync
+    bridge so the delete path can invoke this without spinning up an
+    event loop.
+
+    Strategy: run the async helper on the current loop if one is
+    available; otherwise spin a fresh loop for the duration of the
+    call. Any failure is logged and swallowed — the doc-delete path
+    must never fail because of nav-md cleanup.
+    """
+    try:
+        loop = asyncio.new_event_loop()
+        try:
+            loop.run_until_complete(
+                remove_dataset_nav_doc(tenant_id, kb_id, doc_id),
+            )
+        finally:
+            loop.close()
+    except Exception:
+        logging.exception(
+            "dataset_nav: sync remove failed for kb=%s doc=%s",
+            kb_id,
+            doc_id,
+        )

rag/advanced_rag/knowlege_compile/mind_map_extractor.py

@@ -33,6 +33,7 @@ from common.token_utils import num_tokens_from_string
 @dataclass
 class MindMapResult:
     """Unipartite Mind Graph result class definition."""
+
     output: dict
 
 
@@ -42,14 +43,12 @@ class MindMapExtractor(Extractor):
     _on_error: ErrorHandlerFn
 
     def __init__(
-            self,
-            llm_invoker: CompletionLLM,
-            prompt: str | None = None,
-            input_text_key: str | None = None,
-            on_error: ErrorHandlerFn | None = None,
+        self,
+        llm_invoker: CompletionLLM,
+        prompt: str | None = None,
+        input_text_key: str | None = None,
+        on_error: ErrorHandlerFn | None = None,
     ):
-        """Init method definition."""
-        # TODO: streamline construction
         self._llm = llm_invoker
         self._input_text_key = input_text_key or "input_text"
         self._mind_map_prompt = prompt or MIND_MAP_EXTRACTION_PROMPT
@@ -70,17 +69,10 @@ class MindMapExtractor(Extractor):
             k = self._key(k)
             if k and k not in keyset:
                 keyset.add(k)
-                arr.append(
-                    {
-                        "id": k,
-                        "children": self._be_children(v, keyset)
-                    }
-                )
+                arr.append({"id": k, "children": self._be_children(v, keyset)})
         return arr
 
-    async def __call__(
-            self, sections: list[str], prompt_variables: dict[str, Any] | None = None
-    ) -> MindMapResult:
+    async def __call__(self, sections: list[str], prompt_variables: dict[str, Any] | None = None) -> MindMapResult:
         """Call method definition."""
         if prompt_variables is None:
             prompt_variables = {}
@@ -93,18 +85,14 @@ class MindMapExtractor(Extractor):
         for i in range(len(sections)):
             section_cnt = num_tokens_from_string(sections[i])
             if cnt + section_cnt >= token_count and texts:
-                tasks.append(asyncio.create_task(
-                    self._process_document("".join(texts), prompt_variables, res)
-                ))
+                tasks.append(asyncio.create_task(self._process_document("".join(texts), prompt_variables, res)))
                 texts = []
                 cnt = 0
 
             texts.append(sections[i])
             cnt += section_cnt
         if texts:
-            tasks.append(asyncio.create_task(
-                self._process_document("".join(texts), prompt_variables, res)
-            ))
+            tasks.append(asyncio.create_task(self._process_document("".join(texts), prompt_variables, res)))
         try:
             await asyncio.gather(*tasks, return_exceptions=False)
         except Exception as e:
@@ -119,16 +107,7 @@ class MindMapExtractor(Extractor):
         if len(merge_json) > 1:
             keys = [re.sub(r"\*+", "", k) for k, v in merge_json.items() if isinstance(v, dict)]
             keyset = set(i for i in keys if i)
-            merge_json = {
-                "id": "root",
-                "children": [
-                    {
-                        "id": self._key(k),
-                        "children": self._be_children(v, keyset)
-                    }
-                    for k, v in merge_json.items() if isinstance(v, dict) and self._key(k)
-                ]
-            }
+            merge_json = {"id": "root", "children": [{"id": self._key(k), "children": self._be_children(v, keyset)} for k, v in merge_json.items() if isinstance(v, dict) and self._key(k)]}
         else:
             k = self._key(list(merge_json.keys())[0])
             merge_json = {"id": k, "children": self._be_children(list(merge_json.items())[0][1], {k})}
@@ -176,9 +155,7 @@ class MindMapExtractor(Extractor):
 
         return self._list_to_kv(to_ret)
 
-    async def _process_document(
-            self, text: str, prompt_variables: dict[str, str], out_res
-    ) -> str:
+    async def _process_document(self, text: str, prompt_variables: dict[str, str], out_res) -> str:
         variables = {
             **prompt_variables,
             self._input_text_key: text,

rag/advanced_rag/knowlege_compile/raptor.py

@@ -19,7 +19,6 @@ import logging
 import re
 
 import numpy as np
-import umap
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.mixture import GaussianMixture
 
@@ -54,6 +53,12 @@ class _PsiTreeNode:
     embedding: np.ndarray | None = None
     children: list["_PsiTreeNode"] = field(default_factory=list)
     parent: "_PsiTreeNode | None" = None
+    # Original (leaf-level) chunk ids that contributed to this node. On
+    # a leaf this is a single-element list with the leaf's own id; on an
+    # internal node it's the order-preserving deduped union of its
+    # children's lists. Carried up through the merge tree so each
+    # produced summary knows which source chunks it covers.
+    source_chunk_ids: list[str] = field(default_factory=list)
 
 
 class _PsiUnionFind:
@@ -150,7 +155,7 @@ class _PsiUnionFind:
     @property
     def tree(self) -> list[int]:
         """Return the compact child-to-parent array for constructed nodes."""
-        return self._tree[:self._next_id]
+        return self._tree[: self._next_id]
 
 
 class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
@@ -209,7 +214,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
                 response = re.sub(r"^.*</think>", "", response, flags=re.DOTALL)
                 if response.find("**ERROR**") >= 0:
                     raise Exception(response)
-                await thread_pool_exec(set_llm_cache,self._llm_model.llm_name,system,response,history,gen_conf)
+                await thread_pool_exec(set_llm_cache, self._llm_model.llm_name, system, response, history, gen_conf)
                 return response
             except Exception as exc:
                 last_exc = exc
@@ -305,6 +310,61 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
             labels = np.array([remap[int(lbl)] for lbl in new_labels])
         return labels
 
+    def clustering(self, embeddings, random_state: int, task_id: str = "") -> tuple[int, list[int]]:
+        """Cluster one RAPTOR layer and return contiguous labels."""
+        reduced_embeddings = np.asarray(embeddings, dtype=np.float64)
+        if len(reduced_embeddings) == 0:
+            return 0, []
+
+        # Degrade too much ??
+        n_neighbors = int((len(embeddings) - 1) ** 0.8)
+        import umap
+
+        reduced_embeddings = umap.UMAP(
+            n_neighbors=max(2, n_neighbors),
+            n_components=min(12, len(embeddings) - 2),
+            metric="cosine",
+        ).fit_transform(embeddings)
+        if self._clustering_method == AHC_CLUSTERING_METHOD:
+            logging.info("RAPTOR: using clustering_method=%s before _get_clusters_ahc", self._clustering_method)
+            raw_labels = self._get_clusters_ahc(reduced_embeddings, task_id=task_id)
+            raw_cluster_count = np.unique(raw_labels).size
+            logging.info("RAPTOR AHC: _get_clusters_ahc produced n_clusters=%d", raw_cluster_count)
+            if raw_cluster_count > 1:
+                labels = self._adjust_tree_nodes(reduced_embeddings, raw_labels)
+                adjusted_cluster_count = np.unique(labels).size
+                logging.info("RAPTOR AHC: _adjust_tree_nodes adjusted n_clusters=%d", adjusted_cluster_count)
+            else:
+                labels = raw_labels
+                logging.warning("RAPTOR AHC: _adjust_tree_nodes skipped because _get_clusters_ahc returned one cluster")
+        else:
+            n_clusters = int(self._get_optimal_clusters(reduced_embeddings, random_state, task_id=task_id))
+            if n_clusters <= 1:
+                labels = [0 for _ in range(len(reduced_embeddings))]
+            else:
+                gm = GaussianMixture(n_components=n_clusters, random_state=random_state)
+                gm.fit(reduced_embeddings)
+                probs = gm.predict_proba(reduced_embeddings)
+                labels = []
+                for prob in probs:
+                    candidates = np.where(prob > self._threshold)[0]
+                    labels.append(int(candidates[0]) if len(candidates) else int(np.argmax(prob)))
+
+        normalized_labels: list[int] = []
+        for label in labels:
+            if isinstance(label, np.ndarray):
+                normalized_labels.append(int(label[0]) if len(label) else 0)
+            else:
+                normalized_labels.append(int(label))
+
+        if len(normalized_labels) <= 0:
+            return 0, []
+        unique_labels = np.unique(normalized_labels)
+        if len(unique_labels) <= 1:
+            return 1, [0 for _ in normalized_labels]
+        label_map = {int(old): idx for idx, old in enumerate(unique_labels)}
+        return len(unique_labels), [label_map[label] for label in normalized_labels]
+
     @timeout(60 * 20)
     async def _summarize_texts(self, texts: list[str], callback=None, task_id: str = ""):
         """Summarize a cluster and return text plus embedding when successful."""
@@ -317,11 +377,11 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
                 self._check_task_canceled(task_id, "before LLM call")
 
                 cnt = await self._chat(
-                    "You're a helpful assistant.",
+                    "You're a helpful assistant.\n\nHelp me with the following task.\n\n%s" % self._prompt.format(cluster_content=cluster_content),
                     [
                         {
                             "role": "user",
-                            "content": self._prompt.format(cluster_content=cluster_content),
+                            "content": "Beside the summarization, give a title at the first line of your summarization. Must be in the same language as the paragraphs.",
                         }
                     ],
                     {"max_tokens": max(self._max_token, 512)},  # fix issue:  #10235
@@ -336,7 +396,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
                 self._check_task_canceled(task_id, "before embedding")
 
                 embds = await self._embedding_encode(cnt)
-                return cnt, embds
+                return cnt.split("\n")[0], cnt, embds
         except TaskCanceledException:
             raise
         except Exception as exc:
@@ -406,10 +466,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
             split_groups = [group[labels == center_id].tolist() for center_id in range(fanout)]
             split_groups = [bucket for bucket in split_groups if bucket]
             if len(split_groups) <= 1:
-                split_groups = [
-                    group[start:start + self._psi_bucket_size].tolist()
-                    for start in range(0, len(group), self._psi_bucket_size)
-                ]
+                split_groups = [group[start : start + self._psi_bucket_size].tolist() for start in range(0, len(group), self._psi_bucket_size)]
             groups.extend(split_groups)
 
         buckets = [bucket for bucket in buckets if bucket]
@@ -455,7 +512,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
                 original_children = len(node.children)
                 grouped_children = []
                 for start in range(0, len(node.children), max_children):
-                    batch = node.children[start:start + max_children]
+                    batch = node.children[start : start + max_children]
                     if len(batch) == 1:
                         grouped_children.append(batch[0])
                         batch[0].parent = node
@@ -561,10 +618,21 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
         return self._build_bucketed_psi_structure(nodes, next_index, task_id)
 
     def _build_psi_structure(self, chunks, task_id: str = "") -> tuple[_PsiTreeNode, list[_PsiTreeNode]]:
-        """Build the Psi merge tree from original chunk embeddings."""
+        """Build the Psi merge tree from original chunk embeddings.
+
+        ``chunks`` is expected in the normalized 3-tuple shape
+        ``(text, vec, source_chunk_ids)`` — leaves are seeded with
+        their own source ids, internal nodes get their ids set during
+        layer materialization in ``_build_psi_layers``.
+        """
         leaves = [
-            _PsiTreeNode(index=i, text=text, embedding=np.asarray(embd))
-            for i, (text, embd) in enumerate(chunks)
+            _PsiTreeNode(
+                index=i,
+                text=item[0],
+                embedding=np.asarray(item[1]),
+                source_chunk_ids=list(item[2] if len(item) > 2 else []),
+            )
+            for i, item in enumerate(chunks)
         ]
         if len(leaves) == 1:
             return leaves[0], leaves
@@ -604,7 +672,16 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
             layer_start = len(chunks)
 
             async def summarize_node(node: _PsiTreeNode):
-                """Summarize one Psi internal node if its children have text."""
+                """Summarize one Psi internal node if its children have text.
+
+                Also propagates leaf provenance: the node's
+                ``source_chunk_ids`` becomes the order-preserving deduped
+                union of every child's ``source_chunk_ids``. Because
+                children at this layer have already been processed (leaves
+                first, then bottom-up), each child carries the full set
+                of leaf ids underneath it — so the union here is the
+                complete leaf set this summary covers.
+                """
                 texts = [child.text for child in node.children if child.text]
                 if not texts:
                     logging.warning("RAPTOR Psi node %s skipped because it has no child text to summarize", node.index)
@@ -613,7 +690,15 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
                 if result is None:
                     logging.warning("RAPTOR Psi node %s skipped because summarization failed", node.index)
                     return None
-                node.text, node.embedding = result
+                _, node.text, node.embedding = result
+                merged_ids: list[str] = []
+                seen: set[str] = set()
+                for child in node.children:
+                    for src in child.source_chunk_ids:
+                        if src and src not in seen:
+                            seen.add(src)
+                            merged_ids.append(src)
+                node.source_chunk_ids = merged_ids
                 return node
 
             tasks = [asyncio.create_task(summarize_node(node)) for node in nodes]
@@ -628,7 +713,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
 
             summarized_nodes = [node for node in summarized_nodes if node is not None]
             for node in summarized_nodes:
-                chunks.append((node.text, node.embedding))
+                chunks.append((node.text, node.embedding, list(node.source_chunk_ids)))
 
             if len(chunks) > layer_start:
                 layers.append((layer_start, len(chunks)))
@@ -646,34 +731,119 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
 
         return chunks, layers
 
-    async def __call__(self, chunks, random_state, callback=None, task_id: str = ""):
-        """Build summary chunks and layer boundaries for RAPTOR retrieval."""
+    async def __call__(
+        self,
+        chunks,
+        random_state,
+        callback=None,
+        task_id: str = "",
+        is_tree: bool = False,
+    ):
+        """Build summary chunks and layer boundaries for RAPTOR retrieval.
+
+        ``chunks`` accepts either the legacy 2-tuple shape
+        ``(text, vec)`` or the provenance-carrying 3-tuple shape
+        ``(text, vec, source_chunk_ids)`` where ``source_chunk_ids`` is
+        the list of original chunk ids that produced this entry. Output
+        always uses the 3-tuple shape so every appended summary carries
+        its leaves' ids. ``[]`` is left in the slot for a leaf whose id
+        was missing — see the caller for the normalization rules.
+
+        Return shapes:
+          * ``is_tree=False`` (default) — original behavior: returns
+            ``(chunks, layers)`` where ``chunks`` is the flat list
+            (originals + summaries) and ``layers`` is the per-level
+            index range ``[(start, end), ...]``.
+          * ``is_tree=True`` — returns a hierarchical tree dict via
+            ``_materialize_tree``. Supported for the classic builder
+            only; raises ``NotImplementedError`` for PSI_TREE_BUILDER
+            (PSI's hyperedge-driven summarization doesn't form a strict
+            parent-of relation). Returns ``None`` when there's nothing
+            to materialize.
+        """
         if len(chunks) <= 1:
-            return [], []
-        chunks = [(s, a) for s, a in chunks if s and a is not None and len(a) > 0]
-        if len(chunks) <= 1:
-            return chunks, [(0, len(chunks))]
+            return None if is_tree else ([], [])
+
+        # Normalize input to the 3-tuple shape. Reject empties / bad
+        # vectors at the same time the legacy path used to.
+        def _normalize(item):
+            if len(item) >= 3:
+                text, vec, src = item[0], item[1], item[2]
+            else:
+                text, vec = item[0], item[1]
+                src = []
+            if not text or vec is None or len(vec) <= 0:
+                return None
+            # Defensive: a leaf should carry a list of strings. Drop
+            # falsy entries so we don't propagate empty ids upward.
+            if isinstance(src, (list, tuple)):
+                src = [s for s in src if s]
+            else:
+                src = [src] if src else []
+            return (text, vec, list(src), "")
+
+        normalized = [t for t in (_normalize(c) for c in chunks) if t is not None]
+        if len(normalized) <= 1:
+            return None if is_tree else (normalized, [(0, len(normalized))])
+        chunks = normalized
+
         if self._tree_builder == PSI_TREE_BUILDER:
+            if is_tree:
+                raise NotImplementedError(
+                    "is_tree=True is not supported for PSI_TREE_BUILDER",
+                )
             logging.info("RAPTOR: using %s tree builder for %d chunks", self._tree_builder, len(chunks))
             return await self._build_psi_layers(chunks, callback, task_id)
 
+        # ``parent_child_map`` records each summary's immediate
+        # children so ``_materialize_tree`` can walk back into a tree
+        # when ``is_tree`` is set. Always populated (cheap) so the
+        # tree path is just a return-shape choice at the end.
+        parent_child_map: dict[int, list[int]] = {}
+        n_originals = len(chunks)
+
         layers = [(0, len(chunks))]
         start, end = 0, len(chunks)
 
         @timeout(60 * 20)
         async def summarize(ck_idx: list[int]):
-            """Summarize one classic RAPTOR cluster into the chunk list."""
+            """Summarize one classic RAPTOR cluster into the chunk list.
+
+            On success appends ``(summary_text, summary_vec, src_ids)``
+            where ``src_ids`` is the order-preserving deduped union of
+            the ``source_chunk_ids`` of every chunk indexed in
+            ``ck_idx`` — i.e. the full leaf set that contributed to
+            the cluster, even through nested summaries.
+            """
             nonlocal chunks
 
             texts = [chunks[i][0] for i in ck_idx]
             result = await self._summarize_texts(texts, callback, task_id)
             if result is not None:
-                chunks.append(result)
+                # ``dict.fromkeys`` is the cheapest way to de-dup a
+                # list of strings while preserving first-seen order.
+                merged_ids: list[str] = []
+                seen: set[str] = set()
+                for i in ck_idx:
+                    for src in chunks[i][2]:
+                        if src and src not in seen:
+                            seen.add(src)
+                            merged_ids.append(src)
+                summary_ti, summary_text, summary_vec = result
+                chunks.append((summary_text, summary_vec, merged_ids, summary_ti))
+                # Index of the just-appended summary; map it to its
+                # immediate children for the tree materializer below.
+                parent_child_map[len(chunks) - 1] = list(ck_idx)
 
         while end - start > 1:
             self._check_task_canceled(task_id, "layer processing")
 
-            embeddings = [embd for _, embd in chunks[start:end]]
+            # ``chunks`` is a mix of 3-tuples (layer-0 originals from
+            # _normalize) and 4-tuples (summaries appended by
+            # summarize). Vector is always at index 1 in both shapes,
+            # so use positional access — the older ``_, embd, _, _``
+            # form crashed on layer-0 entries.
+            embeddings = [entry[1] for entry in chunks[start:end]]
             if len(embeddings) == 2:
                 await summarize([start, start + 1])
                 produced = len(chunks) - end
@@ -687,43 +857,34 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
                 end = len(chunks)
                 continue
 
-            n_neighbors = int((len(embeddings) - 1) ** 0.8)
-            reduced_embeddings = umap.UMAP(
-                n_neighbors=max(2, n_neighbors),
-                n_components=min(12, len(embeddings) - 2),
-                metric="cosine",
-            ).fit_transform(embeddings)
-            if self._clustering_method == AHC_CLUSTERING_METHOD:
-                logging.info("RAPTOR: using clustering_method=%s before _get_clusters_ahc", self._clustering_method)
-                raw_labels = self._get_clusters_ahc(reduced_embeddings, task_id=task_id)
-                raw_cluster_count = np.unique(raw_labels).size
-                logging.info("RAPTOR AHC: _get_clusters_ahc produced n_clusters=%d", raw_cluster_count)
-                if raw_cluster_count > 1:
-                    adjusted = self._adjust_tree_nodes(reduced_embeddings, raw_labels)
-                    adjusted_cluster_count = np.unique(adjusted).size
-                    logging.info("RAPTOR AHC: _adjust_tree_nodes adjusted n_clusters=%d", adjusted_cluster_count)
-                else:
-                    adjusted = raw_labels
-                    logging.warning("RAPTOR AHC: _adjust_tree_nodes skipped because _get_clusters_ahc returned one cluster")
-                unique_labels = np.unique(adjusted)
-                label_map = {old: idx for idx, old in enumerate(unique_labels)}
-                lbls = [label_map[int(lbl)] for lbl in adjusted]
-                n_clusters = len(unique_labels)
-            else:
-                n_clusters = self._get_optimal_clusters(reduced_embeddings, random_state, task_id=task_id)
-                if n_clusters == 1:
-                    lbls = [0 for _ in range(len(reduced_embeddings))]
-                else:
-                    gm = GaussianMixture(n_components=n_clusters, random_state=random_state)
-                    gm.fit(reduced_embeddings)
-                    probs = gm.predict_proba(reduced_embeddings)
-                    lbls = [np.where(prob > self._threshold)[0] for prob in probs]
-                    lbls = [lbl[0] if isinstance(lbl, np.ndarray) else lbl for lbl in lbls]
-
-            if n_clusters == 1:
-                lbls = [0 for _ in range(len(reduced_embeddings))]
-            else:
-                lbls = [int(lbl[0]) if isinstance(lbl, np.ndarray) else int(lbl) for lbl in lbls]
+            n_clusters, lbls = self.clustering(
+                embeddings,
+                random_state=random_state,
+                task_id=task_id,
+            )
+            
+            # Loop-termination guarantee. The outer ``while end - start > 1``
+            # relies on each layer strictly shrinking the input count. If
+            # the clusterer degenerates and returns one cluster per input,
+            # every "cluster" is a single chunk, ``summarize()`` produces
+            # one summary per input, and ``produced == end - start`` —
+            # the same count carries into the next iteration and the loop
+            # spins forever, logging "Cluster one layer: N -> N".
+            #
+            # Collapse everything at this level into a single cluster so
+            # the layer produces exactly one summary. The tree gets a
+            # taller-than-usual "single trunk" segment at this depth
+            # instead of an infinite loop; downstream consumers only care
+            # that ``layers`` is monotonically shrinking.
+            if n_clusters >= len(embeddings):
+                logging.warning(
+                    "RAPTOR clustering did not reduce input count "
+                    "(%d inputs → %d clusters); collapsing this layer "
+                    "into a single summary to prevent a non-terminating loop",
+                    len(embeddings), n_clusters,
+                )
+                n_clusters = 1
+                lbls = [0] * len(embeddings)
 
             tasks = []
             for c in range(n_clusters):
@@ -758,4 +919,52 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
             start = end
             end = len(chunks)
 
+        if is_tree:
+            return self._materialize_tree(chunks, layers, parent_child_map, n_originals), []
         return chunks, layers
+
+    @staticmethod
+    def _materialize_tree(chunks, layers, parent_child_map, n_originals):
+        """Walk ``parent_child_map`` from the top layer down to layer-1
+        and emit the user-facing tree dict. See ``__call__``'s
+        ``is_tree=True`` contract for the shape.
+        chunks: [(summary_text, summary_vec, merged_ids, summary_ti)]"""
+        if not layers or len(chunks) == 0:
+            return None
+        top_start, top_end = layers[-1]
+        if top_end <= top_start:
+            return None
+
+        def _title_at(idx: int) -> str:
+            # Summary tuples are (text, vec, merged_ids, summary_ti)
+            # — title is the 4th slot. Layer-0 originals are 3-tuples
+            # and don't appear as tree nodes themselves (they collapse
+            # into source_chunk_ids on their layer-1 parent).
+            return chunks[idx][3] if len(chunks[idx]) >= 4 else ""
+
+        def _desc_at(idx: int) -> str:
+            return chunks[idx][0] if chunks[idx] else ""
+
+        def _build_node(idx: int) -> dict:
+            children_idx = parent_child_map.get(idx, [])
+            # If every immediate child is a layer-0 original, this
+            # node is a "leaf" in the tree contract — collapse to
+            # source_chunk_ids.
+            if children_idx and all(c < n_originals for c in children_idx):
+                ids: list[str] = []
+                seen: set[str] = set()
+                for c in children_idx:
+                    for s in chunks[c][2]:
+                        if s and s not in seen:
+                            seen.add(s)
+                            ids.append(s)
+                return {"title": _title_at(idx), "source_chunk_ids": ids, "description": _desc_at(idx)}
+            return {"children": [_build_node(c) for c in children_idx], "title": _title_at(idx), "description": _desc_at(idx)}
+
+        top_nodes = [_build_node(i) for i in range(top_start, top_end)]
+        if len(top_nodes) == 1:
+            return top_nodes[0]
+        # Multiple top-layer summaries — clustering didn't collapse to
+        # a single root. Wrap in a synthetic root so the caller always
+        # sees one dict.
+        return {"title": "(root)", "children": top_nodes}

rag/flow/extractor/extractor.py

@@ -17,9 +17,16 @@ import logging
 import random
 from copy import deepcopy
 
+from api.db.services.document_service import DocumentService
+from api.db.services.llm_service import LLMBundle
+from common.constants import LLMType
 import xxhash
 
 from agent.component.llm import LLMParam, LLM
+from rag.advanced_rag.knowlege_compile.structure import (
+    compile_structure_from_text,
+    merge_compiled_structures,
+)
 from rag.flow.base import ProcessBase, ProcessParamBase
 from rag.prompts.generator import run_toc_from_text
 
@@ -28,6 +35,7 @@ class ExtractorParam(ProcessParamBase, LLMParam):
     def __init__(self):
         super().__init__()
         self.field_name = ""
+        self.knowledge_compilation = {}
 
     def check(self):
         super().check()
@@ -38,22 +46,25 @@ class Extractor(ProcessBase, LLM):
     component_name = "Extractor"
 
     async def _build_TOC(self, docs):
-        self.callback(0.2,message="Start to generate table of content ...")
-        docs = sorted(docs, key=lambda d:(
-            d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
-            d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0)
-        ))
+        self.callback(0.2, message="Start to generate table of content ...")
+        docs = sorted(
+            docs,
+            key=lambda d: (
+                d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
+                d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0),
+            ),
+        )
         toc = await run_toc_from_text([d["text"] for d in docs], self.chat_mdl)
-        logging.info("------------ T O C -------------\n"+json.dumps(toc, ensure_ascii=False, indent='  '))
+        logging.info("------------ T O C -------------\n" + json.dumps(toc, ensure_ascii=False, indent="  "))
         ii = 0
         while ii < len(toc):
             try:
                 idx = int(toc[ii]["chunk_id"])
                 del toc[ii]["chunk_id"]
                 toc[ii]["ids"] = [docs[idx]["id"]]
-                if ii == len(toc) -1:
+                if ii == len(toc) - 1:
                     break
-                for jj in range(idx+1, int(toc[ii+1]["chunk_id"])+1):
+                for jj in range(idx + 1, int(toc[ii + 1]["chunk_id"]) + 1):
                     toc[ii]["ids"].append(docs[jj]["id"])
             except Exception as e:
                 logging.exception(e)
@@ -71,6 +82,20 @@ class Extractor(ProcessBase, LLM):
             return d
         return None
 
+    async def _knowledge_compile(self, docs):
+        embedding_model = LLMBundle(self._canvas.get_tenant_id(), LLMType.EMBEDDING, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error)
+        self.callback(0.2, message="Start to generate table of content ...")
+        docs = sorted(
+            docs,
+            key=lambda d: (
+                d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
+                d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0),
+            ),
+        )
+        docs = await compile_structure_from_text(docs, self._param.knowledge_compilation, self.chat_mdl, embedding_model, self._canvas._doc_id)
+        info = await merge_compiled_structures(docs, self.chat_mdl, embedding_model, self._canvas.get_tenant_id(), DocumentService.get_knowledgebase_id(self._canvas._doc_id))
+        return info
+
     async def _invoke(self, **kwargs):
         self.set_output("output_format", "chunks")
         self.callback(random.randint(1, 5) / 100.0, "Start to generate.")
@@ -89,10 +114,17 @@ class Extractor(ProcessBase, LLM):
                 for ck in chunks:
                     ck["doc_id"] = self._canvas._doc_id
                     ck["id"] = xxhash.xxh64((ck["text"] + str(ck["doc_id"])).encode("utf-8")).hexdigest()
-                toc =await self._build_TOC(chunks)
+                toc = await self._build_TOC(chunks)
                 chunks.append(toc)
                 self.set_output("chunks", chunks)
                 return
+            if self._param.field_name in ["set", "list", "graph"]:
+                for ck in chunks:
+                    ck["doc_id"] = self._canvas._doc_id
+                    ck["id"] = xxhash.xxh64((ck["text"] + str(ck["doc_id"])).encode("utf-8")).hexdigest()
+                await self._knowledge_compile(chunks)
+                self.set_output("chunks", chunks)
+                return
 
             prog = 0
             for i, ck in enumerate(chunks):
@@ -100,12 +132,11 @@ class Extractor(ProcessBase, LLM):
                 msg, sys_prompt = self._sys_prompt_and_msg([], args)
                 msg.insert(0, {"role": "system", "content": sys_prompt})
                 ck[self._param.field_name] = await self._generate_async(msg)
-                prog += 1./len(chunks)
-                if i % (len(chunks)//100+1) == 1:
-                    self.callback(prog, f"{i+1} / {len(chunks)}")
+                prog += 1.0 / len(chunks)
+                if i % (len(chunks) // 100 + 1) == 1:
+                    self.callback(prog, f"{i + 1} / {len(chunks)}")
             self.set_output("chunks", chunks)
         else:
             msg, sys_prompt = self._sys_prompt_and_msg([], args)
             msg.insert(0, {"role": "system", "content": sys_prompt})
             self.set_output("chunks", [{self._param.field_name: await self._generate_async(msg)}])
-

rag/nlp/search.py

@@ -31,7 +31,9 @@ from common import settings
 
 from common.misc_utils import thread_pool_exec
 
-def index_name(uid): return f"ragflow_{uid}"
+
+def index_name(uid):
+    return f"ragflow_{uid}"
 
 
 class Dealer:
@@ -54,11 +56,10 @@ class Dealer:
         qv, _ = await thread_pool_exec(emb_mdl.encode_queries, txt)
         shape = np.array(qv).shape
         if len(shape) > 1:
-            raise Exception(
-                f"Dealer.get_vector returned array's shape {shape} doesn't match expectation(exact one dimension).")
+            raise Exception(f"Dealer.get_vector returned array's shape {shape} doesn't match expectation(exact one dimension).")
         embedding_data = [get_float(v) for v in qv]
         vector_column_name = f"q_{len(embedding_data)}_vec"
-        return MatchDenseExpr(vector_column_name, embedding_data, 'float', 'cosine', topk, {"similarity": similarity})
+        return MatchDenseExpr(vector_column_name, embedding_data, "float", "cosine", topk, {"similarity": similarity})
 
     async def _existing_doc_ids(self, doc_ids: list[str]) -> set[str]:
         if not doc_ids:
@@ -123,18 +124,14 @@ class Dealer:
             if key in req and req[key] is not None:
                 condition[field] = req[key]
         # TODO(yzc): `available_int` is nullable however infinity doesn't support nullable columns.
-        for key in ["knowledge_graph_kwd", "available_int", "entity_kwd", "from_entity_kwd", "to_entity_kwd",
-                    "removed_kwd"]:
+        for key in ["id", "knowledge_graph_kwd", "available_int", "entity_kwd", "from_entity_kwd", "to_entity_kwd", "removed_kwd"]:
             if key in req and req[key] is not None:
                 condition[key] = req[key]
+        if isinstance(req.get("must_not"), dict):
+            condition["must_not"] = req["must_not"]
         return condition
 
-    async def search(self, req, idx_names: str | list[str],
-               kb_ids: list[str],
-               emb_mdl=None,
-               highlight: bool | list | None = None,
-               rank_feature: dict | None = None
-               ):
+    async def search(self, req, idx_names: str | list[str], kb_ids: list[str], emb_mdl=None, highlight: bool | list | None = None, rank_feature: dict | None = None):
         if highlight is None:
             highlight = False
 
@@ -146,11 +143,33 @@ class Dealer:
         ps = int(req.get("size", topk))
         offset, limit = pg * ps, ps
 
-        src = req.get("fields",
-                      ["docnm_kwd", "content_ltks", "kb_id", "img_id", "title_tks", "important_kwd", "position_int",
-                       "doc_id", "chunk_order_int", "page_num_int", "top_int", "create_timestamp_flt", "knowledge_graph_kwd",
-                       "question_kwd", "question_tks", "doc_type_kwd",
-                       "available_int", "content_with_weight", "mom_id", PAGERANK_FLD, TAG_FLD, "row_id()"])
+        src = req.get(
+            "fields",
+            [
+                "docnm_kwd",
+                "content_ltks",
+                "kb_id",
+                "img_id",
+                "title_tks",
+                "important_kwd",
+                "position_int",
+                "doc_id",
+                "chunk_order_int",
+                "page_num_int",
+                "top_int",
+                "create_timestamp_flt",
+                "knowledge_graph_kwd",
+                "question_kwd",
+                "question_tks",
+                "doc_type_kwd",
+                "available_int",
+                "content_with_weight",
+                "mom_id",
+                PAGERANK_FLD,
+                TAG_FLD,
+                "row_id()",
+            ],
+        )
         kwds = set([])
 
         qst = req.get("question", "")
@@ -173,8 +192,7 @@ class Dealer:
             matchText, keywords = self.qryr.question(qst, min_match=0.3)
             if emb_mdl is None:
                 matchExprs = [matchText]
-                res = await thread_pool_exec(self.dataStore.search, src, highlightFields, filters, matchExprs, orderBy, offset, limit,
-                                            idx_names, kb_ids, rank_feature=rank_feature)
+                res = await thread_pool_exec(self.dataStore.search, src, highlightFields, filters, matchExprs, orderBy, offset, limit, idx_names, kb_ids, rank_feature=rank_feature)
                 total = self.dataStore.get_total(res)
                 logging.debug("Dealer.search TOTAL: {}".format(total))
             else:
@@ -192,8 +210,7 @@ class Dealer:
                 fusionExpr = FusionExpr("weighted_sum", topk, {"weights": "0.05,0.95"})
                 matchExprs = [matchText, matchDense, fusionExpr]
 
-                res = await thread_pool_exec(self.dataStore.search, src, highlightFields, filters, matchExprs, orderBy, offset, limit,
-                                            idx_names, kb_ids, rank_feature=rank_feature)
+                res = await thread_pool_exec(self.dataStore.search, src, highlightFields, filters, matchExprs, orderBy, offset, limit, idx_names, kb_ids, rank_feature=rank_feature)
                 total = self.dataStore.get_total(res)
                 logging.debug("Dealer.search TOTAL: {}".format(total))
 
@@ -205,9 +222,9 @@ class Dealer:
                     else:
                         matchText, _ = self.qryr.question(qst, min_match=0.1)
                         matchDense.extra_options["similarity"] = 0.17
-                        res = await thread_pool_exec(self.dataStore.search, src, highlightFields, filters, [matchText, matchDense, fusionExpr],
-                                                    orderBy, offset, limit, idx_names, kb_ids,
-                                                    rank_feature=rank_feature)
+                        res = await thread_pool_exec(
+                            self.dataStore.search, src, highlightFields, filters, [matchText, matchDense, fusionExpr], orderBy, offset, limit, idx_names, kb_ids, rank_feature=rank_feature
+                        )
                         total = self.dataStore.get_total(res)
                     logging.debug("Dealer.search 2 TOTAL: {}".format(total))
 
@@ -225,22 +242,13 @@ class Dealer:
         keywords = list(kwds)
         highlight = self.dataStore.get_highlight(res, keywords, "content_with_weight")
         aggs = self.dataStore.get_aggregation(res, "docnm_kwd")
-        return self.SearchResult(
-            total=total,
-            ids=ids,
-            query_vector=q_vec,
-            aggregation=aggs,
-            highlight=highlight,
-            field=self.dataStore.get_fields(res, src + ["_score"]),
-            keywords=keywords
-        )
+        return self.SearchResult(total=total, ids=ids, query_vector=q_vec, aggregation=aggs, highlight=highlight, field=self.dataStore.get_fields(res, src + ["_score"]), keywords=keywords)
 
     @staticmethod
     def trans2floats(txt):
         return [get_float(t) for t in txt.split("\t")]
 
-    def insert_citations(self, answer, chunks, chunk_v,
-                         embd_mdl, tkweight=0.1, vtweight=0.9):
+    def insert_citations(self, answer, chunks, chunk_v, embd_mdl, tkweight=0.1, vtweight=0.9):
         assert len(chunks) == len(chunk_v)
         if not chunks:
             return answer, set([])
@@ -256,13 +264,10 @@ class Dealer:
                         i += 1
                     if i < len(pieces):
                         i += 1
-                    pieces_.append("".join(pieces[st: i]) + "\n")
+                    pieces_.append("".join(pieces[st:i]) + "\n")
                 else:
                     # Sentence boundary regex includes Arabic punctuation (، ؛ ؟ ۔)
-                    pieces_.extend(
-                        re.split(
-                            r"([^\|][；。？!！،؛؟۔\n]|[a-z\u0600-\u06FF][.?;!،؛؟][ \n])",
-                            pieces[i]))
+                    pieces_.extend(re.split(r"([^\|][；。？!！،؛؟۔\n]|[a-z\u0600-\u06FF][.?;!،؛؟][ \n])", pieces[i]))
                     i += 1
             pieces = pieces_
         else:
@@ -287,30 +292,21 @@ class Dealer:
         for i in range(len(chunk_v)):
             if len(ans_v[0]) != len(chunk_v[i]):
                 chunk_v[i] = [0.0] * len(ans_v[0])
-                logging.warning(
-                    "The dimension of query and chunk do not match: {} vs. {}".format(len(ans_v[0]), len(chunk_v[i])))
+                logging.warning("The dimension of query and chunk do not match: {} vs. {}".format(len(ans_v[0]), len(chunk_v[i])))
 
-        assert len(ans_v[0]) == len(chunk_v[0]), "The dimension of query and chunk do not match: {} vs. {}".format(
-            len(ans_v[0]), len(chunk_v[0]))
+        assert len(ans_v[0]) == len(chunk_v[0]), "The dimension of query and chunk do not match: {} vs. {}".format(len(ans_v[0]), len(chunk_v[0]))
 
-        chunks_tks = [rag_tokenizer.tokenize(self.qryr.rmWWW(ck)).split()
-                      for ck in chunks]
+        chunks_tks = [rag_tokenizer.tokenize(self.qryr.rmWWW(ck)).split() for ck in chunks]
         cites = {}
         thr = 0.63
         while thr > 0.3 and len(cites.keys()) == 0 and pieces_ and chunks_tks:
             for i, a in enumerate(pieces_):
-                sim, tksim, vtsim = self.qryr.hybrid_similarity(ans_v[i],
-                                                                chunk_v,
-                                                                rag_tokenizer.tokenize(
-                                                                    self.qryr.rmWWW(pieces_[i])).split(),
-                                                                chunks_tks,
-                                                                tkweight, vtweight)
+                sim, tksim, vtsim = self.qryr.hybrid_similarity(ans_v[i], chunk_v, rag_tokenizer.tokenize(self.qryr.rmWWW(pieces_[i])).split(), chunks_tks, tkweight, vtweight)
                 mx = np.max(sim) * 0.99
                 logging.debug("{} SIM: {}".format(pieces_[i], mx))
                 if mx < thr:
                     continue
-                cites[idx[i]] = list(
-                    set([str(ii) for ii in range(len(chunk_v)) if sim[ii] > mx]))[:4]
+                cites[idx[i]] = list(set([str(ii) for ii in range(len(chunk_v)) if sim[ii] > mx]))[:4]
             thr *= 0.8
 
         res = ""
@@ -362,11 +358,9 @@ class Dealer:
                 rank_fea.append(0)
             else:
                 rank_fea.append(nor / np.sqrt(denor) / q_denor)
-        return np.array(rank_fea) * 10. + pageranks
+        return np.array(rank_fea) * 10.0 + pageranks
 
-    async def _knn_scores(self, sres: "Dealer.SearchResult",
-                          idx_names: str | list[str],
-                          kb_ids: list[str]) -> dict[str, float]:
+    async def _knn_scores(self, sres: "Dealer.SearchResult", idx_names: str | list[str], kb_ids: list[str]) -> dict[str, float]:
         """
         Second-pass ES call that returns the cosine similarity between the
         query embedding and each candidate chunk's embedding, filtered to the
@@ -399,10 +393,7 @@ class Dealer:
         )
         return self.dataStore.get_scores(res)
 
-    async def fetch_chunk_vectors(self, chunk_ids: list[str],
-                                  tenant_ids: str | list[str],
-                                  kb_ids: list[str],
-                                  dim: int) -> dict[str, list[float]]:
+    async def fetch_chunk_vectors(self, chunk_ids: list[str], tenant_ids: str | list[str], kb_ids: list[str], dim: int) -> dict[str, list[float]]:
         """
         Citation-time helper: fetch only the embedding vectors for an
         explicit set of chunk ids. Used by callers that need to compute
@@ -440,10 +431,7 @@ class Dealer:
             out[cid] = v
         return out
 
-    def rerank_with_knn(self, sres, query, knn_scores: dict[str, float],
-                        tkweight=0.3, vtweight=0.7,
-                        cfield="content_ltks",
-                        rank_feature: dict | None = None):
+    def rerank_with_knn(self, sres, query, knn_scores: dict[str, float], tkweight=0.3, vtweight=0.7, cfield="content_ltks", rank_feature: dict | None = None):
         """
         Merge ES-side KNN cosine similarity with locally computed term
         similarity using the user-configured weights. Replaces the older
@@ -465,16 +453,12 @@ class Dealer:
             ins_tw.append(tks)
 
         tksim = np.array(self.qryr.token_similarity(keywords, ins_tw), dtype=np.float64)
-        vtsim = np.array([knn_scores.get(chunk_id, 0.0) for chunk_id in sres.ids],
-                         dtype=np.float64)
+        vtsim = np.array([knn_scores.get(chunk_id, 0.0) for chunk_id in sres.ids], dtype=np.float64)
         rank_fea = self._rank_feature_scores(rank_feature, sres)
         sim = tkweight * tksim + vtweight * vtsim + rank_fea
         return sim, tksim, vtsim
 
-    def rerank(self, sres, query, tkweight=0.3,
-               vtweight=0.7, cfield="content_ltks",
-               rank_feature: dict | None = None
-               ):
+    def rerank(self, sres, query, tkweight=0.3, vtweight=0.7, cfield="content_ltks", rank_feature: dict | None = None):
         _, keywords = self.qryr.question(query)
         vector_size = len(sres.query_vector)
         vector_column = f"q_{vector_size}_vec"
@@ -503,16 +487,11 @@ class Dealer:
         ## For rank feature(tag_fea) scores.
         rank_fea = self._rank_feature_scores(rank_feature, sres)
 
-        sim, tksim, vtsim = self.qryr.hybrid_similarity(sres.query_vector,
-                                                        ins_embd,
-                                                        keywords,
-                                                        ins_tw, tkweight, vtweight)
+        sim, tksim, vtsim = self.qryr.hybrid_similarity(sres.query_vector, ins_embd, keywords, ins_tw, tkweight, vtweight)
 
         return sim + rank_fea, tksim, vtsim
 
-    def rerank_by_model(self, rerank_mdl, sres, query, tkweight=0.3,
-                        vtweight=0.7, cfield="content_ltks",
-                        rank_feature: dict | None = None):
+    def rerank_by_model(self, rerank_mdl, sres, query, tkweight=0.3, vtweight=0.7, cfield="content_ltks", rank_feature: dict | None = None):
         _, keywords = self.qryr.question(query)
 
         for i in sres.ids:
@@ -520,7 +499,7 @@ class Dealer:
                 sres.field[i]["important_kwd"] = [sres.field[i]["important_kwd"]]
         ins_tw = []
         for i in sres.ids:
-            #content_ltks = list(OrderedDict.fromkeys(sres.field[i][cfield].split()))
+            # content_ltks = list(OrderedDict.fromkeys(sres.field[i][cfield].split()))
             content_ltks = sres.field[i][cfield].split()
             title_tks = [t for t in sres.field[i].get("title_tks", "").split() if t]
             important_kwd = sres.field[i].get("important_kwd", [])
@@ -540,10 +519,7 @@ class Dealer:
         return tkweight * np.array(tksim) + vtweight * vtsim + rank_fea, tksim, vtsim
 
     def hybrid_similarity(self, ans_embd, ins_embd, ans, inst):
-        return self.qryr.hybrid_similarity(ans_embd,
-                                           ins_embd,
-                                           rag_tokenizer.tokenize(ans).split(),
-                                           rag_tokenizer.tokenize(inst).split())
+        return self.qryr.hybrid_similarity(ans_embd, ins_embd, rag_tokenizer.tokenize(ans).split(), rag_tokenizer.tokenize(inst).split())
 
     @staticmethod
     def _rerank_window(page_size: int, top: int = 0) -> int:
@@ -571,22 +547,22 @@ class Dealer:
         return window
 
     async def retrieval(
-            self,
-            question,
-            embd_mdl,
-            tenant_ids,
-            kb_ids,
-            page,
-            page_size,
-            similarity_threshold=0.2,
-            vector_similarity_weight=0.3,
-            top=1024,
-            doc_ids=None,
-            aggs=True,
-            rerank_mdl=None,
-            highlight=False,
-            rank_feature: dict | None = {PAGERANK_FLD: 10},
-            trace_id=None,
+        self,
+        question,
+        embd_mdl,
+        tenant_ids,
+        kb_ids,
+        page,
+        page_size,
+        similarity_threshold=0.2,
+        vector_similarity_weight=0.3,
+        top=1024,
+        doc_ids=None,
+        aggs=True,
+        rerank_mdl=None,
+        highlight=False,
+        rank_feature: dict | None = {PAGERANK_FLD: 10},
+        trace_id=None,
     ):
         ranks = {"total": 0, "chunks": [], "doc_aggs": {}}
         if not question:
@@ -617,8 +593,7 @@ class Dealer:
             tenant_ids = tenant_ids.split(",")
 
         idx_names = [index_name(tid) for tid in tenant_ids]
-        sres = await self.search(req, idx_names, kb_ids, embd_mdl, highlight,
-                           rank_feature=rank_feature)
+        sres = await self.search(req, idx_names, kb_ids, embd_mdl, highlight, rank_feature=rank_feature)
         # Temporary retrieval-side guard: prune chunks whose parent document no
         # longer exists before reranking and returning results.
         sres = await self._prune_deleted_chunks(sres)
@@ -628,8 +603,7 @@ class Dealer:
 
         term_similarity_weight = 1 - vector_similarity_weight
         logging.debug(
-            "[Search] retrieval weights: trace_id=%s kb_count=%s similarity_threshold=%s "
-            "vector_similarity_weight=%s full_text_weight=%s rerank_enabled=%s",
+            "[Search] retrieval weights: trace_id=%s kb_count=%s similarity_threshold=%s vector_similarity_weight=%s full_text_weight=%s rerank_enabled=%s",
             trace_id,
             len(kb_ids),
             similarity_threshold,
@@ -685,7 +659,7 @@ class Dealer:
             return ranks
 
         # Use stable sort for deterministic ordering when scores are tied
-        sorted_idx = np.argsort(sim_np * -1, kind='stable')
+        sorted_idx = np.argsort(sim_np * -1, kind="stable")
 
         # When vector_similarity_weight is 0, similarity_threshold is not meaningful for term-only scores.
         post_threshold = 0.0 if vector_similarity_weight <= 0 else similarity_threshold
@@ -774,12 +748,17 @@ class Dealer:
         tbl = self.dataStore.sql(sql, fetch_size, format)
         return tbl
 
-    def chunk_list(self, doc_id: str, tenant_id: str,
-                   kb_ids: list[str], max_count=1024,
-                   offset=0,
-                   fields=["docnm_kwd", "content_with_weight", "img_id"],
-                   sort_by_position: bool = False,
-                   retrieve_all: bool = False):
+    def chunk_list(
+        self,
+        doc_id: str,
+        tenant_id: str,
+        kb_ids: list[str],
+        max_count=1024,
+        offset=0,
+        fields=["docnm_kwd", "content_with_weight", "img_id"],
+        sort_by_position: bool = False,
+        retrieve_all: bool = False,
+    ):
         """Return chunks for a document.
 
         By default, preserve the historical max_count cap. When retrieve_all is
@@ -807,8 +786,7 @@ class Dealer:
             limit = bs if retrieve_all else min(bs, max_count - p)
             if limit <= 0:
                 break
-            es_res = self.dataStore.search(fields, [], condition, [], orderBy, p, limit, index_name(tenant_id),
-                                           kb_ids)
+            es_res = self.dataStore.search(fields, [], condition, [], orderBy, p, limit, index_name(tenant_id), kb_ids)
             dict_chunks = self.dataStore.get_fields(es_res, fields)
             for id, doc in dict_chunks.items():
                 doc["id"] = id
@@ -834,15 +812,13 @@ class Dealer:
 
     def tag_content(self, tenant_id: str, kb_ids: list[str], doc, all_tags, topn_tags=3, keywords_topn=30, S=1000):
         idx_nm = index_name(tenant_id)
-        match_txt = self.qryr.paragraph(doc["title_tks"] + " " + doc["content_ltks"], doc.get("important_kwd", []),
-                                        keywords_topn)
+        match_txt = self.qryr.paragraph(doc["title_tks"] + " " + doc["content_ltks"], doc.get("important_kwd", []), keywords_topn)
         res = self.dataStore.search([], [], {}, [match_txt], OrderByExpr(), 0, 0, idx_nm, kb_ids, ["tag_kwd"])
         aggs = self.dataStore.get_aggregation(res, "tag_kwd")
         if not aggs:
             return False
         cnt = np.sum([c for _, c in aggs])
-        tag_fea = sorted([(a, round(0.1 * (c + 1) / (cnt + S) / max(1e-6, all_tags.get(a, 0.0001)))) for a, c in aggs],
-                         key=lambda x: x[1] * -1)[:topn_tags]
+        tag_fea = sorted([(a, round(0.1 * (c + 1) / (cnt + S) / max(1e-6, all_tags.get(a, 0.0001)))) for a, c in aggs], key=lambda x: x[1] * -1)[:topn_tags]
         doc[TAG_FLD] = {a.replace(".", "_"): c for a, c in tag_fea if c > 0}
         return True
 
@@ -857,12 +833,12 @@ class Dealer:
         if not aggs:
             return {}
         cnt = np.sum([c for _, c in aggs])
-        tag_fea = sorted([(a, round(0.1 * (c + 1) / (cnt + S) / max(1e-6, all_tags.get(a, 0.0001)))) for a, c in aggs],
-                         key=lambda x: x[1] * -1)[:topn_tags]
+        tag_fea = sorted([(a, round(0.1 * (c + 1) / (cnt + S) / max(1e-6, all_tags.get(a, 0.0001)))) for a, c in aggs], key=lambda x: x[1] * -1)[:topn_tags]
         return {a.replace(".", "_"): max(1, c) for a, c in tag_fea}
 
     async def retrieval_by_toc(self, query: str, chunks: list[dict], tenant_ids: list[str], chat_mdl, topn: int = 6):
-        from rag.prompts.generator import relevant_chunks_with_toc # moved from the top of the file to avoid circular import
+        from rag.prompts.generator import relevant_chunks_with_toc  # moved from the top of the file to avoid circular import
+
         if not chunks:
             return []
         idx_nms = [index_name(tid) for tid in tenant_ids]
@@ -872,11 +848,9 @@ class Dealer:
                 ranks[ck["doc_id"]] = 0
             ranks[ck["doc_id"]] += ck["similarity"]
             doc_id2kb_id[ck["doc_id"]] = ck["kb_id"]
-        doc_id = sorted(ranks.items(), key=lambda x: x[1] * -1.)[0][0]
+        doc_id = sorted(ranks.items(), key=lambda x: x[1] * -1.0)[0][0]
         kb_ids = [doc_id2kb_id[doc_id]]
-        es_res = self.dataStore.search(["content_with_weight"], [], {"doc_id": doc_id, "toc_kwd": "toc"}, [],
-                                       OrderByExpr(), 0, 128, idx_nms,
-                                       kb_ids)
+        es_res = self.dataStore.search(["content_with_weight"], [], {"doc_id": doc_id, "toc_kwd": "toc"}, [], OrderByExpr(), 0, 128, idx_nms, kb_ids)
         toc = []
         dict_chunks = self.dataStore.get_fields(es_res, ["content_with_weight"])
         for _, doc in dict_chunks.items():
@@ -914,7 +888,7 @@ class Dealer:
                 "term_similarity": sim,
                 "vector": [0.0] * vector_size,
                 "positions": chunk.get("position_int", []),
-                "doc_type_kwd": chunk.get("doc_type_kwd", "")
+                "doc_type_kwd": chunk.get("doc_type_kwd", ""),
             }
             for k in chunk.keys():
                 if k[-4:] == "_vec":
@@ -951,7 +925,8 @@ class Dealer:
             if chunk is None:
                 logging.warning(
                     "Parent chunk '%s' not found in the index; falling back to %d child chunk(s).",
-                    id, len(cks),
+                    id,
+                    len(cks),
                 )
                 chunks.extend(cks)
                 continue
@@ -969,7 +944,7 @@ class Dealer:
                 "term_similarity": np.mean([ck["similarity"] for ck in cks]),
                 "vector": [0.0] * vector_size,
                 "positions": chunk.get("position_int", []),
-                "doc_type_kwd": chunk.get("doc_type_kwd", "")
+                "doc_type_kwd": chunk.get("doc_type_kwd", ""),
             }
             for k in cks[0].keys():
                 if k[-4:] == "_vec":

rag/svr/task_executor.py

@@ -84,7 +84,7 @@ from common.versions import get_ragflow_version
 from api.db.db_models import close_connection
 from rag.app import laws, paper, presentation, manual, qa, table, book, resume, picture, naive, one, audio, email, tag
 from rag.nlp import search, rag_tokenizer, add_positions
-from rag.raptor import (
+from rag.advanced_rag.knowlege_compile.raptor import (
     RAPTOR_TREE_BUILDER,
 )
 from common.token_utils import num_tokens_from_string, truncate
@@ -136,6 +136,8 @@ TASK_TYPE_TO_PIPELINE_TASK_TYPE = {
     "graphrag": PipelineTaskType.GRAPH_RAG,
     "mindmap": PipelineTaskType.MINDMAP,
     "memory": PipelineTaskType.MEMORY,
+    "artifact": PipelineTaskType.ARTIFACT,
+    "skill": PipelineTaskType.SKILL,
 }
 
 UNACKED_ITERATOR = None
@@ -250,6 +252,13 @@ async def collect():
 
     task_type = msg.get("task_type", "")
     task["task_type"] = task_type
+    # Per-doc fan-out task types (today: doc-scoped raptor) carry their
+    # participating doc id list on the Redis message but not on the DB
+    # row. The KB-scoped branch above already does this for FAKE doc
+    # tasks; mirror here so ``ctx.doc_ids`` is populated for the
+    # per-doc path too.
+    if "doc_ids" in msg and not task.get("doc_ids"):
+        task["doc_ids"] = msg.get("doc_ids", []) or []
     if task_type[:8] == "dataflow":
         task["tenant_id"] = msg["tenant_id"]
         task["dataflow_id"] = msg["dataflow_id"]
@@ -1060,7 +1069,7 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si
         """Run RAPTOR and append generated summary chunks for one doc id."""
         nonlocal tk_count, res
         logging.info("RAPTOR: using tree_builder=%s clustering_method=%s for doc %s", tree_builder, clustering_method, did)
-        from rag.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor  # Lazy load, save around 8s
+        from rag.advanced_rag.knowlege_compile.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor  # Lazy load, save around 8s
 
         raptor = Raptor(
             raptor_config.get("max_cluster", 64),
@@ -1532,6 +1541,9 @@ async def do_handle_task(task):
         progress_callback(1, "place holder")
         pass
         return
+    elif task_type == "skill":
+        progress_callback(-1, "Skill generation requires the refactored task executor (TE_RUN_MODE=0).")
+        return
     else:
         # Standard chunking methods
         task["llm_id"] = doc_task_llm_id
@@ -1563,6 +1575,7 @@ async def do_handle_task(task):
         progress_message = "Embedding chunks ({:.2f}s)".format(timer() - start_ts)
         logging.info(progress_message)
         progress_callback(msg=progress_message)
+
         if task["parser_id"].lower() == "naive" and task["parser_config"].get("toc_extraction", False):
             toc_thread = asyncio.create_task(asyncio.to_thread(build_TOC, task, chunks, progress_callback))
 
@@ -1739,8 +1752,11 @@ async def handle_task():
     finally:
         if not task.get("dataflow_id", ""):
             referred_document_id = None
-            if task_type in ["graphrag", "raptor", "mindmap"]:
-                referred_document_id = task["doc_ids"][0]
+            if task_type in ["graphrag", "raptor", "mindmap", "artifact", "skill"]:
+                # KB-level fan-out tasks store the participating doc list in
+                # task["doc_ids"]; the first entry is used as a referent so
+                # the pipeline operation log has something to anchor to.
+                referred_document_id = (task.get("doc_ids") or [None])[0]
             ret = PipelineOperationLogService.record_pipeline_operation(
                 document_id=task["doc_id"], pipeline_id="", task_type=pipeline_task_type, task_id=task_id, referred_document_id=referred_document_id
             )
@@ -1871,7 +1887,6 @@ async def main():
           /____/
     """)
     logging.info(f"RAGFlow ingestion version: {get_ragflow_version()}")
-    logging.info(f"ENABLE_DRY_RUN_COMPARISON: {os.environ.get('ENABLE_DRY_RUN_COMPARISON', '0')}")
     show_configs()
     settings.init_settings()
     settings.check_and_install_torch()

rag/svr/task_executor_refactor/chunk_post_processor.py

@@ -74,8 +74,7 @@ async def extract_keywords(docs: List[Dict], ctx: TaskContext) -> None:
 
         tasks = []
         for doc in docs:
-            tasks.append(
-                asyncio.create_task(doc_keyword_extraction(chat_model, doc, ctx.parser_config["auto_keywords"])))
+            tasks.append(asyncio.create_task(doc_keyword_extraction(chat_model, doc, ctx.parser_config["auto_keywords"])))
         try:
             await asyncio.gather(*tasks, return_exceptions=False)
         except Exception as e:
@@ -116,8 +115,7 @@ async def generate_questions(docs: List[Dict], ctx: TaskContext) -> None:
 
         tasks = []
         for doc in docs:
-            tasks.append(
-                asyncio.create_task(doc_question_proposal(chat_model, doc, ctx.parser_config["auto_questions"])))
+            tasks.append(asyncio.create_task(doc_question_proposal(chat_model, doc, ctx.parser_config["auto_questions"])))
         try:
             await asyncio.gather(*tasks, return_exceptions=False)
         except Exception as e:
@@ -184,18 +182,14 @@ async def generate_metadata(docs: List[Dict], ctx: TaskContext) -> None:
         metadata_conf = build_metadata_config(ctx.parser_config)
 
         async def gen_metadata_task(chat_mdl, d):
-            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "metadata",
-                                   metadata_conf)
+            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "metadata", metadata_conf)
             if not cached:
                 if ctx.has_canceled_func(ctx.id):
                     ctx.progress_cb(-1, msg="Task has been canceled.")
                     return
                 async with chat_limiter:
-                    cached = await gen_metadata(chat_mdl,
-                                                turn2jsonschema(metadata_conf),
-                                                d["content_with_weight"])
-                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, "metadata",
-                              metadata_conf)
+                    cached = await gen_metadata(chat_mdl, turn2jsonschema(metadata_conf), d["content_with_weight"])
+                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, "metadata", metadata_conf)
             if cached:
                 d["metadata_obj"] = cached
 
@@ -256,8 +250,7 @@ async def apply_tags(docs: List[Dict], ctx: TaskContext) -> None:
             if ctx.has_canceled_func(ctx.id):
                 ctx.progress_cb(-1, msg="Task has been canceled.")
                 return
-            if settings.retriever.tag_content(tenant_id, kb_ids, doc, all_tags, topn_tags=topn_tags, S=S) and len(
-                    doc.get(TAG_FLD, [])) > 0:
+            if settings.retriever.tag_content(tenant_id, kb_ids, doc, all_tags, topn_tags=topn_tags, S=S) and len(doc.get(TAG_FLD, [])) > 0:
                 examples.append({"content": doc["content_with_weight"], TAG_FLD: doc[TAG_FLD]})
             else:
                 docs_to_tag.append(doc)
@@ -270,7 +263,7 @@ async def apply_tags(docs: List[Dict], ctx: TaskContext) -> None:
                     return
                 picked_examples = random.choices(examples, k=2) if len(examples) > 2 else examples
                 if not picked_examples:
-                    picked_examples.append({"content": "This is an example", TAG_FLD: {'example': 1}})
+                    picked_examples.append({"content": "This is an example", TAG_FLD: {"example": 1}})
                 async with chat_limiter:
                     cached = await content_tagging(
                         chat_mdl,
@@ -310,3 +303,883 @@ def count_with_key(docs: List[Dict], key: str) -> int:
         Count of docs that have the key.
     """
     return sum(1 for d in docs if d.get(key))
+
+
+# =====================================================================
+# Document post-chunking pipeline
+# ---------------------------------------------------------------------
+# Extracted from ``task_handler`` to keep the handler class small.
+# The public entry point is :func:`run_document_post_chunking_if_last`;
+# everything below is called (transitively) from there:
+#   run_document_post_chunking_if_last
+#     ├─ run_document_structure_compile
+#     │    ├─ run_tree_templates
+#     │    │    ├─ load_chunks_with_vec
+#     │    │    ├─ rechunk_doc_by_tree
+#     │    │    └─ raptor_tree_to_graph
+#     │    └─ (streaming compile via chat models per template)
+#     └─ handler._run_raptor       ← stays on the handler
+#
+# All entries take ``handler`` (``TaskHandler``) as their first arg so
+# they can reach the handler's ``_task_context``, ``_run_raptor``, and
+# ``_load_chunks_for_doc`` without a circular import.
+# =====================================================================
+
+import numpy as np  # noqa: E402
+from typing import Callable, Optional  # noqa: E402
+
+from common.exceptions import TaskCanceledException  # noqa: E402
+from common.misc_utils import thread_pool_exec  # noqa: E402
+from common.token_utils import num_tokens_from_string  # noqa: E402
+from rag.nlp import search  # noqa: E402
+from api.apps.restful_apis.chunk_api import _compilation_template_kind  # noqa: E402
+from api.db.services.document_service import DocumentService  # noqa: E402
+from api.db.services.compilation_template_service import (  # noqa: E402
+    CompilationTemplateService,
+)
+from api.db.services.compilation_template_group_service import (  # noqa: E402
+    CompilationTemplateGroupService,
+)
+from api.db.services.task_service import (  # noqa: E402
+    abort_doc_chunking_counter,
+    clear_doc_chunking_counter,
+    credit_doc_chunking_task,
+    is_doc_chunking_aborted,
+)
+from rag.advanced_rag.knowlege_compile.structure import (  # noqa: E402
+    CHAIN_KINDS,
+    compile_structure_from_text,
+    merge_compiled_structures,
+    validate_and_correct_chain,
+)
+
+
+# ----- tunables ------------------------------------------------------
+# Bound how many source chunks are handed to a single
+# ``compile_structure_from_text`` invocation. The call fans them out
+# across max_workers internally, so a moderate window keeps memory +
+# LLM-context pressure predictable for long docs.
+DOC_STRUCTURE_COMPILE_BATCH_CHUNKS = 4
+
+# Bound how many compiled ES-ready docs may accumulate before we flush
+# them through ``merge_compiled_structures``. The merger does pairwise
+# cosine + LLM duplicate-judging, so it's the more expensive step; we
+# cap the per-flush set to keep the local-dedup buckets tractable.
+DOC_STRUCTURE_MERGE_MAX_DOCS = 512
+
+# Hard wall on the chain-validator LLM correction step. ``list`` and
+# ``timeline`` kinds run this just before each merge flush; anything
+# longer than this is treated as a blocked LLM and the uncorrected
+# docs are flushed instead.
+STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S = 120.0
+
+
+# ----- parser_config helpers -----------------------------------------
+# Duplicated from ``task_handler`` so this module stays free of a
+# reverse import (task_handler → this module via dispatch; the other
+# direction would be circular).
+
+
+def _parser_config_compilation_template_group_ids(parser_config) -> list[str]:
+    def _normalize(raw) -> list[str]:
+        if isinstance(raw, str):
+            raw = [raw]
+        if not isinstance(raw, list):
+            return []
+        ids: list[str] = []
+        seen: set[str] = set()
+        for gid in raw:
+            if not isinstance(gid, str):
+                continue
+            gid = gid.strip()
+            if gid and gid not in seen:
+                seen.add(gid)
+                ids.append(gid)
+        return ids
+
+    if not isinstance(parser_config, dict):
+        return []
+    if "compilation_template_group_id" in parser_config:
+        return _normalize(parser_config.get("compilation_template_group_id"))
+    ext = parser_config.get("ext")
+    if isinstance(ext, dict):
+        return _normalize(ext.get("compilation_template_group_id"))
+    return []
+
+
+def _parser_config_compilation_template_ids(parser_config, tenant_id: str) -> list[str]:
+    template_ids: list[str] = []
+    seen: set[str] = set()
+    for group_id in _parser_config_compilation_template_group_ids(parser_config):
+        for template_id in CompilationTemplateGroupService.resolve_template_ids(
+            group_id,
+            tenant_id,
+        ):
+            if template_id in seen:
+                continue
+            seen.add(template_id)
+            template_ids.append(template_id)
+    return template_ids
+
+
+def _resolve_template_chat_llm_id(parser_cfg: dict, ctx) -> str:
+    """Pick the chat model id for a knowledge-compilation template.
+
+    Resolution order: template ``llm_id`` → doc ``parser_config.llm_id``
+    → ``ctx.llm_id`` (the chunking task's default).
+    """
+    if isinstance(parser_cfg, dict):
+        tid = parser_cfg.get("llm_id")
+        if isinstance(tid, str) and tid.strip():
+            return tid.strip()
+    doc_cfg = getattr(ctx, "parser_config", None) or {}
+    if isinstance(doc_cfg, dict):
+        did = doc_cfg.get("llm_id")
+        if isinstance(did, str) and did.strip():
+            return did.strip()
+    return ctx.llm_id
+
+
+# ----- progress helper -----------------------------------------------
+
+
+def cap_done_progress(progress_cb: Callable) -> Callable:
+    """Wrap a progress callback so any ``prog >= 1`` gets clamped to
+    ``0.99`` — the final ``1.0`` is reserved for the caller who owns
+    the task's terminal state."""
+
+    def capped_progress(*args, **kwargs):
+        args = list(args)
+        if args:
+            prog = args[0]
+            if isinstance(prog, (int, float)) and not isinstance(prog, bool) and prog >= 1:
+                args[0] = 0.99
+        if "prog" in kwargs:
+            prog = kwargs["prog"]
+            if isinstance(prog, (int, float)) and not isinstance(prog, bool) and prog >= 1:
+                kwargs["prog"] = 0.99
+        return progress_cb(*args, **kwargs)
+
+    return capped_progress
+
+
+# ----- tree helpers --------------------------------------------------
+
+
+def raptor_tree_to_graph(tree: Dict) -> Dict:
+    """Project a RAPTOR tree dict (from ``Raptor(is_tree=True)``) onto
+    the ``{entities, relations}`` shape the document-structure graph
+    endpoint already serves for ``page_index``-kind rows."""
+    entities: list[dict] = []
+    relations: list[dict] = []
+
+    def _walk(node: dict, parent_id: Optional[str]) -> None:
+        if not isinstance(node, dict):
+            return
+        title = node.get("title") or ""
+        node_id = title
+        ent: dict = {
+            "name": node_id,
+            "type": "tree_node",
+            "description": node.get("description", title),
+            "mention_count": 1,
+        }
+        src_ids = node.get("source_chunk_ids")
+        if isinstance(src_ids, list) and src_ids:
+            ent["source_chunk_ids"] = [s for s in src_ids if isinstance(s, str) and s]
+        entities.append(ent)
+        if parent_id is not None:
+            relations.append({"from": parent_id, "to": node_id, "type": "child"})
+        for child in node.get("children") or []:
+            _walk(child, node_id)
+
+    _walk(tree, None)
+    return {"entities": entities, "relations": relations}
+
+
+async def load_chunks_with_vec(
+    tenant_id: str,
+    kb_id: str,
+    doc_id: str,
+    vctr_nm: str,
+) -> list[tuple[str, "np.ndarray", str]]:
+    """Page through this doc's chunks pulling content + vector +
+    chunk_id, in the shape ``RaptorService.build_doc_tree`` expects.
+    Mirrors the streaming ``_load_chunks_for_doc`` loader but with the
+    vector field pre-selected."""
+    from common.doc_store.doc_store_base import OrderByExpr
+
+    index_nm = search.index_name(tenant_id)
+    if not settings.docStoreConn.index_exist(index_nm, kb_id):
+        return []
+    select_fields = ["id", "doc_id", "content_with_weight", vctr_nm]
+    order_by = OrderByExpr()
+    order_by.asc("page_num_int")
+    order_by.asc("top_int")
+
+    out: list[tuple[str, "np.ndarray", str]] = []
+    offset = 0
+    PAGE = 500
+    while True:
+        try:
+            res = await thread_pool_exec(
+                settings.docStoreConn.search,
+                select_fields,
+                [],
+                {"doc_id": [doc_id], "available_int": 1},
+                [],
+                order_by,
+                offset,
+                PAGE,
+                index_nm,
+                [kb_id],
+            )
+            field_map = settings.docStoreConn.get_fields(res, select_fields)
+        except Exception:
+            logging.exception(
+                "tree-template: failed to load chunks for doc=%s",
+                doc_id,
+            )
+            break
+        if not field_map:
+            break
+        for row_id, row in field_map.items():
+            if row.get("compile_kwd"):
+                continue
+            text = row.get("content_with_weight") or ""
+            vec = row.get(vctr_nm)
+            if not text or vec is None:
+                continue
+            try:
+                arr = np.asarray(vec, dtype=np.float32)
+            except Exception:
+                continue
+            if arr.size == 0:
+                continue
+            out.append((text, arr, str(row_id)))
+        if len(field_map) < PAGE:
+            break
+        offset += PAGE
+    return out
+
+
+async def rechunk_doc_by_tree(
+    handler,
+    tree: dict,
+    template_id: str,
+    embedding_model,
+) -> None:
+    """Merge each leaf cluster's source chunks into a single
+    replacement chunk and rewrite the tree's leaf-cluster
+    ``source_chunk_ids`` in-place. Original chunks are soft-deleted
+    via ``available_int=0`` and stamped with ``superseded_by_chunk_id``.
+    """
+    from datetime import datetime
+    from common.misc_utils import get_uuid
+
+    ctx = handler._task_context
+
+    cluster_id_map: dict[int, tuple[dict, list[str]]] = {}
+
+    def _is_terminal(node: object) -> bool:
+        return isinstance(node, dict) and not (node.get("children") or [])
+
+    def _walk(node: object) -> None:
+        if not isinstance(node, dict):
+            return
+        children = node.get("children") or []
+        if children and all(_is_terminal(c) for c in children):
+            src_ids: list[str] = []
+            seen: set[str] = set()
+            for c in children:
+                for cid in c.get("source_chunk_ids") or []:
+                    if isinstance(cid, str) and cid and cid not in seen:
+                        seen.add(cid)
+                        src_ids.append(cid)
+            for cid in node.get("source_chunk_ids") or []:
+                if isinstance(cid, str) and cid and cid not in seen:
+                    seen.add(cid)
+                    src_ids.append(cid)
+            if src_ids:
+                cluster_id_map[id(node)] = (node, src_ids)
+        else:
+            for c in children:
+                _walk(c)
+
+    _walk(tree)
+    if not cluster_id_map:
+        return
+
+    all_source_ids = sorted({sid for _, ids in cluster_id_map.values() for sid in ids})
+
+    from common.doc_store.doc_store_base import OrderByExpr
+
+    index_nm = search.index_name(ctx.tenant_id)
+    if not settings.docStoreConn.index_exist(index_nm, ctx.kb_id):
+        return
+
+    vctr_nm = "q_%d_vec" % len(embedding_model.encode(["x"])[0][0])
+    select_fields = [
+        "id",
+        "doc_id",
+        "kb_id",
+        "content_with_weight",
+        "page_num_int",
+        "top_int",
+        "position_int",
+        "docnm_kwd",
+        "title_tks",
+        "title_sm_tks",
+        "available_int",
+    ]
+    try:
+        res = await thread_pool_exec(
+            settings.docStoreConn.search,
+            select_fields,
+            [],
+            {"id": all_source_ids, "available_int": 1},
+            [],
+            OrderByExpr(),
+            0,
+            len(all_source_ids) + 16,
+            index_nm,
+            [ctx.kb_id],
+        )
+        field_map = settings.docStoreConn.get_fields(res, select_fields)
+    except Exception:
+        logging.exception(
+            "rechunk: failed to load source chunks for doc=%s template=%s",
+            ctx.doc_id,
+            template_id,
+        )
+        return
+    if not field_map:
+        return
+
+    chunks_by_id: dict[str, dict] = {str(rid): {**row, "id": str(rid)} for rid, row in field_map.items()}
+
+    merged_rows: list[dict] = []
+    cluster_new_id: dict[int, str] = {}
+
+    for node_id_int, (node, src_ids) in cluster_id_map.items():
+        cluster_chunks = [chunks_by_id[c] for c in src_ids if c in chunks_by_id]
+        if not cluster_chunks:
+            continue
+
+        def _sort_key(c: dict) -> tuple:
+            pages = c.get("page_num_int") or [0]
+            tops = c.get("top_int") or [0]
+            return (
+                min(pages) if pages else 0,
+                min(tops) if tops else 0,
+                c.get("id") or "",
+            )
+
+        cluster_chunks.sort(key=_sort_key)
+
+        merged_content = "\n\n".join((c.get("content_with_weight") or "") for c in cluster_chunks).strip()
+        if not merged_content:
+            continue
+        page_union = sorted({p for c in cluster_chunks for p in (c.get("page_num_int") or [])})
+        top_union = sorted({t for c in cluster_chunks for t in (c.get("top_int") or [])})
+
+        base = dict(cluster_chunks[0])
+        new_id = get_uuid()
+        cluster_new_id[node_id_int] = new_id
+
+        base.update(
+            {
+                "id": new_id,
+                "content_with_weight": merged_content,
+                "content_ltks": rag_tokenizer.tokenize(merged_content),
+                "page_num_int": page_union,
+                "top_int": top_union,
+                "available_int": 1,
+                "rechunk_kwd": "tree",
+                "rechunked_from_template_id": template_id,
+                "rechunked_from_chunk_ids": [c.get("id") for c in cluster_chunks if c.get("id")],
+                "token_num": num_tokens_from_string(merged_content),
+                "create_time": str(datetime.now()).replace("T", " ")[:19],
+                "create_timestamp_flt": datetime.now().timestamp(),
+            }
+        )
+        base["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(base["content_ltks"])
+        merged_rows.append(base)
+
+    if not merged_rows:
+        return
+
+    contents = [r["content_with_weight"] for r in merged_rows]
+    try:
+        vectors, _ = embedding_model.encode(contents)
+    except Exception:
+        logging.exception(
+            "rechunk: embedding failed for doc=%s template=%s",
+            ctx.doc_id,
+            template_id,
+        )
+        return
+    for row, vec in zip(merged_rows, vectors):
+        try:
+            row[vctr_nm] = np.asarray(vec, dtype=np.float32).tolist()
+        except Exception:
+            logging.exception(
+                "rechunk: vector cast failed; skipping row %s",
+                row.get("id"),
+            )
+            row[vctr_nm] = None
+    merged_rows = [r for r in merged_rows if r.get(vctr_nm) is not None]
+    if not merged_rows:
+        return
+
+    try:
+        await thread_pool_exec(
+            settings.docStoreConn.insert,
+            merged_rows,
+            index_nm,
+            ctx.kb_id,
+        )
+    except Exception:
+        logging.exception(
+            "rechunk: insert failed for doc=%s template=%s",
+            ctx.doc_id,
+            template_id,
+        )
+        return
+
+    for node_id_int, new_chunk_id in cluster_new_id.items():
+        node, _ = cluster_id_map[node_id_int]
+        node["source_chunk_ids"] = [new_chunk_id]
+        for child in node.get("children") or []:
+            if isinstance(child, dict):
+                child["source_chunk_ids"] = [new_chunk_id]
+
+    for node_id_int, new_chunk_id in cluster_new_id.items():
+        _, src_ids = cluster_id_map[node_id_int]
+        for cid in src_ids:
+            try:
+                await thread_pool_exec(
+                    settings.docStoreConn.update,
+                    {"id": cid},
+                    {
+                        "available_int": 0,
+                        "superseded_by_chunk_id": new_chunk_id,
+                    },
+                    index_nm,
+                    ctx.kb_id,
+                )
+            except Exception:
+                logging.exception(
+                    "rechunk: soft-delete failed for chunk=%s (merged=%s)",
+                    cid,
+                    new_chunk_id,
+                )
+
+
+async def run_tree_templates(
+    handler,
+    templates: list[tuple[str, dict]],
+    chat_mdl_by_tid: dict[str, "LLMBundle"],
+    embedding_model,
+) -> None:
+    """Run the ``tree``-kind compilation templates for the current
+    doc. Each pair runs RAPTOR with ``is_tree=True`` via
+    ``RaptorService.build_doc_tree`` and persists a single graph row
+    via ``_struct_upsert_graph_json``."""
+    from rag.svr.task_executor_refactor.raptor_service import RaptorService
+    from rag.advanced_rag.knowlege_compile.structure import _struct_upsert_graph_json
+
+    ctx = handler._task_context
+    progress_cb = ctx.progress_cb
+
+    try:
+        doc_id = ctx.doc_id
+    except Exception:
+        doc_id = getattr(ctx, "_task", {}).get("doc_id") if hasattr(ctx, "_task") else None
+    if not doc_id:
+        logging.warning("tree-template: no doc_id on task context; skipping")
+        return
+
+    vctr_nm = "q_%d_vec" % len(embedding_model.encode(["x"])[0][0])
+    chunks = await load_chunks_with_vec(
+        ctx.tenant_id,
+        ctx.kb_id,
+        doc_id,
+        vctr_nm,
+    )
+    if not chunks:
+        progress_cb(msg=f"tree-template: doc {doc_id} has no chunks; skipping")
+        return
+
+    raptor_service = RaptorService(ctx)
+
+    for idx, (template_id, parser_cfg) in enumerate(templates):
+        raptor_cfg = (parser_cfg or {}).get("raptor") or {}
+        raptor_config = {
+            "prompt": raptor_cfg.get("prompt") or "Please write a concise summary of the following texts:\n{cluster_content}",
+            "max_token": int(raptor_cfg.get("max_token") or 512),
+            "threshold": float(raptor_cfg.get("threshold") or 0.1),
+            "random_seed": int(raptor_cfg.get("random_seed") or 0),
+            "max_cluster": int(raptor_cfg.get("max_cluster") or 64),
+            "ext": raptor_cfg.get("ext") or {},
+        }
+        progress_cb(
+            msg=f"tree-template ({idx + 1}/{len(templates)}): building tree for doc={doc_id}",
+        )
+        try:
+            tree = await raptor_service.build_doc_tree(
+                chunks=chunks,
+                raptor_config=raptor_config,
+                chat_mdl=chat_mdl_by_tid[template_id],
+                embd_mdl=embedding_model,
+                tree_builder="raptor",
+                clustering_method="gmm",
+                max_errors=3,
+            )
+        except Exception:
+            logging.exception(
+                "tree-template %s: RAPTOR build failed for doc %s",
+                template_id,
+                doc_id,
+            )
+            continue
+        if tree is None:
+            logging.info(
+                "tree-template %s: no tree produced for doc %s",
+                template_id,
+                doc_id,
+            )
+            continue
+
+        if bool((raptor_cfg or {}).get("rechunk")):
+            try:
+                await rechunk_doc_by_tree(
+                    handler=handler,
+                    tree=tree,
+                    template_id=template_id,
+                    embedding_model=embedding_model,
+                )
+            except Exception:
+                logging.exception(
+                    "tree-template %s: re-chunking failed for doc %s; persisting tree with original chunk ids",
+                    template_id,
+                    doc_id,
+                )
+
+        graph = raptor_tree_to_graph(tree)
+        try:
+            await _struct_upsert_graph_json(
+                graph,
+                ctx.tenant_id,
+                ctx.kb_id,
+                doc_id,
+                compile_kwd="tree",
+                compilation_template_id=template_id,
+            )
+        except Exception:
+            logging.exception(
+                "tree-template %s: graph upsert failed for doc %s",
+                template_id,
+                doc_id,
+            )
+            continue
+
+        try:
+            from rag.advanced_rag.knowlege_compile.dataset_nav import (
+                upsert_dataset_nav_doc,
+            )
+
+            await upsert_dataset_nav_doc(
+                ctx.tenant_id,
+                ctx.kb_id,
+                doc_id,
+                tree,
+            )
+        except Exception:
+            logging.exception(
+                "tree-template %s: dataset_nav upsert failed for doc %s",
+                template_id,
+                doc_id,
+            )
+
+        progress_cb(
+            msg=f"tree-template ({idx + 1}/{len(templates)}): persisted {len(graph['entities'])} node(s), {len(graph['relations'])} edge(s) for doc {doc_id}",
+        )
+
+
+async def run_document_structure_compile(handler, embedding_model: LLMBundle) -> None:
+    """Run document-scoped knowledge compilation for non-artifact
+    templates. Streams the doc's chunks (via
+    ``handler._load_chunks_for_doc``) and fans each batch out to every
+    configured non-artifact template, flushing accumulators through
+    ``merge_compiled_structures`` at :data:`DOC_STRUCTURE_MERGE_MAX_DOCS`.
+    """
+    ctx = handler._task_context
+    template_ids = _parser_config_compilation_template_ids(ctx.parser_config, ctx.tenant_id)
+    if not template_ids:
+        return
+
+    active_templates: list[tuple[str, dict]] = []
+    for template_id in template_ids:
+        template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id)
+        if not template:
+            logging.warning(
+                "document_structure_compile: template %s not found",
+                template_id,
+            )
+            continue
+        parser_cfg = template.get("config") or {}
+        if not isinstance(parser_cfg, dict):
+            logging.warning(
+                "document_structure_compile: template %s config is invalid",
+                template_id,
+            )
+            continue
+        kind = _compilation_template_kind(parser_cfg.get("kind"))
+        if not kind or kind == "artifacts":
+            continue
+        active_templates.append((template_id, parser_cfg))
+
+    if not active_templates:
+        return
+
+    llm_bundle_cache: dict[str, LLMBundle] = {}
+    chat_mdl_by_tid: dict[str, LLMBundle] = {}
+    filtered_templates: list[tuple[str, dict]] = []
+    for template_id, parser_cfg in active_templates:
+        chat_llm_id = _resolve_template_chat_llm_id(parser_cfg, ctx)
+        if chat_llm_id not in llm_bundle_cache:
+            try:
+                cfg = get_model_config_from_provider_instance(
+                    ctx.tenant_id,
+                    LLMType.CHAT,
+                    chat_llm_id,
+                )
+                llm_bundle_cache[chat_llm_id] = LLMBundle(
+                    ctx.tenant_id,
+                    cfg,
+                    lang=ctx.language,
+                )
+            except Exception:
+                logging.exception(
+                    "document_structure_compile: cannot resolve chat model %s for template %s; skipping",
+                    chat_llm_id,
+                    template_id,
+                )
+                continue
+        chat_mdl_by_tid[template_id] = llm_bundle_cache[chat_llm_id]
+        filtered_templates.append((template_id, parser_cfg))
+
+    if not filtered_templates:
+        return
+    active_templates = filtered_templates
+
+    tree_templates: list[tuple[str, dict]] = []
+    non_tree_templates: list[tuple[str, dict]] = []
+    for tid, cfg in active_templates:
+        if _compilation_template_kind((cfg or {}).get("kind")) == "tree":
+            tree_templates.append((tid, cfg))
+        else:
+            non_tree_templates.append((tid, cfg))
+
+    if tree_templates:
+        await run_tree_templates(
+            handler,
+            tree_templates,
+            chat_mdl_by_tid,
+            embedding_model,
+        )
+
+    if not non_tree_templates:
+        return
+    active_templates = non_tree_templates
+
+    progress_cb = ctx.progress_cb
+    total = len(active_templates)
+
+    accumulators: dict[str, list[dict]] = {tid: [] for tid, _ in active_templates}
+    template_kinds: dict[str, str] = {tid: _compilation_template_kind((cfg or {}).get("kind")) for tid, cfg in active_templates}
+    agg_infos: dict[str, dict] = {tid: {"inserted": 0, "updated": 0, "duplicates_dropped": 0} for tid, _ in active_templates}
+    chunks_by_id: dict[str, str] = {}
+
+    async def _flush(template_id: str) -> None:
+        acc = accumulators[template_id]
+        if not acc:
+            return
+        kind = template_kinds.get(template_id, "")
+        if kind in CHAIN_KINDS:
+            try:
+                acc = await asyncio.wait_for(
+                    validate_and_correct_chain(
+                        acc,
+                        chunks_by_id,
+                        chat_mdl_by_tid[template_id],
+                        kind,
+                        callback=progress_cb,
+                    ),
+                    timeout=STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S,
+                )
+                accumulators[template_id] = acc
+            except asyncio.TimeoutError:
+                logging.warning(
+                    "chain validate: timed out after %ss for template %s; using uncorrected docs",
+                    STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S,
+                    template_id,
+                )
+            except Exception:
+                logging.exception(
+                    "chain validate: unexpected failure for template %s; using uncorrected docs",
+                    template_id,
+                )
+        info = await merge_compiled_structures(
+            acc,
+            chat_mdl_by_tid[template_id],
+            embedding_model,
+            ctx.tenant_id,
+            ctx.kb_id,
+            compilation_template_id=template_id,
+            cancel_check=lambda: ctx.has_canceled_func(ctx.id),
+        )
+        acc.clear()
+        if isinstance(info, dict):
+            agg = agg_infos[template_id]
+            for k in ("inserted", "updated", "duplicates_dropped"):
+                agg[k] = agg.get(k, 0) + int(info.get(k, 0) or 0)
+
+    progress_cb(msg=f"Start document knowledge compilation ({total} template(s)) ...")
+
+    batch_no = 0
+    async for batch in handler._load_chunks_for_doc(
+        ctx.tenant_id,
+        ctx.kb_id,
+        ctx.doc_id,
+        batch_size=DOC_STRUCTURE_COMPILE_BATCH_CHUNKS,
+    ):
+        batch_no += 1
+        for chunk in batch:
+            cid = chunk.get("id")
+            if isinstance(cid, str) and cid not in chunks_by_id:
+                text = chunk.get("content_with_weight") or ""
+                chunks_by_id[cid] = text if isinstance(text, str) else ""
+        for idx, (template_id, parser_cfg) in enumerate(active_templates):
+            progress_cb(msg=f"  compile batch {batch_no} ({len(batch)} chunks) for template ({idx + 1}/{total})")
+            docs = await compile_structure_from_text(
+                batch,
+                parser_cfg,
+                chat_mdl_by_tid[template_id],
+                embedding_model,
+                ctx.doc_id,
+                language=ctx.language,
+                callback=progress_cb,
+                compilation_template_id=template_id,
+            )
+            if docs:
+                accumulators[template_id].extend(docs)
+            if len(accumulators[template_id]) >= DOC_STRUCTURE_MERGE_MAX_DOCS:
+                progress_cb(msg=f"  merge flush ({len(accumulators[template_id])} docs) for template ({idx + 1}/{total})")
+                await _flush(template_id)
+
+    for idx, (template_id, _parser_cfg) in enumerate(active_templates):
+        if ctx.has_canceled_func(ctx.id):
+            raise TaskCanceledException(f"Task {ctx.id} was cancelled during document knowledge compilation")
+        await _flush(template_id)
+        agg = agg_infos[template_id]
+        ctx.recording_context.record(f"document_structure_compile:{template_id}", agg)
+        progress_cb(msg=f"Document knowledge compilation done ({idx + 1}/{total}): {agg}")
+
+
+async def run_document_post_chunking_if_last(
+    handler,
+    embedding_model: LLMBundle,
+    vector_size: int,
+    task_start_ts: float,
+    chunks_len: int,
+    token_count: int,
+) -> bool:
+    """Gate: only the last chunking task for a doc runs post-processing.
+    Returns ``True`` if the caller may proceed to its own terminal
+    progress update, ``False`` if the task was cancelled.
+
+    The pass runs :func:`run_document_structure_compile` and
+    ``handler._run_raptor`` concurrently — they read the same chunks
+    but write disjoint ES rows.
+    """
+    ctx = handler._task_context
+    task_id = ctx.id
+    task_doc_id = ctx.doc_id
+
+    if ctx.has_canceled_func(task_id):
+        abort_doc_chunking_counter(task_doc_id)
+        ctx.progress_cb(-1, msg="Task has been canceled.")
+        return False
+
+    chunking_aborted = is_doc_chunking_aborted(task_doc_id)
+    remaining_chunking_tasks = (
+        0 if ctx.write_interceptor
+        else credit_doc_chunking_task(task_doc_id, task_id)
+    )
+    if remaining_chunking_tasks != 0:
+        if chunking_aborted:
+            logging.info(
+                "Chunking for doc %s was aborted before task %s reached post-processing; "
+                "skip document finalizers.",
+                task_doc_id,
+                task_id,
+            )
+        elif remaining_chunking_tasks is not None and remaining_chunking_tasks < 0:
+            logging.warning(
+                "Chunking counter for doc %s is missing or expired after task %s; skip post-processing to avoid duplicate finalizers.",
+                task_doc_id,
+                task_id,
+            )
+        else:
+            logging.info(
+                "Chunk doc(%s), page(%s-%s), chunks(%s), token(%s), elapsed:%.2f; waiting for %s chunking task(s) before post-processing",
+                ctx.name,
+                ctx.from_page,
+                ctx.to_page,
+                chunks_len,
+                token_count,
+                timer() - task_start_ts,
+                remaining_chunking_tasks,
+            )
+        return True
+
+    async def _maybe_run_raptor():
+        raptor_cfg = (ctx.parser_config or {}).get("raptor") or {}
+        if not raptor_cfg.get("use_raptor"):
+            return
+        try:
+            ok_doc, doc_obj = DocumentService.get_by_id(task_doc_id)
+            if ok_doc and doc_obj is not None:
+                ctx.progress_cb(msg="Starting RAPTOR task.")
+                await handler._run_raptor(embedding_model, vector_size, mark_done=False)
+            else:
+                logging.warning(
+                    "raptor: cannot resolve doc %s to queue per-doc task",
+                    task_doc_id,
+                )
+        except Exception:
+            logging.exception(
+                "raptor: failed to queue per-doc task for doc %s",
+                task_doc_id,
+            )
+
+    original_progress_cb = getattr(ctx, "_progress_cb", None)
+    if original_progress_cb is not None:
+        ctx._progress_cb = cap_done_progress(original_progress_cb)
+    try:
+        await asyncio.gather(
+            run_document_structure_compile(handler, embedding_model),
+            _maybe_run_raptor(),
+        )
+    finally:
+        if original_progress_cb is not None:
+            ctx._progress_cb = original_progress_cb
+        clear_doc_chunking_counter(task_doc_id)
+
+    if ctx.has_canceled_func(task_id):
+        abort_doc_chunking_counter(task_doc_id)
+        ctx.progress_cb(-1, msg="Task has been canceled.")
+        return False
+    return True

rag/svr/task_executor_refactor/dataset_skill_generator.py

@@ -0,0 +1,588 @@
+#
+#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+"""Corpus → Skill tree generator.
+
+Extracted from ``rag.svr.task_executor_refactor.task_handler`` where the
+same pipeline previously lived as a set of ``_skill_*`` methods and one
+``_corpus2skill`` orchestrator. The public entry point is
+:func:`run_corpus2skill`; the module-level helpers are internal but
+kept accessible so tests can exercise the individual phases.
+
+Design notes:
+
+* The pipeline is per-KB: given every parsed doc in a KB it produces a
+  hierarchical "skill" tree by summarizing each doc, RAPTOR-clustering
+  the summaries, summarizing each cluster, then repeating until the top
+  fan-out is at or below :data:`SKILL_MAX_TOP_CLUSTERS`.
+* Each node lands in ES twice: one per-node row under
+  ``compile_kwd="skill"`` carrying markdown metadata + a leaf-vs-branch
+  contents section, and one aggregate row under
+  ``compile_kwd="skill_all"`` holding the whole recursive tree as JSON
+  for cheap sidebar reads.
+* The layout mirrors Corpus2Skill's ``SKILL.md`` / ``INDEX.md`` naming
+  so a future on-disk export is a straight projection.
+
+The extraction keeps the callable surface minimal:
+
+    run_corpus2skill(ctx, embedding_model, load_chunks_for_doc)
+
+``load_chunks_for_doc`` is injected rather than imported to keep the
+module decoupled from ``TaskHandler``'s streaming chunk loader — any
+async iterator that yields batches of ``{content_with_weight: str, ...}``
+dicts will do.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import re
+from dataclasses import dataclass, field
+from typing import AsyncIterator, Callable, Dict, Optional
+
+import numpy as np
+import xxhash
+
+from common import settings
+from common.constants import LLMType
+from common.misc_utils import thread_pool_exec
+from common.token_utils import num_tokens_from_string
+from rag.nlp import search
+from rag.svr.task_executor_refactor.task_context import TaskContext
+
+
+# ----- tunables ------------------------------------------------------
+# Stop folding clusters once we've boiled the KB down to ≤ this many
+# top-level nodes. Mirrors Corpus2Skill's default top-of-tree fan-out.
+SKILL_MAX_TOP_CLUSTERS = 8
+# Per-doc summary is built from a budget of this fraction of the chat
+# model's context window. Stops adding chunks once cumulative tokens
+# hit the cap.
+SKILL_DOC_BUDGET_FRACTION = 0.5
+# Concurrency caps for the two LLM-bound stages.
+SKILL_DOC_SUMMARY_CONCURRENCY = 8
+SKILL_LABEL_CONCURRENCY = 10
+# Defensive cap on the clustering loop — a degenerate clustering that
+# keeps returning N clusters for N inputs would otherwise loop forever.
+SKILL_MAX_TREE_ITERATIONS = 12
+# Page size for the streaming chunk reader used during per-doc
+# summarization.
+SKILL_CHUNK_BATCH = 64
+
+
+# A ``load_chunks_for_doc`` callable takes ``(tenant_id, kb_id, doc_id,
+# batch_size)`` and returns an async iterator of chunk-batch lists.
+ChunkLoader = Callable[
+    [str, str, str],  # positional: tenant_id, kb_id, doc_id
+    AsyncIterator[list[dict]],
+]
+
+
+@dataclass
+class SkillNode:
+    """One node in the corpus → skill hierarchy.
+
+    Leaves wrap a single doc; branch nodes wrap a cluster of child
+    nodes plus the LLM summary of their summaries. ``doc_ids`` is the
+    flattened leaf set under this node, so any branch can quickly
+    report ``num_documents``. ``doc_texts`` carries first-page
+    previews keyed by ``doc_id`` only for leaves whose markdown nav
+    file needs the first-line title — branches inherit it as a no-op
+    union so the dict is non-empty across the tree.
+    """
+
+    level: int
+    label: str
+    summary: str
+    vec: np.ndarray
+    doc_ids: list[str]
+    doc_texts: dict[str, str] = field(default_factory=dict)
+    children: list["SkillNode"] = field(default_factory=list)
+    folder_name: str = ""
+
+
+# ----- helpers -------------------------------------------------------
+
+
+def skill_safe_name(text: str, max_len: int = 50) -> str:
+    """Lowercase, hyphen-only, max-len-clamped slug. Mirrors
+    Corpus2Skill ``_safe_name`` for cross-system stability of folder
+    names."""
+    name = (text or "").lower().strip()
+    name = re.sub(r"[^a-z0-9\s-]", "", name)
+    name = re.sub(r"\s+", "-", name)
+    name = name.strip("-")[:max_len]
+    return name
+
+
+async def label_skill_node_one(
+    summary: str,
+    chat_mdl,
+    semaphore: asyncio.Semaphore,
+) -> str:
+    """Generate a single fs-safe label: 2–5 word lowercase
+    hyphenated label, max_tokens=20, sanitized to [a-z0-9-], capped
+    at 50 chars, falls back to "cluster" on any failure.
+    """
+    async with semaphore:
+        try:
+            cnt = await chat_mdl.async_chat(
+                "You generate short filesystem-safe cluster labels. Reply with the label only.",
+                [
+                    {
+                        "role": "user",
+                        "content": (
+                            "Generate a short (2-5 word) filesystem-safe label for this cluster. "
+                            "Use lowercase(MUST be in the same language as 'Summary'), hyphens instead of spaces. No quotes.\n\n"
+                            f"Summary: {(summary or '')[:500]}"
+                        ),
+                    }
+                ],
+                {"max_tokens": 20, "temperature": 0.0},
+            )
+            raw = (cnt or "").strip().lower()
+            label = re.sub(r"[^a-z0-9-]", "-", raw)
+            label = re.sub(r"-+", "-", label).strip("-")[:50]
+            return label or "cluster"
+        except Exception:
+            logging.exception("skill: label generation failed; using fallback")
+            return "cluster"
+
+
+async def doc_summary_for_skill(
+    doc_id: str,
+    raptor,
+    chat_mdl,
+    ctx: TaskContext,
+    load_chunks_for_doc: Callable[..., AsyncIterator[list[dict]]],
+) -> Optional["SkillNode"]:
+    """Concatenate chunks up to half the chat model's context budget,
+    then summarize via RAPTOR's ``_summarize_texts`` (which also
+    returns the embedding). Returns a leaf-shaped :class:`SkillNode`
+    or ``None`` if the doc has no usable chunks."""
+    max_ctx = int(getattr(chat_mdl, "max_length", 4096) or 4096)
+    budget = max(512, int(max_ctx * SKILL_DOC_BUDGET_FRACTION))
+
+    accumulated: list[str] = []
+    running = 0
+    async for batch in load_chunks_for_doc(
+        ctx.tenant_id,
+        ctx.kb_id,
+        doc_id,
+        batch_size=SKILL_CHUNK_BATCH,
+    ):
+        for chunk in batch:
+            text = chunk.get("content_with_weight") or ""
+            if not isinstance(text, str) or not text:
+                continue
+            t_tokens = num_tokens_from_string(text)
+            if running + t_tokens > budget and accumulated:
+                break
+            accumulated.append(text)
+            running += t_tokens
+        else:
+            continue
+        break
+
+    if not accumulated:
+        return None
+
+    result = await raptor._summarize_texts(accumulated, callback=None, task_id="")
+    if result is None:
+        return None
+    title, summary_text, vec = result
+
+    doc_preview = accumulated[0][:600] if accumulated else ""
+    return SkillNode(
+        level=0,
+        label="",  # filled in phase 5
+        summary=summary_text or title,
+        vec=np.asarray(vec),
+        doc_ids=[doc_id],
+        doc_texts={doc_id: doc_preview},
+        children=[],
+        folder_name="",  # filled in phase 5
+    )
+
+
+def build_skill_md(node: "SkillNode") -> str:
+    """SKILL.md (depth 0) / INDEX.md (deeper) text. Mirrors
+    Corpus2Skill's ``_format_skill_md`` (skill_builder.py:193): YAML
+    frontmatter (name / description / level / num_documents), then
+    ``## Overview`` with the full summary, then ``## Contents`` —
+    sub-groups for branches, ``- `doc_id`: <first 120 chars>`` for
+    leaves.
+    """
+    depth = node.level
+    name = node.folder_name or node.label or f"cluster-{depth}"
+    desc = (node.summary or "")[:300].replace("\n", " ").strip()
+    lines: list[str] = [
+        "---",
+        f"name: {name}",
+        "description: >",
+        f"  {desc}",
+        f"level: {depth}",
+        f"num_documents: {len(node.doc_ids)}",
+        "---",
+        "",
+        "## Overview",
+        "",
+        (node.summary or "").strip() or "(no summary)",
+        "",
+        "## Contents",
+        "",
+    ]
+    if node.children:
+        lines.append("### Sub-groups (directories)")
+        lines.append("")
+        for child in node.children:
+            child_name = child.folder_name or child.label or "cluster"
+            summary_snip = (child.summary or "")[:200].replace("\n", " ").strip()
+            lines.append(f"- **{child_name}/** ({len(child.doc_ids)} docs): {summary_snip}")
+        lines.append("")
+    else:
+        lines.append(f"### Documents ({len(node.doc_ids)} items)")
+        lines.append("")
+        for doc_id in node.doc_ids:
+            preview = node.doc_texts.get(doc_id, "")
+            first_line = (preview.split("\n", 1)[0] if preview else "").strip()[:120]
+            lines.append(f"- `{doc_id}`: {first_line}")
+        lines.append("")
+    return "\n".join(lines)
+
+
+def skill_node_es_row(ctx: TaskContext, node: "SkillNode") -> Dict:
+    """Build the ES row for one tree node. Stable id from
+    (kb_id, folder_name) so re-runs upsert cleanly."""
+    kb_id_str = str(ctx.kb_id)
+    row_id = xxhash.xxh64(
+        f"skill:{kb_id_str}:{node.folder_name}".encode("utf-8", "surrogatepass"),
+    ).hexdigest()
+    return {
+        "id": row_id,
+        "kb_id": kb_id_str,
+        "doc_id": kb_id_str,  # KB-scoped sentinel
+        "compile_kwd": "skill",
+        "skill_kwd": node.folder_name,
+        "depth_int": int(node.level),
+        "children_kwd": [c.folder_name for c in node.children],
+        "source_doc_ids": list(node.doc_ids),
+        "md_with_weight": build_skill_md(node),
+        "available_int": 1,
+    }
+
+
+def skill_tree_md_snippet(node: "SkillNode") -> str:
+    """Return only the frontmatter/preamble before the Overview body.
+
+    The one-shot tree browser needs enough metadata to render the skill
+    directory without loading every full node body up front.
+    """
+    md = build_skill_md(node)
+    return md.split("\n## Overview", 1)[0].strip()
+
+
+def skill_tree_node(node: "SkillNode") -> Dict:
+    return {
+        "skill_kwd": node.folder_name,
+        "md_with_weight": skill_tree_md_snippet(node),
+        "children_kwd": [skill_tree_node(child) for child in node.children],
+    }
+
+
+def skill_all_es_row(ctx: TaskContext, roots: list["SkillNode"]) -> Dict:
+    """Build the aggregate tree row loaded by the Skills sidebar."""
+    kb_id_str = str(ctx.kb_id)
+    row_id = xxhash.xxh64(
+        f"skill_all:{kb_id_str}".encode("utf-8", "surrogatepass"),
+    ).hexdigest()
+    return {
+        "id": row_id,
+        "kb_id": kb_id_str,
+        "doc_id": kb_id_str,
+        "compile_kwd": "skill_all",
+        "skill_with_weight": json.dumps(
+            [skill_tree_node(root) for root in roots],
+            ensure_ascii=False,
+            indent=2,
+        ),
+        "available_int": 1,
+    }
+
+
+# ----- main entry ----------------------------------------------------
+
+
+async def run_corpus2skill(
+    ctx: TaskContext,
+    embedding_model,
+    load_chunks_for_doc: Callable[..., AsyncIterator[list[dict]]],
+) -> None:
+    """Build a hierarchical skill tree for the current KB and persist
+    one ES row per node under ``compile_kwd="skill"`` plus a full
+    recursive aggregate row under ``compile_kwd="skill_all"``.
+
+    Always-rebuild semantics for v1: every parsed doc in the KB is
+    re-summarized on each call. (Incremental "only changed docs" is a
+    TODO — needs a per-doc content-hash similar to MAP's
+    ``chunk_hash_kwd``.)
+    """
+    # Local imports so the module doesn't drag in the API service layer
+    # at import time — that's a source of circular-import risk given how
+    # much lives under ``api.db.services``.
+    from api.db.services.document_service import DocumentService
+    from api.db.services.llm_service import LLMBundle
+    from api.db.joint_services.tenant_model_service import (
+        get_tenant_default_model_by_type,
+    )
+    from rag.advanced_rag.knowlege_compile.raptor import (
+        RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor,
+    )
+
+    progress = ctx.progress_cb
+    progress(0.0, "skill: loading documents")
+
+    # ---- Phase 0: chat model + RAPTOR instance for summarization/clustering.
+    chat_model_config = get_tenant_default_model_by_type(ctx.tenant_id, LLMType.CHAT)
+    chat_mdl = LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language)
+
+    raptor = Raptor(
+        max_cluster=128,
+        llm_model=chat_mdl,
+        embd_model=embedding_model,
+        prompt="Please write a concise summary of the following texts:\n{cluster_content}",
+        max_token=256,
+        threshold=0.1,
+        max_errors=3,
+    )
+
+    # ---- Phase 1: per-doc summaries.
+    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_docs = [d for d in (all_docs or []) if d.get("id")]
+    if not eligible_docs:
+        progress(1.0, "skill: no documents in KB")
+        return
+
+    # Phase-1 gate: bail before spinning up N per-doc summarizations.
+    if ctx.has_canceled_func(ctx.id):
+        progress(-1, "skill: task has been canceled")
+        return
+
+    n_docs = len(eligible_docs)
+    progress(0.05, f"skill: summarizing {n_docs} document(s)")
+    doc_sem = asyncio.Semaphore(SKILL_DOC_SUMMARY_CONCURRENCY)
+
+    async def _summarize_doc(d: Dict) -> Optional[SkillNode]:
+        async with doc_sem:
+            try:
+                return await doc_summary_for_skill(
+                    d["id"],
+                    raptor,
+                    chat_mdl,
+                    ctx,
+                    load_chunks_for_doc,
+                )
+            except Exception:
+                logging.exception(
+                    "skill: doc summary failed for doc=%s",
+                    d.get("id"),
+                )
+                return None
+
+    leaf_results = await asyncio.gather(
+        *(_summarize_doc(d) for d in eligible_docs),
+        return_exceptions=False,
+    )
+    leaves: list[SkillNode] = [n for n in leaf_results if n is not None]
+    if not leaves:
+        progress(1.0, "skill: no doc summaries produced")
+        return
+
+    # Post-Phase-1 gate: bail before starting the iterative clustering.
+    if ctx.has_canceled_func(ctx.id):
+        progress(-1, "skill: task has been canceled")
+        return
+
+    # ---- Phase 2-4: iterative clustering until ≤ MAX_TOP.
+    current_layer = leaves
+    level = 0
+    for iteration in range(SKILL_MAX_TREE_ITERATIONS):
+        # Per-iteration gate: caps the wasted LLM cost when the task is
+        # canceled mid-way through a many-layer clustering run.
+        if ctx.has_canceled_func(ctx.id):
+            progress(-1, "skill: task has been canceled")
+            return
+        if len(current_layer) <= SKILL_MAX_TOP_CLUSTERS:
+            break
+        progress(
+            0.3 + 0.4 * iteration / SKILL_MAX_TREE_ITERATIONS,
+            f"skill: clustering layer {level} ({len(current_layer)} nodes)",
+        )
+        try:
+            embeddings = np.asarray([n.vec for n in current_layer])
+            n_clusters, labels = raptor.clustering(
+                embeddings,
+                random_state=0,
+                task_id="",
+            )
+        except Exception:
+            logging.exception("skill: clustering failed at level %d", level)
+            break
+        if n_clusters <= 0 or n_clusters >= len(current_layer):
+            # No reduction → stop to avoid an infinite loop.
+            logging.warning(
+                "skill: clustering did not reduce node count (%d → %d); stopping",
+                len(current_layer),
+                n_clusters,
+            )
+            break
+
+        cluster_buckets: dict[int, list[SkillNode]] = {}
+        for idx, lbl in enumerate(labels):
+            cluster_buckets.setdefault(int(lbl), []).append(current_layer[idx])
+
+        async def _summarize_cluster(children: list[SkillNode]) -> Optional[SkillNode]:
+            texts = [c.summary for c in children if c.summary]
+            if not texts:
+                return None
+            try:
+                res = await raptor._summarize_texts(texts, callback=None, task_id="")
+            except Exception:
+                logging.exception("skill: cluster summary failed")
+                return None
+            if res is None:
+                return None
+            title, summary_text, vec = res
+            merged_doc_ids: list[str] = []
+            seen: set[str] = set()
+            merged_doc_texts: dict[str, str] = {}
+            for c in children:
+                for did in c.doc_ids:
+                    if did and did not in seen:
+                        seen.add(did)
+                        merged_doc_ids.append(did)
+                merged_doc_texts.update(c.doc_texts)
+            return SkillNode(
+                level=level + 1,
+                label="",
+                summary=summary_text or title,
+                vec=np.asarray(vec),
+                doc_ids=merged_doc_ids,
+                doc_texts=merged_doc_texts,
+                children=list(children),
+                folder_name="",
+            )
+
+        parent_results = await asyncio.gather(
+            *(_summarize_cluster(children) for children in cluster_buckets.values()),
+            return_exceptions=False,
+        )
+        next_layer = [p for p in parent_results if p is not None]
+        if not next_layer:
+            logging.warning("skill: no cluster summaries produced at level %d", level)
+            break
+        current_layer = next_layer
+        level += 1
+
+    roots: list[SkillNode] = current_layer
+
+    # ---- Phase 5: label every node (concurrent), then assign folders.
+    all_nodes: list[SkillNode] = []
+
+    def _collect(node: SkillNode) -> None:
+        all_nodes.append(node)
+        for c in node.children:
+            _collect(c)
+
+    for r in roots:
+        _collect(r)
+
+    progress(0.75, f"skill: labelling {len(all_nodes)} cluster node(s)")
+    label_sem = asyncio.Semaphore(SKILL_LABEL_CONCURRENCY)
+    labels_out = await asyncio.gather(
+        *(label_skill_node_one(n.summary, chat_mdl, label_sem) for n in all_nodes),
+        return_exceptions=False,
+    )
+    for n, lbl in zip(all_nodes, labels_out):
+        n.label = lbl or "cluster"
+
+    # Folder naming: roots get ``skill-NN-<label>``; deeper nodes get
+    # ``group-NN-<label>`` keyed by their position in the parent's
+    # children list (mirrors Corpus2Skill's skill_builder).
+    def _assign_folders(node: SkillNode, idx: int, is_root: bool) -> None:
+        slug = skill_safe_name(node.label)
+        prefix = f"skill-{idx:02d}" if is_root else f"group-{idx:02d}"
+        node.folder_name = f"{prefix}-{slug}" if slug else prefix
+        for ci, child in enumerate(node.children):
+            _assign_folders(child, ci, is_root=False)
+
+    for ri, root in enumerate(roots):
+        _assign_folders(root, ri, is_root=True)
+
+    # Final gate before the destructive delete + bulk insert. This is
+    # the most important check — without it a late cancel would still
+    # wipe the KB's existing ``skill``/``skill_all`` rows AND spend a
+    # full bulk-insert round-trip on data the caller no longer wants.
+    if ctx.has_canceled_func(ctx.id):
+        progress(-1, "skill: task has been canceled")
+        return
+
+    # ---- Phase 6: clean + bulk insert.
+    index = search.index_name(ctx.tenant_id)
+    try:
+        await thread_pool_exec(
+            settings.docStoreConn.delete,
+            {"compile_kwd": ["skill", "skill_all"]},
+            index,
+            ctx.kb_id,
+        )
+    except Exception:
+        logging.debug("skill: prior delete failed; relying on id-upsert")
+
+    rows = [skill_node_es_row(ctx, n) for n in all_nodes]
+    rows.append(skill_all_es_row(ctx, roots))
+    if not rows:
+        progress(1.0, "skill: nothing to persist")
+        return
+
+    try:
+        await thread_pool_exec(
+            settings.docStoreConn.insert,
+            rows,
+            index,
+            ctx.kb_id,
+        )
+    except Exception:
+        logging.exception("skill: bulk insert failed (rows=%d)", len(rows))
+        return
+
+    progress(
+        1.0,
+        f"skill: built {len(roots)} top-level skill(s), {len(all_nodes)} total node(s), {len(leaves)} doc(s)",
+    )

rag/svr/task_executor_refactor/dataset_wiki_generator.py

@@ -0,0 +1,805 @@
+#
+#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+"""KB-wide wiki / artifact compilation.
+
+Extracted from ``rag.svr.task_executor_refactor.task_handler`` where the
+same pipeline previously lived as a set of ``_wiki_*`` / ``_persist_wiki_*``
+methods and one ``_run_wiki`` orchestrator. The public entry point is
+:func:`run_wiki`.
+
+The pipeline runs MAP per (doc, template) — each MAP call resumes from
+its own ``artifact_map_extract`` ES rows — then REDUCE / PLAN / REFINE
+KB-wide via ``rag.advanced_rag.knowlege_compile.wiki``. Refined pages
+land in ES twice: once as searchable ``artifact_page`` rows and once as
+``artifact_entity`` / ``artifact_relation`` rows for the dataset
+Artifact tab's canvas graph.
+
+Design notes:
+
+* ``load_chunks_for_doc`` is injected rather than imported to keep the
+  module decoupled from ``TaskHandler``'s streaming chunk loader.
+* The eligibility loop resolves each doc's
+  ``parser_config.compilation_template_group_id`` to a template list
+  via ``CompilationTemplateGroupService.resolve_template_ids``. That
+  helper is duplicated as a small private function here so the module
+  stays free of a task_handler import (which would be circular).
+* The persistence helpers (``persist_wiki_pages_to_es`` etc.) are
+  exposed at module level for testing but are only called from
+  :func:`run_wiki` in production.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import AsyncIterator, Callable, Dict, List, Optional
+
+import xxhash
+
+from common import settings
+from common.constants import LLMType
+from common.misc_utils import thread_pool_exec
+from rag.nlp import search
+from rag.advanced_rag.knowlege_compile.wiki import (
+    wiki_map_from_chunks,
+    wiki_plan_from_reduction,
+    wiki_reduce_from_extracts,
+    wiki_refine_from_plan,
+)
+from rag.svr.task_executor_refactor.task_context import TaskContext
+
+
+# ----- tunables ------------------------------------------------------
+# Artifact-MAP tuning: how many chunks to feed per ``wiki_map_from_chunks``
+# invocation. The function does its own per-call resume-set load + ES
+# persist, so smaller batches mean more (small) ES round-trips but a flat
+# memory footprint. 64 keeps the resume-set re-reads cheap while leaving
+# room for the function's internal split_chunks packing to do real work.
+WIKI_MAP_BATCH_CHUNKS = 64
+
+# Per-node cap on ``source_chunk_ids`` carried by the canvas graph blob.
+# Pages can accumulate hundreds of source chunks; the graph response is
+# meant for fast canvas rendering, not full provenance audit, so we trim
+# each node's list. The full per-page list is still available on the
+# ``artifact_page`` row the UI deep-links into.
+WIKI_GRAPH_MAX_CHUNK_IDS_PER_NODE = 64
+
+# Title + comments stamped on every ``artifact_commit`` row produced by
+# :func:`run_wiki`'s regeneration path (as opposed to the dialog-edit
+# path, which carries the user's own title/comments).
+WIKI_REGEN_COMMIT_TITLE = "Regenerated by artifact compilation"
+WIKI_REGEN_COMMIT_COMMENTS_TEMPLATE = "Auto-update via run_wiki (action={action})"
+
+
+# ----- helpers -------------------------------------------------------
+
+
+def _parser_config_compilation_template_group_id(parser_config) -> str:
+    """Read the single template-group id from a doc's ``parser_config``.
+
+    Duplicated from ``task_handler`` so this module doesn't import from
+    it (avoids a circular import — ``task_handler`` imports this module
+    from its ``task_type == "artifact"`` dispatch branch).
+    """
+    if not isinstance(parser_config, dict):
+        return ""
+    gid = parser_config.get("compilation_template_group_id")
+    if isinstance(gid, str) and gid.strip():
+        return gid.strip()
+    ext = parser_config.get("ext")
+    if isinstance(ext, dict):
+        gid = ext.get("compilation_template_group_id")
+        if isinstance(gid, str) and gid.strip():
+            return gid.strip()
+    return ""
+
+
+def _parser_config_compilation_template_ids(parser_config, tenant_id: str) -> list[str]:
+    """Resolve a doc's ``parser_config`` to its compile-template ids by
+    looking up the configured group. Returns ``[]`` if the doc has no
+    group set or the group cannot be resolved."""
+    from api.db.services.compilation_template_group_service import (
+        CompilationTemplateGroupService,
+    )
+
+    group_id = _parser_config_compilation_template_group_id(parser_config)
+    if not group_id:
+        return []
+    return CompilationTemplateGroupService.resolve_template_ids(group_id, tenant_id)
+
+
+# ----- persistence ---------------------------------------------------
+
+
+async def persist_wiki_pages_to_es(
+    ctx: TaskContext,
+    pages: List[Dict],
+    embd_mdl,
+) -> None:
+    """Insert one ES row per generated artifact page using the
+    knowledge-compilation schema:
+
+      id                  xxh64(kb_id + ":" + slug)
+      compile_kwd         "artifact_page"
+      slug_kwd            page.slug
+      title_kwd           page.title
+      page_type_kwd       page.page_type
+      entity_names_kwd    page.entity_names
+      outlinks_kwd        page.outlinks
+      related_kb_pages_kwd page.related_kb_pages
+      source_chunk_ids    page.source_chunk_ids
+      source_doc_ids      page.source_doc_ids
+      kb_id               ctx.kb_id
+      content_with_weight rendered markdown
+      content_ltks /
+      content_sm_ltks     tokenize(content_md + summary)
+      q_<dim>_vec         embed(summary)
+
+    ``action`` is intentionally not stored — it's a planner artifact
+    and has no meaning post-write.
+    """
+    if not pages:
+        return
+
+    from rag.nlp import rag_tokenizer
+    from common.doc_store.doc_store_base import OrderByExpr
+    from api.db.services.file_commit_service import (
+        FileCommitService as WikiCommitService,
+    )
+
+    index = search.index_name(ctx.tenant_id)
+    kb_id_str = str(ctx.kb_id)
+
+    # Capture the prior rendered content for every slug we're about to
+    # overwrite, so the per-page commit row downstream has a real diff
+    # baseline. Single batch read by slug_kwd IN [...] — one round-trip
+    # regardless of page count. Failures here degrade gracefully.
+    target_slugs: list[str] = [(p.get("slug") or "").strip() for p in pages if isinstance(p.get("slug"), str) and p.get("slug")]
+    prior_by_slug: dict[str, str] = {}
+    if target_slugs:
+        try:
+            res = await thread_pool_exec(
+                settings.docStoreConn.search,
+                ["id", "slug_kwd", "content_with_weight"],
+                [],
+                {"compile_kwd": ["artifact_page"], "slug_kwd": list(target_slugs)},
+                [],
+                OrderByExpr(),
+                0,
+                max(len(target_slugs), 1),
+                index,
+                [ctx.kb_id],
+            )
+            field_map = settings.docStoreConn.get_fields(
+                res,
+                ["id", "slug_kwd", "content_with_weight"],
+            )
+            for row in (field_map or {}).values():
+                s = row.get("slug_kwd")
+                c = row.get("content_with_weight")
+                if isinstance(s, str) and isinstance(c, str):
+                    prior_by_slug[s] = c
+        except Exception:
+            logging.exception(
+                "wiki_persist: prior-content read failed for kb=%s; commit audit will treat all pages as creations",
+                kb_id_str,
+            )
+
+    # Batch the summary embeddings in one model call. Empty strings are
+    # swapped for a single space so the encoder doesn't reject them —
+    # they still yield a vector but contribute nothing meaningful.
+    summaries = [(p.get("summary") or "").strip() for p in pages]
+    embed_inputs = [s if s else " " for s in summaries]
+    try:
+        embeddings, _ = await thread_pool_exec(embd_mdl.encode, embed_inputs)
+    except Exception:
+        logging.exception(
+            "wiki_persist: summary embedding batch failed for kb=%s",
+            kb_id_str,
+        )
+        return
+    try:
+        n_emb = len(embeddings) if embeddings is not None else 0
+    except TypeError:
+        n_emb = 0
+    if n_emb != len(pages):
+        logging.warning(
+            "artifact_persist: embedding count %d != pages %d for kb=%s; aborting",
+            n_emb,
+            len(pages),
+            kb_id_str,
+        )
+        return
+
+    rows: List[Dict] = []
+    for page, vec in zip(pages, embeddings):
+        slug = page.get("slug") or ""
+        if not slug:
+            continue
+        title = page.get("title") or slug
+        summary = page.get("summary") or ""
+        content_md = page.get("content_md_rendered") or page.get("content_md") or page.get("content_md_raw") or ""
+
+        vec_list = vec.tolist() if hasattr(vec, "tolist") else list(vec)
+        if not vec_list:
+            logging.warning(
+                "artifact_persist: empty embedding for slug=%s; skipping",
+                slug,
+            )
+            continue
+
+        text_for_search = (content_md + "\n\n" + summary).strip()
+        content_ltks = rag_tokenizer.tokenize(text_for_search) if text_for_search else ""
+        content_sm_ltks = rag_tokenizer.fine_grained_tokenize(content_ltks) if content_ltks else ""
+
+        row_id = xxhash.xxh64(
+            f"{kb_id_str}:{slug}".encode("utf-8", "surrogatepass"),
+        ).hexdigest()
+
+        rows.append(
+            {
+                "id": row_id,
+                "kb_id": kb_id_str,
+                "doc_id": kb_id_str,  # sentinel; KB-scoped row, real provenance in source_doc_ids
+                "compile_kwd": "artifact_page",
+                "slug_kwd": slug,
+                "title_kwd": title,
+                "page_type_kwd": page.get("page_type") or "concept",
+                "entity_names_kwd": list(page.get("entity_names") or []),
+                "outlinks_kwd": list(page.get("outlinks") or []),
+                "outlinks_int": len(list(page.get("outlinks") or [])),
+                "related_kb_pages_kwd": list(page.get("related_kb_pages") or []),
+                "source_chunk_ids": list(page.get("source_chunk_ids") or []),
+                "source_doc_ids": list(page.get("source_doc_ids") or []),
+                "content_with_weight": content_md,
+                # Summary kept verbatim alongside the rendered body so the
+                # viewer can render it as a distinct (smaller) block above
+                # the main content.
+                "summary_with_weight": summary,
+                "content_ltks": content_ltks,
+                "content_sm_ltks": content_sm_ltks,
+                f"q_{len(vec_list)}_vec": vec_list,
+                "available_int": 1,
+            }
+        )
+
+    if not rows:
+        return
+
+    try:
+        await thread_pool_exec(settings.docStoreConn.insert, rows, index, ctx.kb_id)
+    except Exception:
+        logging.exception(
+            "wiki_persist: bulk insert failed for kb=%s (rows=%d)",
+            kb_id_str,
+            len(rows),
+        )
+        return
+
+    # Audit trail: one ArtifactCommit row per page whose rendered
+    # content actually changed (record_edit silently skips empty diffs).
+    # Best-effort — commit failures log but don't fail the artifact
+    # compile.
+    for page in pages:
+        slug = (page.get("slug") or "").strip()
+        if not slug:
+            continue
+        if not prior_by_slug.get(slug, ""):
+            continue
+        content_md = page.get("content_md_rendered") or page.get("content_md") or page.get("content_md_raw") or ""
+        action = (page.get("action") or "CREATE").upper()
+        try:
+            WikiCommitService.record_page_edit(
+                tenant_id=ctx.tenant_id,
+                kb_id=ctx.kb_id,
+                page_type=page.get("page_type") or "concept",
+                slug=slug,
+                content_before=prior_by_slug.get(slug, ""),
+                content_after=content_md,
+                title=WIKI_REGEN_COMMIT_TITLE,
+                comments=WIKI_REGEN_COMMIT_COMMENTS_TEMPLATE.format(action=action),
+                user_id=None,  # system commit — no human author
+            )
+        except Exception:
+            logging.exception(
+                "wiki_persist: commit record failed for kb=%s slug=%s",
+                kb_id_str,
+                slug,
+            )
+
+
+def build_wiki_page_graph(
+    pages: List[Dict],
+    kb_id: str,
+) -> tuple[List[Dict], List[Dict]]:
+    """Project the REFINE-emitted page list onto per-entity and
+    per-relation ES rows.
+
+    Returns:
+        (entity_rows, relation_rows)
+
+    Both lists are ES-ready docs (one per node / per surviving edge)
+    using the standard artifact envelope. They are BM25-only (no
+    ``q_<dim>_vec``) — entities carry ``content_ltks`` derived from
+    ``slug + " " + summary`` so name/summary lookups hit the lexical
+    index.
+
+    ``by_slug`` is kept internally only to filter dangling outlinks
+    (a target slug not present as a node in this KB).
+    """
+    from rag.nlp import rag_tokenizer
+
+    by_slug: Dict[str, Dict] = {}
+    entity_rows: List[Dict] = []
+    for p in pages or []:
+        slug = (p.get("slug") or "").strip()
+        if not slug:
+            continue
+        outlinks_raw = p.get("outlinks") or []
+        # ``weight`` is the page's outlink count. Drives node size on
+        # the canvas. Computed on the raw outlink list before dangling-
+        # target filtering so visual weight reflects what the writer
+        # actually emitted.
+        weight = len(outlinks_raw) if isinstance(outlinks_raw, list) else 0
+
+        # Per-node provenance: union of source chunks REFINE attributed
+        # to this page. Dedup preserves first-seen order; cap the list
+        # to keep the graph blob small.
+        raw_chunk_ids = p.get("source_chunk_ids") or []
+        seen_chunk_ids: dict[str, None] = {}
+        for cid in raw_chunk_ids:
+            if isinstance(cid, str) and cid and cid not in seen_chunk_ids:
+                seen_chunk_ids[cid] = None
+                if len(seen_chunk_ids) >= WIKI_GRAPH_MAX_CHUNK_IDS_PER_NODE:
+                    break
+        capped_chunk_ids = list(seen_chunk_ids.keys())
+
+        page_type = p.get("page_type") or "concept"
+        description = p.get("summary") or ""
+        name = p.get("title") or slug
+        aliases = list(p.get("entity_names") or [])
+
+        by_slug[slug] = {
+            "slug": slug,
+            "name": name,
+            "aliases": aliases,
+            "description": description,
+            "type": page_type,
+            "weight": weight,
+            "source_chunk_ids": capped_chunk_ids,
+        }
+
+        # Per-entity ES row. content_ltks is built from slug + summary
+        # so BM25 hits both the deep-link key and human prose.
+        content_text = (slug + " " + description).strip()
+        entity_payload = {
+            "slug": slug,
+            "name": name,
+            "aliases": aliases,
+            "description": description,
+            "type": page_type,
+            "weight": weight,
+        }
+        entity_rows.append(
+            {
+                "id": xxhash.xxh64(
+                    f"artifact_entity:{kb_id}:{slug}".encode("utf-8", "surrogatepass"),
+                ).hexdigest(),
+                "kb_id": kb_id,
+                "doc_id": kb_id,  # KB-scoped sentinel
+                "available_int": 1,
+                "compile_kwd": "artifact_entity",
+                "type_kwd": "artifact_" + page_type,
+                "slug_kwd": slug,
+                "weight_int": int(weight),
+                "source_chunk_ids": capped_chunk_ids,
+                "content_ltks": rag_tokenizer.tokenize(content_text) if content_text else "",
+                "content_with_weight": json.dumps(entity_payload, ensure_ascii=False),
+            }
+        )
+
+    relation_rows: List[Dict] = []
+    for p in pages or []:
+        src = (p.get("slug") or "").strip()
+        if not src or src not in by_slug:
+            continue
+        for raw_target in p.get("outlinks") or []:
+            if isinstance(raw_target, str):
+                tgt = raw_target.strip()
+            elif isinstance(raw_target, dict):
+                tgt = str(raw_target.get("slug") or "").strip()
+            else:
+                tgt = ""
+            if not tgt or tgt == src or tgt not in by_slug:
+                continue
+            relation_payload = {"from": src, "to": tgt}
+            relation_rows.append(
+                {
+                    "id": xxhash.xxh64(
+                        f"artifact_relation:{kb_id}:{src}:{tgt}".encode("utf-8", "surrogatepass"),
+                    ).hexdigest(),
+                    "kb_id": kb_id,
+                    "doc_id": kb_id,
+                    "available_int": 1,
+                    "compile_kwd": "artifact_relation",
+                    "type_kwd": "artifact_relation",
+                    "from_kwd": src,
+                    "to_kwd": tgt,
+                    "content_with_weight": json.dumps(relation_payload, ensure_ascii=False),
+                }
+            )
+
+    return entity_rows, relation_rows
+
+
+async def persist_wiki_page_graph_to_es(
+    ctx: TaskContext,
+    pages: List[Dict],
+) -> None:
+    """Materialize and store the per-entity / per-relation ES rows
+    derived from artifact pages.
+
+    Writes two row types — both delete-then-insert for idempotent
+    re-runs:
+
+    1. ``compile_kwd="artifact_entity"`` — one row per page node,
+       BM25-only via ``content_ltks``.
+    2. ``compile_kwd="artifact_relation"`` — one row per surviving
+       edge (dangling outlinks dropped by the builder).
+
+    Also sweeps any leftover legacy ``artifact_page_graph`` blob so
+    the index doesn't accumulate stale state.
+    """
+    kb_id_str = str(ctx.kb_id)
+    entity_rows, relation_rows = build_wiki_page_graph(pages or [], kb_id_str)
+
+    index = search.index_name(ctx.tenant_id)
+
+    async def _replace_bucket(kwd: str, rows: List[Dict]) -> None:
+        try:
+            await thread_pool_exec(
+                settings.docStoreConn.delete,
+                {"compile_kwd": kwd},
+                index,
+                ctx.kb_id,
+            )
+        except Exception:
+            logging.debug(
+                "%s: prior delete failed; relying on id-upsert",
+                kwd,
+            )
+        if not rows:
+            return
+        try:
+            await thread_pool_exec(
+                settings.docStoreConn.insert,
+                rows,
+                index,
+                ctx.kb_id,
+            )
+        except Exception:
+            logging.exception(
+                "%s: insert failed for kb=%s (%d rows)",
+                kwd,
+                kb_id_str,
+                len(rows),
+            )
+
+    async def _sweep_legacy_blob() -> None:
+        try:
+            await thread_pool_exec(
+                settings.docStoreConn.delete,
+                {"compile_kwd": "artifact_page_graph"},
+                index,
+                ctx.kb_id,
+            )
+        except Exception:
+            logging.debug(
+                "artifact_page_graph: legacy blob sweep failed for kb=%s",
+                kb_id_str,
+            )
+
+    await asyncio.gather(
+        _replace_bucket("artifact_entity", entity_rows),
+        _replace_bucket("artifact_relation", relation_rows),
+        _sweep_legacy_blob(),
+    )
+
+
+# ----- main entry ----------------------------------------------------
+
+
+async def run_wiki(
+    ctx: TaskContext,
+    embedding_model,
+    load_chunks_for_doc: Callable[..., AsyncIterator[list[dict]]],
+) -> None:
+    """KB-wide artifact compilation task.
+
+    Runs after the user clicks the "Artifact" button in the dataset
+    generate menu. Iterates every doc in the KB whose parser_config
+    has a compilation template group resolving to an artifacts-kind
+    child, runs MAP per-doc (which uses ES-stored resume rows to skip
+    chunks already processed in a previous run), then runs REDUCE /
+    PLAN / REFINE KB-wide and persists pages.
+
+    Batching: each MAP call uses ``batch_size_cap=8`` and
+    ``window_fraction=0.5`` — i.e. roll over to a new batch when the
+    current batch reaches 8 chunks OR its accumulated token count
+    exceeds 50% of the chat model's ``max_length``.
+    """
+    # Local imports so this module doesn't drag in the API service
+    # layer at import time — that's a source of circular-import risk
+    # given how much lives under ``api.db.services``.
+    from api.db.services.document_service import DocumentService
+    from api.db.services.knowledgebase_service import KnowledgebaseService
+    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,
+        get_model_config_from_provider_instance,
+    )
+    from api.apps.restful_apis.chunk_api import _compilation_template_kind
+
+    progress = ctx.progress_cb
+    progress(0.0, "Loading documents for wiki compilation...")
+
+    # 1. Resolve KB metadata for PLAN.
+    ok, kb = KnowledgebaseService.get_by_id(ctx.kb_id)
+    if not ok:
+        progress(-1, f"KB {ctx.kb_id} not found.")
+        return
+    kb_name = kb.name
+    kb_description = kb.description
+
+    # 2. Pick docs eligible for artifact compilation (those whose
+    # configured template group resolves to at least one artifacts-kind
+    # child). The frontend Artifact button targets the KB, but the
+    # per-doc opt-in is what gates inclusion.
+    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 = []
+    for d in all_docs or []:
+        pc = d.get("parser_config") or {}
+        for template_id in _parser_config_compilation_template_ids(pc, ctx.tenant_id):
+            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 == "artifacts":
+                eligible.append((d, template_id))
+                break
+    if not eligible:
+        progress(1.0, "No documents are configured for wiki compilation.")
+        return
+
+    # 3. Resolve chat models. MAP is per-(doc, template) so each pair
+    # uses its template's own ``llm_id``. REDUCE / PLAN / REFINE are
+    # KB-wide and need exactly one model — we pick the first eligible
+    # template's ``llm_id`` as the canonical KB 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 = get_model_config_from_provider_instance(
+                    ctx.tenant_id,
+                    LLMType.CHAT,
+                    key,
+                )
+        except Exception:
+            logging.exception(
+                "wiki: chat model resolution failed for llm_id=%s (kb=%s); falling back to tenant default",
+                key,
+                ctx.kb_id,
+            )
+            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
+
+    def _stage_cb(prefix: str):
+        def _cb(*args, **kwargs):
+            try:
+                if args and isinstance(args[0], (int, float)):
+                    msg = args[1] if len(args) > 1 else kwargs.get("msg", "")
+                    progress(msg=f"{prefix} {msg}")
+                else:
+                    msg = kwargs.get("msg") or (args[0] if args else "")
+                    progress(msg=f"{prefix} {msg}")
+            except Exception:
+                logging.exception("wiki: progress callback failed")
+
+        return _cb
+
+    # 4. MAP per eligible doc. Each MAP call's own resume mechanism
+    # (artifact_map_extract rows keyed by chunk_id) skips chunks that
+    # were already processed in a prior run — this is the incremental
+    # behavior the user asked for.
+    #
+    # ``kb_chat_llm_id`` is captured from the first eligible template
+    # and used as the canonical chat model for REDUCE/PLAN/REFINE.
+    # ``kb_writer_example`` follows the same first-template-wins rule:
+    # the REFINE writer's page-structure section is pulled from the
+    # first eligible artifact template's ``parser_config.example``
+    # override (None falls back to the built-in ``WIKI_TEMPLATE_EXAMPLE``).
+    kb_chat_llm_id: Optional[str] = None
+    kb_writer_example: Optional[str] = None
+    n_docs = len(eligible)
+    for i, (doc, template_id) in enumerate(eligible):
+        doc_id = doc["id"]
+        progress(
+            0.05 + 0.6 * (i / n_docs),
+            f"MAP {i + 1}/{n_docs}: {doc.get('name', doc_id)}",
+        )
+
+        template = CompilationTemplateService.get_saved(template_id, ctx.tenant_id)
+        if not template:
+            logging.warning(
+                "artifact: template %s not found for doc %s; skipping",
+                template_id,
+                doc_id,
+            )
+            continue
+        parser_cfg = template.get("config") or {}
+
+        map_llm_id = (parser_cfg.get("llm_id") or "").strip() if isinstance(parser_cfg, dict) else ""
+        map_chat_mdl = _bundle_for(map_llm_id)
+        if kb_chat_llm_id is None:
+            # First eligible template wins — canonical for KB-wide
+            # REDUCE/PLAN/REFINE further down.
+            kb_chat_llm_id = map_llm_id or None
+            tmpl_example = parser_cfg.get("example") if isinstance(parser_cfg, dict) else None
+            if isinstance(tmpl_example, str) and tmpl_example.strip():
+                kb_writer_example = tmpl_example
+
+        # Stream the doc's chunks in batches and call MAP per batch so
+        # peak memory stays bounded for long docs.
+        agg = {"entities": 0, "concepts": 0, "claims": 0, "relations": 0}
+        agg_delta = {"new": 0, "changed": 0, "unchanged": 0, "deleted": 0}
+        doc_had_delta = False
+        saw_any = False
+        batch_no = 0
+        async for batch in load_chunks_for_doc(
+            ctx.tenant_id,
+            ctx.kb_id,
+            doc_id,
+            batch_size=WIKI_MAP_BATCH_CHUNKS,
+        ):
+            saw_any = True
+            batch_no += 1
+            try:
+                phase1 = await wiki_map_from_chunks(
+                    chunks=batch,
+                    chat_mdl=map_chat_mdl,
+                    embd_mdl=embedding_model,
+                    doc_id=doc_id,
+                    tenant_id=ctx.tenant_id,
+                    kb_id=ctx.kb_id,
+                    language=ctx.language,
+                    callback=_stage_cb(f"[wiki MAP {i + 1}/{n_docs} b{batch_no}]"),
+                    parser_config=parser_cfg,
+                    batch_size_cap=8,
+                    window_fraction=0.5,
+                )
+            except Exception:
+                logging.exception(
+                    "wiki: MAP failed for doc %s batch %d",
+                    doc_id,
+                    batch_no,
+                )
+                continue
+            for k in agg.keys():
+                agg[k] += len(phase1.get(k) or [])
+            meta = phase1.get("_meta") or {}
+            if isinstance(meta, dict):
+                for k in agg_delta.keys():
+                    agg_delta[k] += int(meta.get(k, 0) or 0)
+                if meta.get("had_delta"):
+                    doc_had_delta = True
+
+        if not saw_any:
+            logging.info("wiki: no chunks for doc %s; skipping", doc_id)
+            continue
+        logging.info(
+            "wiki: MAP doc=%s entities=%d concepts=%d claims=%d relations=%d (batches=%d, new=%d changed=%d unchanged=%d deleted=%d, delta=%s)",
+            doc_id,
+            agg["entities"],
+            agg["concepts"],
+            agg["claims"],
+            agg["relations"],
+            batch_no,
+            agg_delta["new"],
+            agg_delta["changed"],
+            agg_delta["unchanged"],
+            agg_delta["deleted"],
+            doc_had_delta,
+        )
+
+    # 5. REDUCE / PLAN / REFINE KB-wide. Each phase has its own
+    # input_hash gate (REDUCE keys off the MAP-state hash, PLAN off
+    # REDUCE's hash, REFINE off PLAN's hash) so re-runs without an
+    # upstream delta short-circuit at the cache layer.
+    kb_chat_mdl = _bundle_for(kb_chat_llm_id)
+    try:
+        progress(0.65, "Reducing extracts KB-wide...")
+        await wiki_reduce_from_extracts(
+            chat_mdl=kb_chat_mdl,
+            embd_mdl=embedding_model,
+            tenant_id=ctx.tenant_id,
+            kb_id=ctx.kb_id,
+            callback=_stage_cb("[wiki REDUCE]"),
+        )
+
+        progress(0.75, "Planning wiki pages...")
+        await wiki_plan_from_reduction(
+            chat_mdl=kb_chat_mdl,
+            embd_mdl=embedding_model,
+            tenant_id=ctx.tenant_id,
+            kb_id=ctx.kb_id,
+            kb_name=kb_name,
+            kb_description=kb_description,
+            callback=_stage_cb("[wiki PLAN]"),
+        )
+
+        progress(0.85, "Refining pages...")
+        pages = await wiki_refine_from_plan(
+            chat_mdl=kb_chat_mdl,
+            embd_mdl=embedding_model,
+            tenant_id=ctx.tenant_id,
+            kb_id=ctx.kb_id,
+            callback=_stage_cb("[wiki REFINE]"),
+            example=kb_writer_example,
+        )
+    except Exception:
+        logging.exception("wiki: REDUCE/PLAN/REFINE failed for kb %s", ctx.kb_id)
+        progress(-1, "Wiki pipeline failed during REDUCE/PLAN/REFINE.")
+        return
+
+    # 6. Persist searchable artifact_page rows.
+    try:
+        await persist_wiki_pages_to_es(ctx, pages or [], embedding_model)
+    except Exception:
+        logging.exception("wiki: ES 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, 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).")

rag/svr/task_executor_refactor/raptor_service.py

@@ -21,12 +21,14 @@ Provides [`RaptorService`](rag/svr/task_executor_refactor/raptor_service.py:48)
 """
 
 import copy
+import json
 import logging
 import os
 from datetime import datetime
 from typing import Dict, List, Optional, Set, Tuple
 
 import numpy as np
+import xxhash
 
 from api.db.services.document_service import DocumentService
 from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID
@@ -48,6 +50,31 @@ from rag.utils.raptor_utils import (
 from rag.svr.task_executor_refactor.task_context import TaskContext
 
 
+def _sum_tree_text_tokens(tree) -> int:
+    """Count tokens across every ``title`` string in the RAPTOR tree.
+
+    Mirrors the legacy ``tk_count`` semantic (sum over summary texts)
+    so the orchestrator's downstream logging / billing keeps working
+    when the tree path replaces the per-summary rows. Walks the dict
+    iteratively to avoid recursion-limit issues on deep trees.
+    """
+    if not isinstance(tree, dict):
+        return 0
+    total = 0
+    stack = [tree]
+    while stack:
+        node = stack.pop()
+        if not isinstance(node, dict):
+            continue
+        title = node.get("title")
+        if isinstance(title, str) and title:
+            total += num_tokens_from_string(title)
+        children = node.get("children")
+        if isinstance(children, list):
+            stack.extend(children)
+    return total
+
+
 class RaptorService:
     """Service for RAPTOR summary generation.
 
@@ -106,15 +133,11 @@ class RaptorService:
         # Determine scope
         if raptor_config.get("scope", "file") == "file":
             res, tk_count = await self._run_file_level_raptor(
-                raptor_config, tree_builder, clustering_method,
-                chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
-                max_errors, res, tk_count, cleanup_raptor_chunks
+                raptor_config, tree_builder, clustering_method, chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id, max_errors, res, tk_count, cleanup_raptor_chunks
             )
         else:
             res, tk_count = await self._run_dataset_level_raptor(
-                raptor_config, tree_builder, clustering_method,
-                chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
-                max_errors, res, tk_count, cleanup_raptor_chunks
+                raptor_config, tree_builder, clustering_method, chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id, max_errors, res, tk_count, cleanup_raptor_chunks
             )
 
         return res, tk_count, cleanup_raptor_chunks
@@ -135,15 +158,11 @@ class RaptorService:
             }
         return doc_info_by_id
 
-    async def _run_file_level_raptor(
-        self, raptor_config, tree_builder, clustering_method,
-        chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
-        max_errors, res, tk_count, cleanup_raptor_chunks
-    ):
+    async def _run_file_level_raptor(self, raptor_config, tree_builder, clustering_method, chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id, max_errors, res, tk_count, cleanup_raptor_chunks):
         """Run RAPTOR at file level (per document)."""
         ctx = self._task_context
         fake_doc_id = GRAPH_RAPTOR_FAKE_DOC_ID
-        if self._task_context.write_interceptor: # dry run mode
+        if self._task_context.write_interceptor:  # dry run mode
             dataset_methods = set()
         else:
             dataset_methods = await self._get_raptor_chunk_methods(fake_doc_id, ctx.tenant_id, ctx.kb_id)
@@ -155,7 +174,7 @@ class RaptorService:
 
         for x, doc_id in enumerate(doc_ids):
             if self._should_skip_raptor(doc_id, doc_info_by_id, raptor_config):
-                self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
+                self._task_context.progress_cb(prog=(x + 1.0) / len(doc_ids))
                 continue
             if self._task_context.write_interceptor:
                 existing_methods = set()
@@ -164,12 +183,10 @@ class RaptorService:
             if tree_builder in existing_methods:
                 has_file_level_target = True
                 if existing_methods != {tree_builder}:
-                    self._schedule_raptor_cleanup(
-                        doc_id, tree_builder, cleanup_raptor_chunks
-                    )
+                    self._schedule_raptor_cleanup(doc_id, tree_builder, cleanup_raptor_chunks)
                     self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} will remove old RAPTOR summaries after insert.")
                 self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} already has {tree_builder} RAPTOR chunks, skipping.")
-                self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
+                self._task_context.progress_cb(prog=(x + 1.0) / len(doc_ids))
                 continue
 
             if existing_methods:
@@ -180,36 +197,25 @@ class RaptorService:
                 continue
 
             before_generate = len(res)
-            new_chunks, new_tk_count = await self._generate_raptor(
-                chunks, doc_id, raptor_config, chat_mdl, embd_mdl,
-                tree_builder, clustering_method, max_errors, doc_info_by_id
-            )
+            new_chunks, new_tk_count = await self._generate_raptor(chunks, doc_id, raptor_config, chat_mdl, embd_mdl, tree_builder, clustering_method, max_errors, doc_info_by_id)
             res.extend(new_chunks)
             tk_count += new_tk_count
 
             if len(res) > before_generate:
                 has_file_level_target = True
                 if existing_methods:
-                    self._schedule_raptor_cleanup(
-                        doc_id, tree_builder, cleanup_raptor_chunks
-                    )
-            self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
+                    self._schedule_raptor_cleanup(doc_id, tree_builder, cleanup_raptor_chunks)
+            self._task_context.progress_cb(prog=(x + 1.0) / len(doc_ids))
 
         if remove_dataset_summaries:
             if has_file_level_target:
-                self._schedule_raptor_cleanup(
-                    fake_doc_id, None, cleanup_raptor_chunks
-                )
+                self._schedule_raptor_cleanup(fake_doc_id, None, cleanup_raptor_chunks)
             else:
                 self._task_context.progress_cb(msg="[RAPTOR] kept dataset-level summaries because no file-level summaries were built.")
 
         return res, tk_count
 
-    async def _run_dataset_level_raptor(
-        self, raptor_config, tree_builder, clustering_method,
-        chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
-        max_errors, res, tk_count, cleanup_raptor_chunks
-    ):
+    async def _run_dataset_level_raptor(self, raptor_config, tree_builder, clustering_method, chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id, max_errors, res, tk_count, cleanup_raptor_chunks):
         """Run RAPTOR at dataset level (all documents combined)."""
         ctx = self._task_context
         fake_doc_id = GRAPH_RAPTOR_FAKE_DOC_ID
@@ -230,9 +236,7 @@ class RaptorService:
                 migrated_file_docs += 1
 
         if migrated_file_docs:
-            self._task_context.progress_cb(
-                msg=f"[RAPTOR] will remove file-level summaries for {migrated_file_docs} docs after dataset-level build succeeds."
-            )
+            self._task_context.progress_cb(msg=f"[RAPTOR] will remove file-level summaries for {migrated_file_docs} docs after dataset-level build succeeds.")
 
         if self._task_context.write_interceptor:
             existing_methods = set()
@@ -240,9 +244,7 @@ class RaptorService:
             existing_methods = await self._get_raptor_chunk_methods(fake_doc_id, ctx.tenant_id, ctx.kb_id)
         if tree_builder in existing_methods:
             if existing_methods != {tree_builder}:
-                self._schedule_raptor_cleanup(
-                    fake_doc_id, tree_builder, cleanup_raptor_chunks
-                )
+                self._schedule_raptor_cleanup(fake_doc_id, tree_builder, cleanup_raptor_chunks)
                 self._task_context.progress_cb(msg="[RAPTOR] will remove old dataset-level RAPTOR summaries after insert.")
             for doc_id in file_cleanup_doc_ids:
                 self._schedule_raptor_cleanup(doc_id, None, cleanup_raptor_chunks)
@@ -262,10 +264,7 @@ class RaptorService:
             return res, tk_count
 
         before_generate = len(res)
-        new_chunks, new_tk_count = await self._generate_raptor(
-            chunks, fake_doc_id, raptor_config, chat_mdl, embd_mdl,
-            tree_builder, clustering_method, max_errors, doc_info_by_id
-        )
+        new_chunks, new_tk_count = await self._generate_raptor(chunks, fake_doc_id, raptor_config, chat_mdl, embd_mdl, tree_builder, clustering_method, max_errors, doc_info_by_id)
         res.extend(new_chunks)
         tk_count += new_tk_count
 
@@ -273,15 +272,11 @@ class RaptorService:
             for doc_id in file_cleanup_doc_ids:
                 self._schedule_raptor_cleanup(doc_id, None, cleanup_raptor_chunks)
             if migrate_dataset_summaries:
-                self._schedule_raptor_cleanup(
-                    fake_doc_id, tree_builder, cleanup_raptor_chunks
-                )
+                self._schedule_raptor_cleanup(fake_doc_id, tree_builder, cleanup_raptor_chunks)
 
         return res, tk_count
 
-    def _should_skip_raptor(
-        self, doc_id: str, doc_info_by_id: Dict, raptor_config: Dict
-    ) -> bool:
+    def _should_skip_raptor(self, doc_id: str, doc_info_by_id: Dict, raptor_config: Dict) -> bool:
         """Check if RAPTOR should be skipped for a document."""
         ctx = self._task_context
         doc_info = doc_info_by_id.get(doc_id, {})
@@ -297,67 +292,64 @@ class RaptorService:
             return True
         return False
 
-    def _load_doc_chunks(self, doc_id: str, vctr_nm: str) -> List[Tuple[str, np.ndarray]]:
-        """Load chunks for a single document."""
+    def _load_doc_chunks(self, doc_id: str, vctr_nm: str) -> List[Tuple[str, np.ndarray, str]]:
+        """Load chunks for a single document.
+
+        Returns ``(content, vector, chunk_id)`` triples so downstream
+        RAPTOR can attach ``source_chunk_ids`` provenance onto every
+        summary it produces. ``chunk_id`` may be an empty string if the
+        retriever didn't surface one — defensive against legacy rows.
+        """
         ctx = self._task_context
-        chunks = []
+        chunks: List[Tuple[str, np.ndarray, str]] = []
         skipped_chunks = 0
 
-        fields = ["content_with_weight", vctr_nm]
-        for d in settings.retriever.chunk_list(
-            doc_id, ctx.tenant_id, [str(ctx.kb_id)],
-            fields=fields,
-            sort_by_position=True
-        ):
+        # ``id`` is included so the source-chunk provenance survives
+        # through summarization; the retriever otherwise drops it when
+        # ``fields`` is provided.
+        fields = ["id", "content_with_weight", vctr_nm]
+        for d in settings.retriever.chunk_list(doc_id, ctx.tenant_id, [str(ctx.kb_id)], fields=fields, sort_by_position=True):
             if vctr_nm not in d or d[vctr_nm] is None:
                 skipped_chunks += 1
                 logging.warning(f"RAPTOR: Chunk missing vector field '{vctr_nm}' in doc {doc_id}, skipping")
                 continue
-            chunks.append((d["content_with_weight"], np.array(d[vctr_nm])))
+            chunks.append((d["content_with_weight"], np.array(d[vctr_nm]), str(d.get("id") or "")))
 
         if skipped_chunks > 0:
-            self._task_context.progress_cb(
-                msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}' for doc {doc_id}."
-            )
+            self._task_context.progress_cb(msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}' for doc {doc_id}.")
         if not chunks:
             logging.warning(f"RAPTOR: No valid chunks with vectors found for doc {doc_id}")
             self._task_context.progress_cb(msg=f"[WARN] No valid chunks with vectors found for doc {doc_id}, skipping")
 
         return chunks
 
-    def _load_all_doc_chunks(
-        self, doc_ids: List[str], vctr_nm: str, skipped_doc_ids: Set[str]
-    ) -> List[Tuple[str, np.ndarray]]:
-        """Load chunks for all documents."""
+    def _load_all_doc_chunks(self, doc_ids: List[str], vctr_nm: str, skipped_doc_ids: Set[str]) -> List[Tuple[str, np.ndarray, str]]:
+        """Load chunks for all documents — returns provenance-carrying
+        ``(content, vector, chunk_id)`` triples. See ``_load_doc_chunks``
+        for the per-doc variant."""
         ctx = self._task_context
-        chunks = []
+        chunks: List[Tuple[str, np.ndarray, str]] = []
         skipped_chunks = 0
 
-        fields = ["content_with_weight", vctr_nm]
+        fields = ["id", "content_with_weight", vctr_nm]
         for doc_id in doc_ids:
             if doc_id in skipped_doc_ids:
                 continue
-            for d in settings.retriever.chunk_list(
-                doc_id, ctx.tenant_id, [str(ctx.kb_id)],
-                fields=fields,
-                sort_by_position=True
-            ):
+            for d in settings.retriever.chunk_list(doc_id, ctx.tenant_id, [str(ctx.kb_id)], fields=fields, sort_by_position=True):
                 if vctr_nm not in d or d[vctr_nm] is None:
                     skipped_chunks += 1
                     logging.warning(f"RAPTOR: Chunk missing vector field '{vctr_nm}' in doc {doc_id}, skipping")
                     continue
-                chunks.append((d["content_with_weight"], np.array(d[vctr_nm])))
+                chunks.append((d["content_with_weight"], np.array(d[vctr_nm]), str(d.get("id") or "")))
 
         if skipped_chunks > 0:
-            self._task_context.progress_cb(
-                msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}'."
-            )
+            self._task_context.progress_cb(msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}'.")
 
         return chunks
 
     async def _generate_raptor(
         self,
-        chunks: List[Tuple[str, np.ndarray]],
+        chunks: List[Tuple[str, np.ndarray, str]],
         doc_id: str,
         raptor_config: Dict,
         chat_mdl,
@@ -366,10 +358,19 @@ class RaptorService:
         clustering_method: str,
         max_errors: int,
         doc_info_by_id: Dict,
+        is_tree: bool = False,
     ) -> Tuple[List[Dict], int]:
-        """Run RAPTOR and generate summary chunks."""
+        """Run RAPTOR and generate summary chunks.
+
+        ``chunks`` is the provenance-carrying triple shape produced by
+        ``_load_doc_chunks`` / ``_load_all_doc_chunks``:
+        ``(content, vector, chunk_id)``. Each leaf is wrapped into the
+        ``(text, vec, [chunk_id])`` shape RAPTOR expects so every
+        summary it produces carries the order-preserving deduped union
+        of the leaf ids underneath it.
+        """
         ctx = self._task_context
-        from rag.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor
+        from rag.advanced_rag.knowlege_compile.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor
 
         raptor_ext_config = raptor_config.get("ext") or {}
         assert chunks, "_generate_raptor must not be called with empty chunks"
@@ -389,13 +390,173 @@ class RaptorService:
             psi_bucket_size=raptor_ext_config.get("psi_bucket_size", 1024),
         )
 
-        original_length = len(chunks)
-        processed_chunks, layers = await raptor(
-            chunks, raptor_config["random_seed"], self._task_context.progress_cb, ctx.id
-        )
+        # Seed each leaf with its own id as the start of its
+        # ``source_chunk_ids`` provenance trail. The id may be empty
+        # for malformed retriever rows; ``Raptor.__call__`` filters
+        # those out of the union on the inbound normalize step.
+        raptor_input = [(content, vctr, [chunk_id] if chunk_id else []) for content, vctr, chunk_id in chunks]
 
         effective_doc_name = ctx.name if doc_id == GRAPH_RAPTOR_FAKE_DOC_ID else doc_info_by_id.get(doc_id, {}).get("name") or ctx.name
 
+        # Default path: ask RAPTOR for a single hierarchical tree dict
+        # and persist it as ONE non-searchable ES row. PSI's
+        # hyperedge-driven summarization can't form a strict
+        # parent-of relation, so __call__(is_tree=True) raises
+        # NotImplementedError there — catch and fall through to the
+        # legacy per-summary materialization below for that case.
+        original_length = len(chunks)
+        try:
+            processed_chunks, layers = await raptor(
+                raptor_input,
+                raptor_config["random_seed"],
+                self._task_context.progress_cb,
+                ctx.id,
+                is_tree=is_tree,
+            )
+        except NotImplementedError:
+            return await self._generate_raptor_legacy_rows(
+                raptor,
+                raptor_input,
+                raptor_config,
+                doc_id,
+                effective_doc_name,
+                tree_builder,
+                vctr_nm,
+            )
+
+        if processed_chunks is None:
+            return [], 0
+        doc = {
+            "doc_id": doc_id,
+            "kb_id": [str(ctx.kb_id)],
+            "docnm_kwd": effective_doc_name,
+            "title_tks": rag_tokenizer.tokenize(effective_doc_name),
+            "raptor_kwd": "raptor",
+            "extra": {"raptor_method": tree_builder},
+            "create_time": str(datetime.now()).replace("T", " ")[:19],
+            "create_timestamp_flt": datetime.now().timestamp(),
+        }
+        if ctx.pagerank:
+            doc[PAGERANK_FLD] = int(ctx.pagerank)
+
+        if not is_tree:
+            # Build index→layer mapping
+            chunk_layer = {}
+            for layer_idx, (layer_start, layer_end) in enumerate(layers):
+                if layer_idx == 0:
+                    continue
+                for ci in range(layer_start, layer_end):
+                    chunk_layer[ci] = layer_idx
+
+            res = []
+            tk_count = 0
+            for idx, (content, vctr, _, _) in enumerate(processed_chunks[original_length:], start=original_length):
+                d = copy.deepcopy(doc)
+                d["id"] = make_raptor_summary_chunk_id(content, doc_id)
+                d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
+                d["create_timestamp_flt"] = datetime.now().timestamp()
+                d[vctr_nm] = vctr.tolist()
+                d["content_with_weight"] = content
+                d["content_ltks"] = rag_tokenizer.tokenize(content)
+                d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
+                d["raptor_layer_int"] = chunk_layer.get(idx, 1)
+                res.append(d)
+                tk_count += num_tokens_from_string(content)
+            return res, tk_count
+
+        row_id = xxhash.xxh64(
+            f"raptor_tree:{doc_id}:{tree_builder}".encode("utf-8", "surrogatepass"),
+        ).hexdigest()
+        row = {
+            **doc,
+            "id": row_id,
+            "raptor_kwd": "raptor_tree",
+            "content_with_weight": json.dumps(processed_chunks, ensure_ascii=False),
+            "available_int": 0,
+        }
+        return [row], _sum_tree_text_tokens(processed_chunks)
+
+    async def build_doc_tree(
+        self,
+        chunks: List[Tuple[str, np.ndarray, str]],
+        raptor_config: Dict,
+        chat_mdl,
+        embd_mdl,
+        tree_builder: str,
+        clustering_method: str,
+        max_errors: int,
+    ) -> Optional[Dict]:
+        """Build a RAPTOR tree dict for one document — no ES IO.
+
+        Used by the ``tree``-kind compilation template, which wraps the
+        returned tree into a per-template structure-graph row. Returns
+        None when the input has no chunks, the PSI builder is selected
+        (which can't form a strict tree), or RAPTOR itself fails.
+        """
+        if not chunks:
+            return None
+        from rag.advanced_rag.knowlege_compile.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor
+
+        raptor_ext_config = raptor_config.get("ext") or {}
+        raptor = Raptor(
+            raptor_config.get("max_cluster", 64),
+            chat_mdl,
+            embd_mdl,
+            raptor_config["prompt"],
+            raptor_config["max_token"],
+            raptor_config["threshold"],
+            max_errors=max_errors,
+            tree_builder=tree_builder,
+            clustering_method=clustering_method,
+            psi_exact_max_leaves=raptor_ext_config.get("psi_exact_max_leaves", 4096),
+            psi_bucket_size=raptor_ext_config.get("psi_bucket_size", 1024),
+        )
+
+        raptor_input = [(content, vctr, [chunk_id] if chunk_id else []) for content, vctr, chunk_id in chunks]
+        try:
+            tree, _ = await raptor(
+                raptor_input,
+                raptor_config["random_seed"],
+                self._task_context.progress_cb,
+                self._task_context.id,
+                is_tree=True,
+            )
+        except NotImplementedError:
+            # PSI builder — not supported in tree mode; surface as None
+            # so the compilation-template path can skip the doc cleanly.
+            logging.warning(
+                "build_doc_tree: PSI builder doesn't support is_tree; skipping",
+            )
+            return None
+        return tree if isinstance(tree, dict) else None
+
+    async def _generate_raptor_legacy_rows(
+        self,
+        raptor,
+        raptor_input,
+        raptor_config,
+        doc_id,
+        effective_doc_name,
+        tree_builder,
+        vctr_nm,
+    ) -> Tuple[List[Dict], int]:
+        """Legacy per-summary materialization, kept only for PSI builds.
+
+        PSI's hyperedge summaries don't map to a strict tree, so the
+        ``is_tree=True`` default in ``_generate_raptor`` raises and
+        falls through here. Same shape this function produced before
+        the tree migration — one ES row per appended summary, marked
+        ``raptor_kwd="raptor"``.
+        """
+        ctx = self._task_context
+        original_length = len(raptor_input)
+        processed_chunks, layers = await raptor(
+            raptor_input,
+            raptor_config["random_seed"],
+            self._task_context.progress_cb,
+            ctx.id,
+        )
+
         doc = {
             "doc_id": doc_id,
             "kb_id": [str(ctx.kb_id)],
@@ -407,7 +568,6 @@ class RaptorService:
         if ctx.pagerank:
             doc[PAGERANK_FLD] = int(ctx.pagerank)
 
-        # Build index→layer mapping
         chunk_layer = {}
         for layer_idx, (layer_start, layer_end) in enumerate(layers):
             if layer_idx == 0:
@@ -417,7 +577,12 @@ class RaptorService:
 
         res = []
         tk_count = 0
-        for idx, (content, vctr) in enumerate(processed_chunks[original_length:], start=original_length):
+        for idx, item in enumerate(processed_chunks[original_length:], start=original_length):
+            if len(item) >= 3:
+                content, vctr, source_chunk_ids = item[0], item[1], item[2] or []
+            else:
+                content, vctr = item[0], item[1]
+                source_chunk_ids = []
             d = copy.deepcopy(doc)
             d["id"] = make_raptor_summary_chunk_id(content, doc_id)
             d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
@@ -427,6 +592,8 @@ class RaptorService:
             d["content_ltks"] = rag_tokenizer.tokenize(content)
             d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
             d["raptor_layer_int"] = chunk_layer.get(idx, 1)
+            if source_chunk_ids:
+                d["source_chunk_ids"] = list(source_chunk_ids)
             res.append(d)
             tk_count += num_tokens_from_string(content)
 
@@ -445,16 +612,17 @@ class RaptorService:
         from common.doc_store.doc_store_base import OrderByExpr
 
         async def search_fields(fields: list, condition: dict, order_by=None):
-            res = await thread_pool_exec(
-                settings.docStoreConn.search,
-                fields, [], condition, [], order_by or OrderByExpr(),
-                0, 10000, search.index_name(tenant_id), [kb_id]
-            )
+            res = await thread_pool_exec(settings.docStoreConn.search, fields, [], condition, [], order_by or OrderByExpr(), 0, 10000, search.index_name(tenant_id), [kb_id])
             return settings.docStoreConn.get_fields(res, fields)
 
         try:
+            # Accept both ``raptor`` (legacy per-summary rows, PSI
+            # builder still produces these) and ``raptor_tree`` (new
+            # single-row tree blob) so existing-method detection stays
+            # accurate across the migration.
             primary = await search_fields(
-                ["raptor_kwd", "extra"], {"doc_id": doc_id, "raptor_kwd": ["raptor"]}
+                ["raptor_kwd", "extra"],
+                {"doc_id": doc_id, "raptor_kwd": ["raptor", "raptor_tree"]},
             )
             if collect_raptor_chunk_ids(primary):
                 return collect_raptor_methods(primary)
@@ -469,3 +637,188 @@ class RaptorService:
         except Exception:
             logging.exception("Failed to check RAPTOR chunks for doc %s", doc_id)
             raise
+
+    @staticmethod
+    def _build_raptor_graph(rows: List[Dict]) -> Dict:
+        """Project loaded RAPTOR summary rows onto the canvas graph shape.
+
+        Each row contributes one entity::
+
+            {
+              "id":          xxh128(content)           # 32-char hex
+              "name":        first 16 whitespace tokens
+              "description": content_with_weight
+              "source_chunk_ids": row.source_chunk_ids
+            }
+
+        Relations: full bipartite layer-by-layer fan-out — every node at
+        layer K gets an edge to every node at layer K-1 (because we only
+        loaded ``content_with_weight`` + ``raptor_layer_int`` we don't
+        have the specific parent linkage). Self-edges and dangling
+        targets are dropped (the latter only matters if the layer-int
+        values are non-contiguous).
+        """
+        # Build entities. Dedup by id so two identical-content summaries
+        # collapse to one node — the canvas can't render multiple nodes
+        # at the same id anyway, and identical content is a defensible
+        # collapse.
+        by_id: Dict[str, Dict] = {}
+        by_layer: Dict[int, List[str]] = {}
+
+        for row in rows:
+            content = row.get("content_with_weight")
+            if not isinstance(content, str) or not content.strip():
+                continue
+            try:
+                layer = int(row.get("raptor_layer_int") or 0)
+            except (TypeError, ValueError):
+                layer = 0
+            if layer <= 0:
+                # Layer 0 would be the original leaf chunks; RAPTOR
+                # summaries start at layer 1. Anything claiming layer 0
+                # here is malformed; skip.
+                continue
+
+            name = " ".join(content.split()[:16])
+            nid = xxhash.xxh128(
+                content.encode("utf-8", "surrogatepass"),
+            ).hexdigest()  # 32-char hex
+            if nid in by_id:
+                continue
+            source_chunk_ids = row.get("source_chunk_ids") or []
+            if not isinstance(source_chunk_ids, list):
+                source_chunk_ids = []
+            by_id[nid] = {
+                "id": nid,
+                "name": name,
+                "description": content,
+                "source_chunk_ids": list(source_chunk_ids),
+            }
+            by_layer.setdefault(layer, []).append(nid)
+
+        # Layered fan-out from parent (higher layer) → child (lower layer).
+        relations: List[Dict] = []
+        layers_sorted = sorted(by_layer.keys())
+        for layer in layers_sorted:
+            child_layer = layer - 1
+            if child_layer not in by_layer:
+                continue
+            for parent in by_layer[layer]:
+                for child in by_layer[child_layer]:
+                    if parent == child:
+                        continue
+                    relations.append({"from": parent, "to": child})
+
+        return {"entities": list(by_id.values()), "relations": relations}
+
+    async def _persist_raptor_graph_to_es(self, doc_id: str) -> None:
+        """Load the just-inserted RAPTOR summaries for ``doc_id`` and
+        persist a single graph row that the dataset structure-graph
+        endpoint can surface as a tree.
+
+        Loads only ``content_with_weight`` + ``raptor_layer_int`` +
+        ``source_chunk_ids`` (per
+        the smallest-payload contract) and writes one row with::
+
+            compile_kwd:                  "raptor_graph"
+            compilation_template_kind_kwd:"raptor"
+            doc_id:                       <doc_id>
+
+        The row id is deterministic per ``(kb_id, doc_id)`` so re-runs
+        delete-and-replace cleanly through the same primary key.
+        ``knowledge_graph_kwd`` is intentionally NOT set — that field
+        belongs to the KG feature; this row is identified via
+        ``compile_kwd`` so the two paths stay semantically distinct.
+        """
+        from common.doc_store.doc_store_base import OrderByExpr
+
+        ctx = self._task_context
+        tenant_id = ctx.tenant_id
+        kb_id_str = str(ctx.kb_id)
+        index_nm = search.index_name(tenant_id)
+        select_fields = ["content_with_weight", "raptor_layer_int", "source_chunk_ids"]
+        try:
+            res = await thread_pool_exec(
+                settings.docStoreConn.search,
+                select_fields,
+                [],
+                {"raptor_kwd": ["raptor"], "doc_id": [doc_id]},
+                [],
+                OrderByExpr(),
+                0,
+                10000,
+                index_nm,
+                [kb_id_str],
+            )
+            field_map = settings.docStoreConn.get_fields(res, select_fields)
+        except Exception:
+            logging.exception(
+                "raptor_graph: load failed for kb=%s doc=%s",
+                kb_id_str,
+                doc_id,
+            )
+            return
+
+        rows = list((field_map or {}).values())
+        if not rows:
+            logging.info(
+                "raptor_graph: no summaries to render for kb=%s doc=%s",
+                kb_id_str,
+                doc_id,
+            )
+            return
+
+        graph = self._build_raptor_graph(rows)
+        if not graph["entities"]:
+            logging.info(
+                "raptor_graph: projection produced no entities for kb=%s doc=%s",
+                kb_id_str,
+                doc_id,
+            )
+            return
+
+        row_id = xxhash.xxh64(
+            f"raptor_graph:{kb_id_str}:{doc_id}".encode("utf-8", "surrogatepass"),
+        ).hexdigest()
+        row = {
+            "id": row_id,
+            "kb_id": kb_id_str,
+            "doc_id": doc_id,
+            "compile_kwd": "raptor_graph",
+            "compilation_template_kind_kwd": "raptor",
+            "content_with_weight": json.dumps(graph, ensure_ascii=False),
+            "available_int": 0,
+        }
+        try:
+            await thread_pool_exec(
+                settings.docStoreConn.delete,
+                {"compile_kwd": "raptor_graph", "doc_id": [doc_id]},
+                index_nm,
+                ctx.kb_id,
+            )
+        except Exception:
+            logging.debug(
+                "raptor_graph: prior delete failed for kb=%s doc=%s; relying on id-upsert",
+                kb_id_str,
+                doc_id,
+            )
+        try:
+            await thread_pool_exec(
+                settings.docStoreConn.insert,
+                [row],
+                index_nm,
+                ctx.kb_id,
+            )
+            logging.info(
+                "raptor_graph: stored %d entities / %d relations for kb=%s doc=%s",
+                len(graph["entities"]),
+                len(graph["relations"]),
+                kb_id_str,
+                doc_id,
+            )
+        except Exception:
+            logging.exception(
+                "raptor_graph: insert failed for kb=%s doc=%s",
+                kb_id_str,
+                doc_id,
+            )

rag/svr/task_executor_refactor/raptor_utils.py

@@ -38,11 +38,7 @@ async def get_raptor_chunk_field_map(doc_id: str, tenant_id: str, kb_id: str) ->
 
     async def search_fields(fields: list[str], condition: dict, order_by=None):
         """Search chunk fields in the current knowledge base."""
-        res = await thread_pool_exec(
-            settings.docStoreConn.search,
-            fields, [], condition, [], order_by or OrderByExpr(),
-            0, RAPTOR_METHOD_SEARCH_LIMIT, nlp_search.index_name(tenant_id), [kb_id]
-        )
+        res = await thread_pool_exec(settings.docStoreConn.search, fields, [], condition, [], order_by or OrderByExpr(), 0, RAPTOR_METHOD_SEARCH_LIMIT, nlp_search.index_name(tenant_id), [kb_id])
         return settings.docStoreConn.get_fields(res, fields)
 
     primary = await search_fields(["raptor_kwd", "extra"], {"doc_id": doc_id, "raptor_kwd": ["raptor"]})
@@ -65,11 +61,17 @@ async def delete_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, keep_met
     if keep_method is None:
         logging.info(
             "delete_raptor_chunks: removing all RAPTOR summaries (doc=%s tenant=%s kb=%s)",
-            doc_id, tenant_id, kb_id,
+            doc_id,
+            tenant_id,
+            kb_id,
         )
+        # Sweep both row types — legacy per-summary (``raptor``, still
+        # used by the PSI builder) and the new single tree blob
+        # (``raptor_tree``) — so re-runs always start from a clean
+        # slate regardless of which path produced the prior state.
         await thread_pool_exec(
             settings.docStoreConn.delete,
-            {"doc_id": doc_id, "raptor_kwd": ["raptor"]},
+            {"doc_id": doc_id, "raptor_kwd": ["raptor", "raptor_tree"]},
             nlp_search.index_name(tenant_id),
             kb_id,
         )
@@ -80,13 +82,20 @@ async def delete_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, keep_met
     if not chunk_ids:
         logging.debug(
             "delete_raptor_chunks: no stale RAPTOR chunks to remove (doc=%s tenant=%s kb=%s keep=%s)",
-            doc_id, tenant_id, kb_id, keep_method,
+            doc_id,
+            tenant_id,
+            kb_id,
+            keep_method,
         )
         return 0
 
     logging.info(
         "delete_raptor_chunks: removing %d stale RAPTOR chunks (doc=%s tenant=%s kb=%s keep=%s)",
-        len(chunk_ids), doc_id, tenant_id, kb_id, keep_method,
+        len(chunk_ids),
+        doc_id,
+        tenant_id,
+        kb_id,
+        keep_method,
     )
     await thread_pool_exec(
         settings.docStoreConn.delete,

rag/svr/task_executor_refactor/task_handler.py

@@ -23,20 +23,26 @@ for handling document processing tasks with refactored, testable methods.
 import asyncio
 import logging
 import json
+
+# Wiki / artifact compilation pipeline lives in
+# ``rag.svr.task_executor_refactor.dataset_wiki_generator`` — see the
+# ``task_type == "artifact"`` branch of ``TaskHandler.run`` for the
+# dispatch call.
+# Document-structure compilation helpers (CHAIN_KINDS,
+# compile_structure_from_text, merge_compiled_structures,
+# validate_and_correct_chain) moved to ``chunk_post_processor``.
 import xxhash
 
 from timeit import default_timer as timer
-from typing import Callable, Dict, List, Optional
+from typing import AsyncIterator, Callable, Dict, List, Optional
 
 from api.db.services.document_service import DocumentService
 from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.compilation_template_group_service import CompilationTemplateGroupService
 from api.db.joint_services.memory_message_service import handle_save_to_memory_task
-from api.db.joint_services.tenant_model_service import (
-    get_tenant_default_model_by_type,
-    get_model_config_from_provider_instance
-)
+from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance
 from api.db.services.llm_service import LLMBundle
-from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID
+from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID, abort_doc_chunking_counter
 from common.constants import LLMType
 from common.exceptions import TaskCanceledException
 from common.connection_utils import timeout
@@ -57,6 +63,96 @@ from rag.prompts.generator import run_toc_from_text
 from common import settings
 
 
+def _parser_config_compilation_template_group_ids(parser_config) -> list[str]:
+    """Read template-group ids from a doc's parser_config.
+
+    Templates were previously referenced as a list
+    (``compilation_template_ids``); after the template-group refactor
+    a doc instead points at one or more groups, and the orchestrator
+    resolves each group's child templates at runtime. Old
+    ``compilation_template_ids`` data is intentionally ignored per
+    the migration spec.
+    """
+
+    def _normalize(raw) -> list[str]:
+        if isinstance(raw, str):
+            raw = [raw]
+        if not isinstance(raw, list):
+            return []
+        ids: list[str] = []
+        seen: set[str] = set()
+        for gid in raw:
+            if not isinstance(gid, str):
+                continue
+            gid = gid.strip()
+            if gid and gid not in seen:
+                seen.add(gid)
+                ids.append(gid)
+        return ids
+
+    if not isinstance(parser_config, dict):
+        return []
+    if "compilation_template_group_id" in parser_config:
+        return _normalize(parser_config.get("compilation_template_group_id"))
+    ext = parser_config.get("ext")
+    if isinstance(ext, dict):
+        return _normalize(ext.get("compilation_template_group_id"))
+    return []
+
+
+def _parser_config_compilation_template_ids(parser_config, tenant_id: str) -> list[str]:
+    """Resolve a doc's parser_config to compile-template ids by
+    looking up configured groups. Returns ``[]`` if the doc has no
+    group set or no group can be resolved.
+    """
+    template_ids: list[str] = []
+    seen: set[str] = set()
+    for group_id in _parser_config_compilation_template_group_ids(parser_config):
+        for template_id in CompilationTemplateGroupService.resolve_template_ids(group_id, tenant_id):
+            if template_id in seen:
+                continue
+            seen.add(template_id)
+            template_ids.append(template_id)
+    return template_ids
+
+
+def _resolve_template_chat_llm_id(parser_cfg: dict, ctx) -> str:
+    """Pick the chat model id for a knowledge-compilation template.
+
+    Resolution order:
+      1. The template's own ``llm_id`` (what the user picked in the
+         compilation-template panel).
+      2. The doc's ``parser_config.llm_id`` (the doc-level chunking
+         model).
+      3. ``ctx.llm_id`` (the chunking task's default).
+    """
+    if isinstance(parser_cfg, dict):
+        tid = parser_cfg.get("llm_id")
+        if isinstance(tid, str) and tid.strip():
+            return tid.strip()
+    doc_cfg = getattr(ctx, "parser_config", None) or {}
+    if isinstance(doc_cfg, dict):
+        did = doc_cfg.get("llm_id")
+        if isinstance(did, str) and did.strip():
+            return did.strip()
+    return ctx.llm_id
+
+
+# Document-structure compilation tunables
+# (DOC_STRUCTURE_COMPILE_BATCH_CHUNKS, DOC_STRUCTURE_MERGE_MAX_DOCS,
+# STRUCTURE_CHAIN_CORRECTION_TIMEOUT_S) moved to
+# ``chunk_post_processor``.
+
+# Wiki / artifact tunables (``WIKI_MAP_BATCH_CHUNKS``,
+# ``WIKI_GRAPH_MAX_CHUNK_IDS_PER_NODE``, commit-title / comments
+# templates) moved to ``dataset_wiki_generator``.
+
+# The corpus → skill compilation pipeline lives in
+# ``rag.svr.task_executor_refactor.dataset_skill_generator``. Its entry
+# point is :func:`run_corpus2skill`; this handler invokes it from the
+# ``task_type == "skill"`` branch of ``run`` below.
+
+
 class TaskHandler:
     """Main task handler for document processing.
 
@@ -85,32 +181,52 @@ class TaskHandler:
         self._task_context = ctx
         self._billing_hook = billing_hook
 
+    @staticmethod
+    def _is_standard_chunking_task(task_type: str) -> bool:
+        task_type = (task_type or "").lower()
+        return task_type not in {
+            "memory",
+            "raptor",
+            "graphrag",
+            "mindmap",
+            "artifact",
+            "skill",
+            "evaluation",
+            "reembedding",
+            "clone",
+        } and not task_type.startswith("dataflow")
+
     async def handle_task(self) -> None:
         try:
             await self.handle()
+        except Exception:
+            if self._is_standard_chunking_task(self._task_context.task_type):
+                abort_doc_chunking_counter(self._task_context.doc_id)
+            raise
         finally:
             task_id = self._task_context.id
             task_tenant_id = self._task_context.tenant_id
             task_dataset_id = self._task_context.kb_id
             task_doc_id = self._task_context.doc_id
             if self._task_context.has_canceled_func(task_id):
-                try:
-                    exists = await thread_pool_exec(
-                        settings.docStoreConn.index_exist,
-                        search.index_name(task_tenant_id),
-                        task_dataset_id,
-                    )
-                    if exists:
-                        ret = await thread_pool_exec(
-                            settings.docStoreConn.delete,
-                            {"doc_id": task_doc_id},
+                if self._is_standard_chunking_task(self._task_context.task_type):
+                    abort_doc_chunking_counter(task_doc_id)
+                    try:
+                        exists = await thread_pool_exec(
+                            settings.docStoreConn.index_exist,
                             search.index_name(task_tenant_id),
                             task_dataset_id,
                         )
-                        self._task_context.recording_context.save_func_return_value("docStoreConn.delete", ret)
-                except Exception as e:
-                    logging.exception(
-                        f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}")
+                        if exists:
+                            ret = await thread_pool_exec(
+                                settings.docStoreConn.delete,
+                                {"doc_id": task_doc_id},
+                                search.index_name(task_tenant_id),
+                                task_dataset_id,
+                            )
+                            self._task_context.recording_context.save_func_return_value("docStoreConn.delete", ret)
+                    except Exception as e:
+                        logging.exception(f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}")
 
     @timeout(60 * 60 * 3, 1)
     async def handle(self) -> None:
@@ -134,7 +250,7 @@ class TaskHandler:
             ctx.progress_cb(-1, msg="Task has been canceled.")
             return
 
-        # Language defaults to "Chinese" via TaskContext._DEFAULTS — safe to bind model directly.
+        # Language defaults to "Chinese" via TaskContext._DEFAULTS 鈥?safe to bind model directly.
         # Bind embedding model (matching original do_handle_task order: bind + init_kb before routing)
         result = await self._bind_embedding_model()
         if result is None:
@@ -154,12 +270,30 @@ class TaskHandler:
                 return
 
             # Route to appropriate handler
-            if task_type == "raptor":
-                await self._run_raptor(embedding_model, vector_size)
-            elif task_type == "graphrag":
+            if task_type == "graphrag":
                 await self._run_graphrag(embedding_model)
             elif task_type == "mindmap":
                 ctx.progress_cb(1, "place holder")
+            elif task_type == "artifact":
+                from rag.svr.task_executor_refactor.dataset_wiki_generator import (
+                    run_wiki,
+                )
+
+                await run_wiki(
+                    self._task_context,
+                    embedding_model,
+                    self._load_chunks_for_doc,
+                )
+            elif task_type == "skill":
+                from rag.svr.task_executor_refactor.dataset_skill_generator import (
+                    run_corpus2skill,
+                )
+
+                await run_corpus2skill(
+                    self._task_context,
+                    embedding_model,
+                    self._load_chunks_for_doc,
+                )
             elif task_type == "evaluation":
                 await self._run_evaluation()
             elif task_type == "reembedding":
@@ -167,8 +301,7 @@ class TaskHandler:
             elif task_type == "clone":
                 await self._run_clone()
             else:
-                await self._run_standard_chunking(embedding_model)
-
+                await self._run_standard_chunking(embedding_model, vector_size)
 
     def _init_kb(self, vector_size: int) -> None:
         """Initialize knowledge base index."""
@@ -214,18 +347,14 @@ class TaskHandler:
 
         try:
             if task_embedding_id:
-                embd_model_config = get_model_config_from_provider_instance(
-                    task_tenant_id, LLMType.EMBEDDING, task_embedding_id
-                )
+                embd_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.EMBEDDING, task_embedding_id)
             else:
-                embd_model_config = get_tenant_default_model_by_type(
-                    task_tenant_id, LLMType.EMBEDDING
-                )
+                embd_model_config = get_tenant_default_model_by_type(task_tenant_id, LLMType.EMBEDDING)
             embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language)
             vts, _ = embedding_model.encode(["ok"])
             return embedding_model, len(vts[0])
         except Exception as e:
-            error_message = f'Fail to bind embedding model: {str(e)}'
+            error_message = f"Fail to bind embedding model: {str(e)}"
             ctx.progress_cb(-1, msg=error_message)
             logging.exception(error_message)
             raise
@@ -234,6 +363,7 @@ class TaskHandler:
         self,
         embedding_model: LLMBundle,
         vector_size: int,
+        mark_done: bool = True,
     ) -> None:
         """Run RAPTOR summary generation."""
         ctx = self._task_context
@@ -248,19 +378,21 @@ class TaskHandler:
 
         kb_parser_config = kb.parser_config
         if not kb_parser_config.get("raptor", {}).get("use_raptor", False):
-            kb_parser_config.update({
-                "raptor": {
-                    "use_raptor": True,
-                    "prompt": "Please summarize the following paragraphs. Be careful with the numbers, do not make things up. Paragraphs as following:\n      {cluster_content}\nThe above is the content you need to summarize.",
-                    "max_token": 256,
-                    "threshold": 0.1,
-                    "max_cluster": 64,
-                    "random_seed": 0,
-                    "scope": "file",
-                    "clustering_method": "gmm",
-                    "tree_builder": "raptor",
-                },
-            })
+            kb_parser_config.update(
+                {
+                    "raptor": {
+                        "use_raptor": True,
+                        "prompt": "Please summarize the following paragraphs. Be careful with the numbers, do not make things up. Paragraphs as following:\n      {cluster_content}\nThe above is the content you need to summarize.",
+                        "max_token": 256,
+                        "threshold": 0.1,
+                        "max_cluster": 64,
+                        "random_seed": 0,
+                        "scope": "file",
+                        "clustering_method": "gmm",
+                        "tree_builder": "raptor",
+                    },
+                }
+            )
             if ctx.write_interceptor:
                 update_result = ctx.write_interceptor.intercept("KnowledgebaseService.update_by_id")
             else:
@@ -271,11 +403,8 @@ class TaskHandler:
                 return
 
         # Bind LLM for raptor
-        chat_model_config = get_model_config_from_provider_instance(
-            task_tenant_id, LLMType.CHAT, kb_task_llm_id
-        )
+        chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
         with LLMBundle(task_tenant_id, chat_model_config, lang=ctx.language) as chat_model:
-
             # Run RAPTOR
             raptor_service = RaptorService(ctx=ctx)
 
@@ -285,7 +414,7 @@ class TaskHandler:
                     chat_mdl=chat_model,
                     embd_mdl=embedding_model,
                     vector_size=vector_size,
-                    doc_ids=ctx.doc_ids,
+                    doc_ids=ctx.doc_ids or [ctx.doc_id],
                 )
 
             ctx.recording_context.record("raptor_chunks", chunks)
@@ -293,7 +422,7 @@ class TaskHandler:
 
             # Insert RAPTOR chunks
             if chunks:
-                task_doc_id = (ctx.doc_ids or [GRAPH_RAPTOR_FAKE_DOC_ID])[0]
+                task_doc_id = (ctx.doc_ids or [ctx.doc_id] or [GRAPH_RAPTOR_FAKE_DOC_ID])[0]
                 chunk_service = ChunkService(ctx=ctx)
                 insert_result = await chunk_service.insert_chunks(ctx.id, task_tenant_id, task_dataset_id, chunks)
                 if insert_result:
@@ -304,27 +433,43 @@ class TaskHandler:
                 # Cleanup stale RAPTOR chunks
                 cleaned_chunks = 0
                 for cleanup_doc_id, keep_method in raptor_cleanup_chunks:
-                    ret = await self._delete_raptor_chunks(
-                        cleanup_doc_id, task_tenant_id, task_dataset_id, keep_method
-                    )
+                    ret = await self._delete_raptor_chunks(cleanup_doc_id, task_tenant_id, task_dataset_id, keep_method)
                     cleaned_chunks += ret
 
                 if cleaned_chunks:
                     ctx.progress_cb(msg=f"Cleaned up {cleaned_chunks} stale RAPTOR chunks.")
 
+                # Build the per-doc RAPTOR tree graph from the just-
+                # inserted summaries. Each chunk in ``chunks`` carries
+                # the doc_id it was written under (real doc id for
+                # scope="file"; GRAPH_RAPTOR_FAKE_DOC_ID for the
+                # dataset-scope path). We materialize one graph row per
+                # distinct doc_id so the dataset structure-graph
+                # endpoint can surface a RAPTOR tab per document.
+                # Failure here is best-effort — the summaries are
+                # already persisted; the tab just won't render.
+                raptor_doc_ids = {str(c.get("doc_id")) for c in chunks if c.get("doc_id")}
+                for raptor_doc_id in raptor_doc_ids:
+                    try:
+                        await raptor_service._persist_raptor_graph_to_es(raptor_doc_id)
+                    except Exception:
+                        logging.exception(
+                            "raptor_graph: build failed for kb=%s doc=%s",
+                            task_dataset_id,
+                            raptor_doc_id,
+                        )
+
                 # Update document stats
                 if ctx.write_interceptor:
                     ctx.write_interceptor.intercept("DocumentService.increment_chunk_num")
                 else:
                     DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, len(chunks), 0)
 
-            ctx.recording_context.record("task_status", "completed")
-            ctx.progress_cb(prog=1.0, msg="RAPTOR done")
+            if mark_done:
+                ctx.recording_context.record("task_status", "completed")
+                ctx.progress_cb(prog=1.0, msg="RAPTOR done")
 
-    async def _run_graphrag(
-        self,
-        embedding_model: LLMBundle
-    ) -> None:
+    async def _run_graphrag(self, embedding_model: LLMBundle) -> None:
         """Run GraphRAG."""
         ctx = self._task_context
         task_tenant_id = ctx.tenant_id
@@ -339,29 +484,31 @@ class TaskHandler:
 
         kb_parser_config = kb.parser_config
         if not kb_parser_config.get("graphrag", {}).get("use_graphrag", False):
-            kb_parser_config.update({
-                "graphrag": {
-                    "use_graphrag": True,
-                    "entity_types": [
-                        "organization",
-                        "person",
-                        "geo",
-                        "event",
-                        "category",
-                    ],
-                    "method": "light",
-                    "batch_chunk_token_size": 4096,
-                    "retry_attempts": 2,
-                    "retry_backoff_seconds": 2.0,
-                    "retry_backoff_max_seconds": 60.0,
-                    "build_subgraph_timeout_per_chunk_seconds": 300,
-                    "build_subgraph_min_timeout_seconds": 600,
-                    "merge_timeout_seconds": 180,
-                    "resolution_timeout_seconds": 1800,
-                    "community_timeout_seconds": 1800,
-                    "lock_acquire_timeout_seconds": 600,
+            kb_parser_config.update(
+                {
+                    "graphrag": {
+                        "use_graphrag": True,
+                        "entity_types": [
+                            "organization",
+                            "person",
+                            "geo",
+                            "event",
+                            "category",
+                        ],
+                        "method": "light",
+                        "batch_chunk_token_size": 4096,
+                        "retry_attempts": 2,
+                        "retry_backoff_seconds": 2.0,
+                        "retry_backoff_max_seconds": 60.0,
+                        "build_subgraph_timeout_per_chunk_seconds": 300,
+                        "build_subgraph_min_timeout_seconds": 600,
+                        "merge_timeout_seconds": 180,
+                        "resolution_timeout_seconds": 1800,
+                        "community_timeout_seconds": 1800,
+                        "lock_acquire_timeout_seconds": 600,
+                    }
                 }
-            })
+            )
             if ctx.write_interceptor:
                 update_result = ctx.write_interceptor.intercept("KnowledgebaseService.update_by_id")
             else:
@@ -372,11 +519,8 @@ class TaskHandler:
 
         graphrag_conf = kb_parser_config.get("graphrag", {})
         start_ts = timer()
-        chat_model_config = get_model_config_from_provider_instance(
-            task_tenant_id, LLMType.CHAT, kb_task_llm_id
-        )
+        chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id)
         with LLMBundle(task_tenant_id, chat_model_config, lang=task_language) as chat_model:
-
             with_resolution = graphrag_conf.get("resolution", False)
             with_community = graphrag_conf.get("community", False)
 
@@ -399,7 +543,20 @@ class TaskHandler:
 
     async def _run_standard_chunking(
         self,
-        embedding_model: LLMBundle
+        embedding_model: LLMBundle,
+        vector_size: int,
+    ) -> None:
+        ctx = self._task_context
+        try:
+            await self._run_standard_chunking_impl(embedding_model, vector_size)
+        except Exception:
+            abort_doc_chunking_counter(ctx.doc_id)
+            raise
+
+    async def _run_standard_chunking_impl(
+        self,
+        embedding_model: LLMBundle,
+        vector_size: int,
     ) -> None:
         """Run standard chunking pipeline."""
         ctx = self._task_context
@@ -409,7 +566,7 @@ class TaskHandler:
         task_doc_id = ctx.doc_id
         task_start_ts = timer()
         doc_task_llm_id = ctx.parser_config.get("llm_id") or ctx.llm_id
-        ctx.raw_task['llm_id'] = doc_task_llm_id
+        ctx.raw_task["llm_id"] = doc_task_llm_id
 
         # Build chunks
         start_ts = timer()
@@ -419,9 +576,7 @@ class TaskHandler:
         bucket, name = File2DocumentService.get_storage_address(doc_id=ctx.doc_id)
         binary = await self._get_storage_binary(bucket, name)
         if binary is None:
-            raise FileNotFoundError(
-                f"Can not find file <{ctx.name}> from minio. Could you try it again."
-            )
+            raise FileNotFoundError(f"Can not find file <{ctx.name}> from minio. Could you try it again.")
 
         chunks = await chunk_service.build_chunks(binary)
         ctx.recording_context.record("chunks", chunks)
@@ -431,7 +586,18 @@ class TaskHandler:
         logging.info("Build document {}: {:.2f}s".format(ctx.name, timer() - start_ts))
 
         if not chunks:
-            ctx.progress_cb(1., msg=f"No chunk built from {ctx.name}")
+            ctx.progress_cb(msg=f"No chunk built from {ctx.name}")
+            if not await self._run_document_post_chunking_if_last(
+                embedding_model,
+                vector_size,
+                task_start_ts,
+                0,
+                0,
+            ):
+                return
+            task_time_cost = timer() - task_start_ts
+            ctx.recording_context.record("task_status", "completed")
+            ctx.progress_cb(prog=1.0, msg="Task done ({:.2f}s)".format(task_time_cost))
             return
 
         ctx.progress_cb(msg="Generate {} chunks".format(len(chunks)))
@@ -440,9 +606,7 @@ class TaskHandler:
         start_ts = timer()
         embedding_service = EmbeddingService(ctx=ctx)
         try:
-            token_count, vector_size = await embedding_service.embed_chunks(
-                chunks, embedding_model, ctx.parser_config
-            )
+            token_count, vector_size = await embedding_service.embed_chunks(chunks, embedding_model, ctx.parser_config)
         except TaskCanceledException:
             raise
         except Exception as e:
@@ -457,11 +621,6 @@ class TaskHandler:
         logging.info(progress_message)
         ctx.progress_cb(msg=progress_message)
 
-        # Build TOC if needed
-        toc_thread = None
-        if ctx.parser_id.lower() == "naive" and ctx.parser_config.get("toc_extraction", False):
-            toc_thread = asyncio.create_task(asyncio.to_thread(self._build_toc, ctx, chunks, ctx.progress_cb))
-
         # Insert chunks
         chunk_count = len(set([chunk["id"] for chunk in chunks]))
         start_ts = timer()
@@ -469,15 +628,15 @@ class TaskHandler:
         chunk_service = ChunkService(ctx=ctx)
 
         if ctx.has_canceled_func(task_id):
+            abort_doc_chunking_counter(task_doc_id)
             ctx.progress_cb(-1, msg="Task has been canceled.")
             return
 
-        insert_result = await chunk_service.insert_chunks(
-            task_id, task_tenant_id, task_dataset_id, chunks
-        )
+        insert_result = await chunk_service.insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks)
 
         if not insert_result:
             ctx.recording_context.record("insertion_result", "failed")
+            abort_doc_chunking_counter(task_doc_id)
             return
         ctx.recording_context.record("insertion_result", "success")
 
@@ -487,12 +646,8 @@ class TaskHandler:
 
         ctx.progress_cb(msg="Indexing done ({:.2f}s).".format(timer() - start_ts))
 
-        toc_chunk = await self._process_toc_thread(toc_thread)
-        if toc_chunk:
-            ctx.recording_context.record("toc_chunk", [toc_chunk])
-            await post_processor.insert_toc_chunk(toc_chunk, chunk_service)
-
         if ctx.has_canceled_func(task_id):
+            abort_doc_chunking_counter(task_doc_id)
             ctx.progress_cb(-1, msg="Task has been canceled.")
             return
 
@@ -502,17 +657,44 @@ class TaskHandler:
         else:
             DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, chunk_count, 0)
 
+        if not await self._run_document_post_chunking_if_last(
+            embedding_model,
+            vector_size,
+            task_start_ts,
+            len(chunks),
+            token_count,
+        ):
+            return
+
         task_time_cost = timer() - task_start_ts
         ctx.recording_context.record("task_status", "completed")
         ctx.progress_cb(prog=1.0, msg="Task done ({:.2f}s)".format(task_time_cost))
 
-        logging.info(
-            "Chunk doc({}), page({}-{}), chunks({}), token({}), elapsed:{:.2f}".format(
-                ctx.name, ctx.from_page, ctx.to_page,
-                len(chunks), token_count, task_time_cost
-            )
+        logging.info("Chunk doc({}), page({}-{}), chunks({}), token({}), elapsed:{:.2f}".format(ctx.name, ctx.from_page, ctx.to_page, len(chunks), token_count, task_time_cost))
+
+    async def _run_document_post_chunking_if_last(
+        self,
+        embedding_model: LLMBundle,
+        vector_size: int,
+        task_start_ts: float,
+        chunks_len: int,
+        token_count: int,
+    ) -> bool:
+        """Thin delegator. The pipeline lives in
+        ``rag.svr.task_executor_refactor.chunk_post_processor``.
+        """
+        from rag.svr.task_executor_refactor.chunk_post_processor import (
+            run_document_post_chunking_if_last,
         )
 
+        return await run_document_post_chunking_if_last(
+            self,
+            embedding_model,
+            vector_size,
+            task_start_ts,
+            chunks_len,
+            token_count,
+        )
 
     async def _process_toc_thread(self, toc_thread):
         try:
@@ -527,30 +709,105 @@ class TaskHandler:
     @classmethod
     async def _get_storage_binary(cls, bucket: str, name: str) -> bytes:
         from common import settings
+
         """Get binary from storage."""
         return await thread_pool_exec(settings.STORAGE_IMPL.get, bucket, name)
 
+    @staticmethod
+    async def _load_chunks_for_doc(
+        tenant_id: str,
+        kb_id: str,
+        doc_id: str,
+        batch_size: int = 500,
+    ) -> AsyncIterator[List[Dict]]:
+        """Stream a document's chunks from the doc store one batch at a time.
+
+        Async generator that yields successive batches of up to ``batch_size``
+        chunks. Order is pushed to the doc store via
+        ``OrderByExpr().asc("page_num_int").asc("top_int")`` so callers do
+        not need to re-sort. Rows with a ``compile_kwd`` marker (artifact
+        pages, structure entities, etc.) are filtered out defensively.
+
+        Memory is bounded by ``batch_size``: at most one page is materialised
+        at a time, so long documents do not balloon the worker's heap.
+        """
+        from common.doc_store.doc_store_base import OrderByExpr
+
+        index_nm = search.index_name(tenant_id)
+        if not settings.docStoreConn.index_exist(index_nm, kb_id):
+            return
+
+        select_fields = [
+            "id",
+            "doc_id",
+            "content_with_weight",
+            "page_num_int",
+            "top_int",
+        ]
+        order_by = OrderByExpr()
+        order_by.asc("page_num_int")
+        order_by.asc("top_int")
+
+        offset = 0
+        while True:
+            try:
+                res = await thread_pool_exec(
+                    settings.docStoreConn.search,
+                    select_fields,
+                    [],
+                    {"doc_id": [doc_id], "available_int": 1},
+                    [],
+                    order_by,
+                    offset,
+                    batch_size,
+                    index_nm,
+                    [kb_id],
+                )
+                field_map = settings.docStoreConn.get_fields(res, select_fields)
+            except Exception:
+                logging.exception("load_chunks_for_doc: failed to load chunks for doc=%s", doc_id)
+                return
+            if not field_map:
+                return
+
+            batch: List[Dict] = []
+            for row_id, row in field_map.items():
+                if row.get("compile_kwd"):
+                    continue
+                batch.append(
+                    {
+                        "id": row_id,
+                        "doc_id": row.get("doc_id") or doc_id,
+                        "content_with_weight": row.get("content_with_weight") or "",
+                        "page_num_int": row.get("page_num_int", 0),
+                        "top_int": row.get("top_int", 0),
+                    }
+                )
+            if batch:
+                yield batch
+            if len(field_map) < batch_size:
+                return
+            offset += batch_size
+
     @classmethod
     def _build_toc(cls, ctx: TaskContext, docs: List[Dict], progress_cb: Callable) -> Optional[Dict]:
         """Build table of contents."""
         progress_cb(msg="Start to generate table of content ...")
-        chat_model_config = get_model_config_from_provider_instance(
-            ctx.tenant_id, LLMType.CHAT, ctx.llm_id
-        )
+        chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
         with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_mdl:
-
-            docs = sorted(docs, key=lambda d: (
-                d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
-                d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0)
-            ))
+            docs = sorted(
+                docs,
+                key=lambda d: (
+                    d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
+                    d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0),
+                ),
+            )
 
             # NOTE: asyncio.run() creates a new event loop in the worker thread
             # (this method is called via asyncio.to_thread), which is the
             # intended pattern for bridging sync -> async in a thread context.
-            toc: list[dict] = asyncio.run(
-                run_toc_from_text([d["content_with_weight"] for d in docs], chat_mdl, progress_cb)
-            )
-            logging.info("------------ T O C -------------\n" + json.dumps(toc, ensure_ascii=False, indent='  '))
+            toc: list[dict] = asyncio.run(run_toc_from_text([d["content_with_weight"] for d in docs], chat_mdl, progress_cb))
+            logging.info("------------ T O C -------------\n" + json.dumps(toc, ensure_ascii=False, indent="  "))
 
             for ii, item in enumerate(toc):
                 try:
@@ -578,19 +835,17 @@ class TaskHandler:
 
             if toc:
                 import copy
+
                 d = copy.deepcopy(docs[-1])
                 d["content_with_weight"] = json.dumps(toc, ensure_ascii=False)
                 d["toc_kwd"] = "toc"
                 d["available_int"] = 0
                 d["page_num_int"] = [100000000]
-                d["id"] = xxhash.xxh64(
-                    (d["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
+                d["id"] = xxhash.xxh64((d["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
                 return d
             return None
 
-    async def _delete_raptor_chunks(
-        self, doc_id: str, tenant_id: str, kb_id: str, keep_method: Optional[str]
-    ) -> int:
+    async def _delete_raptor_chunks(self, doc_id: str, tenant_id: str, kb_id: str, keep_method: Optional[str]) -> int:
         """Delete RAPTOR chunks."""
         if self._task_context.write_interceptor:
             return self._task_context.write_interceptor.intercept("delete_raptor_chunks")

rag/utils/es_conn.py

@@ -65,12 +65,7 @@ class ESConnection(ESConnectionBase):
     """
 
     def _es_search_once(self, index_names: list[str], query: dict, track_total_hits: bool):
-        return self.es.search(
-            index=index_names,
-            body=query,
-            timeout="600s",
-            track_total_hits=track_total_hits
-        )
+        return self.es.search(index=index_names, body=query, timeout="600s", track_total_hits=track_total_hits)
 
     def _search_with_search_after(self, index_names: list[str], query: dict, offset: int, limit: int):
         q_base = copy.deepcopy(query)
@@ -139,17 +134,18 @@ class ESConnection(ESConnectionBase):
         return template_res
 
     def search(
-            self, select_fields: list[str],
-            highlight_fields: list[str],
-            condition: dict,
-            match_expressions: list[MatchExpr],
-            order_by: OrderByExpr,
-            offset: int,
-            limit: int,
-            index_names: str | list[str],
-            knowledgebase_ids: list[str],
-            agg_fields: list[str] | None = None,
-            rank_feature: dict | None = None
+        self,
+        select_fields: list[str],
+        highlight_fields: list[str],
+        condition: dict,
+        match_expressions: list[MatchExpr],
+        order_by: OrderByExpr,
+        offset: int,
+        limit: int,
+        index_names: str | list[str],
+        knowledgebase_ids: list[str],
+        agg_fields: list[str] | None = None,
+        rank_feature: dict | None = None,
     ):
         """
         Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl.html
@@ -166,19 +162,22 @@ class ESConnection(ESConnectionBase):
                 if v == 0:
                     bool_query.filter.append(Q("range", available_int={"lt": 1}))
                 else:
-                    bool_query.filter.append(
-                        Q("bool", must_not=Q("range", available_int={"lt": 1})))
+                    bool_query.filter.append(Q("bool", must_not=Q("range", available_int={"lt": 1})))
                 continue
             if k == "id":
                 if not v:
                     continue
                 if isinstance(v, list):
-                    bool_query.filter.append(
-                        Q("bool", should=[Q("terms", id=v), Q("terms", _id=v)], minimum_should_match=1))
+                    bool_query.filter.append(Q("bool", should=[Q("terms", id=v), Q("terms", _id=v)], minimum_should_match=1))
                 elif isinstance(v, str) or isinstance(v, int):
-                    bool_query.filter.append(
-                        Q("bool", should=[Q("term", id=v), Q("term", _id=v)], minimum_should_match=1))
+                    bool_query.filter.append(Q("bool", should=[Q("term", id=v), Q("term", _id=v)], minimum_should_match=1))
                 continue
+            if k == "must_not":
+                if isinstance(v, dict):
+                    for kk, vv in v.items():
+                        if kk == "exists":
+                            bool_query.must_not.append(Q("exists", field=vv))
+                    continue
             if not v:
                 continue
             if isinstance(v, list):
@@ -186,17 +185,18 @@ class ESConnection(ESConnectionBase):
             elif isinstance(v, str) or isinstance(v, int):
                 bool_query.filter.append(Q("term", **{k: v}))
             else:
-                raise Exception(
-                    f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
+                raise Exception(f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
 
         s = Search()
         vector_similarity_weight = 0.5
         for m in match_expressions:
             if isinstance(m, FusionExpr) and m.method == "weighted_sum" and "weights" in m.fusion_params:
-                assert len(match_expressions) == 3 and isinstance(match_expressions[0], MatchTextExpr) and isinstance(
-                    match_expressions[1],
-                    MatchDenseExpr) and isinstance(
-                    match_expressions[2], FusionExpr)
+                assert (
+                    len(match_expressions) == 3
+                    and isinstance(match_expressions[0], MatchTextExpr)
+                    and isinstance(match_expressions[1], MatchDenseExpr)
+                    and isinstance(match_expressions[2], FusionExpr)
+                )
                 weights = m.fusion_params["weights"]
                 vector_similarity_weight = get_float(weights.split(",")[1])
         for m in match_expressions:
@@ -204,24 +204,22 @@ class ESConnection(ESConnectionBase):
                 minimum_should_match = m.extra_options.get("minimum_should_match", 0.0)
                 if isinstance(minimum_should_match, float):
                     minimum_should_match = str(int(minimum_should_match * 100)) + "%"
-                bool_query.must.append(Q("query_string", fields=m.fields,
-                                         type="best_fields", query=m.matching_text,
-                                         minimum_should_match=minimum_should_match,
-                                         boost=1))
+                bool_query.must.append(Q("query_string", fields=m.fields, type="best_fields", query=m.matching_text, minimum_should_match=minimum_should_match, boost=1))
                 bool_query.boost = 1.0 - vector_similarity_weight
 
             elif isinstance(m, MatchDenseExpr):
-                assert (bool_query is not None)
+                assert bool_query is not None
                 similarity = 0.0
                 if "similarity" in m.extra_options:
                     similarity = m.extra_options["similarity"]
-                s = s.knn(m.vector_column_name,
-                          m.topn,
-                          m.topn * 2,
-                          query_vector=list(m.embedding_data),
-                          filter=bool_query.to_dict(),
-                          similarity=similarity,
-                          )
+                s = s.knn(
+                    m.vector_column_name,
+                    m.topn,
+                    m.topn * 2,
+                    query_vector=list(m.embedding_data),
+                    filter=bool_query.to_dict(),
+                    similarity=similarity,
+                )
 
         if bool_query and rank_feature:
             for fld, sc in rank_feature.items():
@@ -239,31 +237,25 @@ class ESConnection(ESConnectionBase):
             for field, order in order_by.fields:
                 order = "asc" if order == 0 else "desc"
                 if field in ["page_num_int", "top_int"]:
-                    order_info = {"order": order, "unmapped_type": "float",
-                                  "mode": "avg", "numeric_type": "double"}
+                    order_info = {"order": order, "unmapped_type": "float", "mode": "avg", "numeric_type": "double"}
                 elif field.endswith("_int") or field.endswith("_flt"):
                     order_info = {"order": order, "unmapped_type": "float"}
                 elif field == "id":
-                    continue # id as "text", not a "keyword", order by it will cause error
+                    continue  # id as "text", not a "keyword", order by it will cause error
                 else:
                     order_info = {"order": order, "unmapped_type": "keyword"}
                 orders.append({field: order_info})
             s = s.sort(*orders)
         if agg_fields:
             for fld in agg_fields:
-                s.aggs.bucket(f'aggs_{fld}', 'terms', field=fld, size=1000000)
+                s.aggs.bucket(f"aggs_{fld}", "terms", field=fld, size=1000000)
 
         has_dense = any(isinstance(m, MatchDenseExpr) for m in match_expressions)
         has_explicit_sort = bool(order_by and order_by.fields)
-        use_search_after = (
-            limit > 0
-            and (offset + limit > MAX_RESULT_WINDOW)
-            and has_explicit_sort
-            and not has_dense
-        )
+        use_search_after = limit > 0 and (offset + limit > MAX_RESULT_WINDOW) and has_explicit_sort and not has_dense
 
         if limit > 0 and not use_search_after:
-            s = s[offset:offset + limit]
+            s = s[offset : offset + limit]
         # Filter _source to only requested fields for efficiency, and add vector
         # fields to "fields" param so they appear in hit.fields when ES 9.x
         # exclude_source_vectors is enabled (dense_vector not in _source).
@@ -295,7 +287,7 @@ class ESConnection(ESConnectionBase):
                 continue
             except Exception as e:
                 # Only log debug for NotFoundError(accepted when metadata index doesn't exist)
-                if 'NotFound' in str(e):
+                if "NotFound" in str(e):
                     self.logger.debug(f"ESConnection.search {str(index_names)} query: " + str(q) + " - " + str(e))
                 else:
                     self.logger.exception(f"ESConnection.search {str(index_names)} query: " + str(q) + str(e))
@@ -314,16 +306,14 @@ class ESConnection(ESConnectionBase):
             d_copy["kb_id"] = knowledgebase_id
             # Use id as _id for uniqueness, also keep "id" as a regular field for sorting
             meta_id = d_copy.get("id", "")
-            operations.append(
-                {"index": {"_index": index_name, "_id": meta_id}})
+            operations.append({"index": {"_index": index_name, "_id": meta_id}})
             operations.append(d_copy)
 
         res = []
         for _ in range(ATTEMPT_TIME):
             try:
                 res = []
-                r = self.es.bulk(index=index_name, operations=operations,
-                                 refresh="wait_for", timeout="60s")
+                r = self.es.bulk(index=index_name, operations=operations, refresh="wait_for", timeout="60s")
                 if re.search(r"False", str(r["errors"]), re.IGNORECASE):
                     return res
 
@@ -359,10 +349,9 @@ class ESConnection(ESConnectionBase):
                     if "feas" != k.split("_")[-1]:
                         continue
                     try:
-                        self.es.update(index=index_name, id=chunk_id, script=f"ctx._source.remove(\"{k}\");")
+                        self.es.update(index=index_name, id=chunk_id, script=f'ctx._source.remove("{k}");')
                     except Exception:
-                        self.logger.exception(
-                            f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
+                        self.logger.exception(f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
                 try:
                     if remove_field is not None:
                         self.es.update(
@@ -375,9 +364,7 @@ class ESConnection(ESConnectionBase):
                         params = {}
                         for kk, vv in remove_dict.items():
                             scripts.append(
-                                f"if (ctx._source.containsKey('{kk}') && ctx._source.{kk} != null) "
-                                f"{{ int i = ctx._source.{kk}.indexOf(params.p_{kk}); "
-                                f"if (i >= 0) {{ ctx._source.{kk}.remove(i); }} }}"
+                                f"if (ctx._source.containsKey('{kk}') && ctx._source.{kk} != null) {{ int i = ctx._source.{kk}.indexOf(params.p_{kk}); if (i >= 0) {{ ctx._source.{kk}.remove(i); }} }}"
                             )
                             params[f"p_{kk}"] = vv
                         if scripts:
@@ -391,9 +378,7 @@ class ESConnection(ESConnectionBase):
                     if remove_field is not None or remove_dict is not None or doc_part:
                         return True
                 except Exception as e:
-                    self.logger.exception(
-                        f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception: " + str(
-                            e))
+                    self.logger.exception(f"ESConnection.update(index={index_name}, id={chunk_id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception: " + str(e))
                     break
             return False
 
@@ -405,13 +390,18 @@ class ESConnection(ESConnectionBase):
             if k == "exists":
                 bool_query.filter.append(Q("exists", field=v))
                 continue
+            if k == "must_not":
+                if isinstance(v, dict):
+                    for kk, vv in v.items():
+                        if kk == "exists":
+                            bool_query.must_not.append(Q("exists", field=vv))
+                    continue
             if isinstance(v, list):
                 bool_query.filter.append(Q("terms", **{k: v}))
             elif isinstance(v, str) or isinstance(v, int):
                 bool_query.filter.append(Q("term", **{k: v}))
             else:
-                raise Exception(
-                    f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
+                raise Exception(f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
         scripts = []
         params = {}
         for k, v in new_value.items():
@@ -441,11 +431,8 @@ class ESConnection(ESConnectionBase):
                 scripts.append(f"ctx._source.{k}=params.pp_{k};")
                 params[f"pp_{k}"] = json.dumps(v, ensure_ascii=False)
             else:
-                raise Exception(
-                    f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.")
-        ubq = UpdateByQuery(
-            index=index_name).using(
-            self.es).query(bool_query)
+                raise Exception(f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.")
+        ubq = UpdateByQuery(index=index_name).using(self.es).query(bool_query)
         ubq = ubq.script(source="".join(scripts), params=params)
         ubq = ubq.params(refresh=True)
         ubq = ubq.params(slices=5)
@@ -563,10 +550,7 @@ class ESConnection(ESConnectionBase):
         self.logger.debug("ESConnection.delete query: " + json.dumps(qry.to_dict()))
         for _ in range(ATTEMPT_TIME):
             try:
-                res = self.es.delete_by_query(
-                    index=index_name,
-                    body=Search().query(qry).to_dict(),
-                    refresh=True)
+                res = self.es.delete_by_query(index=index_name, body=Search().query(qry).to_dict(), refresh=True)
                 return res["deleted"]
             except ConnectionTimeout:
                 self.logger.exception("ES request timeout")

rag/utils/opensearch_conn.py

@@ -26,8 +26,7 @@ from opensearchpy import UpdateByQuery, Q, Search, Index
 from opensearchpy import ConnectionTimeout
 from common.decorator import singleton
 from common.file_utils import get_project_base_directory
-from common.doc_store.doc_store_base import DocStoreConnection, MatchExpr, OrderByExpr, MatchTextExpr, MatchDenseExpr, \
-    FusionExpr
+from common.doc_store.doc_store_base import DocStoreConnection, MatchExpr, OrderByExpr, MatchTextExpr, MatchDenseExpr, FusionExpr
 from rag.nlp import is_english, rag_tokenizer
 from common.constants import PAGERANK_FLD, TAG_FLD
 from common import settings
@@ -58,7 +57,7 @@ if (nw <= 0.0) {
 }
 """
 
-logger = logging.getLogger('ragflow.opensearch_conn')
+logger = logging.getLogger("ragflow.opensearch_conn")
 
 
 @singleton
@@ -70,10 +69,9 @@ class OSConnection(DocStoreConnection):
             try:
                 self.os = OpenSearch(
                     settings.OS["hosts"].split(","),
-                    http_auth=(settings.OS["username"], settings.OS[
-                        "password"]) if "username" in settings.OS and "password" in settings.OS else None,
+                    http_auth=(settings.OS["username"], settings.OS["password"]) if "username" in settings.OS and "password" in settings.OS else None,
                     verify_certs=False,
-                    timeout=600
+                    timeout=600,
                 )
                 if self.os:
                     self.info = self.os.info()
@@ -118,8 +116,7 @@ class OSConnection(DocStoreConnection):
         warning, leave it off, and search() keeps doing vector-only.
         """
         self.hybrid_search_enabled = False
-        self._hybrid_pipeline = os.environ.get("OS_HYBRID_PIPELINE") \
-            or settings.OS.get("hybrid_search_pipeline") or "ragflow_hybrid_pipeline"
+        self._hybrid_pipeline = os.environ.get("OS_HYBRID_PIPELINE") or settings.OS.get("hybrid_search_pipeline") or "ragflow_hybrid_pipeline"
 
         version_number = self.info.get("version", {}).get("number", "")
         try:
@@ -127,34 +124,32 @@ class OSConnection(DocStoreConnection):
         except (ValueError, AttributeError):
             version = (0, 0)
         if version < self.HYBRID_MIN_VERSION:
-            logger.warning(f"OpenSearch {version_number or 'unknown'} does not support the "
-                           f"normalization-processor (requires >= {self.HYBRID_MIN_VERSION[0]}."
-                           f"{self.HYBRID_MIN_VERSION[1]}); hybrid search is disabled and "
-                           f"queries fall back to vector-only.")
+            logger.warning(
+                f"OpenSearch {version_number or 'unknown'} does not support the "
+                f"normalization-processor (requires >= {self.HYBRID_MIN_VERSION[0]}."
+                f"{self.HYBRID_MIN_VERSION[1]}); hybrid search is disabled and "
+                f"queries fall back to vector-only."
+            )
             return
 
         weights = settings.OS.get("hybrid_search_weights", [0.5, 0.5])
         pipeline_body = {
             "description": "RAGFlow hybrid search normalization pipeline (BM25 + KNN).",
-            "phase_results_processors": [
-                {"normalization-processor": {
-                    "normalization": {"technique": "min_max"},
-                    "combination": {"technique": "arithmetic_mean",
-                                    "parameters": {"weights": weights}}}}
-            ],
+            "phase_results_processors": [{"normalization-processor": {"normalization": {"technique": "min_max"}, "combination": {"technique": "arithmetic_mean", "parameters": {"weights": weights}}}}],
         }
         try:
-            self.os.transport.perform_request(
-                "PUT", f"/_search/pipeline/{self._hybrid_pipeline}", body=pipeline_body)
+            self.os.transport.perform_request("PUT", f"/_search/pipeline/{self._hybrid_pipeline}", body=pipeline_body)
             self.hybrid_search_enabled = True
-            logger.info(f"OpenSearch hybrid search enabled via pipeline "
-                        f"'{self._hybrid_pipeline}' (weights {weights}).")
+            logger.info(f"OpenSearch hybrid search enabled via pipeline '{self._hybrid_pipeline}' (weights {weights}).")
         except Exception:
-            logger.warning(f"Could not create OpenSearch search pipeline '{self._hybrid_pipeline}'; "
-                           f"hybrid search is disabled and queries fall back to vector-only. "
-                           f"Creating a search pipeline needs the "
-                           f"'cluster:admin/search/pipeline/put' privilege (relevant on "
-                           f"locked-down or managed OpenSearch).", exc_info=True)
+            logger.warning(
+                f"Could not create OpenSearch search pipeline '{self._hybrid_pipeline}'; "
+                f"hybrid search is disabled and queries fall back to vector-only. "
+                f"Creating a search pipeline needs the "
+                f"'cluster:admin/search/pipeline/put' privilege (relevant on "
+                f"locked-down or managed OpenSearch).",
+                exc_info=True,
+            )
 
     """
     Database operations
@@ -177,8 +172,8 @@ class OSConnection(DocStoreConnection):
             return True
         try:
             from opensearchpy.client import IndicesClient
-            return IndicesClient(self.os).create(index=indexName,
-                                                 body=self.mapping)
+
+            return IndicesClient(self.os).create(index=indexName, body=self.mapping)
         except Exception:
             logger.exception("OSConnection.createIndex error %s" % (indexName))
 
@@ -215,6 +210,7 @@ class OSConnection(DocStoreConnection):
                 mappings = {**mappings, "dynamic": True}
 
             from opensearchpy.client import IndicesClient
+
             body = {
                 "settings": doc_meta_mapping["settings"],
                 "mappings": mappings,
@@ -316,17 +312,18 @@ class OSConnection(DocStoreConnection):
     """
 
     def search(
-            self, select_fields: list[str],
-            highlight_fields: list[str],
-            condition: dict,
-            match_expressions: list[MatchExpr],
-            order_by: OrderByExpr,
-            offset: int,
-            limit: int,
-            index_names: str | list[str],
-            knowledgebase_ids: list[str],
-            agg_fields: list[str] = [],
-            rank_feature: dict | None = None
+        self,
+        select_fields: list[str],
+        highlight_fields: list[str],
+        condition: dict,
+        match_expressions: list[MatchExpr],
+        order_by: OrderByExpr,
+        offset: int,
+        limit: int,
+        index_names: str | list[str],
+        knowledgebase_ids: list[str],
+        agg_fields: list[str] = [],
+        rank_feature: dict | None = None,
     ):
         """
         Refers to https://github.com/opensearch-project/opensearch-py/blob/main/guides/dsl.md
@@ -345,9 +342,22 @@ class OSConnection(DocStoreConnection):
                 if v == 0:
                     bqry.filter.append(Q("range", available_int={"lt": 1}))
                 else:
-                    bqry.filter.append(
-                        Q("bool", must_not=Q("range", available_int={"lt": 1})))
+                    bqry.filter.append(Q("bool", must_not=Q("range", available_int={"lt": 1})))
                 continue
+            if k == "id":
+                if not v:
+                    continue
+                if isinstance(v, list):
+                    bqry.filter.append(Q("bool", should=[Q("terms", id=v), Q("ids", values=v)], minimum_should_match=1))
+                elif isinstance(v, str) or isinstance(v, int):
+                    bqry.filter.append(Q("bool", should=[Q("term", id=v), Q("ids", values=[v])], minimum_should_match=1))
+                continue
+            if k == "must_not":
+                if isinstance(v, dict):
+                    for kk, vv in v.items():
+                        if kk == "exists":
+                            bqry.must_not.append(Q("exists", field=vv))
+                    continue
             if not v:
                 continue
             if isinstance(v, list):
@@ -355,16 +365,18 @@ class OSConnection(DocStoreConnection):
             elif isinstance(v, str) or isinstance(v, int):
                 bqry.filter.append(Q("term", **{k: v}))
             else:
-                raise Exception(
-                    f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
+                raise Exception(f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
 
         s = Search()
         vector_similarity_weight = 0.5
         for m in match_expressions:
             if isinstance(m, FusionExpr) and m.method == "weighted_sum" and "weights" in m.fusion_params:
-                assert len(match_expressions) == 3 and isinstance(match_expressions[0], MatchTextExpr) and isinstance(match_expressions[1],
-                                                                                                        MatchDenseExpr) and isinstance(
-                    match_expressions[2], FusionExpr)
+                assert (
+                    len(match_expressions) == 3
+                    and isinstance(match_expressions[0], MatchTextExpr)
+                    and isinstance(match_expressions[1], MatchDenseExpr)
+                    and isinstance(match_expressions[2], FusionExpr)
+                )
                 weights = m.fusion_params["weights"]
                 vector_similarity_weight = float(weights.split(",")[1])
         knn_query = {}
@@ -374,10 +386,7 @@ class OSConnection(DocStoreConnection):
                 minimum_should_match = m.extra_options.get("minimum_should_match", 0.0)
                 if isinstance(minimum_should_match, float):
                     minimum_should_match = str(int(minimum_should_match * 100)) + "%"
-                bqry.must.append(Q("query_string", fields=m.fields,
-                                   type="best_fields", query=m.matching_text,
-                                   minimum_should_match=minimum_should_match,
-                                   boost=1))
+                bqry.must.append(Q("query_string", fields=m.fields, type="best_fields", query=m.matching_text, minimum_should_match=minimum_should_match, boost=1))
                 bqry.boost = 1.0 - vector_similarity_weight
 
             # Elasticsearch has the encapsulation of KNN_search in python sdk
@@ -385,7 +394,7 @@ class OSConnection(DocStoreConnection):
             # the following codes implement KNN_search in OpenSearch using DSL
             # Besides, Opensearch's DSL for KNN_search query syntax differs from that in Elasticsearch, I also made some adaptions for it
             elif isinstance(m, MatchDenseExpr):
-                assert (bqry is not None)
+                assert bqry is not None
                 similarity = 0.0
                 if "similarity" in m.extra_options:
                     similarity = m.extra_options["similarity"]
@@ -419,8 +428,7 @@ class OSConnection(DocStoreConnection):
             for field, order in order_by.fields:
                 order = "asc" if order == 0 else "desc"
                 if field in ["page_num_int", "top_int"]:
-                    order_info = {"order": order, "unmapped_type": "float",
-                                  "mode": "avg", "numeric_type": "double"}
+                    order_info = {"order": order, "unmapped_type": "float", "mode": "avg", "numeric_type": "double"}
                 elif field.endswith("_int") or field.endswith("_flt"):
                     order_info = {"order": order, "unmapped_type": "float"}
                 else:
@@ -429,10 +437,10 @@ class OSConnection(DocStoreConnection):
             s = s.sort(*orders)
 
         for fld in agg_fields:
-            s.aggs.bucket(f'aggs_{fld}', 'terms', field=fld, size=1000000)
+            s.aggs.bucket(f"aggs_{fld}", "terms", field=fld, size=1000000)
 
         if limit > 0:
-            s = s[offset:offset + limit]
+            s = s[offset : offset + limit]
         q = s.to_dict()
         logger.debug(f"OSConnection.search {str(index_names)} query: " + json.dumps(q))
 
@@ -455,13 +463,15 @@ class OSConnection(DocStoreConnection):
 
         for i in range(ATTEMPT_TIME):
             try:
-                res = self.os.search(index=index_names,
-                                     body=q,
-                                     timeout=600,
-                                     # search_type="dfs_query_then_fetch",
-                                     track_total_hits=True,
-                                     _source=True,
-                                     **search_kwargs)
+                res = self.os.search(
+                    index=index_names,
+                    body=q,
+                    timeout=600,
+                    # search_type="dfs_query_then_fetch",
+                    track_total_hits=True,
+                    _source=True,
+                    **search_kwargs,
+                )
                 if str(res.get("timed_out", "")).lower() == "true":
                     raise Exception("OpenSearch Timeout.")
                 logger.debug(f"OSConnection.search {str(index_names)} res: " + str(res))
@@ -477,8 +487,11 @@ class OSConnection(DocStoreConnection):
     def get(self, chunkId: str, indexName: str, knowledgebaseIds: list[str]) -> dict | None:
         for i in range(ATTEMPT_TIME):
             try:
-                res = self.os.get(index=(indexName),
-                                  id=chunkId, _source=True, )
+                res = self.os.get(
+                    index=(indexName),
+                    id=chunkId,
+                    _source=True,
+                )
                 if str(res.get("timed_out", "")).lower() == "true":
                     raise Exception("Es Timeout.")
                 chunk = res["_source"]
@@ -505,16 +518,14 @@ class OSConnection(DocStoreConnection):
             # doc-meta read path (DocMetadataService filters on / sorts by the
             # "id" field) can find it, mirroring ESConnection.insert().
             meta_id = d_copy.get("id", "")
-            operations.append(
-                {"index": {"_index": indexName, "_id": meta_id}})
+            operations.append({"index": {"_index": indexName, "_id": meta_id}})
             operations.append(d_copy)
 
         res = []
         for _ in range(ATTEMPT_TIME):
             try:
                 res = []
-                r = self.os.bulk(index=(indexName), body=operations,
-                                 refresh="wait_for", timeout=60)
+                r = self.os.bulk(index=(indexName), body=operations, refresh="wait_for", timeout=60)
                 if re.search(r"False", str(r["errors"]), re.IGNORECASE):
                     return res
 
@@ -556,9 +567,7 @@ class OSConnection(DocStoreConnection):
                         params = {}
                         for kk, vv in remove_dict.items():
                             scripts.append(
-                                f"if (ctx._source.containsKey('{kk}') && ctx._source.{kk} != null) "
-                                f"{{ int i = ctx._source.{kk}.indexOf(params.p_{kk}); "
-                                f"if (i >= 0) {{ ctx._source.{kk}.remove(i); }} }}"
+                                f"if (ctx._source.containsKey('{kk}') && ctx._source.{kk} != null) {{ int i = ctx._source.{kk}.indexOf(params.p_{kk}); if (i >= 0) {{ ctx._source.{kk}.remove(i); }} }}"
                             )
                             params[f"p_{kk}"] = vv
                         if scripts:
@@ -572,8 +581,7 @@ class OSConnection(DocStoreConnection):
                     if remove_field is not None or remove_dict is not None or doc_part:
                         return True
                 except Exception as e:
-                    logger.exception(
-                        f"OSConnection.update(index={indexName}, id={id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
+                    logger.exception(f"OSConnection.update(index={indexName}, id={id}, doc={json.dumps(condition, ensure_ascii=False)}) got exception")
                     if re.search(r"(timeout|connection)", str(e).lower()):
                         continue
                     break
@@ -592,8 +600,7 @@ class OSConnection(DocStoreConnection):
             elif isinstance(v, str) or isinstance(v, int):
                 bqry.filter.append(Q("term", **{k: v}))
             else:
-                raise Exception(
-                    f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
+                raise Exception(f"Condition `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str or list.")
         scripts = []
         params = {}
         for k, v in newValue.items():
@@ -623,11 +630,8 @@ class OSConnection(DocStoreConnection):
                 scripts.append(f"ctx._source.{k}=params.pp_{k};")
                 params[f"pp_{k}"] = json.dumps(v, ensure_ascii=False)
             else:
-                raise Exception(
-                    f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.")
-        ubq = UpdateByQuery(
-            index=indexName).using(
-            self.os).query(bqry)
+                raise Exception(f"newValue `{str(k)}={str(v)}` value type is {str(type(v))}, expected to be int, str.")
+        ubq = UpdateByQuery(index=indexName).using(self.os).query(bqry)
         ubq = ubq.script(source="".join(scripts), params=params)
         ubq = ubq.params(refresh=True)
         ubq = ubq.params(slices=5)
@@ -734,10 +738,7 @@ class OSConnection(DocStoreConnection):
         for _ in range(ATTEMPT_TIME):
             try:
                 # print(Search().query(qry).to_dict(), flush=True)
-                res = self.os.delete_by_query(
-                    index=indexName,
-                    body=Search().query(qry).to_dict(),
-                    refresh=True)
+                res = self.os.delete_by_query(index=indexName, body=Search().query(qry).to_dict(), refresh=True)
                 return res["deleted"]
             except Exception as e:
                 logger.warning("OSConnection.delete got exception: " + str(e))
@@ -820,8 +821,7 @@ class OSConnection(DocStoreConnection):
             txts = []
             for t in re.split(r"[.?!;\n]", txt):
                 for w in keywords:
-                    t = re.sub(r"(^|[ .?/'\"\(\)!,:;-])(%s)([ .?/'\"\(\)!,:;-])" % re.escape(w), r"\1<em>\2</em>\3", t,
-                               flags=re.IGNORECASE | re.MULTILINE)
+                    t = re.sub(r"(^|[ .?/'\"\(\)!,:;-])(%s)([ .?/'\"\(\)!,:;-])" % re.escape(w), r"\1<em>\2</em>\3", t, flags=re.IGNORECASE | re.MULTILINE)
                 if not re.search(r"<em>[^<>]+</em>", t, flags=re.IGNORECASE | re.MULTILINE):
                     continue
                 txts.append(t)
@@ -847,14 +847,8 @@ class OSConnection(DocStoreConnection):
         replaces = []
         for r in re.finditer(r" ([a-z_]+_l?tks)( like | ?= ?)'([^']+)'", sql):
             fld, v = r.group(1), r.group(3)
-            match = " MATCH({}, '{}', 'operator=OR;minimum_should_match=30%') ".format(
-                fld, rag_tokenizer.fine_grained_tokenize(rag_tokenizer.tokenize(v)))
-            replaces.append(
-                ("{}{}'{}'".format(
-                    r.group(1),
-                    r.group(2),
-                    r.group(3)),
-                 match))
+            match = " MATCH({}, '{}', 'operator=OR;minimum_should_match=30%') ".format(fld, rag_tokenizer.fine_grained_tokenize(rag_tokenizer.tokenize(v)))
+            replaces.append(("{}{}'{}'".format(r.group(1), r.group(2), r.group(3)), match))
 
         for p, r in replaces:
             sql = sql.replace(p, r, 1)
@@ -862,8 +856,7 @@ class OSConnection(DocStoreConnection):
 
         for i in range(ATTEMPT_TIME):
             try:
-                res = self.os.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format,
-                                        request_timeout="2s")
+                res = self.os.sql.query(body={"query": sql, "fetch_size": fetch_size}, format=format, request_timeout="2s")
                 return res
             except ConnectionTimeout:
                 logger.exception("OSConnection.sql timeout")

rag/utils/redis_conn.py

@@ -216,6 +216,14 @@ class RedisDB:
             self.__open__()
         return False
 
+    def set_if_absent(self, k, v, exp=3600):
+        try:
+            return bool(self.REDIS.set(k, v, exp, nx=True))
+        except Exception as e:
+            logging.warning("RedisDB.set_if_absent " + str(k) + " got exception: " + str(e))
+            self.__open__()
+        return False
+
     def sadd(self, key: str, member: str):
         try:
             self.REDIS.sadd(key, member)