mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-13 16:38:26 +08:00
3603 lines
132 KiB
Python
3603 lines
132 KiB
Python
#
|
|
# 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="artifact_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 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,
|
|
run_chunked_pipeline as _run_chunked_pipeline,
|
|
stable_row_id as _stable_row_id,
|
|
)
|
|
|
|
|
|
# Global pipeline-rev — bumping this constant invalidates every cached
|
|
# artifact_map_extract / artifact_reduce_result / artifact_compilation_plan
|
|
# / artifact_page_draft / artifact_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 = "artifact_map_extract"
|
|
DEFAULT_WIKI_MAP_WORKERS = 6
|
|
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" - discription: {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 artifact_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 ``artifact_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,
|
|
)
|
|
try:
|
|
res = await asyncio.wait_for(
|
|
gen_json(WIKI_MAP_SYSTEM, user_prompt, chat_mdl, gen_conf={"temperature": 0.1}),
|
|
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 ``artifact_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 6.
|
|
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)
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Migrated from D:/git/arkon/app/ai/mrp/reducer.py, steps 2.1-2.4.
|
|
# KB reconciliation (arkon 2.5-2.6) and the planning LLM call (arkon 2.7) are
|
|
# deferred to the PLAN phase — they belong with the planner, not the dedup.
|
|
#
|
|
# Scope difference from arkon: arkon REDUCE runs per source document. Here it
|
|
# runs per knowledge base — one set of canonical entities/concepts for the
|
|
# entire KB. Inputs come from ES (every artifact_map_extract row in this KB across
|
|
# all docs); the result lives in ES under artifact_reduce_result.
|
|
|
|
WIKI_REDUCE_COMPILE_KWD = "artifact_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 artifact_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_compute_map_input_hash(tenant_id: str, kb_id: str) -> str:
|
|
"""xxh64 fingerprint of the **current** ``artifact_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 = "",
|
|
) -> None:
|
|
"""Upsert the single non-searchable artifact_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.
|
|
"""
|
|
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],
|
|
"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 ``artifact_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
|
|
``artifact_reduce_result`` row per KB. Subsequent calls with
|
|
``force_rerun=False`` (default) return the cached row immediately; pass
|
|
``force_rerun=True`` after new ``artifact_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 artifact_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)
|
|
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)
|
|
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)
|
|
|
|
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)
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Migrated from D:/git/arkon/app/ai/mrp/reducer.py, steps 2.5-2.7 + 2.8 persist.
|
|
# Scope: per KB (one Compilation Plan covering the entire knowledge base),
|
|
# matching the REDUCE phase above.
|
|
#
|
|
# Flow:
|
|
# 1. Resume — return cached artifact_compilation_plan ES row when present.
|
|
# 2. Load REDUCE output from artifact_reduce_result.
|
|
# 3. KB reconciliation — batch-embed entity/concept query texts and run a
|
|
# per-item KNN against existing artifact_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 artifact_compilation_plan row per KB.
|
|
#
|
|
# Differences vs arkon: KB-scoped instead of per-source; no `source` pages
|
|
# emitted (chunk_ids attribution is enough); plan status defaults to
|
|
# "approved" so REFINE can consume immediately (review workflow deferred).
|
|
|
|
WIKI_PLAN_COMPILE_KWD = "artifact_compilation_plan"
|
|
WIKI_PAGE_COMPILE_KWD = "artifact_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_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 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=artifact_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)
|
|
)
|
|
|
|
try:
|
|
res = await asyncio.wait_for(
|
|
gen_json(WIKI_PLAN_RECONCILE_SYSTEM, user_prompt, chat_mdl, gen_conf={"temperature": 0.0}),
|
|
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,
|
|
) -> 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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
try:
|
|
res = await asyncio.wait_for(
|
|
gen_json(WIKI_PLAN_PLANNING_SYSTEM, user_prompt, chat_mdl, gen_conf={"temperature": 0.1}),
|
|
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"] = []
|
|
for page in res["pages"]:
|
|
if not isinstance(page, dict):
|
|
continue
|
|
topic = page.get("topic")
|
|
if not isinstance(topic, str) or not topic.strip():
|
|
fallback = page.get("title") or page.get("slug") or page.get("page_type") or "General"
|
|
page["topic"] = str(fallback).strip() or "General"
|
|
if "estimated_page_count" not in res:
|
|
res["estimated_page_count"] = len(res["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 = "",
|
|
) -> None:
|
|
"""Upsert the single non-searchable artifact_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.
|
|
"""
|
|
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],
|
|
"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 ``artifact_reduce_result`` for this KB, reconciles every
|
|
canonical entity/concept against existing ``artifact_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="artifact_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 artifact_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)
|
|
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)
|
|
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)
|
|
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)
|
|
|
|
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)
|
|
# ---------------------------------------------------------------------------
|
|
#
|
|
# Migrated from D:/git/arkon/app/ai/mrp/writer.py (simple writer path) and
|
|
# merger.py (merge_page_content).
|
|
#
|
|
# Scope: per KB. Consumes the artifact_compilation_plan row written by PLAN,
|
|
# writes one artifact_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 artifact_page row (with embedding) so PLAN reconciliation finds it
|
|
# on the next REDUCE→PLAN cycle.
|
|
#
|
|
# Resume: per-slug artifact_page_draft rows act as a cache; a re-entry skips
|
|
# slugs already cached unless force_rerun=True.
|
|
#
|
|
# Differences vs arkon: no full_text — source context is the union of the
|
|
# evidence chunks fetched from ES by id. Image-marker handling and the
|
|
# complex tool-using writer are deliberately deferred.
|
|
|
|
WIKI_DRAFT_COMPILE_KWD = "artifact_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"
|
|
"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.
|
|
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"\[\[([^\[\]\|]+?)\]\]")
|
|
|
|
|
|
def _wiki_transform_links(content_md: str, kb_id: str) -> tuple[str, list[str]]:
|
|
"""Rewrite ``[[slug]]`` / ``[[slug|display]]`` wikilinks to standard
|
|
markdown links whose href encodes ``(kb_id, slug)`` so a renderer can
|
|
fetch the target page from ES.
|
|
|
|
Returns ``(rewritten_md, unique_outlinks)`` — outlinks are slug strings
|
|
in first-seen order. The href format is ``artifact/{kb_id}/{slug}`` which is
|
|
relative; clients are expected to map this to whatever route serves the
|
|
page (e.g. ``/api/v1/artifact/{kb_id}/{slug}``).
|
|
"""
|
|
kb_id_str = str(kb_id)
|
|
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 _piped(m: re.Match) -> str:
|
|
slug = m.group(1).strip()
|
|
text = m.group(2).strip()
|
|
_track(slug)
|
|
return f"[{text}](artifact/{kb_id_str}/{slug})"
|
|
|
|
def _simple(m: re.Match) -> str:
|
|
slug = m.group(1).strip()
|
|
_track(slug)
|
|
return f"[{slug}](artifact/{kb_id_str}/{slug})"
|
|
|
|
rewritten = _WIKILINK_PIPE_RE.sub(_piped, content_md or "")
|
|
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")
|
|
try:
|
|
raw = await asyncio.wait_for(
|
|
chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": temperature}),
|
|
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),
|
|
)
|
|
|
|
return await _wiki_chat_text(
|
|
chat_mdl,
|
|
_build_refine_writer_system(example),
|
|
user_prompt,
|
|
temperature=0.15,
|
|
llm_timeout=llm_timeout,
|
|
)
|
|
|
|
|
|
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 = "",
|
|
) -> None:
|
|
"""Upsert one non-searchable wiki_page_draft row (resume cache).
|
|
|
|
``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.
|
|
"""
|
|
from common import settings
|
|
from rag.nlp import search as _rag_search
|
|
|
|
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)
|
|
row = {
|
|
"id": _wiki_draft_row_id(kb_id, slug),
|
|
"doc_id": str(kb_id),
|
|
"compile_kwd": WIKI_DRAFT_COMPILE_KWD,
|
|
"artifact_slug_kwd": slug,
|
|
"source_id": [str(kb_id)],
|
|
"input_hash_kwd": plan_input_hash,
|
|
"content_with_weight": content_with_weight,
|
|
"available_int": 0, # non-searchable
|
|
}
|
|
try:
|
|
try:
|
|
await thread_pool_exec(
|
|
settings.docStoreConn.delete,
|
|
{"compile_kwd": WIKI_DRAFT_COMPILE_KWD, "artifact_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", "artifact_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("artifact_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]
|
|
|
|
# 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)
|
|
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 artifact_page persistence has moved to the task
|
|
# handler (TaskHandler._persist_wiki_pages_to_es) so the ES
|
|
# 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,
|
|
)
|
|
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
|
|
|
|
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",
|
|
]
|