Files
ragflow/rag/advanced_rag/harness/tools/navigation.py
Kevin Hu d5d04ad639 Feat: compilation result navigation in agentic search (#17002)
### Summary

Compilation result navigation in agentic search.
2026-07-16 20:19:32 +08:00

608 lines
24 KiB
Python

"""Navigation tools: document catalog and concept mindmap.
Both answer from a document's *compiled structure* rather than going straight to
retrieval: they read the compiled entities+relations for the documents in scope,
let the chat model answer from that outline, and pull the underlying chunks back
via each relevant entity's ``source_chunk_ids``.
Knowledge compilation writes every template kind into the same graph rows
(``{"entities": [...], "relations": [...]}``), so the two tools share one
implementation and differ only in which *kinds* they read:
* ``catalog_navigate`` — the document's layout: tree / TOC, page index, RAPTOR.
* ``mindmap_navigate`` — the concept mindmap.
Both take the same ``keywords`` the search tools do — keywords drive query
expansion and the keyword-sentence narrowing; without them navigation would hand
back full, un-narrowed chunks.
"""
import json
import logging
import re
import json_repair
_LOG = logging.getLogger(__name__)
# Compiled-structure kinds that describe a document's *layout*: a tree/outline
# or a page index. ``_compilation_template_kind`` in the API folds page_index
# and knowledge_graph into "timeline"; RAPTOR is its own bucket and is
# inherently tree-like, so it counts too.
_CATALOG_KINDS = {"tree", "timeline", "raptor", "page_index", "pageindex"}
# Compiled-structure kinds that describe the document's *concepts*.
_MINDMAP_KINDS = {"mindmap", "mind_map"}
# Cap how much compiled structure we render into the prompt.
_MAX_ENTITIES = 300
_MAX_RELATIONS = 300
_MAX_EVIDENCE_CHUNKS = 24
_NAV_SYSTEM = """You are given {noun} of one or more documents — an outline of entities and their relations — and a question.
Decide whether that outline alone already answers the question.
Rules:
1. Answer ONLY from the outline below. Do not invent facts.
2. Set "is_sufficient" to true only when the outline genuinely answers the question; otherwise false with an empty answer.
3. Always fill "relevant_entities" with the exact `name` values of the entities most related to the question (up to 10), even when the outline is not sufficient — they are used to pull the underlying source text.
Output ONLY JSON, no prose, no code fences:
{{"is_sufficient": true/false, "answer": "<answer, or empty>", "relevant_entities": ["<entity name>", ...]}}"""
def _normalize_kind(kind) -> str:
"""Mirror the API's kind normalization (page_index/knowledge_graph -> timeline)."""
if not isinstance(kind, str):
return ""
normalized = kind.strip().lower().replace("-", "_")
if normalized in {"pageindex", "page_index", "knowledge_graph"}:
return "timeline"
return normalized
async def _load_compiled_structure(tools, doc_id: str, kinds: set) -> dict:
"""Read a document's compiled graphs for the given template ``kinds``.
Mirrors the fetch in ``GET /datasets/<id>/documents/<doc_id>/structure/graph``:
one query for the template-authored graph rows, a second for the RAPTOR row
(which carries ``compile_kwd`` but no ``knowledge_graph_kwd``).
Returns ``{"entities": [...], "relations": [...]}`` for the matching buckets.
"""
from common import settings
from common.doc_store.doc_store_base import OrderByExpr
from common.misc_utils import thread_pool_exec
from rag.nlp import search
resolved = await thread_pool_exec(tools._resolve_doc_tenant, doc_id)
if not resolved:
return {"entities": [], "relations": []}
kb_id, tenant_id = resolved
index_name = search.index_name(tenant_id)
fields = [
"content_with_weight",
"compile_kwd",
"compilation_template_ids",
"compilation_template_kind_kwd",
]
async def _query(condition: dict, limit: int) -> dict:
try:
res = await thread_pool_exec(
settings.docStoreConn.search,
fields,
[],
condition,
[],
OrderByExpr(),
0,
limit,
index_name,
[kb_id],
)
return settings.docStoreConn.get_fields(res, fields) or {}
except Exception:
_LOG.exception("catalog_navigate: failed reading compiled structure for doc=%s", doc_id)
return {}
rows = dict(await _query({"doc_id": [doc_id], "knowledge_graph_kwd": ["graph"]}, 1000))
rows.update(await _query({"doc_id": [doc_id], "compile_kwd": ["raptor_graph"]}, 16))
entities: list[dict] = []
relations: list[dict] = []
for row in rows.values():
try:
graph = json.loads(row.get("content_with_weight") or "{}")
except Exception:
continue
if not isinstance(graph, dict):
continue
compile_kwd = row.get("compile_kwd") or ""
kind = _normalize_kind(row.get("compilation_template_kind_kwd") or compile_kwd)
if compile_kwd == "raptor_graph":
kind = "raptor"
if kind not in kinds:
continue
entities.extend(graph.get("entities") or [])
relations.extend(graph.get("relations") or [])
return {"entities": entities, "relations": relations}
def _render_structure(entities: list[dict], relations: list[dict]) -> str:
"""Render the compiled structure as a compact outline for the prompt."""
lines: list[str] = []
if entities:
lines.append("Entities:")
for e in entities[:_MAX_ENTITIES]:
name = (e.get("name") or "").strip()
if not name:
continue
typ = (e.get("type") or "other").strip()
desc = " ".join((e.get("discription") or "").split())
lines.append(f"- {name} ({typ})" + (f": {desc}" if desc else ""))
if relations:
lines.append("\nRelations:")
for r in relations[:_MAX_RELATIONS]:
src, tgt = (r.get("from") or "").strip(), (r.get("to") or "").strip()
if not src or not tgt:
continue
lines.append(f"- {src} -[{(r.get('type') or 'related').strip()}]-> {tgt}")
return "\n".join(lines)
def _entity_chunk_ids(entities: list[dict], wanted_names: list[str]) -> list[str]:
"""Union the source_chunk_ids of the entities the model called relevant.
Matches on ``name`` or any alias, case-insensitively. Relations carry no
source_chunk_ids, so entities are the only evidence anchor.
"""
wanted = {n.strip().lower() for n in wanted_names if isinstance(n, str) and n.strip()}
ids: list[str] = []
seen: set[str] = set()
for e in entities:
names = {(e.get("name") or "").strip().lower()}
names.update((a or "").strip().lower() for a in (e.get("aliases") or []))
if not (names & wanted):
continue
for cid in e.get("source_chunk_ids") or []:
if isinstance(cid, str) and cid and cid not in seen:
seen.add(cid)
ids.append(cid)
return ids
async def _load_chunks_by_ids(tools, doc_id: str, chunk_ids: list[str]) -> list[dict]:
"""Fetch chunks by their ids from the doc store."""
if not chunk_ids:
return []
from common import settings
from common.doc_store.doc_store_base import OrderByExpr
from common.misc_utils import thread_pool_exec
from rag.nlp import search
resolved = await thread_pool_exec(tools._resolve_doc_tenant, doc_id)
if not resolved:
return []
kb_id, tenant_id = resolved
fields = ["content_with_weight", "docnm_kwd", "doc_id"]
try:
res = await thread_pool_exec(
settings.docStoreConn.search,
fields,
[],
{"id": chunk_ids[:_MAX_EVIDENCE_CHUNKS]},
[],
OrderByExpr(),
0,
_MAX_EVIDENCE_CHUNKS,
search.index_name(tenant_id),
[kb_id],
)
rows = settings.docStoreConn.get_fields(res, fields) or {}
except Exception:
_LOG.exception("catalog_navigate: failed loading evidence chunks for doc=%s", doc_id)
return []
chunks = []
for cid, row in rows.items():
chunks.append(
{
"chunk_id": cid,
"content_with_weight": row.get("content_with_weight") or "",
"docnm_kwd": row.get("docnm_kwd") or "",
"doc_id": row.get("doc_id") or doc_id,
}
)
return chunks
def _doc_aggs(chunks: list[dict]) -> list[dict]:
aggs, seen = [], set()
for c in chunks:
did = c.get("doc_id")
if did and did not in seen:
seen.add(did)
aggs.append({"doc_id": did, "doc_name": c.get("docnm_kwd") or ""})
return aggs
async def _ask_structure(tools, topic: str, entities: list[dict], relations: list[dict], noun: str, label: str) -> tuple[str, list[str]]:
"""Ask the chat model to answer ``topic`` from the rendered outline.
Returns ``(answer, relevant_entity_names)`` — ``answer`` is empty unless the
model judged the outline sufficient; the names are always returned so the
caller can pull the underlying source chunks.
"""
verdict = {}
try:
from rag.prompts.generator import form_message, message_fit_in
user = f"Question:\n{topic}\n\n{noun.capitalize()}:\n{_render_structure(entities, relations)}\n\nOutput JSON:"
_, msg = message_fit_in(form_message(_NAV_SYSTEM.format(noun=f"the {noun}"), user), tools.chat_mdl.max_length)
ans = await tools.chat_mdl.async_chat(msg[0]["content"], msg[1:], {"temperature": 0.2})
if isinstance(ans, tuple):
ans = ans[0]
cleaned = re.sub(r"^.*</think>", "", ans, flags=re.DOTALL)
cleaned = re.sub(r"```(?:json)?\s*|\s*```", "", cleaned).strip()
verdict = json_repair.loads(cleaned) or {}
if not isinstance(verdict, dict):
verdict = {}
except Exception:
_LOG.exception(f"[{label}] Could not read the outline with the model.")
is_sufficient = bool(verdict.get("is_sufficient"))
answer = (verdict.get("answer") or "").strip() if is_sufficient else ""
relevant = [n for n in (verdict.get("relevant_entities") or []) if isinstance(n, str)]
_LOG.info(
"[%s] The %s %s the question; %d relevant entity(ies): %s",
label,
noun,
"answers" if is_sufficient else "does not fully answer",
len(relevant),
", ".join(relevant[:10]) or "none",
)
return answer, relevant
async def _navigate_compiled(
tools,
topic: str,
keywords: str,
doc_scope: list[str] | None,
kinds: set,
label: str,
noun: str,
) -> dict:
"""Answer ``topic`` from the compiled structure of the docs in ``doc_scope``.
Shared by :func:`catalog_navigate` and :func:`mindmap_navigate` — compilation
writes every template kind into the same graph rows, so the tools differ only
in which ``kinds`` they read and how they describe themselves.
Falls back to a plain hybrid search when there is no doc scope, no compiled
structure of these kinds, or no source chunks behind the relevant entities.
:returns: ``{"answer": str, "chunks": [...], "doc_aggs": [...]}``
"""
from rag.advanced_rag.harness.tools.search import _narrow_by_keywords, hybrid_search
if not doc_scope:
_LOG.info(f"[{label}] No document scope given — falling back to a normal search.")
return await hybrid_search(tools, query=topic, keywords=keywords)
entities: list[dict] = []
relations: list[dict] = []
per_doc: list[tuple[str, list[dict]]] = []
for doc_id in doc_scope:
structure = await _load_compiled_structure(tools, doc_id, kinds)
if structure["entities"] or structure["relations"]:
per_doc.append((doc_id, structure["entities"]))
entities.extend(structure["entities"])
relations.extend(structure["relations"])
if not entities and not relations:
_LOG.info(f"[{label}] These documents have no {noun} — falling back to a normal search.")
return await hybrid_search(tools, query=topic, keywords=keywords, doc_scope=doc_scope)
_LOG.info("[%s] Read an outline of %d entity(ies) and %d relation(s).", label, len(entities), len(relations))
answer, relevant = await _ask_structure(tools, topic, entities, relations, noun, label)
# Pull the source text behind the relevant entities, per originating doc.
chunks: list[dict] = []
for doc_id, doc_entities in per_doc:
ids = _entity_chunk_ids(doc_entities, relevant)
if ids:
chunks.extend(await _load_chunks_by_ids(tools, doc_id, ids))
if not chunks:
_LOG.info(f"[{label}] No source text behind those entities — falling back to a normal search.")
return await hybrid_search(tools, query=topic, keywords=keywords, doc_scope=doc_scope)
before = len(chunks)
narrowed = _narrow_by_keywords(chunks, keywords)
if narrowed:
chunks = narrowed
_LOG.info("[%s] Pulled %d source passage(s) behind those entities, kept %d after keyword filtering.", label, before, len(chunks))
else:
_LOG.info("[%s] Keyword filtering removed all %d passage(s); restoring unfiltered set.", label, before)
return {"answer": answer, "chunks": chunks, "doc_aggs": _doc_aggs(chunks)}
async def catalog_navigate(tools, topic: str, keywords: str = "", doc_scope: list[str] | None = None) -> dict:
"""Answer from the documents' compiled catalog (tree / TOC, page index, RAPTOR).
:returns: ``{"answer": str, "chunks": [...], "doc_aggs": [...]}``
"""
_LOG.info(f'[Catalog navigation] Looking through the document catalog for "{topic}" (keywords: {keywords})')
return await _navigate_compiled(
tools,
topic,
keywords,
doc_scope,
kinds=_CATALOG_KINDS,
label="Catalog navigation",
noun="document catalog",
)
async def mindmap_navigate(tools, topic: str, keywords: str = "", doc_scope: list[str] | None = None) -> dict:
"""Answer from the documents' compiled concept mindmap.
``topic`` (not ``concept``) — the parameter name must match the registered
``_navigate_schema``, otherwise every LLM call raises a TypeError.
:returns: ``{"answer": str, "chunks": [...], "doc_aggs": [...]}``
"""
_LOG.info(f'[Mindmap navigation] Following the concept mindmap for "{topic}" (keywords: {keywords})')
return await _navigate_compiled(
tools,
topic,
keywords,
doc_scope,
kinds=_MINDMAP_KINDS,
label="Mindmap navigation",
noun="concept mindmap",
)
# ── Knowledge-graph exploration ─────────────────────────────────────────────
#
# Unlike catalog/mindmap (which read the merged "graph" JSON of one doc), the KG
# store keeps one searchable row per entity/relation, so graph_explore *searches*
# its way to a small subgraph: seed entities by the question, hop out over
# relations to the 2nd-degree neighbours, then answer from that subgraph.
_KG_SEEDS = 3 # entities matched directly to the question
_KG_HOPS = 1 # relation hops out from the seeds (1 => "2nd degree")
_KG_NEIGHBORS = 10 # neighbour entity rows resolved per hop
_KG_REL_LIMIT = 32 # relations fetched per endpoint filter
async def _kg_scopes(tools, doc_scope: list[str] | None):
"""Resolve the (kb_id, tenant_id, doc_ids|None) groups to search.
With a ``doc_scope`` the graph is limited to those docs (grouped by their
KB); otherwise the whole bound KB graph is explored.
"""
from common.misc_utils import thread_pool_exec
if doc_scope:
by_kb: dict[tuple, list[str]] = {}
for doc_id in doc_scope:
resolved = await thread_pool_exec(tools._resolve_doc_tenant, doc_id)
if resolved:
by_kb.setdefault(resolved, []).append(doc_id)
return [(kb, tenant, docs) for (kb, tenant), docs in by_kb.items()]
return [(kb.id, kb.tenant_id, None) for kb in getattr(tools, "kbs", []) or []]
async def _kg_search(tools, kb_id: str, tenant_id: str, doc_ids, kind: str, text: str = "", top_n: int = 8, extra: dict | None = None) -> list[dict]:
"""Search the compiled KG rows of one KB and return the raw field maps."""
from common import settings
from common.doc_store.doc_store_base import MatchTextExpr, OrderByExpr
from common.misc_utils import thread_pool_exec
from rag.nlp import search
condition: dict = {"knowledge_graph_kwd": [kind]}
if doc_ids:
condition["doc_id"] = list(doc_ids)
if extra:
condition.update(extra)
fields = ["content_with_weight", "source_chunk_ids", "doc_id", "docnm_kwd", "from_entity_kwd", "to_entity_kwd"]
exprs = []
if text:
if getattr(tools, "embed_mdl", None):
try:
exprs.append(await settings.retriever.get_vector(text, tools.embed_mdl, top_n, 0.1))
except Exception:
_LOG.exception("[Graph exploration] vector build failed; using keyword match")
if not exprs:
exprs.append(MatchTextExpr(["content_ltks", "content_sm_ltks"], text, top_n))
try:
res = await thread_pool_exec(
settings.docStoreConn.search,
fields,
[],
condition,
exprs,
OrderByExpr(),
0,
top_n,
search.index_name(tenant_id),
[kb_id],
)
rows = settings.docStoreConn.get_fields(res, fields) or {}
except Exception:
_LOG.exception("[Graph exploration] KG search failed (kind=%s)", kind)
return {}
return rows
def _kg_parse_entity(row: dict) -> dict | None:
try:
payload = json.loads(row.get("content_with_weight") or "{}")
except Exception:
payload = {}
name = (payload.get("name") or payload.get("term") or payload.get("title") or "").strip()
if not name:
return None
aliases = [str(a).strip() for a in (payload.get("aliases") or []) if str(a).strip()]
return {
"name": name,
"type": (payload.get("type") or "other"),
"discription": (payload.get("discription") or payload.get("description") or ""),
"aliases": aliases,
"source_chunk_ids": list(row.get("source_chunk_ids") or []),
"doc_id": row.get("doc_id") or "",
"docnm_kwd": row.get("docnm_kwd") or "",
}
def _kg_parse_relation(row: dict) -> dict | None:
src = (row.get("from_entity_kwd") or "").strip()
tgt = (row.get("to_entity_kwd") or "").strip()
if not src or not tgt:
return None
typ = "related"
try:
payload = json.loads(row.get("content_with_weight") or "{}")
typ = payload.get("type") or payload.get("relation") or "related"
except Exception:
pass
return {
"from": src,
"to": tgt,
"type": typ,
"source_chunk_ids": list(row.get("source_chunk_ids") or []),
"doc_id": row.get("doc_id") or "",
}
def _collect_evidence_ids(entities: list[dict], relations: list[dict], relevant_names: list[str]) -> dict:
"""Group the source_chunk_ids of the relevant entities AND relations by doc.
An entity is relevant when its name/alias was named by the model; a relation
is relevant when either endpoint is.
"""
wanted = {n.strip().lower() for n in relevant_names if isinstance(n, str) and n.strip()}
by_doc: dict[str, list[str]] = {}
seen: set[tuple[str, str]] = set()
def _add(doc_id: str, ids):
for cid in ids or []:
if not (isinstance(cid, str) and cid):
continue
key = (doc_id, cid)
if key in seen:
continue
seen.add(key)
by_doc.setdefault(doc_id, []).append(cid)
for e in entities:
names = {(e.get("name") or "").strip().lower(), *[(a or "").strip().lower() for a in (e.get("aliases") or [])]}
if names & wanted:
_add(e.get("doc_id") or "", e.get("source_chunk_ids"))
for r in relations:
if {(r.get("from") or "").strip().lower(), (r.get("to") or "").strip().lower()} & wanted:
_add(r.get("doc_id") or "", r.get("source_chunk_ids"))
return by_doc
async def graph_explore(tools, query: str, keywords: str = "", doc_scope: list[str] | None = None) -> dict:
"""Explore the compiled knowledge graph to answer ``query``.
Searches seed entities for the question, hops out over their relations to the
2nd-degree neighbours, asks the chat model to answer from that subgraph, then
pulls the source passages behind the entities/relations it found relevant and
narrows them with ``keywords``. Falls back to a plain hybrid search when the
KB has no compiled graph or no source text is behind the relevant nodes.
:returns: ``{"answer": str, "chunks": [...], "doc_aggs": [...]}``
"""
from rag.advanced_rag.harness.tools.search import _narrow_by_keywords, hybrid_search
_LOG.info(f'[Graph exploration] Exploring the knowledge graph for "{query}" (keywords: {keywords})')
scopes = await _kg_scopes(tools, doc_scope)
if not scopes:
_LOG.info("[Graph exploration] No knowledge base in scope — falling back to a normal search.")
return await hybrid_search(tools, query=query, keywords=keywords, doc_scope=doc_scope)
text = f"{query} {keywords}".strip()
entities: list[dict] = []
relations: list[dict] = []
ent_names: set[str] = set()
def _add_entities(new: list[dict], scope_key: str = "") -> list[str]:
added = []
for e in new:
key = f"{scope_key}:{e['name'].lower()}"
if key in ent_names:
continue
ent_names.add(key)
entities.append(e)
added.append(e["name"])
return added
for kb_id, tenant_id, doc_ids in scopes:
seed_rows = await _kg_search(tools, kb_id, tenant_id, doc_ids, "entity", text=text, top_n=_KG_SEEDS)
seeds = [e for e in (_kg_parse_entity(r) for r in seed_rows.values()) if e]
frontier = _add_entities(seeds, kb_id)
_LOG.info("[Graph exploration] Seeded %d entity(ies): %s", len(frontier), ", ".join(frontier) or "none")
for _hop in range(_KG_HOPS):
if not frontier:
break
# Relations touching the current frontier (as source OR target).
rel_rows: dict = {}
rel_rows.update(await _kg_search(tools, kb_id, tenant_id, doc_ids, "relation", top_n=_KG_REL_LIMIT, extra={"from_entity_kwd": frontier}))
rel_rows.update(await _kg_search(tools, kb_id, tenant_id, doc_ids, "relation", top_n=_KG_REL_LIMIT, extra={"to_entity_kwd": frontier}))
hop_relations = [r for r in (_kg_parse_relation(x) for x in rel_rows.values()) if r]
relations.extend(hop_relations)
neighbour_names = {n for r in hop_relations for n in (r["from"], r["to"]) if n.lower() not in ent_names}
if not neighbour_names:
break
neigh_rows = await _kg_search(tools, kb_id, tenant_id, doc_ids, "entity", text=" ".join(neighbour_names), top_n=_KG_NEIGHBORS)
wanted = {n.lower() for n in neighbour_names}
neighbours = [e for e in (_kg_parse_entity(r) for r in neigh_rows.values()) if e and e["name"].lower() in wanted]
frontier = _add_entities(neighbours, kb_id)
_LOG.info("[Graph exploration] Hop %d reached %d neighbour entity(ies).", _hop + 1, len(frontier))
if not entities and not relations:
_LOG.info("[Graph exploration] No compiled knowledge graph here — falling back to a normal search.")
return await hybrid_search(tools, query=query, keywords=keywords, doc_scope=doc_scope)
_LOG.info("[Graph exploration] Built a subgraph of %d entity(ies) and %d relation(s).", len(entities), len(relations))
answer, relevant = await _ask_structure(tools, query, entities, relations, "knowledge graph", "Graph exploration")
evidence = _collect_evidence_ids(entities, relations, relevant)
chunks: list[dict] = []
for doc_id, ids in evidence.items():
if doc_id and ids:
chunks.extend(await _load_chunks_by_ids(tools, doc_id, ids))
if not chunks:
_LOG.info("[Graph exploration] No source text behind those nodes — falling back to a normal search.")
fallback = await hybrid_search(tools, query=query, keywords=keywords, doc_scope=doc_scope)
fallback["answer"] = answer
return fallback
before = len(chunks)
chunks = _narrow_by_keywords(chunks, keywords)
_LOG.info("[Graph exploration] Pulled %d source passage(s) behind those nodes, kept %d after keyword filtering.", before, len(chunks))
return {"answer": answer, "chunks": chunks, "doc_aggs": _doc_aggs(chunks)}