ragflow/rag/advanced_rag/knowlege_compile/wiki.py

#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
"""WIKI compilation pipeline — MAP phase.

  - Chunks come from ES (or any pre-chunked list passed in by the caller).
    No outline-driven chunking; per-chunk byte offsets are not tracked.
  - The LLM goes through ``rag.prompts.generator.gen_json`` (json_repair-backed).
    Embeddings go through ``LLMBundle.encode`` via ``thread_pool_exec`` (kept
    in the signature for symmetry with the downstream REDUCE / REFINE phases
    even though MAP itself does not embed).
  - Citation anchor is the source chunk id (``source_chunk_ids`` list per item), not
    a byte position. The LLM is prompted to tag each extracted item with the
    ``[CHUNK_ID …]`` of the chunk it came from.
  - Resume: per-chunk extracts are persisted to ES under
    ``compile_kwd="wiki_map_extract"`` with ``available_int=0`` and no vector
    / token-list fields, so retrievers ignore them but downstream phases can
    fetch them by ``doc_id`` + ``source_chunk_ids``. Re-running MAP for the same
    ``doc_id`` skips chunks that already have an extract row.

Public entry: ``wiki_map_from_chunks``.
"""

import asyncio
import json
import logging
import re
from typing import Callable, Optional
from urllib.parse import urlsplit
from common.misc_utils import thread_pool_exec
from common.token_utils import num_tokens_from_string
from rag.prompts.generator import gen_json, message_fit_in

import xxhash as _xxhash

from ._common import (
    build_chunk_batches as _build_chunk_batches,
    bulk_dedup_items as _bulk_dedup_items,
    ensure_llm_bundle as _ensure_llm_bundle,
    knowledge_compile_gen_conf as _knowledge_compile_gen_conf,
    run_chunked_pipeline as _run_chunked_pipeline,
    stable_row_id as _stable_row_id,
)


# Global pipeline-rev — bumping this constant invalidates every cached
# wiki_map_extract / wiki_reduce_result / wiki_compilation_plan
# / wiki_page_draft / wiki_page row on the next re-run. Use it
# when a prompt or extraction schema changes in a way that should
# invalidate prior caches.
_WIKI_PIPELINE_REV = "v1"


def _chunk_hash(content: str) -> str:
    """xxh64 of a chunk's ``content_with_weight`` mixed with the global
    pipeline rev. The mix-in means a prompt / schema bump invalidates
    every cached row without us having to touch each row individually.
    """
    body = (content or "") + "|" + _WIKI_PIPELINE_REV
    return _xxhash.xxh64(body.encode("utf-8", "surrogatepass")).hexdigest()


# Tiny parser_config helpers shared with the structure pipeline. Pulled in
# here so the MAP entity/relation schemas and rules can be driven from the
# same ``parser_config`` shape that ``compile_structure_from_text`` uses.
from .structure import (
    _struct_get,
    _struct_localize,
)


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

WIKI_MAP_COMPILE_KWD = "wiki_map_extract"
DEFAULT_WIKI_MAP_WORKERS = 20
DEFAULT_WIKI_MAP_TIMEOUT = 600


WIKI_MAP_SYSTEM = (
    "You are a knowledge extraction engine. Extract structured knowledge from the "
    "provided document section. Return ONLY valid JSON matching the schema exactly. "
    "Never include any text outside the JSON object. If a category has no items, use []."
    "Keep the chunks' original language (Chinese/English etc.) for generated data."
)


_DEFAULT_ENTITY_SCHEMA_BODY = (
    '      "name": "string — entity canonical name as it appears in text",\n'
    '      "type": "string — one of: person|org|product|regulation|location|system|equipment|other",\n'
    '      "aliases": ["string"],\n'
    '      "source_chunk_id": "string — exact value from the chunk_id list above"'
)

_DEFAULT_RELATION_SCHEMA_BODY = (
    '      "from": "string — source entity/concept name",\n'
    '      "to": "string — target entity/concept name",\n'
    '      "type": "string — e.g. owns|part_of|caused_by|regulates|uses|located_in|other",\n'
    '      "source_chunk_id": "string — exact value from the chunk_id list above"'
)


WIKI_MAP_USER_TEMPLATE = """\
## Document context
Document id: {doc_id}
Batch contains {chunk_count} packed chunk(s). Each chunk is introduced by a
``[CHUNK_ID <id>]`` line. The chunk_id values to choose from are:
{chunk_id_list}

## Packed chunks
{packed_chunks}

---

Extract all knowledge from every chunk and return a single JSON object with this
exact schema:

{{
  "entities": [
    {{
      "name": "string - entity canonical name as it appears in text",
      "type": "string - {entity_type_rules}",
      "aliases": ["string"],
      "source_chunk_id": "string - exact value from the chunk_id list above"
    }}
  ],
  "concepts": [
    {{
      "term": "string - {concept_term}",
      "definition_excerpt": "string - {concept_definition_excerpt}",
      "source_chunk_id": "string - exact value from the chunk_id list above"
    }}
  ],
  "claims": [
    {{
      "statement": "string - {claim_statement}",
      "subject": "string - {claim_subject}",
      "confidence": "explicit",
      "source_chunk_id": "string - exact value from the chunk_id list above"
    }}
  ],
  "relations": [
    {{
      "from": "string - source entity/concept name",
      "to": "string - target entity/concept name",
      "type": "string - {relation_type_rules}",
      "source_chunk_id": "string - exact value from the chunk_id list above"
    }}
  ],
  "topics": ["string"]
}}

Rules:
- ``source_chunk_id`` MUST be one of the chunk_id values listed above (they
  look like ``C1``, ``C2``, …); do NOT invent new ids. Pick the chunk where
  the item is primarily stated.
- The ``[CHUNK_ID …]`` header lines AND the ``C1``/``C2``/… chunk tags are
  prompt scaffolding — they are NOT part of the document content. Do NOT
  extract them (or any other identifier-looking strings from the headers)
  as entities, concepts, claims, or relations. Entity ``name`` / concept
  ``term`` values must come from the human-readable chunk body only.
- NEVER use bare hexadecimal hashes (such as ``a3f1b2c4d5e6f7a8``),
  UUIDs, database row ids, or any other opaque identifier-looking token
  as an entity ``name`` or concept ``term``. If you cannot find a
  human-readable name for a candidate entity in the chunk body, drop it.
- Concrete examples of values that are ALWAYS WRONG:
    BAD entity: {{"name": "C1", "type": "product", "aliases": ["C1"]}}
    BAD entity: {{"name": "C3", "type": "location"}}
    BAD concept: {{"term": "C2"}}
    BAD entity: {{"name": "d523a888c5b2a167", "type": "location"}}
    BAD entity: {{"name": "41a5271858ca11f1bbb9047c16ec874f", "type": "product"}}
  ``C1`` / ``C2`` / etc. are CHUNK TAGS, not products or locations. The
  hex hashes are DATABASE IDS, not entities. If your candidate ``name``
  matches any of these shapes, do not include the item in the output.
- ``confidence`` is ``"explicit"`` (directly stated) or ``"inferred"`` (implied
  by the text).
- Be exhaustive — include all named entities, defined terms, and factual claims.
- For ``concepts``, extract BOTH (a) named terms with definitions AND (b)
  coherent thematic sub-topics that could become their own wiki page.
- Extract ``claims`` LIBERALLY: every factual sentence about an entity is a
  claim. Definitions, attributes, ownership, locations, dates, actions,
  events, financial figures, regulations cited — all qualify. If you
  extract an entity, you should usually extract one or more claims that
  mention it. An empty ``claims`` array is almost always wrong unless the
  chunks are pure boilerplate.
- ``relations`` only fire when the text states an explicit link between two
  named entities/concepts (``A owns B``, ``A is part of B``, ``A regulates B``).
  Otherwise leave ``relations`` empty.
- Return empty arrays ``[]`` for categories with no findings.
- Return ONLY the JSON object, no markdown fences, no commentary.
{custom_rules}"""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_EXTRACT_LIST_KEYS = ("entities", "concepts", "claims", "relations")


def _wiki_empty_extract() -> dict:
    return {
        "entities": [],
        "concepts": [],
        "claims": [],
        "relations": [],
        "topics": [],
    }


# ---- parser_config-driven schema rendering --------------------------------
#
# artifact MAP's prompt previously hardcoded the entity & relation field set
# (name/type/aliases for entities; from/to/type for relations). The shape
# is now driven by ``parser_config.output.entities.fields`` and
# ``parser_config.output.relations.fields`` — the same YAML-style config
# used by ``compile_structure_from_text``. When no fields are configured
# (or no parser_config is passed) we fall back to the original artifact
# defaults so existing call sites keep working.


def _wiki_render_schema_body(fields, language: str, default_body: str, *, indent: int = 6) -> str:
    """Render the JSON body for one item in the entity/relation schema.

    Produces one line per field of the form::

        "<field>": <placeholder>

    where ``<placeholder>`` carries a ``string — <description>`` hint for
    string fields (so the LLM sees the user's intent) or a typed example
    for list/int/float/bool fields. Always appends ``source_chunk_id`` as
    the final field — chunk attribution is structural and not user-tunable.

    Falls back to ``default_body`` when ``fields`` is empty or only contains
    invalid entries.
    """
    if not fields:
        return default_body

    pad = " " * indent
    lines: list[str] = []
    seen: set[str] = set()
    for f in fields:
        if not isinstance(f, dict):
            continue
        name = f.get("name") or ""
        name = name.strip() if isinstance(name, str) else ""
        if not name or name in seen or name == "source_chunk_id":
            # Skip duplicates and any user-supplied source_chunk_id — we
            # always append our canonical one below.
            continue
        seen.add(name)

        ftype = f.get("type", "str")
        desc = _struct_localize(f.get("description", ""), language)
        if ftype == "list":
            placeholder = '["string"]'
        elif ftype == "int":
            placeholder = "0"
        elif ftype == "float":
            placeholder = "0.0"
        elif ftype == "bool":
            placeholder = "false"
        else:
            if desc:
                # Strip newlines and curly braces from the description so it
                # doesn't break the prompt's JSON layout or str.format.
                safe = desc.replace("\n", " ").replace("{", "(").replace("}", ")").strip()
                placeholder = f'"string — {safe}"'
            else:
                placeholder = '"string"'
        lines.append(f'{pad}"{name}": {placeholder}')

    if not lines:
        return default_body

    lines.append(f'{pad}"source_chunk_id": "string — exact value from the chunk_id list above"')
    return ",\n".join(lines)


def _wiki_build_custom_rules(parser_config, language: str) -> str:
    """Concatenate user-provided entity/relation rules into bullet-style
    sections appended to the prompt's Rules section.

    Returns an empty string when no parser_config rules are present so the
    template's trailing triple-quote closes cleanly without an extra blank
    line.
    """
    if not isinstance(parser_config, dict):
        return ""

    guideline = _struct_get(parser_config, "guideline", default={}) or {}
    rules_e = _struct_localize(_struct_get(guideline, "rules_for_entities"), language)
    rules_r = _struct_localize(_struct_get(guideline, "rules_for_relations"), language)

    sections: list[str] = []
    if rules_e:
        sections.append("## Entity extraction rules (from knowledge base config):\n" + rules_e)
    if rules_r:
        sections.append("## Relation extraction rules (from knowledge base config):\n" + rules_r)

    if not sections:
        return ""
    return "\n" + "\n\n".join(sections) + "\n"


def _wiki_template_fields(parser_config, section: str) -> list:
    if not isinstance(parser_config, dict):
        return []
    cfg = _struct_get(parser_config, section, default={}) or {}
    fields = _struct_get(cfg, "fields", default=[]) or []
    return fields if isinstance(fields, list) else []


def _wiki_type_rules(fields: list) -> str:
    lines: list[str] = []
    for field in fields:
        if not isinstance(field, dict):
            continue
        typ = field.get("type")
        typ = typ.strip() if isinstance(typ, str) else ""
        if not typ:
            continue
        description = field.get("description")
        description = description.strip() if isinstance(description, str) else ""
        rule = field.get("rule")
        rule = rule.strip() if isinstance(rule, str) else ""
        lines.append(f"type: {typ}")
        if description:
            lines.append(f"  - description: {description}")
        if rule:
            lines.append(f"  - rule: {rule}")
    return "\n".join(lines)


def _wiki_pipe_join(fields: list, key: str) -> str:
    values: list[str] = []
    for field in fields:
        if not isinstance(field, dict):
            continue
        value = field.get(key)
        value = value.strip() if isinstance(value, str) else ""
        if value:
            values.append(value)
    return "|".join(values)


def _wiki_colon_join(fields: list, left_key: str, right_key: str) -> str:
    values: list[str] = []
    for field in fields:
        if not isinstance(field, dict):
            continue
        left = field.get(left_key)
        left = left.strip() if isinstance(left, str) else ""
        right = field.get(right_key)
        right = right.strip() if isinstance(right, str) else ""
        if left or right:
            values.append(f"{left}:{right}")
    return "\n".join(values)


def _wiki_named_field_description(fields: list, name: str) -> str:
    for field in fields:
        if not isinstance(field, dict):
            continue
        field_name = field.get("name")
        field_name = field_name.strip().lower() if isinstance(field_name, str) else ""
        if field_name == name:
            description = field.get("description")
            description = description.strip() if isinstance(description, str) else ""
            if description:
                return description
        legacy = field.get(name)
        legacy = legacy.strip() if isinstance(legacy, str) else ""
        if legacy:
            return legacy
    return ""


def _wiki_template_custom_rules(parser_config) -> str:
    if not isinstance(parser_config, dict):
        return ""
    rules = parser_config.get("global_rules")
    return rules.strip() if isinstance(rules, str) else ""


def _wiki_build_user_prompt(
    *,
    parser_config,
    language: str,
    doc_id,
    chunk_count: int,
    chunk_id_list: str,
    packed_chunks: str,
) -> str:
    """Fill ``WIKI_MAP_USER_TEMPLATE`` with the dynamic entity / relation
    schema bodies plus optional rules drawn from ``parser_config``."""
    ent_fields = _wiki_template_fields(parser_config, "entity")
    rel_fields = _wiki_template_fields(parser_config, "relation")
    concept_fields = _wiki_template_fields(parser_config, "concept")
    claim_fields = _wiki_template_fields(parser_config, "claim")
    entity_type_rules = _wiki_type_rules(ent_fields)
    relation_type_rules = _wiki_type_rules(rel_fields)
    concept_term = _wiki_pipe_join(concept_fields, "term")
    concept_definition_excerpt = _wiki_colon_join(concept_fields, "term", "definition_excerpt")
    claim_statement = _wiki_named_field_description(claim_fields, "statement")
    claim_subject = _wiki_named_field_description(claim_fields, "subject")
    custom_rules = _wiki_template_custom_rules(parser_config)

    if isinstance(parser_config, dict):
        output = _struct_get(parser_config, "output", default={}) or {}
        entities_cfg = _struct_get(output, "entities", default={}) or {}
        relations_cfg = _struct_get(output, "relations", default={}) or {}
        legacy_ent_fields = _struct_get(entities_cfg, "fields", default=[]) or []
        legacy_rel_fields = _struct_get(relations_cfg, "fields", default=[]) or []
        if not entity_type_rules and legacy_ent_fields:
            entity_type_rules = _wiki_render_schema_body(
                legacy_ent_fields,
                language,
                _DEFAULT_ENTITY_SCHEMA_BODY,
            )
        if not relation_type_rules and legacy_rel_fields:
            relation_type_rules = _wiki_render_schema_body(
                legacy_rel_fields,
                language,
                _DEFAULT_RELATION_SCHEMA_BODY,
            )

    if not entity_type_rules:
        entity_type_rules = "person|org|product|regulation|location|system|equipment|other"
    if not relation_type_rules:
        relation_type_rules = "include|ordered|owns|part_of|caused_by|regulates|uses|located_in|other"
    if not concept_term:
        concept_term = "named term or topic"
    if not concept_definition_excerpt:
        concept_definition_excerpt = "short definition excerpt from the source text"
    if not claim_statement:
        claim_statement = "factual statement"
    if not claim_subject:
        claim_subject = "entity or concept that the claim is about"
    if not custom_rules:
        custom_rules = _wiki_build_custom_rules(parser_config, language)

    return WIKI_MAP_USER_TEMPLATE.format(
        doc_id=doc_id,
        chunk_count=chunk_count,
        chunk_id_list=chunk_id_list,
        packed_chunks=packed_chunks,
        entity_type_rules=entity_type_rules,
        relation_type_rules=relation_type_rules,
        concept_term=concept_term,
        concept_definition_excerpt=concept_definition_excerpt,
        claim_statement=claim_statement,
        claim_subject=claim_subject,
        custom_rules=custom_rules,
    )


def _wiki_pick_chunk_text(chunk: dict) -> str:
    text = chunk.get("text") or chunk.get("content_with_weight") or chunk.get("content") or ""
    return text if isinstance(text, str) else ""


# Matches a bare 16-char lowercase-hex token bounded by non-word chars on
# both sides — the shape of an xxh64 hexdigest (chunk_id / row id).
_HEX16_TOKEN_RE = re.compile(r"(?<![0-9a-zA-Z])[0-9a-f]{16}(?![0-9a-zA-Z])")
# Similar for the 32-char doc-id / uuid-without-dashes pattern.
_HEX32_TOKEN_RE = re.compile(r"(?<![0-9a-zA-Z])[0-9a-f]{32}(?![0-9a-zA-Z])")


def _wiki_scrub_known_ids(text: str, ids_to_remove) -> str:
    """Defensive scrub: strip any literal occurrence of known ES ids from
    chunk text before sending to the extraction LLM.

    Some chunkers embed the chunk_id / doc_id into the body (e.g. as a
    header, footer, or breadcrumb). Without this scrub the extraction LLM
    grabs the hash and reports it as an entity name (commonly mis-typed as
    "location"). We belt-and-brace by removing:

      1. Every literal id passed in ``ids_to_remove`` (chunk_ids of the
         batch + the doc_id).
      2. Any standalone 16-hex or 32-hex token still left over after (1).
    """
    if not text:
        return text
    out = text
    for h in ids_to_remove or ():
        if h and isinstance(h, str) and h in out:
            out = out.replace(h, "")
    out = _HEX16_TOKEN_RE.sub("", out)
    out = _HEX32_TOKEN_RE.sub("", out)
    return out


def _wiki_format_batch_prompt(packed: list[dict]) -> tuple[str, list[str]]:
    """Render the [CHUNK_ID …]-labelled body and return (body_text, label_order)."""
    parts: list[str] = []
    labels: list[str] = []
    for entry in packed:
        labels.append(entry["label"])
        parts.append(f"[CHUNK_ID {entry['label']}]\n{entry['text']}")
    return "\n\n".join(parts), labels


def _wiki_unwrap_extract(res) -> dict:
    """Coerce LLM JSON to the canonical 5-key shape with defaulted lists."""
    out = _wiki_empty_extract()
    if not isinstance(res, dict):
        return out
    for k in _EXTRACT_LIST_KEYS:
        v = res.get(k)
        if isinstance(v, list):
            out[k] = [item for item in v if isinstance(item, dict)]
    topics = res.get("topics")
    if isinstance(topics, list):
        out["topics"] = [t for t in topics if isinstance(t, str) and t.strip()]
    return out


# Matches strings the LLM should NEVER use as an entity / concept / claim name:
#   - chunk tag scaffolding: C1, C2, c0001, …
#   - bare hexadecimal hashes (xxh64 16-char, doc-id 32-char)
#   - UUIDs with or without dashes
# Anything matching is dropped post-extraction as defensive filtering.
_WIKI_IDENTIFIER_LIKE_RE = re.compile(
    r"""^\s*(
        [Cc]\d{1,5}                       # chunk tag like C1, c0001
        | [0-9a-fA-F]{16}                 # xxh64 hexdigest
        | [0-9a-fA-F]{32}                 # md5 / doc_id-shaped
        | [0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}  # UUID
    )\s*$""",
    re.VERBOSE,
)


def _wiki_looks_like_identifier(s) -> bool:
    """True when ``s`` looks like a chunk tag, hash, or UUID rather than a name."""
    if not isinstance(s, str):
        return False
    return bool(_WIKI_IDENTIFIER_LIKE_RE.match(s))


def _wiki_item_has_identifier_name(key: str, item: dict) -> bool:
    """Return True when an extracted item's display name is identifier-shaped.

    Different list keys carry the name field under different keys:
      - entities  → ``name``
      - concepts  → ``term``
      - claims    → ``subject``
      - relations → ``from`` and ``to`` (drop if either is identifier-shaped)
    """
    if key == "entities":
        return _wiki_looks_like_identifier(item.get("name", ""))
    if key == "concepts":
        return _wiki_looks_like_identifier(item.get("term", ""))
    if key == "claims":
        return _wiki_looks_like_identifier(item.get("subject", ""))
    if key == "relations":
        return _wiki_looks_like_identifier(item.get("from", "")) or _wiki_looks_like_identifier(item.get("to", ""))
    return False


def _wiki_resolve_chunk_ids(
    extract: dict,
    label_to_id: dict[str, str],
) -> tuple[dict, dict[str, dict]]:
    """Split a batch extract by source chunk id.

    Returns:
        merged: the input extract with ``source_chunk_id`` rewritten to
                ``chunk_ids=[real_id]`` per item, dropping items whose label
                does not match any in ``label_to_id``.
        per_chunk: {real_chunk_id: extract-shaped dict containing only the
                   items attributed to that chunk}. Includes empty extracts
                   for every label in ``label_to_id`` so resume knows the
                   chunk was processed even when nothing was extracted.
    """
    per_chunk: dict[str, dict] = {real_id: _wiki_empty_extract() for real_id in label_to_id.values()}
    merged = _wiki_empty_extract()
    merged["topics"] = list(extract.get("topics") or [])

    dropped = 0
    dropped_identifier = 0
    for key in _EXTRACT_LIST_KEYS:
        for item in extract.get(key) or []:
            label = item.get("source_chunk_id")
            real = label_to_id.get(label) if isinstance(label, str) else None
            if real is None:
                dropped += 1
                continue
            # Drop items whose display name is identifier-shaped — the LLM
            # occasionally grabs prompt scaffolding (C1, C2, …) or leftover
            # hash tokens and reports them as entities/concepts/claims. The
            # prompt forbids this but post-filtering is the bulletproof guard.
            if _wiki_item_has_identifier_name(key, item):
                dropped_identifier += 1
                continue
            new_item = {k: v for k, v in item.items() if k != "source_chunk_id"}
            new_item["chunk_ids"] = [real]
            merged[key].append(new_item)
            per_chunk[real][key].append(new_item)

    if dropped:
        logging.debug(f"wiki_map: dropped {dropped} item(s) with unrecognized source_chunk_id")
    if dropped_identifier:
        logging.info(
            "wiki_map: dropped %d item(s) whose name looked like a prompt-scaffolding tag or hash",
            dropped_identifier,
        )

    return merged, per_chunk


def _wiki_merge_extracts(extracts: list[dict]) -> dict:
    """Concat the 5 lists across multiple batch extracts (no entity-level
    dedup — that is the REDUCE phase's job)."""
    out = _wiki_empty_extract()
    seen_topics: set[str] = set()
    for ex in extracts:
        if not isinstance(ex, dict):
            continue
        for key in _EXTRACT_LIST_KEYS:
            out[key].extend(ex.get(key) or [])
        for t in ex.get("topics") or []:
            if t not in seen_topics:
                seen_topics.add(t)
                out["topics"].append(t)
    return out


def _wiki_build_resume_doc(
    chunk_id: str,
    doc_id: str,
    per_chunk_extract: dict,
    chunk_hash: str = "",
) -> dict:
    """Build the non-searchable ES doc that records a per-chunk MAP extract.

    Intentionally omits ``q_<dim>_vec`` / ``content_ltks`` / ``content_sm_ltks``
    so retrievers cannot surface this row; also sets ``available_int=0`` which
    most ragflow retrievers already filter on.

    ``chunk_hash`` fingerprints the chunk's content as of extraction time.
    The incremental MAP re-run reads it back and compares against the
    current chunk's hash to decide whether to re-extract.
    """
    content_with_weight = json.dumps(per_chunk_extract, ensure_ascii=False)
    doc_id_str = str(doc_id)
    return {
        "id": _stable_row_id(content_with_weight, doc_id_str, chunk_id),
        "doc_id": doc_id_str,
        "compile_kwd": WIKI_MAP_COMPILE_KWD,
        "source_chunk_ids": [chunk_id],
        "chunk_hash_kwd": chunk_hash,
        "content_with_weight": content_with_weight,
        "available_int": 0,
    }


async def _wiki_load_resume_map(
    doc_id: str,
    tenant_id: str,
    kb_id: str,
) -> dict[str, str]:
    """Query ES for chunks that already have a wiki_map_extract row for
    this doc. Returns ``{chunk_id → chunk_hash}``.

    ``chunk_hash`` may be empty for legacy rows that predate the field —
    callers treat empty as "definitely re-MAP" (no hash to compare).
    """
    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)
    condition = {
        "compile_kwd": [WIKI_MAP_COMPILE_KWD],
        "doc_id": [str(doc_id)],
    }
    select_fields = ["id", "source_chunk_ids", "chunk_hash_kwd"]
    try:
        res = await thread_pool_exec(
            settings.docStoreConn.search,
            select_fields,
            [],
            condition,
            [],
            OrderByExpr(),
            0,
            10000,
            index,
            [kb_id],
        )
        field_map = settings.docStoreConn.get_fields(res, select_fields)
    except Exception:
        logging.exception("wiki_map: failed to query resume map; will re-extract all chunks")
        return {}

    seen: dict[str, str] = {}
    for row in field_map.values():
        src = row.get("source_chunk_ids") or []
        hh = row.get("chunk_hash_kwd")
        if not isinstance(hh, str):
            hh = ""
        if isinstance(src, list):
            for cid in src:
                if isinstance(cid, str) and cid:
                    # First-write-wins is fine: if a doc has two rows for
                    # the same chunk_id (legacy / dirty state), we treat
                    # the first as the canonical and let the changed-hash
                    # path or the deletion sweep clean it up later.
                    seen.setdefault(cid, hh)
    return seen


async def _wiki_delete_map_rows(
    doc_id: str,
    chunk_ids: list[str],
    tenant_id: str,
    kb_id: str,
) -> int:
    """Delete ``wiki_map_extract`` rows for ``(doc_id, chunk_id)`` pairs.

    Used by the incremental MAP path:
      * stale rows whose chunk content has changed → re-extracted next.
      * rows whose chunk_id is gone from the doc (chunk deleted upstream).

    Returns the number of distinct ``chunk_ids`` we attempted to drop;
    the backend may delete more (e.g. duplicate rows) — we don't try to
    track that precisely.
    """
    if not chunk_ids:
        return 0
    from common import settings
    from rag.nlp import search as _rag_search

    index = _rag_search.index_name(tenant_id)
    condition = {
        "compile_kwd": [WIKI_MAP_COMPILE_KWD],
        "doc_id": [str(doc_id)],
        "source_chunk_ids": list(chunk_ids),
    }
    try:
        await thread_pool_exec(
            settings.docStoreConn.delete,
            condition,
            index,
            kb_id,
        )
    except Exception:
        logging.exception(
            "wiki_map: failed to delete %d stale extract row(s) for doc %s",
            len(chunk_ids),
            doc_id,
        )
        return 0
    return len(chunk_ids)


async def _wiki_persist_extracts(
    per_chunk: dict[str, dict],
    doc_id: str,
    tenant_id: str,
    kb_id: str,
    chunk_hashes: Optional[dict[str, str]] = None,
) -> None:
    """Write one non-searchable ES doc per source chunk_id.

    ``chunk_hashes`` (``{chunk_id → chunk_hash}``) is stamped onto each
    row so the next incremental run can decide whether to re-MAP.
    Missing entries default to '' (treated as "definitely re-MAP" by the
    resume-map comparator).
    """
    if not per_chunk:
        return
    from common import settings
    from rag.nlp import search as _rag_search

    index = _rag_search.index_name(tenant_id)
    hashes = chunk_hashes or {}
    docs = [
        _wiki_build_resume_doc(
            chunk_id,
            doc_id,
            extract,
            chunk_hash=hashes.get(chunk_id, ""),
        )
        for chunk_id, extract in per_chunk.items()
        if chunk_id
    ]
    if not docs:
        return
    try:
        await thread_pool_exec(settings.docStoreConn.insert, docs, index, kb_id)
    except Exception:
        logging.exception("wiki_map: failed to persist %d resume docs", len(docs))


# ---------------------------------------------------------------------------
# Per-batch extraction
# ---------------------------------------------------------------------------


async def _wiki_extract_one_batch(
    packed: list[dict],
    doc_id: str,
    chat_mdl,
    language: str,
    llm_timeout: int,
    parser_config: Optional[dict] = None,
) -> dict:
    """Single LLM call for one packed batch. Returns the raw (label-tagged)
    extract dict.

    The entity / relation schemas and the extra rules sections of the
    prompt are rendered from ``parser_config`` when supplied (mirroring
    ``compile_structure_from_text``); when omitted, the built-in defaults
    are used."""
    body, labels = _wiki_format_batch_prompt(packed)
    user_prompt = _wiki_build_user_prompt(
        parser_config=parser_config,
        language=language,
        doc_id=doc_id,
        chunk_count=len(packed),
        chunk_id_list="\n".join(f"- {label}" for label in labels),
        packed_chunks=body,
    )
    request_conf = _knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.1})
    try:
        res = await asyncio.wait_for(
            gen_json(
                WIKI_MAP_SYSTEM,
                user_prompt,
                chat_mdl,
                gen_conf=request_conf,
            ),
            timeout=llm_timeout,
        )
    except asyncio.TimeoutError:
        logging.warning("wiki_map: batch extraction timed out after %ds (%d chunks)", llm_timeout, len(packed))
        return _wiki_empty_extract()
    except Exception:
        logging.exception("wiki_map: batch extraction failed (%d chunks)", len(packed))
        return _wiki_empty_extract()
    _ = language  # reserved for future localization
    return _wiki_unwrap_extract(res)


async def _wiki_process_batch(
    packed: list[dict],
    batch_idx: int,
    total_batches: int,
    doc_id: str,
    tenant_id: str,
    kb_id: str,
    chat_mdl,
    language: str,
    llm_timeout: int,
    semaphore: Optional[asyncio.Semaphore],
    callback: Optional[Callable],
    parser_config: Optional[dict] = None,
    chunk_hashes: Optional[dict[str, str]] = None,
) -> dict:
    """Run one batch end-to-end: LLM extract → split by source_chunk_id →
    persist resume docs → return the merged batch extract.

    ``chunk_hashes`` is the ``{chunk_id → chunk_hash}`` map captured at
    the top of ``wiki_map_from_chunks``; threaded through so the
    persisted resume rows record the right hash and the next
    incremental run can compare cleanly.
    """
    if not packed:
        return _wiki_empty_extract()

    label_to_id = {entry["label"]: entry["chunk_id"] for entry in packed}

    async def _run() -> dict:
        raw_extract = await _wiki_extract_one_batch(
            packed,
            doc_id,
            chat_mdl,
            language,
            llm_timeout,
            parser_config=parser_config,
        )
        merged, per_chunk = _wiki_resolve_chunk_ids(raw_extract, label_to_id)
        await _wiki_persist_extracts(
            per_chunk,
            doc_id,
            tenant_id,
            kb_id,
            chunk_hashes=chunk_hashes,
        )
        if callback:
            try:
                n_items = sum(len(merged.get(k) or []) for k in _EXTRACT_LIST_KEYS)
                callback(
                    (batch_idx + 1) / max(1, total_batches),
                    f"Wiki MAP {batch_idx + 1}/{total_batches}: {n_items} items from {len(packed)} chunks",
                )
            except Exception:
                logging.debug("wiki_map: progress callback failed", exc_info=True)
        return merged

    if semaphore is not None:
        async with semaphore:
            return await _run()
    return await _run()


# ---------------------------------------------------------------------------
# Public entry
# ---------------------------------------------------------------------------


async def wiki_map_from_chunks(
    chunks: list[dict],
    chat_mdl,
    embd_mdl,
    doc_id: str,
    tenant_id: str,
    kb_id: str,
    language: str = "en",
    max_workers: int = DEFAULT_WIKI_MAP_WORKERS,
    llm_timeout: int = DEFAULT_WIKI_MAP_TIMEOUT,
    callback: Optional[Callable] = None,
    parser_config: Optional[dict] = None,
    batch_size_cap: Optional[int] = None,
    window_fraction: Optional[float] = None,
) -> dict:
    """Phase 1 (MAP) of the artifact compilation pipeline.

    Packs the provided RAGFlow chunks into batches via ``split_chunks``, runs
    one ``gen_json`` extraction call per batch in parallel (bounded by
    ``max_workers``), then splits each batch's output back to per-chunk
    extracts and persists them to ES as non-searchable ``wiki_map_extract``
    rows so subsequent runs can skip chunks already processed.

    Args:
        chunks: list of dicts; each must expose ``id`` and ``text`` (with
            ``content_with_weight`` / ``content`` accepted as fallbacks).
        chat_mdl: LLMBundle for chat (used via ``gen_json``).
        embd_mdl: LLMBundle for embeddings — accepted for downstream symmetry
            with REDUCE/REFINE but **not used in MAP itself**.
        doc_id: source document id; stamped onto every resume doc and on every
            extracted item via ``chunk_ids``.
        tenant_id, kb_id: address the doc-store index for resume reads + writes.
        language: reserved for future prompt localization.
        max_workers: maximum concurrent batches. Defaults to 20, matching the
            task-scoped Wiki LLM pool used by the task executor.
        llm_timeout: seconds per batch extraction call.
        callback: optional ``(progress: float, msg: str)`` progress callback.
        parser_config: optional YAML-style config (same shape that
            ``compile_structure_from_text`` accepts).
            ``source_chunk_id`` field is always appended so chunk
            attribution survives regardless of the user's schema. When
            omitted, the built-in default artifact schema is used.

    Returns:
        ``{"entities", "concepts", "claims", "relations", "topics"}`` where
        every item (except ``topics`` strings) carries a
        ``chunk_ids=[<source chunk id>]`` field. No entity-level dedup is
        performed here — that is the REDUCE phase's responsibility.
    """
    _ = embd_mdl  # noqa: F841 — accepted for symmetry with downstream phases
    if not chunks:
        # Even with zero chunks we still want to sweep any orphaned MAP
        # rows that point at chunks the doc no longer has — otherwise
        # deletions never propagate.
        prior_resume_map = await _wiki_load_resume_map(doc_id, tenant_id, kb_id)
        if prior_resume_map:
            await _wiki_delete_map_rows(
                doc_id,
                list(prior_resume_map.keys()),
                tenant_id,
                kb_id,
            )
            logging.info(
                "wiki_map: doc %s now has zero chunks; swept %d stale extract row(s)",
                doc_id,
                len(prior_resume_map),
            )
        out = _wiki_empty_extract()
        out["_meta"] = {
            "doc_id": str(doc_id),
            "new": 0,
            "changed": 0,
            "deleted": len(prior_resume_map),
            "unchanged": 0,
            "had_delta": bool(prior_resume_map),
        }
        return out

    # Incremental decision per current chunk:
    #
    #   * Compute the fresh chunk hash for every chunk in this call.
    #   * Load the prior resume map (chunk_id → hash from the last MAP).
    #   * NEW       — chunk_id not in prior     → MAP this chunk.
    #   * UNCHANGED — chunk_id in prior, hash matches → skip (resume).
    #   * CHANGED   — chunk_id in prior, hash differs → delete prior
    #                 row, then MAP this chunk.
    #   * DELETED   — chunk_id only in prior     → delete prior row
    #                 (chunk was removed upstream).
    #
    # The "resume set" handed to ``_build_chunk_batches`` is just the
    # UNCHANGED ids — those are the only ones the packer should skip.
    current_chunk_hashes: dict[str, str] = {}
    for chunk in chunks:
        cid = chunk.get("id") or chunk.get("chunk_id")
        if not isinstance(cid, str) or not cid:
            continue
        text = _wiki_pick_chunk_text(chunk) or ""
        current_chunk_hashes[cid] = _chunk_hash(text)

    prior_resume_map = await _wiki_load_resume_map(doc_id, tenant_id, kb_id)
    unchanged_ids: set[str] = set()
    changed_ids: list[str] = []
    new_ids: list[str] = []
    for cid, h in current_chunk_hashes.items():
        prior_h = prior_resume_map.get(cid)
        if prior_h is None:
            new_ids.append(cid)
        elif prior_h and prior_h == h:
            unchanged_ids.add(cid)
        else:
            # Empty stored hash = legacy row written before chunk_hash_kwd
            # existed → re-MAP. Differing hash = content changed → re-MAP.
            changed_ids.append(cid)
    deleted_ids = [cid for cid in prior_resume_map if cid not in current_chunk_hashes]

    if changed_ids or deleted_ids:
        await _wiki_delete_map_rows(
            doc_id,
            list(set(changed_ids) | set(deleted_ids)),
            tenant_id,
            kb_id,
        )

    if unchanged_ids or changed_ids or deleted_ids or new_ids:
        logging.info(
            "wiki_map: doc %s — new=%d changed=%d unchanged=%d deleted=%d",
            doc_id,
            len(new_ids),
            len(changed_ids),
            len(unchanged_ids),
            len(deleted_ids),
        )

    # The packer's "resume" set is the UNCHANGED ids only — NEW and
    # CHANGED both need re-extraction.
    resume_set = unchanged_ids

    # Defensive scrub: chunkers sometimes embed the chunk_id / doc_id into
    # the body (e.g. as a header). Without this the extraction LLM tends to
    # grab the hash as an "entity" — see _wiki_scrub_known_ids.
    all_known_ids: list[str] = []
    for chunk in chunks:
        cid = chunk.get("id") or chunk.get("chunk_id")
        if isinstance(cid, str) and cid:
            all_known_ids.append(cid)
    if doc_id:
        all_known_ids.append(str(doc_id))

    prompt_overhead = num_tokens_from_string(WIKI_MAP_SYSTEM + WIKI_MAP_USER_TEMPLATE)
    packed_batches, _info = _build_chunk_batches(
        chunks,
        chat_mdl,
        prompt_overhead_tokens=prompt_overhead,
        resume_chunk_ids=resume_set,
        scrub_text=lambda t: _wiki_scrub_known_ids(t, all_known_ids),
        chunk_text_picker=_wiki_pick_chunk_text,
        batch_size_cap=batch_size_cap,
        window_fraction=window_fraction,
    )
    if not packed_batches:
        return _wiki_empty_extract()

    async def _process_one(batch: list[dict], bi: int, total: int) -> dict:
        # ``_run_chunked_pipeline`` already wraps each task in the engine's
        # semaphore, so pass ``semaphore=None`` here to avoid nesting.
        return await _wiki_process_batch(
            packed=batch,
            batch_idx=bi,
            total_batches=total,
            doc_id=doc_id,
            tenant_id=tenant_id,
            kb_id=kb_id,
            chat_mdl=chat_mdl,
            language=language,
            llm_timeout=llm_timeout,
            semaphore=None,
            callback=callback,
            parser_config=parser_config,
            chunk_hashes=current_chunk_hashes,
        )

    merged = await _run_chunked_pipeline(
        packed_batches,
        process_batch=_process_one,
        aggregate=_wiki_merge_extracts,
        max_workers=max_workers,
        callback=callback,
        log_prefix="wiki_map",
    )
    logging.info(
        "wiki_map: doc %s — entities=%d concepts=%d claims=%d relations=%d topics=%d",
        doc_id,
        len(merged["entities"]),
        len(merged["concepts"]),
        len(merged["claims"]),
        len(merged["relations"]),
        len(merged["topics"]),
    )
    # Surface the incremental decisions to the orchestrator. ``had_delta``
    # is the most useful summary: REDUCE/PLAN/REFINE can short-circuit
    # KB-wide when no doc's MAP touched any rows on this run.
    merged["_meta"] = {
        "doc_id": str(doc_id),
        "new": len(new_ids),
        "changed": len(changed_ids),
        "unchanged": len(unchanged_ids),
        "deleted": len(deleted_ids),
        "had_delta": bool(new_ids or changed_ids or deleted_ids),
    }
    return merged


# ---------------------------------------------------------------------------
# REDUCE phase (KB-scoped)
# ---------------------------------------------------------------------------
#

WIKI_REDUCE_COMPILE_KWD = "wiki_reduce_result"
DEFAULT_WIKI_REDUCE_MERGE_THRESHOLD = 0.95
DEFAULT_WIKI_REDUCE_AMBIGUOUS_LOW = 0.75
DEFAULT_WIKI_REDUCE_AMBIGUOUS_BATCH = 50
DEFAULT_WIKI_REDUCE_TIMEOUT = 60


# System prompt for the LLM disambiguation batch. The shared engine
# (``_common.bulk_dedup_items``) defaults to the same wording via
# ``_common.DEFAULT_DISAMBIGUATE_SYSTEM``; we keep the local alias so the
# constant name stays usable by call sites and external imports.
WIKI_REDUCE_DISAMBIGUATE_SYSTEM = "You are a named-entity resolution assistant. Return only JSON."


# --- ES I/O ----------------------------------------------------------------


async def _wiki_load_all_map_extracts(tenant_id: str, kb_id: str) -> dict:
    """Aggregate every wiki_map_extract row in this KB into one merged dict.

    Pages through ES if the KB has more than the per-call cap. Returns a dict
    in the same shape as wiki_map_from_chunks' return value.
    """
    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)
    condition = {"compile_kwd": [WIKI_MAP_COMPILE_KWD]}
    select_fields = ["id", "content_with_weight"]

    PAGE_SIZE = 1000
    offset = 0
    merged = _wiki_empty_extract()
    seen_topics: set[str] = set()

    while True:
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                condition,
                [],
                OrderByExpr(),
                offset,
                PAGE_SIZE,
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields)
        except Exception:
            logging.exception("wiki_reduce: failed to page wiki_map_extract rows")
            break

        if not field_map:
            break

        for row in field_map.values():
            content = row.get("content_with_weight")
            if not isinstance(content, str) or not content:
                continue
            try:
                payload = json.loads(content)
            except Exception:
                logging.debug("wiki_reduce: skipping unparseable extract row")
                continue
            if not isinstance(payload, dict):
                continue
            for key in _EXTRACT_LIST_KEYS:
                items = payload.get(key)
                if isinstance(items, list):
                    merged[key].extend(item for item in items if isinstance(item, dict))
            topics = payload.get("topics")
            if isinstance(topics, list):
                for t in topics:
                    if isinstance(t, str) and t and t not in seen_topics:
                        seen_topics.add(t)
                        merged["topics"].append(t)

        if len(field_map) < PAGE_SIZE:
            break
        offset += PAGE_SIZE

    return merged


async def _wiki_all_map_doc_ids(tenant_id: str, kb_id: str) -> list[str]:
    """Distinct ``doc_id`` across every ``wiki_map_extract`` row in this KB.

    These are the documents that fed the current compilation. Stamped onto
    the KB-wide aggregate rows (REDUCE / PLAN) as ``source_doc_ids`` so a
    document delete can reference-count them — the aggregate is dropped only
    once its last contributing document is gone.
    """
    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)
    condition = {"compile_kwd": [WIKI_MAP_COMPILE_KWD]}
    select_fields = ["id", "doc_id"]

    PAGE_SIZE = 1000
    offset = 0
    doc_ids: list[str] = []
    seen: set[str] = set()
    while True:
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                condition,
                [],
                OrderByExpr(),
                offset,
                PAGE_SIZE,
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields)
        except Exception:
            logging.exception("wiki: failed to scan MAP doc ids for kb=%s (offset=%d)", kb_id, offset)
            break
        if not field_map:
            break
        for row in field_map.values():
            raw = row.get("doc_id")
            candidates = raw if isinstance(raw, list) else [raw]
            for d in candidates:
                if isinstance(d, str) and d and d not in seen:
                    seen.add(d)
                    doc_ids.append(d)
        if len(field_map) < PAGE_SIZE:
            break
        offset += PAGE_SIZE
    return doc_ids


async def _wiki_compute_map_input_hash(tenant_id: str, kb_id: str) -> str:
    """xxh64 fingerprint of the **current** ``wiki_map_extract`` rows for
    this KB — used by REDUCE / PLAN to cache-bust when MAP changed.

    Built from ``sorted((chunk_id, chunk_hash))`` so:
      * adding a new chunk → new pair appears → hash flips.
      * editing a chunk → MAP row deleted + re-inserted with new hash → flips.
      * deleting a chunk → MAP row gone → its pair drops → flips.
      * everything stable → identical hash.

    Empty / missing ``chunk_hash_kwd`` (legacy rows) defaults to '' so a
    legacy KB still produces a stable hash; once those rows are touched
    by an incremental MAP run, the hash naturally upgrades.

    Pages through ES in windows of ``PAGE_SIZE`` rows — single-shot
    "give me everything" reads hit doc-store limits on KBs with many
    chunks. The accumulated ``pairs`` are sorted once at the end so the
    fingerprint is independent of page order.
    """
    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)
    condition = {"compile_kwd": [WIKI_MAP_COMPILE_KWD]}
    select_fields = ["id", "source_chunk_ids", "chunk_hash_kwd"]

    PAGE_SIZE = 128
    offset = 0
    pairs: list[tuple[str, str]] = []
    while True:
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                condition,
                [],
                OrderByExpr(),
                offset,
                PAGE_SIZE,
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields)
        except Exception:
            logging.exception(
                "wiki: failed to compute MAP input hash for kb=%s (offset=%d)",
                kb_id,
                offset,
            )
            # Partial scan → cannot trust the resulting hash; return ""
            # so REDUCE / PLAN fall through to a full re-run rather than
            # cache-hitting against an incomplete fingerprint.
            return ""
        if not field_map:
            break
        for row in field_map.values():
            hh = row.get("chunk_hash_kwd")
            if not isinstance(hh, str):
                hh = ""
            src = row.get("source_chunk_ids") or []
            if isinstance(src, list):
                for cid in src:
                    if isinstance(cid, str) and cid:
                        pairs.append((cid, hh))
        if len(field_map) < PAGE_SIZE:
            break
        offset += PAGE_SIZE

    pairs.sort()
    body = "|".join(f"{cid}:{hh}" for cid, hh in pairs) + "|" + _WIKI_PIPELINE_REV
    return _xxhash.xxh64(body.encode("utf-8", "surrogatepass")).hexdigest()


async def _wiki_load_reduce_resume(
    tenant_id: str,
    kb_id: str,
) -> Optional[tuple[dict, str]]:
    """Return ``(cached_result, stored_input_hash)`` or None."""
    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)
    condition = {"compile_kwd": [WIKI_REDUCE_COMPILE_KWD]}
    select_fields = ["id", "content_with_weight", "input_hash_kwd"]
    try:
        res = await thread_pool_exec(
            settings.docStoreConn.search,
            select_fields,
            [],
            condition,
            [],
            OrderByExpr(),
            0,
            1,
            index,
            [kb_id],
        )
        field_map = settings.docStoreConn.get_fields(res, select_fields)
    except Exception:
        logging.exception("wiki_reduce: failed to load resume cache")
        return None
    if not field_map:
        return None
    row = next(iter(field_map.values()))
    content = row.get("content_with_weight")
    if not isinstance(content, str) or not content:
        return None
    try:
        cached = json.loads(content)
    except Exception:
        logging.debug("wiki_reduce: cached result unparseable; ignoring")
        return None
    if not isinstance(cached, dict):
        return None
    stored_hash = row.get("input_hash_kwd")
    if not isinstance(stored_hash, str):
        stored_hash = ""
    return cached, stored_hash


async def _wiki_persist_reduce(
    reduced: dict,
    tenant_id: str,
    kb_id: str,
    input_hash: str = "",
    source_doc_ids: Optional[list[str]] = None,
) -> None:
    """Upsert the single non-searchable wiki_reduce_result row for this KB.

    ``input_hash`` records the MAP-state fingerprint this reduction was
    computed from; the next call compares it before re-running.
    ``source_doc_ids`` is the set of documents that fed this reduction, used
    for delete-time reference counting.
    """
    from common import settings
    from rag.nlp import search as _rag_search

    index = _rag_search.index_name(tenant_id)
    kb_id_str = str(kb_id)
    content_with_weight = json.dumps(reduced, ensure_ascii=False)
    # Stable id per KB so a re-run upserts the same row.
    row_id = _stable_row_id(WIKI_REDUCE_COMPILE_KWD, kb_id_str)
    doc = {
        "id": row_id,
        "doc_id": kb_id_str,  # sentinel — KB-scoped row, not a real document
        "compile_kwd": WIKI_REDUCE_COMPILE_KWD,
        "source_id": [kb_id_str],
        "source_doc_ids": list(source_doc_ids or []),
        "input_hash_kwd": input_hash,
        "content_with_weight": content_with_weight,
        "available_int": 0,
    }
    try:
        # Best-effort delete then insert so re-runs replace cleanly.
        try:
            await thread_pool_exec(
                settings.docStoreConn.delete,
                {"compile_kwd": WIKI_REDUCE_COMPILE_KWD},
                index,
                kb_id,
            )
        except Exception:
            logging.debug("wiki_reduce: prior result delete failed; will overwrite by id")
        await thread_pool_exec(settings.docStoreConn.insert, [doc], index, kb_id)
    except Exception:
        logging.exception("wiki_reduce: failed to persist result row")


# --- public entry ----------------------------------------------------------


async def wiki_reduce_from_extracts(
    chat_mdl,
    embd_mdl,
    tenant_id: str,
    kb_id: str,
    merge_threshold: float = DEFAULT_WIKI_REDUCE_MERGE_THRESHOLD,
    ambiguous_low: float = DEFAULT_WIKI_REDUCE_AMBIGUOUS_LOW,
    ambiguous_batch_size: int = DEFAULT_WIKI_REDUCE_AMBIGUOUS_BATCH,
    llm_timeout: int = DEFAULT_WIKI_REDUCE_TIMEOUT,
    force_rerun: bool = False,
    callback: Optional[Callable] = None,
) -> dict:
    """Phase 2 (REDUCE/Dedup) — KB-scoped.

    Loads every ``wiki_map_extract`` row in this KB (across all documents) and
    produces a single canonical dict of entities/concepts via:
        1. Exact dedup by ``(normalize(name), type)`` for entities and by
           ``normalize(term)`` for concepts.
        2. Embedding dedup of entity names: vectorized pairwise cosine over
           ``embd_mdl.encode(...)`` output. Pairs of the same type with
           similarity ≥ ``merge_threshold`` auto-merge; pairs in
           ``[ambiguous_low, merge_threshold)`` go to step 3.
        3. LLM disambiguation: batches of ambiguous pairs are sent to
           ``chat_mdl`` via ``gen_json``; true verdicts collapse via union-find.
        4. Apply merges: sum ``mention_count``, union ``aliases`` and
           ``chunk_ids`` per canonical entity.

    The result is persisted to ES as a single non-searchable
    ``wiki_reduce_result`` row per KB. Subsequent calls with
    ``force_rerun=False`` (default) return the cached row immediately; pass
    ``force_rerun=True`` after new ``wiki_map_extract`` rows have been added.

    Args:
        chat_mdl, embd_mdl: ragflow LLMBundle instances.
        tenant_id, kb_id: address the doc-store index.
        merge_threshold: cosine ≥ this auto-merges. Default 0.90.
        ambiguous_low: cosine in [ambiguous_low, merge_threshold) goes to LLM.
        ambiguous_batch_size: max pairs per LLM disambiguation call.
        llm_timeout: seconds per LLM disambiguation batch.
        force_rerun: bypass the cached wiki_reduce_result.
        callback: optional ``(progress: float, msg: str)`` callback.

    Returns the canonical extract dict::

        {
          "entities":  [{"name","type","aliases","mention_count","chunk_ids"}, ...],
          "concepts":  [{"term","definition_excerpt","mention_count","chunk_ids"}, ...],
          "claims":    [...],   # pass-through from MAP
          "relations": [...],   # pass-through from MAP
          "topics":    [...],   # pass-through from MAP
        }
    """
    # Incremental gate: the current MAP-state fingerprint is the union
    # of every MAP row's (chunk_id, chunk_hash). If a cached REDUCE row
    # exists AND its stored input_hash equals the current fingerprint,
    # the upstream chunks haven't changed → cached output is still
    # correct. ``force_rerun=True`` bypasses both checks for the
    # legacy / admin "rebuild from scratch" path.
    current_input_hash = await _wiki_compute_map_input_hash(tenant_id, kb_id)
    reduce_source_doc_ids = await _wiki_all_map_doc_ids(tenant_id, kb_id)
    if not force_rerun:
        cached_pair = await _wiki_load_reduce_resume(tenant_id, kb_id)
        if cached_pair is not None:
            cached, stored_hash = cached_pair
            if stored_hash and stored_hash == current_input_hash:
                if callback:
                    try:
                        callback(1.0, "wiki REDUCE: cache hit (input unchanged)")
                    except Exception:
                        pass
                return cached
            # Cache present but stale (no hash, or hash mismatch). Fall
            # through to a full re-reduce and write a fresh stamp.

    if callback:
        try:
            callback(0.05, "wiki REDUCE: loading MAP extracts")
        except Exception:
            pass

    raw = await _wiki_load_all_map_extracts(tenant_id, kb_id)
    raw_entities = raw.get("entities") or []
    raw_concepts = raw.get("concepts") or []
    logging.info(
        "wiki_reduce: kb=%s loaded raw entities=%d concepts=%d claims=%d relations=%d",
        kb_id,
        len(raw_entities),
        len(raw_concepts),
        len(raw.get("claims") or []),
        len(raw.get("relations") or []),
    )

    if not raw_entities and not raw_concepts:
        # Nothing to reduce; persist an empty result so resume can short-circuit.
        empty = _wiki_empty_extract()
        await _wiki_persist_reduce(empty, tenant_id, kb_id, input_hash=current_input_hash, source_doc_ids=reduce_source_doc_ids)
        return empty

    if callback:
        try:
            callback(0.25, "wiki REDUCE: dedup (exact + embedding + LLM)")
        except Exception:
            pass

    # Entities: full three-phase dedup keyed by (normalized name, type).
    canonical_entities = await _bulk_dedup_items(
        raw_entities,
        name_key="name",
        type_key="type",
        chat_mdl=chat_mdl,
        embd_mdl=embd_mdl,
        merge_threshold=merge_threshold,
        ambiguous_low=ambiguous_low,
        ambiguous_batch_size=ambiguous_batch_size,
        disambiguate_system_prompt=WIKI_REDUCE_DISAMBIGUATE_SYSTEM,
        llm_timeout=llm_timeout,
    )

    # Concepts: exact-dedup only (current behaviour); keep the longest
    # definition_excerpt across the group via aggregate_extra.
    def _concept_extras(group: list[dict]) -> dict:
        best_def = max(
            ((c.get("definition_excerpt") or "") for c in group if isinstance(c, dict)),
            key=lambda s: len(s) if isinstance(s, str) else 0,
            default="",
        )
        return {"definition_excerpt": best_def}

    canonical_concepts = await _bulk_dedup_items(
        raw_concepts,
        name_key="term",
        type_key=None,
        aggregate_extra=_concept_extras,
    )

    logging.info(
        "wiki_reduce: after dedup entities=%d concepts=%d",
        len(canonical_entities),
        len(canonical_concepts),
    )

    reduced = {
        "entities": canonical_entities,
        "concepts": canonical_concepts,
        "claims": list(raw.get("claims") or []),
        "relations": list(raw.get("relations") or []),
        "topics": list(raw.get("topics") or []),
    }

    if callback:
        try:
            callback(0.9, "wiki REDUCE: persisting result")
        except Exception:
            pass
    await _wiki_persist_reduce(reduced, tenant_id, kb_id, input_hash=current_input_hash, source_doc_ids=reduce_source_doc_ids)

    logging.info(
        "wiki_reduce: kb=%s done — entities=%d concepts=%d claims=%d relations=%d topics=%d",
        kb_id,
        len(reduced["entities"]),
        len(reduced["concepts"]),
        len(reduced["claims"]),
        len(reduced["relations"]),
        len(reduced["topics"]),
    )

    if callback:
        try:
            callback(1.0, "wiki REDUCE: done")
        except Exception:
            pass

    return reduced


# ---------------------------------------------------------------------------
# PLAN phase (KB-scoped)
# ---------------------------------------------------------------------------
# Flow:
#   1. Resume — return cached wiki_compilation_plan ES row when present.
#   2. Load REDUCE output from wiki_reduce_result.
#   3. KB reconciliation — batch-embed entity/concept query texts and run a
#      per-item KNN against existing wiki_page rows in this KB. Classify
#      UPDATE / MAYBE / CREATE by similarity. Batched LLM resolves MAYBE.
#   4. Planning call — one gen_json call producing the Compilation Plan JSON.
#   5. Attach raw items as side context for REFINE (no extra ES round-trips).
#   6. Persist as a single non-searchable wiki_compilation_plan row per KB.
#

WIKI_PLAN_COMPILE_KWD = "wiki_compilation_plan"
WIKI_PAGE_COMPILE_KWD = "wiki_page"
DEFAULT_WIKI_PLAN_UPDATE_THRESHOLD = 0.95
DEFAULT_WIKI_PLAN_MAYBE_THRESHOLD = 0.60
DEFAULT_WIKI_PLAN_TIMEOUT = 600  # ~10 min — the planning call emits one big
# JSON plan and reasoning models can spend a
# long time thinking before emitting tokens.
# Override via the ``llm_timeout`` arg to
# ``wiki_plan_from_reduction``.
DEFAULT_WIKI_PLAN_RECONCILE_BATCH = 50
_WIKI_PLAN_MAX_OUTPUT_TOKENS = 4096
_WIKI_PLAN_OUTPUT_SAFETY_TOKENS = 256
_WIKI_PLAN_PAGE_TOKEN_ESTIMATE = 48
_WIKI_PLAN_ITEMS_PER_BATCH = 40
_WIKI_PLAN_MAX_CONCURRENT_BATCHES = 4


WIKI_PLAN_PLANNING_SYSTEM = (
    "You are a knowledge compilation planner. Given extracted entities and their "
    "relationship to an existing knowledge base, produce a compilation plan. "
    "Return ONLY valid JSON."
    "Keep the user's original language (Chinese/English etc.) for generated data."
)


WIKI_PLAN_RECONCILE_SYSTEM = "You are a knowledge base assistant. Return only a JSON boolean array.Keep the user's original language (Chinese/English etc.) for generated data."


WIKI_PLAN_USER_TEMPLATE = """\
## Knowledge base context
Name: {kb_name}
Description: {kb_description}

## Extracted entities (with mention counts)
{entities_summary}

## Extracted concepts (with mention counts)
{concepts_summary}

## Extracted topics
{topics_summary}

## KB reconciliation results
{kb_reconciliation}

Produce a JSON compilation plan:

{{
  "pages": [
    {{
      "action": "CREATE",
      "slug": "concept/example-name",
      "title": "Example Page Title",
      "page_type": "entity | concept | topic",
      "topic": "short canonical topic name",
      "entity_names": ["entity or concept name covered by this page"],
      "related_kb_pages": ["existing-slug-1"],
      "priority": 1
    }}
  ],
  "estimated_page_count": 5,
  "compilation_notes": "any important notes for the compiler"
}}

Rules:
- action must be "CREATE" or "UPDATE".
- For UPDATE, slug MUST be an existing wiki page slug from the KB
  reconciliation list above.
- page_type is one of: entity | concept | topic. Do NOT use "source".
- topic is required for every page. Prefer a topic from the extracted
  topics implied by the entities/concepts. If none fits, create a short
  canonical topic name in the user's language. For topic pages, topic should
  usually match the page title.

# Slug format (CRITICAL — every slug must follow this shape exactly)
- The slug is ``<page_type>/<short-descriptive-name>``. The separator
  between the type and the name MUST be a forward slash ``/``. Do NOT use a
  hyphen here.
- The descriptive part is lowercase, English/Latin only (transliterate
  non-English names), and uses hyphens to join multi-word names. Keep it
  short — 1 to 4 words is ideal.
- The descriptive part MUST be unique to that page's specific subject. Do
  NOT prefix every slug with the same KB-wide topic word. If the KB is
  about logistics, do NOT emit ``concept/logistics-channels``,
  ``concept/logistics-warehousing``, ``concept/logistics-fleet`` — emit
  ``concept/distribution-channels``, ``concept/warehousing``,
  ``concept/fleet-management`` instead.
- Do NOT append numeric suffixes (``-1``, ``-2``, ``-v2``) or random hex
  tags to make slugs distinct. If two candidate slugs collide, rename one
  to use a different descriptive word.

Examples of GOOD slugs:
  - ``entity/jane-doe``               (entity page about a person)
  - ``entity/acme-corp``               (entity page about a company)
  - ``concept/fire-safety``            (concept page about a topic)
  - ``concept/expense-approval``       (concept page about a process)
  - ``topic/water-treatment``          (topic page grouping related items)

Examples of BAD slugs (do NOT produce):
  - ``concept-fire-safety``            (missing the ``/`` between type and name)
  - ``concept/logistics-channels-1``   (numeric suffix to distinguish pages)
  - ``concept/logistics-channels-abc`` (random hex tag)
  - ``logistics/concept-channels``     (type and topic order swapped)
  - ``concept/example-name``           (just duplicate the sample)

# Other rules
- Group closely related small entities onto the same page (max 3-4 per page).
  BUT if a primary entity is described through several distinct thematic
  sections that appear as concepts above, prefer a separate ``concept`` page
  for EACH such section instead of collapsing them onto the entity page.
- priority 1 = highest importance (process first).
- entity_names must match the names in the entities / concepts lists above.
- Target approximately {target_page_count} total pages (feel free to deviate
  by ±50% if the KB content warrants it).
- Return no more than {max_page_count} page objects. This is a hard limit;
  never continue the JSON beyond this number.
- Return ONLY the JSON object.
"""


# --- helpers ---------------------------------------------------------------


def _wiki_target_page_count(total_items: int) -> int:
    """Item-count-based heuristic: clamp(8, total // 3, 60)."""
    if total_items <= 0:
        return 8
    return max(8, min(60, total_items // 3))


def _wiki_format_entity_for_plan(entity: dict, reconciliation: dict) -> str:
    aliases = ", ".join((entity.get("aliases") or [])[:3])
    rec = reconciliation.get(entity.get("name", ""), {})
    action = rec.get("action", "CREATE")
    slug = rec.get("page_slug", "")
    kb_info = f"→ {action} {slug}".rstrip()
    line = f"  - {entity.get('name', '')} ({entity.get('type', '')}, {entity.get('mention_count', 0)} mentions"
    if aliases:
        line += f", aliases: {aliases}"
    line += f") {kb_info}"
    return line


def _wiki_format_concept_for_plan(concept: dict, reconciliation: dict) -> str:
    rec = reconciliation.get(concept.get("term", ""), {})
    action = rec.get("action", "CREATE")
    slug = rec.get("page_slug", "")
    kb_info = f"→ {action} {slug}".rstrip()
    return f"  - {concept.get('term', '')} ({concept.get('mention_count', 0)} mentions) {kb_info}"


async def _wiki_reconcile_with_kb(
    canonical_entities: list[dict],
    canonical_concepts: list[dict],
    embd_mdl,
    tenant_id: str,
    kb_id: str,
    update_threshold: float,
    maybe_threshold: float,
) -> dict[str, dict]:
    """Per-entity / per-concept KNN against compile_kwd=wiki_page rows in this KB.

    Returns ``{name_or_term: {"action", "page_slug", "page_title", "page_id",
    "similarity"}}``. When no artifact pages exist (first run before REFINE), every
    item maps to ``action="CREATE"``.
    """
    from common import settings
    from common.doc_store.doc_store_base import MatchDenseExpr, OrderByExpr
    from rag.nlp import search as _rag_search

    items: list[tuple[str, str, dict]] = []  # (kind, key, source_dict)
    for e in canonical_entities:
        name = e.get("name")
        if isinstance(name, str) and name:
            items.append(("entity", name, e))
    for c in canonical_concepts:
        term = c.get("term")
        if isinstance(term, str) and term:
            items.append(("concept", term, c))

    reconciliation: dict[str, dict] = {}
    if not items:
        return reconciliation

    # Embed all query texts in one batch.
    query_texts: list[str] = []
    for kind, key, src in items:
        if kind == "concept":
            defn = src.get("definition_excerpt") or ""
            text = f"{key}: {defn[:200]}" if defn else key
        else:
            text = key
        query_texts.append(text[:4000])

    try:
        embeddings, _ = await thread_pool_exec(embd_mdl.encode, query_texts)
        vectors = list(embeddings)
    except Exception:
        logging.exception("wiki_plan: reconciliation embedding failed — all items will be CREATE")
        for _, key, _ in items:
            reconciliation[key] = {
                "action": "CREATE",
                "page_slug": None,
                "page_title": None,
                "page_id": None,
                "similarity": 0.0,
            }
        return reconciliation

    if len(vectors) != len(items):
        logging.error(
            "wiki_plan: reconciliation embedding count mismatch (%d vs %d); CREATE all",
            len(vectors),
            len(items),
        )
        for _, key, _ in items:
            reconciliation[key] = {
                "action": "CREATE",
                "page_slug": None,
                "page_title": None,
                "page_id": None,
                "similarity": 0.0,
            }
        return reconciliation

    index = _rag_search.index_name(tenant_id)
    condition = {"compile_kwd": [WIKI_PAGE_COMPILE_KWD]}

    select_fields = ["id", "slug_kwd", "title_kwd", "page_type_kwd", "_score"]
    for (_kind, key, _src), vec in zip(items, vectors):
        vec_list = list(vec) if not hasattr(vec, "tolist") else vec.tolist()
        if not vec_list:
            reconciliation[key] = {
                "action": "CREATE",
                "page_slug": None,
                "page_title": None,
                "page_id": None,
                "similarity": 0.0,
            }
            continue
        match_expr = MatchDenseExpr(
            vector_column_name=f"q_{len(vec_list)}_vec",
            embedding_data=vec_list,
            embedding_data_type="float",
            distance_type="cosine",
            topn=1,
            extra_options={"similarity": update_threshold},
        )
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                condition,
                [match_expr],
                OrderByExpr(),
                0,
                1,
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields)
        except Exception:
            logging.exception("wiki_plan: KNN failed for %r", key)
            reconciliation[key] = {
                "action": "CREATE",
                "page_slug": None,
                "page_title": None,
                "page_id": None,
                "similarity": 0.0,
            }
            continue

        if not field_map:
            reconciliation[key] = {
                "action": "CREATE",
                "page_slug": None,
                "page_title": None,
                "page_id": None,
                "similarity": 0.0,
            }
            continue

        top_id, top_row = next(iter(field_map.items()))
        # Pull similarity from the search result if exposed; fall back to threshold floor.
        sim = 0.0
        try:
            sim = float(getattr(top_row, "_score", None))
        except Exception:
            sim = 0.0
        if sim <= 0.0:
            sim = float(top_row.get("similarity", maybe_threshold))

        slug = top_row.get("slug_kwd")
        title = top_row.get("title_kwd")
        if sim >= update_threshold:
            action = "UPDATE"
        else:
            action = "MAYBE"
        reconciliation[key] = {
            "action": action,
            "page_slug": slug,
            "page_title": title,
            "page_id": top_id,
            "similarity": sim,
        }

    return reconciliation


async def _wiki_resolve_maybe_items(
    reconciliation: dict[str, dict],
    chat_mdl,
    batch_size: int,
    llm_timeout: int,
) -> None:
    """Flip MAYBE → UPDATE | CREATE via batched LLM calls. Mutates in place."""
    maybe_items = [(k, v) for k, v in reconciliation.items() if v.get("action") == "MAYBE"]
    if not maybe_items:
        return

    for batch_start in range(0, len(maybe_items), batch_size):
        batch = maybe_items[batch_start : batch_start + batch_size]
        lines = []
        for k, (name, rec) in enumerate(batch):
            title = rec.get("page_title") or rec.get("page_slug") or ""
            slug = rec.get("page_slug") or ""
            sim = rec.get("similarity", 0.0)
            lines.append(f'{k + 1}. Entity: "{name}" — existing wiki page: "{title}" (slug: {slug}, similarity: {sim:.2f})')

        user_prompt = (
            "For each pair below, decide whether the entity refers to the same "
            "real-world concept as the existing wiki page (true = UPDATE existing "
            "page, false = CREATE new page).\n"
            f"Return a JSON array of exactly {len(batch)} booleans. "
            "Return ONLY the JSON array.\n\n" + "\n".join(lines)
        )

        request_conf = _knowledge_compile_gen_conf(chat_mdl, {"temperature": 0.0})
        try:
            res = await asyncio.wait_for(
                gen_json(
                    WIKI_PLAN_RECONCILE_SYSTEM,
                    user_prompt,
                    chat_mdl,
                    gen_conf=request_conf,
                ),
                timeout=llm_timeout,
            )
        except asyncio.TimeoutError:
            logging.warning("wiki_plan: MAYBE resolution timed out (%d pairs); defaulting CREATE", len(batch))
            for name, _ in batch:
                reconciliation[name]["action"] = "CREATE"
            continue
        except Exception:
            logging.exception("wiki_plan: MAYBE resolution failed (%d pairs); defaulting CREATE", len(batch))
            for name, _ in batch:
                reconciliation[name]["action"] = "CREATE"
            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("wiki_plan: MAYBE LLM returned unexpected shape %r; CREATE all", type(res))
            for name, _ in batch:
                reconciliation[name]["action"] = "CREATE"
            continue

        for k, (name, _) in enumerate(batch):
            verdict = decisions[k] if k < len(decisions) else False
            reconciliation[name]["action"] = "UPDATE" if verdict else "CREATE"


async def _wiki_planning_call(
    canonical_entities: list[dict],
    canonical_concepts: list[dict],
    raw_topics: list,
    reconciliation: dict[str, dict],
    chat_mdl,
    kb_name: str | None,
    kb_description: str | None,
    target_page_count: int,
    llm_timeout: int,
    _batch_depth: int = 0,
) -> dict:
    """Single LLM call → Compilation Plan JSON."""
    # Sort by mention count descending so the planner sees the most important
    # items first; cap to keep the prompt size reasonable.
    sorted_entities = sorted(
        canonical_entities,
        key=lambda x: x.get("mention_count", 0),
        reverse=True,
    )
    sorted_concepts = sorted(
        canonical_concepts,
        key=lambda x: x.get("mention_count", 0),
        reverse=True,
    )

    model_context = int(getattr(chat_mdl, "max_length", 8192) or 8192)
    output_tokens = min(
        _WIKI_PLAN_MAX_OUTPUT_TOKENS,
        max(1024, int(model_context * 0.4)),
    )
    output_page_capacity = max(
        1,
        (output_tokens - _WIKI_PLAN_OUTPUT_SAFETY_TOKENS) // _WIKI_PLAN_PAGE_TOKEN_ESTIMATE,
    )
    max_page_count = min(
        output_page_capacity,
        max(target_page_count + 8, target_page_count * 2),
    )

    all_items = [("entity", item) for item in sorted_entities] + [("concept", item) for item in sorted_concepts]
    if _batch_depth == 0 and len(all_items) > _WIKI_PLAN_ITEMS_PER_BATCH:
        batches = [all_items[offset : offset + _WIKI_PLAN_ITEMS_PER_BATCH] for offset in range(0, len(all_items), _WIKI_PLAN_ITEMS_PER_BATCH)]
        total_items = len(all_items)
        semaphore = asyncio.Semaphore(_WIKI_PLAN_MAX_CONCURRENT_BATCHES)

        async def _plan_batch(batch):
            async with semaphore:
                batch_entities = [item for kind, item in batch if kind == "entity"]
                batch_concepts = [item for kind, item in batch if kind == "concept"]
                batch_keys = {item.get("name") if kind == "entity" else item.get("term") for kind, item in batch}
                batch_reconciliation = {key: value for key, value in reconciliation.items() if key in batch_keys}
                batch_target = max(1, round(target_page_count * len(batch) / total_items))
                return await _wiki_planning_call(
                    batch_entities,
                    batch_concepts,
                    raw_topics,
                    batch_reconciliation,
                    chat_mdl,
                    kb_name,
                    kb_description,
                    batch_target,
                    llm_timeout,
                    _batch_depth=1,
                )

        batch_plans = await asyncio.gather(*(_plan_batch(batch) for batch in batches))
        merged_pages = []
        seen_slugs = set()
        for batch_plan in batch_plans:
            for page in batch_plan.get("pages") or []:
                slug = page.get("slug")
                if slug and slug not in seen_slugs:
                    seen_slugs.add(slug)
                    merged_pages.append(page)
        logging.info(
            "wiki_plan: batched planning items=%d batches=%d merged_pages=%d",
            total_items,
            len(batches),
            len(merged_pages),
        )
        return {
            "pages": merged_pages,
            "estimated_page_count": len(merged_pages),
            "compilation_notes": "planned in batches",
        }

    entities_summary = "\n".join(_wiki_format_entity_for_plan(e, reconciliation) for e in sorted_entities[:200]) or "  (none)"
    concepts_summary = "\n".join(_wiki_format_concept_for_plan(c, reconciliation) for c in sorted_concepts[:200]) or "  (none)"
    topics_summary = "\n".join(f"  - {t.strip()}" for t in raw_topics[:200] if isinstance(t, str) and t.strip()) or "  (none)"

    kb_lines: list[str] = []
    for name, rec in reconciliation.items():
        if rec.get("action") == "UPDATE" and rec.get("page_slug"):
            kb_lines.append(f"  - UPDATE: {name} → {rec['page_slug']} (sim={rec.get('similarity', 0.0):.2f})")
    kb_reconciliation = "\n".join(kb_lines) if kb_lines else "  (all items are new)"

    user_prompt = WIKI_PLAN_USER_TEMPLATE.format(
        kb_name=kb_name or "(unspecified)",
        kb_description=kb_description or "(no description)",
        entities_summary=entities_summary,
        concepts_summary=concepts_summary,
        topics_summary=topics_summary,
        kb_reconciliation=kb_reconciliation,
        target_page_count=target_page_count,
        max_page_count=max_page_count,
    )

    request_conf = _knowledge_compile_gen_conf(
        chat_mdl,
        {"temperature": 0.1, "max_tokens": output_tokens},
    )
    try:
        res = await asyncio.wait_for(
            gen_json(
                WIKI_PLAN_PLANNING_SYSTEM,
                user_prompt,
                chat_mdl,
                gen_conf=request_conf,
            ),
            timeout=llm_timeout,
        )
    except asyncio.TimeoutError:
        logging.warning("wiki_plan: planning LLM call timed out after %ds", llm_timeout)
        return {"pages": [], "estimated_page_count": 0, "compilation_notes": "planning timeout"}
    except Exception:
        logging.exception("wiki_plan: planning LLM call failed")
        return {"pages": [], "estimated_page_count": 0, "compilation_notes": "planning failed"}

    if not isinstance(res, dict):
        return {"pages": [], "estimated_page_count": 0, "compilation_notes": "planner returned non-object"}
    if "pages" not in res or not isinstance(res.get("pages"), list):
        res["pages"] = []

    valid_pages = []
    for page in res["pages"]:
        if not isinstance(page, dict):
            continue
        action = page.get("action")
        slug = page.get("slug")
        title = page.get("title")
        page_type = page.get("page_type")
        topic = page.get("topic")
        if action not in {"CREATE", "UPDATE"}:
            continue
        if not isinstance(slug, str) or not re.fullmatch(r"(?:entity|concept|topic)/[a-z0-9]+(?:-[a-z0-9]+)*", slug):
            logging.warning("wiki_plan: dropped invalid planner slug %r", slug)
            continue
        if not isinstance(title, str) or not title.strip():
            logging.warning("wiki_plan: dropped page with missing title slug=%s", slug)
            continue
        if page_type not in {"entity", "concept", "topic"}:
            logging.warning("wiki_plan: dropped page with invalid page_type slug=%s", slug)
            continue
        if slug.split("/", 1)[0] != page_type:
            logging.warning("wiki_plan: dropped page with mismatched page_type slug=%s page_type=%s", slug, page_type)
            continue
        if action == "UPDATE" and not any(rec.get("action") == "UPDATE" and rec.get("page_slug") == slug for rec in reconciliation.values()):
            logging.warning("wiki_plan: dropped UPDATE for unreconciled slug=%s", slug)
            continue
        if not isinstance(topic, str) or not topic.strip():
            logging.warning("wiki_plan: dropped page with missing topic slug=%s", slug)
            continue
        valid_pages.append(page)
        if len(valid_pages) >= max_page_count:
            break
    res["pages"] = valid_pages
    res["estimated_page_count"] = len(valid_pages)
    res.setdefault("compilation_notes", "")
    return res


# --- ES I/O ---------------------------------------------------------------


async def _wiki_load_reduce_result(tenant_id: str, kb_id: str) -> Optional[dict]:
    """Load the cached REDUCE output for this KB."""
    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)
    condition = {"compile_kwd": [WIKI_REDUCE_COMPILE_KWD]}
    select_fields = ["id", "content_with_weight"]
    try:
        res = await thread_pool_exec(
            settings.docStoreConn.search,
            select_fields,
            [],
            condition,
            [],
            OrderByExpr(),
            0,
            1,
            index,
            [kb_id],
        )
        field_map = settings.docStoreConn.get_fields(res, select_fields)
    except Exception:
        logging.exception("wiki_plan: failed to load wiki_reduce_result")
        return None
    if not field_map:
        return None
    row = next(iter(field_map.values()))
    content = row.get("content_with_weight")
    if not isinstance(content, str) or not content:
        return None
    try:
        cached = json.loads(content)
    except Exception:
        logging.debug("wiki_plan: wiki_reduce_result unparseable; ignoring")
        return None
    return cached if isinstance(cached, dict) else None


async def _wiki_load_reduce_input_hash(tenant_id: str, kb_id: str) -> str:
    """Read just the ``input_hash_kwd`` off the REDUCE row (without
    deserializing the body). Used by PLAN's incremental gate so we can
    short-circuit without re-running the planner.
    """
    pair = await _wiki_load_reduce_resume(tenant_id, kb_id)
    if pair is None:
        return ""
    _cached, stored_hash = pair
    return stored_hash


async def _wiki_load_plan_resume(
    tenant_id: str,
    kb_id: str,
) -> Optional[tuple[dict, str]]:
    """Return ``(cached_plan, stored_input_hash)`` or None.

    The stored hash is whatever REDUCE's ``input_hash_kwd`` was when this
    plan was last written. PLAN's cache check compares it to the
    current REDUCE input hash to decide whether to re-plan.
    """
    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)
    condition = {"compile_kwd": [WIKI_PLAN_COMPILE_KWD]}
    select_fields = ["id", "content_with_weight", "input_hash_kwd"]
    try:
        res = await thread_pool_exec(
            settings.docStoreConn.search,
            select_fields,
            [],
            condition,
            [],
            OrderByExpr(),
            0,
            1,
            index,
            [kb_id],
        )
        field_map = settings.docStoreConn.get_fields(res, select_fields)
    except Exception:
        logging.exception("wiki_plan: failed to load cached plan")
        return None
    if not field_map:
        return None
    row = next(iter(field_map.values()))
    content = row.get("content_with_weight")
    if not isinstance(content, str) or not content:
        return None
    try:
        cached = json.loads(content)
    except Exception:
        logging.debug("wiki_plan: cached plan unparseable; ignoring")
        return None
    if not isinstance(cached, dict):
        return None
    stored_hash = row.get("input_hash_kwd")
    if not isinstance(stored_hash, str):
        stored_hash = ""
    return cached, stored_hash


async def _wiki_persist_plan(
    plan: dict,
    tenant_id: str,
    kb_id: str,
    input_hash: str = "",
    source_doc_ids: Optional[list[str]] = None,
) -> None:
    """Upsert the single non-searchable wiki_compilation_plan row for this KB.

    ``input_hash`` records the REDUCE-state fingerprint this plan was
    derived from; the next call compares it before re-planning.
    ``source_doc_ids`` is the set of documents that fed this plan, used for
    delete-time reference counting.
    """
    from common import settings
    from rag.nlp import search as _rag_search

    index = _rag_search.index_name(tenant_id)
    kb_id_str = str(kb_id)
    content_with_weight = json.dumps(plan, ensure_ascii=False)
    row_id = _stable_row_id(WIKI_PLAN_COMPILE_KWD, kb_id_str)
    doc = {
        "id": row_id,
        "doc_id": kb_id_str,  # sentinel — KB-scoped row, not a real document
        "compile_kwd": WIKI_PLAN_COMPILE_KWD,
        "source_id": [kb_id_str],
        "source_doc_ids": list(source_doc_ids or []),
        "input_hash_kwd": input_hash,
        "content_with_weight": content_with_weight,
        "available_int": 0,
    }
    try:
        try:
            await thread_pool_exec(
                settings.docStoreConn.delete,
                {"compile_kwd": WIKI_PLAN_COMPILE_KWD},
                index,
                kb_id,
            )
        except Exception:
            logging.debug("wiki_plan: prior plan delete failed; relying on id-based upsert")
        await thread_pool_exec(settings.docStoreConn.insert, [doc], index, kb_id)
    except Exception:
        logging.exception("wiki_plan: failed to persist plan row")


# --- public entry ---------------------------------------------------------


async def wiki_plan_from_reduction(
    chat_mdl,
    embd_mdl,
    tenant_id: str,
    kb_id: str,
    kb_name: Optional[str] = None,
    kb_description: Optional[str] = None,
    update_threshold: float = DEFAULT_WIKI_PLAN_UPDATE_THRESHOLD,
    maybe_threshold: float = DEFAULT_WIKI_PLAN_MAYBE_THRESHOLD,
    reconcile_batch_size: int = DEFAULT_WIKI_PLAN_RECONCILE_BATCH,
    llm_timeout: int = DEFAULT_WIKI_PLAN_TIMEOUT,
    force_rerun: bool = False,
    callback: Optional[Callable] = None,
) -> dict:
    """Phase 3 (PLAN) — KB-scoped.

    Loads the cached ``wiki_reduce_result`` for this KB, reconciles every
    canonical entity/concept against existing ``wiki_page`` rows in the same
    KB (top-1 KNN, with MAYBE matches resolved by a batched LLM call), then
    asks the LLM for one Compilation Plan JSON. The plan is persisted under
    ``compile_kwd="wiki_compilation_plan"`` with ``_status="approved"`` so
    REFINE can consume it immediately.

    Args:
        chat_mdl, embd_mdl: ragflow LLMBundle instances.
        tenant_id, kb_id: address the doc-store index.
        kb_name / kb_description: optional KB-level metadata that biases the
            planner's slug and tone choices.
        update_threshold: cosine ≥ this → UPDATE the existing page outright.
        maybe_threshold: cosine in [maybe_threshold, update_threshold) → ask LLM.
        reconcile_batch_size: max pairs per LLM MAYBE-resolution call.
        llm_timeout: seconds per LLM call (both MAYBE resolution and planning).
        force_rerun: bypass the cached wiki_compilation_plan.
        callback: optional ``(progress: float, msg: str)`` callback.

    Returns the plan dict with this shape (plus underscore-prefixed side
    context fields for REFINE)::

        {
          "pages":               [{action, slug, title, page_type, entity_names, related_kb_pages, priority}, ...],
          "estimated_page_count": int,
          "compilation_notes":   str,
          "_status":             "approved",
          "_entities":           [...],   # canonical entities from REDUCE
          "_concepts":           [...],
          "_claims":             [...],
          "_relations":          [...],
          "_topics":             [...],
          "_reconciliation":     {name: {action, page_slug, page_id, similarity}, ...},
        }
    """
    # Incremental gate: PLAN keys off REDUCE's input_hash. If the cached
    # plan was stamped with the same hash REDUCE is currently exposing,
    # nothing upstream has changed and the plan is still valid.
    current_reduce_hash = await _wiki_load_reduce_input_hash(tenant_id, kb_id)
    plan_source_doc_ids = await _wiki_all_map_doc_ids(tenant_id, kb_id)
    if not force_rerun:
        cached_pair = await _wiki_load_plan_resume(tenant_id, kb_id)
        if cached_pair is not None:
            cached, stored_hash = cached_pair
            if stored_hash and stored_hash == current_reduce_hash:
                if callback:
                    try:
                        callback(1.0, "wiki PLAN: cache hit (REDUCE unchanged)")
                    except Exception:
                        pass
                return cached

    if callback:
        try:
            callback(0.05, "wiki PLAN: loading REDUCE result")
        except Exception:
            pass

    reduced = await _wiki_load_reduce_result(tenant_id, kb_id)
    if reduced is None:
        logging.warning("wiki_plan: no wiki_reduce_result found for kb=%s — returning empty plan", kb_id)
        empty = {
            "pages": [],
            "estimated_page_count": 0,
            "compilation_notes": "no REDUCE result available",
            "_status": "approved",
            "_entities": [],
            "_concepts": [],
            "_claims": [],
            "_relations": [],
            "_topics": [],
            "_reconciliation": {},
        }
        await _wiki_persist_plan(empty, tenant_id, kb_id, input_hash=current_reduce_hash, source_doc_ids=plan_source_doc_ids)
        return empty

    canonical_entities = reduced.get("entities") or []
    canonical_concepts = reduced.get("concepts") or []
    raw_claims = reduced.get("claims") or []
    raw_relations = reduced.get("relations") or []
    raw_topics = reduced.get("topics") or []

    total_items = len(canonical_entities) + len(canonical_concepts)
    logging.info(
        "wiki_plan: kb=%s reducing-input entities=%d concepts=%d (total=%d)",
        kb_id,
        len(canonical_entities),
        len(canonical_concepts),
        total_items,
    )

    if total_items == 0:
        empty = {
            "pages": [],
            "estimated_page_count": 0,
            "compilation_notes": "no canonical items",
            "_status": "approved",
            "_entities": canonical_entities,
            "_concepts": canonical_concepts,
            "_claims": raw_claims,
            "_relations": raw_relations,
            "_topics": raw_topics,
            "_reconciliation": {},
        }
        await _wiki_persist_plan(empty, tenant_id, kb_id, input_hash=current_reduce_hash, source_doc_ids=plan_source_doc_ids)
        return empty

    if callback:
        try:
            callback(0.25, "wiki PLAN: KB reconciliation")
        except Exception:
            pass

    reconciliation = await _wiki_reconcile_with_kb(
        canonical_entities=canonical_entities,
        canonical_concepts=canonical_concepts,
        embd_mdl=embd_mdl,
        tenant_id=tenant_id,
        kb_id=kb_id,
        update_threshold=update_threshold,
        maybe_threshold=maybe_threshold,
    )

    if callback:
        n_maybe = sum(1 for v in reconciliation.values() if v.get("action") == "MAYBE")
        try:
            callback(0.55, f"wiki PLAN: resolving {n_maybe} MAYBE items")
        except Exception:
            pass

    await _wiki_resolve_maybe_items(
        reconciliation,
        chat_mdl,
        batch_size=reconcile_batch_size,
        llm_timeout=llm_timeout,
    )

    if callback:
        try:
            callback(0.75, "wiki PLAN: planning LLM call")
        except Exception:
            pass

    target = _wiki_target_page_count(total_items)
    plan = await _wiki_planning_call(
        canonical_entities=canonical_entities,
        canonical_concepts=canonical_concepts,
        raw_topics=raw_topics,
        reconciliation=reconciliation,
        chat_mdl=chat_mdl,
        kb_name=kb_name,
        kb_description=kb_description,
        target_page_count=target,
        llm_timeout=llm_timeout,
    )

    plan["_status"] = "approved"
    plan["_entities"] = canonical_entities
    plan["_concepts"] = canonical_concepts
    plan["_claims"] = raw_claims
    plan["_relations"] = raw_relations
    plan["_topics"] = raw_topics
    plan["_reconciliation"] = reconciliation

    if callback:
        try:
            callback(0.9, "wiki PLAN: persisting plan")
        except Exception:
            pass
    await _wiki_persist_plan(plan, tenant_id, kb_id, input_hash=current_reduce_hash, source_doc_ids=plan_source_doc_ids)

    logging.info(
        "wiki_plan: kb=%s done — pages=%d (target=%d) updates=%d creates=%d",
        kb_id,
        len(plan.get("pages") or []),
        target,
        sum(1 for v in reconciliation.values() if v.get("action") == "UPDATE"),
        sum(1 for v in reconciliation.values() if v.get("action") == "CREATE"),
    )

    if callback:
        try:
            callback(1.0, "wiki PLAN: done")
        except Exception:
            pass

    return plan


# ---------------------------------------------------------------------------
# REFINE phase (KB-scoped)
# ---------------------------------------------------------------------------
#
#
# Scope: per KB. Consumes the wiki_compilation_plan row written by PLAN,
# writes one wiki_page per planned page in parallel under a semaphore.
# UPDATE actions LLM-merge new vs existing content with a 70 % shrink-check
# fallback to the new content. Each written page is persisted to ES as a
# searchable wiki_page row (with embedding) so PLAN reconciliation finds it
# on the next REDUCE→PLAN cycle.
#
# Resume: per-slug wiki_page_draft rows act as a cache; a re-entry skips
# slugs already cached unless force_rerun=True.


WIKI_DRAFT_COMPILE_KWD = "wiki_page_draft"
DEFAULT_WIKI_REFINE_WORKERS = 4
DEFAULT_WIKI_REFINE_TIMEOUT = 300
WIKI_REFINE_SOURCE_BUDGET_CHARS = 60_000
WIKI_MERGE_BODY_SHRINK_THRESHOLD = 0.7
WIKI_MERGE_TIMEOUT = 600


WIKI_TEMPLATE_EXAMPLE = (
    "Each page must be a proper encyclopedic article, NOT a flat bullet list:\n"
    "1. Opening paragraph (2-4 sentences defining what this is). No heading.\n"
    "2. Sections with H2 headings, each starting with prose before sub-bullets.\n"
    "   Put every heading on its own line and separate every paragraph with a blank line.\n"
    "3. Bold key terms on first use; link them with [[ ]] wikilinks.\n"
    "4. Examples or implications where the source provides them.\n"
    "5. ## See also section at the end with wikilinks to highly related pages(less than 12).\n\n"
    "Page structure could be as following:\n(Not provided)"
)

# Writer system prompt as a template: the ``{template_example}``
# placeholder is filled in at request time so each artifact compilation
# template can override the page-structure section without touching the
# rest of the writer's guidance. Use ``_build_refine_writer_system`` to
# materialize a concrete prompt; ``WIKI_REFINE_WRITER_SYSTEM`` is
# kept as the default-filled value for back-compat with any code that
# still imports it.
WIKI_REFINE_WRITER_SYSTEM_TEMPLATE = (
    "You are an enterprise knowledge compilation writer. Your job is to write a single, "
    "high-quality wiki page by reading the SOURCE TEXT provided and using the "
    "evidence checklist as guidance for what to cover.\n\n"
    "# Mindset: COMPILE, do NOT summarize\n"
    "You are not writing an executive summary. You are extracting structured "
    "knowledge and rewriting it into a reusable wiki page. The output should "
    "contain MORE information density than a summary — organized differently, "
    "but not condensed. A summary loses specifics. A wiki page preserves them "
    "in a queryable structure.\n\n"
    "# What to KEEP from the source (do not lose these)\n"
    "- Specific numbers: thresholds, dosages, timeframes, dimensions, percentages.\n"
    "- Named regulations, laws, articles, code references.\n"
    "- Equipment names, model numbers, product specs.\n"
    "- Procedure steps in order, with actual actions.\n"
    "- Worked examples and exceptions.\n"
    "- Named parties, roles, contact paths, escalation chains.\n"
    "- Definitions verbatim or near-verbatim if the source is authoritative.\n"
    "- Cause-effect statements ('X causes Y because Z') — preserve all three parts.\n\n"
    "# What to DROP\n"
    "- Marketing language, mission statements, ceremonial filler.\n"
    "- Source-specific framing: 'This document explains…', 'In Section 3 below…'.\n"
    "- Repeated boilerplate, tables of contents, cover-page metadata.\n"
    "- Prose that just rephrases what was already said.\n\n"
    "# Language\n"
    "Write in the SAME LANGUAGE as the source text. Never translate content.\n\n"
    "# Page structure — CRITICAL\n"
    "{template_example}\n\n"
    "# What NOT to do\n"
    "- Do NOT dump raw bullet points from the source as the entire content.\n"
    "- Do NOT omit the opening prose paragraph.\n"
    "- Do NOT include Citations / Footnotes sections.\n"
    "- Do NOT use [^N] footnote markers.\n"
    "- Do NOT translate the content language.\n\n"
    "# Wikilinks\n"
    "- Use [[slug]] or [[slug|display text]] to cross-link.\n"
    "- CRITICAL: You may ONLY link to slugs from the 'Available pages' list.\n"
    "  Do NOT invent or hallucinate slugs.\n\n"
    "# Minimum depth\n"
    "- concept/topic pages: at least 200 words of actual prose+structure.\n"
    "- entity pages: at least 100 words.\n"
)


def _build_refine_writer_system(example: str | None) -> str:
    """Return the writer system prompt with the configured page-structure
    example (or ``WIKI_TEMPLATE_EXAMPLE`` when ``example`` is empty /
    whitespace-only). Used by the REFINE phase to let each compilation
    template override just the page-structure section.

    The default-filled form is also exposed as
    ``WIKI_REFINE_WRITER_SYSTEM`` for callers that don't have an
    override to apply.
    """
    body = (example or "").strip() or WIKI_TEMPLATE_EXAMPLE
    return WIKI_REFINE_WRITER_SYSTEM_TEMPLATE.format(template_example=body)


WIKI_REFINE_WRITER_SYSTEM = _build_refine_writer_system(None)


WIKI_REFINE_WRITER_USER_TEMPLATE = """\
## Task
{action} the following wiki page.

## Page specification
- Slug: {slug}
- Title: {title}
- Type: {page_type}

## Available pages (ONLY use these slugs for [[wikilinks]])
{all_plan_slugs}

{existing_section}

## Source document text
Read this carefully. Extract all relevant facts for this page's topic.

{source_context}

## Evidence checklist ({evidence_count} items)
The following items were pre-extracted and should be covered in the page.
Use them as a checklist — make sure you don't miss any of these facts.
But also look for additional relevant information in the source text above.

{evidence_blocks}

## Instructions
Write the complete wiki page in markdown based on the source text above.
Put every heading on its own line and separate every paragraph with a blank line. Do not return the page as one line.
Cross-link to other pages using [[slug]] or [[slug|display text]] — ONLY
use slugs from the "Available pages" list. Do NOT invent new slugs.
Do NOT include Citations or Footnotes sections.
MUST be in the language as the same as the source document text is.

Return ONLY the markdown content, no other text.
"""


WIKI_REFINE_MERGE_SYSTEM = (
    "You are a wiki page merger. You receive two versions of the same wiki page:\n"
    "- EXISTING: the current version in the knowledge base.\n"
    "- INCOMING: a new version generated from a different source document.\n\n"
    "Your job is to produce a SINGLE unified page that preserves ALL factual "
    "content from BOTH versions. Rules:\n\n"
    "1. KEEP all facts, numbers, procedures, names from both versions.\n"
    "2. REMOVE exact duplicates — if both versions state the same fact, keep it once.\n"
    "3. ORGANIZE coherently — clear H2 sections, opening paragraph, ## See also.\n"
    "4. PRESERVE [[wikilinks]] from both versions.\n"
    "5. Write in the SAME LANGUAGE as the existing content.\n"
    "6. Do NOT summarize or condense — the merged page should be AT LEAST as long "
    "as the longer of the two inputs.\n"
    "7. Do NOT add any facts not present in either version.\n\n"
    "Return ONLY the merged markdown content, no other text."
)


# --- helpers ---------------------------------------------------------------


_REFINE_THINK_PREFIX_RE = re.compile(r"^.*</think>", re.DOTALL)


def _wiki_strip_think(raw: str) -> str:
    """Strip a leading ``...</think>`` block that some LLMs emit."""
    if not isinstance(raw, str):
        return ""
    return _REFINE_THINK_PREFIX_RE.sub("", raw).strip()


def _wiki_assemble_evidence(
    plan_item: dict,
    claims: list[dict],
    entity_by_name: dict[str, dict] | None = None,
    concept_by_term: dict[str, dict] | None = None,
) -> list[dict]:
    """Find claims whose `subject` matches any `entity_name` in the plan item.

    Match is case-insensitive: exact match on the full normalized subject, or
    whole-word substring match for multi-word subjects. Each returned
    evidence item carries chunk_ids[] for downstream source-context loading.

    Fallback: if no claim attributes this page (a common case when the MAP
    LLM extracted entities but no claims for them), synthesize a single
    evidence stub from the canonical entity/concept records — that way
    provenance (chunk_ids / source_doc_ids) and the source-context fetch
    still resolve to the chunks that produced the entity/concept itself.
    Pass ``entity_by_name`` / ``concept_by_term`` (lowercased-key lookups
    over ``plan["_entities"]`` / ``plan["_concepts"]``) to enable the
    fallback.
    """
    raw_names = [n.strip() for n in (plan_item.get("entity_names") or []) if isinstance(n, str) and n.strip()]
    if not raw_names:
        return []

    names_lower = [n.lower() for n in raw_names]
    patterns = [re.compile(rf"\b{re.escape(n)}\b", re.IGNORECASE) for n in raw_names]

    evidence: list[dict] = []
    for claim in claims:
        if not isinstance(claim, dict):
            continue
        subj_raw = (claim.get("subject") or "").strip()
        if not subj_raw:
            continue
        subj_lower = subj_raw.lower()

        matched = subj_lower in names_lower or any(p.search(subj_raw) for p in patterns)
        if not matched:
            continue

        chunk_ids = claim.get("chunk_ids") or []
        evidence.append(
            {
                "statement": claim.get("statement", ""),
                "subject": claim.get("subject", ""),
                "confidence": claim.get("confidence", "explicit"),
                "chunk_ids": [c for c in chunk_ids if isinstance(c, str) and c],
            }
        )

    if evidence:
        return evidence

    # ---- Fallback: derive evidence from entity/concept chunk_ids. -------
    if not entity_by_name and not concept_by_term:
        return []

    fallback_chunk_ids: list[str] = []
    matched_names: list[str] = []
    for name, name_lc in zip(raw_names, names_lower):
        hit = None
        if entity_by_name:
            hit = entity_by_name.get(name_lc)
        if hit is None and concept_by_term:
            hit = concept_by_term.get(name_lc)
        if not hit:
            continue
        for cid in hit.get("chunk_ids") or []:
            if isinstance(cid, str) and cid and cid not in fallback_chunk_ids:
                fallback_chunk_ids.append(cid)
        matched_names.append(name)

    if not fallback_chunk_ids:
        return []

    # Marker ``_synthetic`` keeps this item out of the writer prompt — it
    # exists only to carry chunk_ids forward for provenance and source-context
    # fetching. _wiki_format_evidence_blocks filters it out.
    return [
        {
            "statement": "",
            "subject": matched_names[0] if matched_names else raw_names[0],
            "confidence": "inferred",
            "chunk_ids": fallback_chunk_ids,
            "_synthetic": True,
        }
    ]


def _wiki_format_evidence_blocks(evidence: list[dict]) -> str:
    # Filter out synthetic stubs (entity-fallback chunk-id carriers) — they
    # don't represent real claims and shouldn't appear in the writer's
    # evidence checklist.
    real_evidence = [ev for ev in (evidence or []) if not ev.get("_synthetic")]
    if not real_evidence:
        return "(no pre-extracted evidence — extract facts directly from the source document text above)"
    lines: list[str] = []
    for i, ev in enumerate(real_evidence, 1):
        confidence = (ev.get("confidence") or "explicit").upper()
        subject = ev.get("subject") or ""
        statement = ev.get("statement") or ""
        lines.append(f"{i}. [{confidence}] {subject}\n   {statement}")
    return "\n\n".join(lines)


def _wiki_collect_evidence_chunk_ids(evidence: list[dict]) -> list[str]:
    seen: list[str] = []
    for ev in evidence:
        for cid in ev.get("chunk_ids") or []:
            if isinstance(cid, str) and cid and cid not in seen:
                seen.append(cid)
    return seen


async def _wiki_load_chunks_by_id(
    chunk_ids: list[str],
    tenant_id: str,
    kb_id: str,
) -> dict[str, str]:
    """Fetch chunks from ES by id. Returns ``{chunk_id: content_with_weight}``.

    Tries a batch search first (``condition={"id": [chunk_ids]}``) and falls
    back to ``docStoreConn.get(chunk_id, …)`` per missing id. The fallback
    exists because the ``id`` filter is translated differently by each
    backend (ES uses a bool/should over the ``id`` keyword and the reserved
    ``_id``; Infinity translates to SQL ``id IN (...)``; OpenSearch routes
    to the ``ids`` query). A chunk that was inserted without populating the
    queryable ``id`` field — or whose dataset/table doesn't expose it the
    same way — silently drops out of the batch query but is still
    addressable via the primary-key ``get()``. The fallback heals both.
    """
    if not chunk_ids:
        return {}
    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)
    select_fields = ["id", "content_with_weight"]
    out: dict[str, str] = {}
    unique_ids = [cid for cid in dict.fromkeys(chunk_ids) if isinstance(cid, str) and cid]
    if not unique_ids:
        return {}

    BATCH = 500
    for i in range(0, len(unique_ids), BATCH):
        batch_ids = unique_ids[i : i + BATCH]
        condition = {"id": batch_ids}
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                condition,
                [],
                OrderByExpr(),
                0,
                len(batch_ids),
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields)
        except Exception:
            logging.exception("wiki_refine: batch chunk fetch failed (%d ids)", len(batch_ids))
            field_map = {}
        for cid, row in field_map.items():
            content = row.get("content_with_weight")
            if isinstance(content, str) and content:
                out[cid] = content

    # Per-id fallback for anything the batch missed — robust against backend-
    # specific quirks in how ``condition={"id": [list]}`` is translated.
    missing = [cid for cid in unique_ids if cid not in out]
    if missing:
        logging.warning(
            "wiki_refine: batch chunk fetch missed %d/%d id(s) in kb=%s; falling back to per-id get() (first missing: %s)",
            len(missing),
            len(unique_ids),
            kb_id,
            missing[0],
        )

        def _get_one(cid: str):
            try:
                return cid, settings.docStoreConn.get(cid, index, [kb_id])
            except Exception:
                logging.exception("wiki_refine: per-id get failed for %s", cid)
                return cid, None

        # Run the per-id gets concurrently to keep latency reasonable.
        results = await asyncio.gather(*[thread_pool_exec(_get_one, cid) for cid in missing], return_exceptions=False)

        recovered = 0
        for cid, doc in results:
            if not isinstance(doc, dict):
                continue
            content = doc.get("content_with_weight")
            if isinstance(content, str) and content:
                out[cid] = content
                recovered += 1

        if recovered:
            logging.info(
                "wiki_refine: per-id fallback recovered %d/%d missing chunk(s)",
                recovered,
                len(missing),
            )

    final_missing = [cid for cid in unique_ids if cid not in out]
    if final_missing:
        logging.warning(
            "wiki_refine: %d chunk(s) still unresolved after fallback in kb=%s (first: %s) — check that the chunk_ids exist in the doc-store and that the row's kb_id matches the request.",
            len(final_missing),
            kb_id,
            final_missing[0],
        )

    return out


async def _wiki_build_source_context(
    evidence: list[dict],
    tenant_id: str,
    kb_id: str,
    budget: int = WIKI_REFINE_SOURCE_BUDGET_CHARS,
) -> str:
    """Concatenate evidence chunks into a labelled source-context block.

    Budget is char-based. Evidence chunks come first (preserve their order of
    appearance in the evidence list); if total exceeds budget the tail is
    truncated with a marker.
    """
    chunk_ids = _wiki_collect_evidence_chunk_ids(evidence)
    if not chunk_ids:
        return "(no source chunks available)"

    chunk_map = await _wiki_load_chunks_by_id(chunk_ids, tenant_id, kb_id)
    if not chunk_map:
        print("chunk_map::::::::::::::", chunk_ids, tenant_id, kb_id, flush=True)
        return "(source chunks could not be loaded)"

    parts: list[str] = []
    total = 0
    truncated = 0
    for cid in chunk_ids:
        content = chunk_map.get(cid)
        if not content:
            continue
        block = f"[CHUNK {cid}]\n{content}"
        if total + len(block) + 2 > budget:
            remaining = budget - total
            if remaining > 1000:
                parts.append(block[:remaining] + "\n\n[…chunk truncated…]")
                total += remaining
            truncated += 1
            continue
        parts.append(block)
        total += len(block) + 2

    if truncated:
        parts.append(f"\n\n[…{truncated} chunk(s) omitted to fit context budget…]")

    return "\n\n".join(parts)


# --- artifactlink rewriting and doc-id collection ------------------------------

_WIKILINK_PIPE_RE = re.compile(r"\[\[([^\[\]\|]+?)\|([^\[\]]+?)\]\]")
_WIKILINK_SIMPLE_RE = re.compile(r"\[\[([^\[\]\|]+?)\]\]")
_WIKI_MARKDOWN_LINK_RE = re.compile(r"\[([^\]]+)\]\(([^)]+)\)")


def _wiki_transform_links(
    content_md: str,
    kb_id: str,
    page_titles: dict[str, str] | None = None,
    valid_slugs: set[str] | None = None,
) -> tuple[str, list[str]]:
    """Normalize wiki links and return ``(rendered_md, unique_outlinks)``.

    Both the canonical ``[[slug]]`` form and Markdown links emitted by an LLM
    are accepted. Artifact links are rewritten to the relative
    ``artifact/{kb_id}/{slug}`` form. When an artifact link uses its slug as
    the label, prefer the planned page title; otherwise derive readable text
    from the slug.
    """
    kb_id_str = str(kb_id)
    page_titles = page_titles or {}
    if valid_slugs is not None:
        valid_slugs = {str(slug).strip() for slug in valid_slugs if str(slug).strip()}
    seen: set[str] = set()
    outlinks: list[str] = []

    def _track(slug: str) -> None:
        s = slug.strip()
        if s and s not in seen:
            seen.add(s)
            outlinks.append(s)

    def _display_text(label: str, slug: str) -> str:
        label = label.strip()
        if label not in {slug, slug.rsplit("/", 1)[-1]}:
            return label
        planned_title = page_titles.get(slug)
        if planned_title:
            return planned_title
        readable = slug.rsplit("/", 1)[-1].replace("-", " ").replace("_", " ").strip()
        return readable.title() or label

    def _is_valid(slug: str) -> bool:
        return valid_slugs is None or slug in valid_slugs

    def _wiki_slug(href: str) -> str | None:
        parsed = urlsplit(href)
        if parsed.scheme or parsed.netloc:
            if parsed.netloc != "artifact":
                return None
            path = parsed.path
        else:
            path = parsed.path
        parts = path.strip("/").split("/")
        if parts and parts[0] == "artifact":
            parts = parts[1:]
        if len(parts) < 2 or parts[0] != kb_id_str:
            return None
        return "/".join(parts[1:])

    def _markdown_artifact(m: re.Match) -> str:
        slug = _wiki_slug(m.group(2))
        if not slug:
            return m.group(0)
        if not _is_valid(slug):
            return _display_text(m.group(1), slug)
        _track(slug)
        return f"[{_display_text(m.group(1), slug)}](artifact/{kb_id_str}/{slug})"

    def _piped(m: re.Match) -> str:
        slug = m.group(1).strip()
        text = m.group(2).strip()
        if not _is_valid(slug):
            return text
        _track(slug)
        return f"[{text}](artifact/{kb_id_str}/{slug})"

    def _simple(m: re.Match) -> str:
        slug = m.group(1).strip()
        if not _is_valid(slug):
            return _display_text(slug, slug)
        _track(slug)
        return f"[{_display_text(slug, slug)}](artifact/{kb_id_str}/{slug})"

    rewritten = _WIKI_MARKDOWN_LINK_RE.sub(_markdown_artifact, content_md or "")
    rewritten = _WIKILINK_PIPE_RE.sub(_piped, rewritten)
    rewritten = _WIKILINK_SIMPLE_RE.sub(_simple, rewritten)
    return rewritten, outlinks


async def _wiki_collect_doc_ids(
    chunk_ids: list[str],
    tenant_id: str,
    kb_id: str,
) -> list[str]:
    """Look up ``doc_id`` for each chunk by id. Returns the unique list in
    first-seen order (subset of the source chunks' parents).

    Defensive: handles both string and list shapes of the ``doc_id`` field
    (different doc-store connectors normalize scalar keyword fields
    differently). Logs when nothing comes back so the empty-source_doc_ids
    failure mode is diagnosable.
    """
    if not chunk_ids:
        return []
    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)
    select_fields = ["id", "doc_id"]
    out: list[str] = []
    seen: set[str] = set()
    total_rows_seen = 0

    def _accept(did) -> None:
        if isinstance(did, str):
            if did and did not in seen:
                seen.add(did)
                out.append(did)
        elif isinstance(did, (list, tuple)):
            for d in did:
                if isinstance(d, str) and d and d not in seen:
                    seen.add(d)
                    out.append(d)

    BATCH = 500
    for i in range(0, len(chunk_ids), BATCH):
        batch_ids = chunk_ids[i : i + BATCH]
        condition = {"id": batch_ids}
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                condition,
                [],
                OrderByExpr(),
                0,
                len(batch_ids),
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields)
        except Exception:
            logging.exception("wiki_refine: failed to fetch doc_ids for %d chunks", len(batch_ids))
            continue
        total_rows_seen += len(field_map)
        for row in field_map.values():
            _accept(row.get("doc_id"))

    if chunk_ids and not out:
        logging.warning(
            "wiki_refine: doc_id resolution returned 0 for %d chunk(s) (rows_found=%d, kb=%s); first chunk_id=%s",
            len(chunk_ids),
            total_rows_seen,
            kb_id,
            chunk_ids[0],
        )
    return out


async def _wiki_get_existing_page(
    slug: str,
    tenant_id: str,
    kb_id: str,
) -> Optional[dict]:
    """Fetch a wiki_page row by slug from this KB. Returns ``{id, content_md,
    content_md_raw, title, page_type}`` or None. ``content_md_raw`` is the
    pre-link-transform markdown — what the merger should consume."""
    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)
    condition = {
        "compile_kwd": [WIKI_PAGE_COMPILE_KWD],
        "slug_kwd": [slug],
    }
    select_fields = [
        "id",
        "content_with_weight",
        "title_kwd",
        "page_type_kwd",
    ]
    try:
        res = await thread_pool_exec(
            settings.docStoreConn.search,
            select_fields,
            [],
            condition,
            [],
            OrderByExpr(),
            0,
            1,
            index,
            [kb_id],
        )
        field_map = settings.docStoreConn.get_fields(res, select_fields)
    except Exception:
        logging.exception("wiki_refine: failed to fetch existing page for slug=%s", slug)
        return None
    if not field_map:
        return None
    row_id, row = next(iter(field_map.items()))
    rendered = row.get("content_with_weight") or ""
    return {
        "id": row_id,
        "content_md": rendered,
        "content_md_raw": rendered,
        "title": row.get("title_kwd") or "",
        "page_type": row.get("page_type_kwd") or "concept",
    }


async def _wiki_chat_text(
    chat_mdl,
    system_prompt: str,
    user_prompt: str,
    temperature: float,
    llm_timeout: int,
) -> str:
    """Single chat call returning the raw text. Trims to chat_mdl.max_length
    via message_fit_in and strips a leading </think> block."""
    msg = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    try:
        _, msg = message_fit_in(msg, chat_mdl.max_length)
    except Exception:
        logging.exception("wiki_refine: message_fit_in failed; sending untrimmed")
    request_conf = _knowledge_compile_gen_conf(chat_mdl, {"temperature": temperature})
    try:
        raw = await asyncio.wait_for(
            chat_mdl.async_chat(
                msg[0]["content"],
                msg[1:],
                request_conf,
            ),
            timeout=llm_timeout,
        )
    except asyncio.TimeoutError:
        logging.warning("wiki_refine: chat call timed out after %ds", llm_timeout)
        return ""
    except Exception:
        logging.exception("wiki_refine: chat call failed")
        return ""
    if isinstance(raw, tuple):
        raw = raw[0]
    return _wiki_strip_think(raw or "")


async def _wiki_write_page_simple(
    plan_item: dict,
    evidence: list[dict],
    existing_md: Optional[str],
    source_context: str,
    all_plan_slugs: list[str],
    chat_mdl,
    llm_timeout: int,
    example: Optional[str] = None,
) -> str:
    """Single LLM call → markdown content.

    ``example`` is the per-template ``parser_config.example`` override
    for the writer's page-structure section. Falsy / whitespace-only
    values fall through to ``WIKI_TEMPLATE_EXAMPLE``.
    """
    own_slug = plan_item.get("slug") or ""
    available = [s for s in all_plan_slugs if s and s != own_slug]
    slugs_block = "\n".join(f"- [[{s}]]" for s in available) if available else "(none — this is the only page)"

    if existing_md:
        existing_section = f"## Existing page content (UPDATE — integrate new evidence into this)\n\n{existing_md}\n"
    else:
        existing_section = ""

    user_prompt = WIKI_REFINE_WRITER_USER_TEMPLATE.format(
        action=plan_item.get("action", "CREATE"),
        slug=own_slug,
        title=plan_item.get("title", own_slug),
        page_type=plan_item.get("page_type", "concept"),
        all_plan_slugs=slugs_block,
        existing_section=existing_section,
        source_context=source_context,
        evidence_count=len(evidence),
        evidence_blocks=_wiki_format_evidence_blocks(evidence),
    )

    content = await _wiki_chat_text(
        chat_mdl,
        _build_refine_writer_system(example),
        user_prompt,
        temperature=0.15,
        llm_timeout=llm_timeout,
    )
    return content


async def _wiki_merge_page_content(
    existing_md: str,
    new_md: str,
    slug: str,
    chat_mdl,
    shrink_threshold: float = WIKI_MERGE_BODY_SHRINK_THRESHOLD,
    llm_timeout: int = WIKI_MERGE_TIMEOUT,
) -> str:
    """LLM-merge existing vs new. Falls back to ``new_md`` on shrink-check
    failure or LLM error."""
    if not existing_md or len(existing_md.strip()) < 50:
        return new_md
    if existing_md.strip() == (new_md or "").strip():
        return new_md
    if not new_md:
        return existing_md

    user_prompt = (
        f"Merge these two versions of wiki page `{slug}`:\n\n"
        f"## EXISTING VERSION\n\n{existing_md}\n\n"
        "---\n\n"
        f"## INCOMING VERSION\n\n{new_md}\n\n"
        "---\n\n"
        "Produce the merged page now. Return ONLY the markdown content."
    )
    merged = await _wiki_chat_text(
        chat_mdl,
        WIKI_REFINE_MERGE_SYSTEM,
        user_prompt,
        temperature=0.1,
        llm_timeout=llm_timeout,
    )
    if not merged:
        return new_md

    max_input_len = max(len(existing_md), len(new_md))
    min_acceptable = int(max_input_len * shrink_threshold)
    if len(merged) < min_acceptable:
        logging.warning(
            "wiki_refine: merge rejected for slug=%s (merged=%d chars < %d threshold; max input=%d). Falling back to new content.",
            slug,
            len(merged),
            min_acceptable,
            max_input_len,
        )
        return new_md
    return merged


def _wiki_extract_summary(content_md: str, max_chars: int = 300) -> str:
    """First non-heading paragraph of the markdown, capped at ``max_chars``."""
    if not isinstance(content_md, str) or not content_md.strip():
        return ""
    buf: list[str] = []
    for line in content_md.splitlines():
        s = line.strip()
        if not s or s.startswith("#"):
            if buf:
                break
            continue
        buf.append(s)
        if len(" ".join(buf)) >= max_chars:
            break
    return " ".join(buf)[:max_chars]


def _wiki_draft_row_id(kb_id: str, slug: str) -> str:
    return _stable_row_id(WIKI_DRAFT_COMPILE_KWD, kb_id, slug)


async def _wiki_persist_draft(
    page: dict,
    tenant_id: str,
    kb_id: str,
    plan_input_hash: str = "",
    embd_mdl=None,
) -> None:
    """Upsert one wiki_page_draft row (resume cache + searchable page).

    ``plan_input_hash`` is the PLAN's ``input_hash_kwd`` at the time this
    draft was produced. The next REFINE re-entry compares it against the
    current PLAN hash to decide whether the cached draft is still
    valid; a mismatch forces a rewrite for that slug.

    When ``embd_mdl`` is provided the row is made searchable: the title/body are
    tokenized (``title_tks`` / ``content_ltks`` / ``content_sm_ltks``) and a
    ``q_<dim>_vec`` page embedding is attached, with ``available_int=1`` so the
    agent's ``wiki_query`` tool can retrieve it. Without an embedder the row stays
    a non-searchable resume cache. ``content_with_weight`` is left as the page
    JSON either way, so ``_wiki_load_refine_resume`` still restores the draft.
    """
    from common import settings
    from rag.nlp import search as _rag_search
    from rag.nlp import rag_tokenizer

    slug = page.get("slug") or ""
    if not slug:
        return
    index = _rag_search.index_name(tenant_id)
    content_with_weight = json.dumps(page, ensure_ascii=False)
    draft_doc_ids = [d for d in (page.get("source_doc_ids") or []) if isinstance(d, str) and d]
    row = {
        "id": _wiki_draft_row_id(kb_id, slug),
        "doc_id": str(kb_id),
        "compile_kwd": WIKI_DRAFT_COMPILE_KWD,
        "wiki_slug_kwd": slug,
        "source_id": [str(kb_id)],
        "source_doc_ids": draft_doc_ids,
        "input_hash_kwd": plan_input_hash,
        "content_with_weight": content_with_weight,
        "available_int": 0,  # non-searchable unless made searchable below
    }

    # Make the draft searchable when an embedder is available. content_with_weight
    # is deliberately left untouched (the page JSON) — the tokenized fields drive
    # BM25 and q_<dim>_vec drives dense retrieval.
    if embd_mdl is not None:
        title = str(page.get("title") or slug)
        body = str(page.get("content_md_rendered") or page.get("content_md") or page.get("content_md_raw") or "")
        summary = str(page.get("summary") or "")
        content_ltks = rag_tokenizer.tokenize(body)
        row.update(
            {
                "docnm_kwd": title,
                "title_kwd": title,
                "title_tks": rag_tokenizer.tokenize(title),
                "content_ltks": content_ltks,
                "content_sm_ltks": rag_tokenizer.fine_grained_tokenize(content_ltks),
            }
        )
        try:
            emb_text = (summary or f"{title}\n{body}").strip()[:2048] or title
            vectors, _ = await thread_pool_exec(embd_mdl.encode, [emb_text])
            vec = vectors[0]
            vec_list = vec.tolist() if hasattr(vec, "tolist") else list(vec)
            if vec_list:
                row[f"q_{len(vec_list)}_vec"] = vec_list
                row["available_int"] = 1
        except Exception:
            logging.exception("wiki_refine: draft embedding failed slug=%s; row stays non-searchable", slug)

    try:
        try:
            await thread_pool_exec(
                settings.docStoreConn.delete,
                {"compile_kwd": WIKI_DRAFT_COMPILE_KWD, "wiki_slug_kwd": slug},
                index,
                kb_id,
            )
        except Exception:
            logging.debug("wiki_refine: prior draft delete failed; relying on id upsert")
        await thread_pool_exec(settings.docStoreConn.insert, [row], index, kb_id)
    except Exception:
        logging.exception("wiki_refine: failed to persist draft slug=%s", slug)


async def _wiki_load_refine_resume(
    tenant_id: str,
    kb_id: str,
) -> dict[str, tuple[dict, str]]:
    """Load all cached wiki_page_draft rows for this KB.

    Returns ``{slug: (page, stored_plan_input_hash)}``. The hash lets
    REFINE invalidate drafts whose upstream plan has shifted on a
    re-run; legacy rows without the field show up as ``""`` and are
    treated as always-stale.
    """
    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)
    condition = {"compile_kwd": [WIKI_DRAFT_COMPILE_KWD]}
    select_fields = ["id", "wiki_slug_kwd", "content_with_weight", "input_hash_kwd"]

    PAGE_SIZE = 500
    offset = 0
    out: dict[str, tuple[dict, str]] = {}
    while True:
        try:
            res = await thread_pool_exec(
                settings.docStoreConn.search,
                select_fields,
                [],
                condition,
                [],
                OrderByExpr(),
                offset,
                PAGE_SIZE,
                index,
                [kb_id],
            )
            field_map = settings.docStoreConn.get_fields(res, select_fields)
        except Exception:
            logging.exception("wiki_refine: failed to page draft cache")
            break
        if not field_map:
            break
        for row in field_map.values():
            slug = row.get("wiki_slug_kwd")
            content = row.get("content_with_weight")
            if not isinstance(slug, str) or not isinstance(content, str):
                continue
            try:
                cached = json.loads(content)
            except Exception:
                continue
            if isinstance(cached, dict):
                stored_hash = row.get("input_hash_kwd")
                if not isinstance(stored_hash, str):
                    stored_hash = ""
                out[slug] = (cached, stored_hash)
        if len(field_map) < PAGE_SIZE:
            break
        offset += PAGE_SIZE
    return out


# --- public entry ---------------------------------------------------------


async def wiki_refine_from_plan(
    chat_mdl,
    embd_mdl,
    tenant_id: str,
    kb_id: str,
    max_workers: int = DEFAULT_WIKI_REFINE_WORKERS,
    llm_timeout: int = DEFAULT_WIKI_REFINE_TIMEOUT,
    source_budget_chars: int = WIKI_REFINE_SOURCE_BUDGET_CHARS,
    merge_shrink_threshold: float = WIKI_MERGE_BODY_SHRINK_THRESHOLD,
    force_rerun: bool = False,
    callback: Optional[Callable] = None,
    example: Optional[str] = None,
) -> list[dict]:
    """Phase 4 (REFINE) — KB-scoped.

    Reads the cached ``wiki_compilation_plan`` for this KB and writes one
    wiki page per planned entry. Writers run in parallel under
    ``asyncio.Semaphore(max_workers)``. UPDATE pages are LLM-merged against
    their existing content (sanity-checked at ``merge_shrink_threshold``).
    Each finished page is persisted as a searchable ``wiki_page`` row in ES,
    plus a non-searchable ``wiki_page_draft`` row for resume.

    Args:
        chat_mdl, embd_mdl: ragflow LLMBundle instances.
        tenant_id, kb_id: address the doc-store index.
        max_workers: max concurrent writers (default 4).
        llm_timeout: seconds per writer LLM call (default 300).
        source_budget_chars: max chars of source-chunk context per writer call.
        merge_shrink_threshold: a merged body shorter than this fraction of
            the longest input falls back to the new content.
        force_rerun: ignore the wiki_page_draft cache and re-write everything.
        callback: optional ``(progress: float, msg: str)`` callback.

    Returns the list of page dicts (one per planned entry). Each page dict
    has ``slug, title, page_type, topic, action, content_md, summary,
    entity_names, related_kb_pages, source_chunk_ids``.
    """
    # Defensive: some callers accidentally pass the result of
    # ``LLMBundle.encode()`` (a tuple) instead of the bundle itself. Earlier
    # phases often hit a resume cache so this surfaces here for the first
    # time. ``_ensure_llm_bundle`` unwraps the tuple if possible, otherwise
    # logs and returns ``None`` so we bail cleanly.
    embd_mdl = _ensure_llm_bundle(embd_mdl, "encode", label="wiki_refine: embd_mdl")
    if embd_mdl is None:
        return []
    chat_mdl = _ensure_llm_bundle(chat_mdl, "async_chat", label="wiki_refine: chat_mdl")
    if chat_mdl is None:
        return []

    if callback:
        try:
            callback(0.02, "wiki REFINE: loading plan")
        except Exception:
            pass

    plan_pair = await _wiki_load_plan_resume(tenant_id, kb_id)
    if plan_pair is None:
        logging.warning("wiki_refine: no wiki_compilation_plan found for kb=%s", kb_id)
        return []
    plan, plan_input_hash = plan_pair
    if not isinstance(plan, dict):
        logging.warning("wiki_refine: cached plan is not a dict for kb=%s", kb_id)
        return []

    pages_spec = plan.get("pages") or []
    if not pages_spec:
        logging.info("wiki_refine: plan has no pages for kb=%s", kb_id)
        return []
    # Sort by priority then dedupe by slug, keeping the first (highest-priority)
    # entry. The planning LLM sometimes emits the same slug multiple times,
    # which both wastes writer calls and bloats every prompt's "Available
    # pages" list with duplicates.
    sorted_spec = sorted(
        [p for p in pages_spec if isinstance(p, dict) and p.get("slug")],
        key=lambda p: float(p.get("priority", 99)),
    )
    seen_slugs: set[str] = set()
    pages_spec = []
    duplicates_dropped = 0
    for p in sorted_spec:
        s = p.get("slug")
        if not s:
            continue
        if s in seen_slugs:
            duplicates_dropped += 1
            continue
        seen_slugs.add(s)
        pages_spec.append(p)
    if duplicates_dropped:
        logging.info(
            "wiki_refine: dropped %d duplicate slug entr(ies) from plan for kb=%s",
            duplicates_dropped,
            kb_id,
        )

    all_claims = plan.get("_claims") or []
    # ``all_plan_slugs`` is implicitly deduped now (pages_spec is unique).
    all_plan_slugs = [p["slug"] for p in pages_spec]
    page_titles = {str(p["slug"]): str(p.get("title") or "").strip() for p in pages_spec if p.get("slug") and str(p.get("title") or "").strip()}

    # Build canonical entity/concept lookups for evidence fallback. When MAP
    # produced no claims (a real failure mode we've seen on Chinese / dense
    # technical content), provenance still resolves via the chunk_ids on
    # the entities and concepts themselves. The lookups index every name
    # variant (canonical + aliases) so the planner LLM picking an alias
    # spelling still hits the right canonical record.
    entity_by_name: dict[str, dict] = {}
    for e in plan.get("_entities") or []:
        if not isinstance(e, dict):
            continue
        canon = (e.get("name") or "").strip()
        if canon:
            entity_by_name.setdefault(canon.lower(), e)
        for alias in e.get("aliases") or []:
            if isinstance(alias, str) and alias.strip():
                entity_by_name.setdefault(alias.strip().lower(), e)

    concept_by_term: dict[str, dict] = {}
    for c in plan.get("_concepts") or []:
        if not isinstance(c, dict):
            continue
        term = (c.get("term") or "").strip()
        if term:
            concept_by_term.setdefault(term.lower(), c)
        # Concepts in REDUCE output rarely carry aliases, but accept them if
        # present so a future MAP schema change is forward-compatible.
        for alias in c.get("aliases") or []:
            if isinstance(alias, str) and alias.strip():
                concept_by_term.setdefault(alias.strip().lower(), c)

    # Resume cache — only honour drafts whose stored PLAN input_hash
    # matches the current plan's. Mismatch (or missing on legacy rows)
    # forces that slug to be rewritten. ``force_rerun`` still nukes
    # everything for the admin "rebuild from scratch" path.
    cached: dict[str, dict] = {}
    stale_drafts = 0
    if not force_rerun:
        all_drafts = await _wiki_load_refine_resume(tenant_id, kb_id)
        for slug, (page, stored_hash) in all_drafts.items():
            if plan_input_hash and stored_hash and stored_hash == plan_input_hash:
                cached[slug] = page
            else:
                stale_drafts += 1
        if cached or stale_drafts:
            logging.info(
                "wiki_refine: resume — %d fresh, %d stale draft(s) for kb=%s",
                len(cached),
                stale_drafts,
                kb_id,
            )

    pending = [p for p in pages_spec if p.get("slug") not in cached]
    total = max(1, len(pending))

    if callback:
        try:
            callback(0.1, f"wiki REFINE: writing {len(pending)} page(s) (cached={len(cached)})")
        except Exception:
            pass

    semaphore = asyncio.Semaphore(max_workers) if max_workers and max_workers > 0 else None
    completed = 0
    completed_lock = asyncio.Lock()

    async def _write_one(plan_item: dict) -> Optional[dict]:
        nonlocal completed
        slug = plan_item.get("slug") or ""
        action = (plan_item.get("action") or "CREATE").upper()
        title = plan_item.get("title") or slug
        page_type = plan_item.get("page_type") or "concept"

        async def _run() -> Optional[dict]:
            nonlocal completed
            try:
                evidence = _wiki_assemble_evidence(
                    plan_item,
                    all_claims,
                    entity_by_name=entity_by_name,
                    concept_by_term=concept_by_term,
                )
                source_chunk_ids = _wiki_collect_evidence_chunk_ids(evidence)
                source_context = await _wiki_build_source_context(
                    evidence,
                    tenant_id,
                    kb_id,
                    budget=source_budget_chars,
                )

                # Use the raw [[slug]] form for the writer and merger so the
                # LLM sees a stable, well-known artifactlink notation; we render
                # to clickable links once at persist time.
                existing_md_raw: Optional[str] = None
                if action == "UPDATE":
                    existing = await _wiki_get_existing_page(slug, tenant_id, kb_id)
                    if existing:
                        existing_md_raw = existing.get("content_md_raw") or existing.get("content_md")

                content_md_raw = await _wiki_write_page_simple(
                    plan_item,
                    evidence,
                    existing_md_raw,
                    source_context,
                    all_plan_slugs,
                    chat_mdl,
                    llm_timeout,
                    example=example,
                )
                if not content_md_raw:
                    content_md_raw = f"# {title}\n\n(Page generation produced no content.)"

                if existing_md_raw:
                    content_md_raw = await _wiki_merge_page_content(
                        existing_md_raw,
                        content_md_raw,
                        slug,
                        chat_mdl,
                        shrink_threshold=merge_shrink_threshold,
                    )

                # Render artifactlinks once, here, after all LLM transforms.
                content_md_rendered, outlinks = _wiki_transform_links(
                    content_md_raw,
                    kb_id,
                    page_titles=page_titles,
                    valid_slugs=set(all_plan_slugs),
                )
                source_doc_ids = await _wiki_collect_doc_ids(source_chunk_ids, tenant_id, kb_id)
                summary = _wiki_extract_summary(content_md_rendered) or title

                topic = plan_item.get("topic")
                if not isinstance(topic, str) or not topic.strip():
                    topic = title or slug

                page = {
                    "slug": slug,
                    "title": title,
                    "page_type": page_type,
                    "topic": topic.strip(),
                    "action": action,
                    # Rendered content (with clickable artifact/{kb_id}/{slug} links) is
                    # what callers and the UI consume; the raw [[slug]] form is
                    # preserved for LLM-facing re-reads and the merger.
                    "content_md": content_md_rendered,
                    "content_md_rendered": content_md_rendered,
                    "content_md_raw": content_md_raw,
                    "outlinks": outlinks,
                    "summary": summary,
                    "entity_names": plan_item.get("entity_names") or [],
                    "related_kb_pages": plan_item.get("related_kb_pages") or [],
                    "source_chunk_ids": source_chunk_ids,
                    "source_doc_ids": source_doc_ids,
                    "kb_id": str(kb_id),
                }
            except Exception:
                logging.exception("wiki_refine: writer failed for slug=%s", slug)
                return None

            # Searchable wiki_page persistence has moved to the task
            # handler so the doc-storage schema can be controlled in one
            # place at the ingest layer.
            # REFINE now just builds the page dict and resume cache.
            try:
                await _wiki_persist_draft(
                    page,
                    tenant_id,
                    kb_id,
                    plan_input_hash=plan_input_hash,
                    embd_mdl=embd_mdl,
                )
            except Exception:
                logging.exception("wiki_refine: persist_draft failed for slug=%s", slug)

            if callback:
                async with completed_lock:
                    completed += 1
                    done = completed
                progress = 0.1 + 0.85 * (done / total)
                try:
                    callback(progress, f"wiki REFINE: {done}/{total} pages written ({slug})")
                except Exception:
                    pass
            return page

        if semaphore is not None:
            async with semaphore:
                return await _run()
        return await _run()

    tasks = [asyncio.create_task(_write_one(p)) for p in pending]
    if tasks:
        try:
            new_pages = await asyncio.gather(*tasks, return_exceptions=False)
        except Exception:
            for t in tasks:
                t.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
            raise
    else:
        new_pages = []

    results: list[dict] = []
    # Cached pages first (in plan order), then freshly written ones.
    for p in pages_spec:
        slug = p.get("slug")
        if not slug:
            continue
        if slug in cached:
            results.append(cached[slug])
        else:
            # Look up the freshly produced page (None on writer failure).
            for np in new_pages:
                if np and np.get("slug") == slug:
                    results.append(np)
                    break

    # A planned page can still fail during REFINE, and cached pages may carry
    # links from an older plan. Re-render against the pages that actually
    # survived this run so no dangling artifact link reaches persistence.
    actual_slugs = {str(p.get("slug")).strip() for p in results if p.get("slug")}
    actual_titles = {str(p.get("slug")).strip(): str(p.get("title") or "").strip() for p in results if p.get("slug") and str(p.get("title") or "").strip()}
    for page in results:
        raw_content = page.get("content_md_raw") or page.get("content_md") or ""
        rendered, outlinks = _wiki_transform_links(
            raw_content,
            kb_id,
            page_titles=actual_titles,
            valid_slugs=actual_slugs,
        )
        page["content_md"] = rendered
        page["content_md_rendered"] = rendered
        page["outlinks"] = outlinks
        page["summary"] = _wiki_extract_summary(rendered) or page.get("title") or page.get("slug") or ""
        try:
            await _wiki_persist_draft(
                page,
                tenant_id,
                kb_id,
                plan_input_hash=plan_input_hash,
                embd_mdl=embd_mdl,
            )
        except Exception:
            logging.exception("wiki_refine: persist cleaned draft failed for slug=%s", page.get("slug"))

    logging.info(
        "wiki_refine: kb=%s done — pages written=%d (cached=%d new=%d)",
        kb_id,
        len(results),
        len(cached),
        sum(1 for p in new_pages if p),
    )

    if callback:
        try:
            callback(1.0, "wiki REFINE: done")
        except Exception:
            pass

    return results


__all__ = [
    "WIKI_MAP_COMPILE_KWD",
    "WIKI_REDUCE_COMPILE_KWD",
    "WIKI_PLAN_COMPILE_KWD",
    "WIKI_PAGE_COMPILE_KWD",
    "WIKI_DRAFT_COMPILE_KWD",
    "wiki_map_from_chunks",
    "wiki_reduce_from_extracts",
    "wiki_plan_from_reduction",
    "wiki_refine_from_plan",
]