From 770809af3aaf0865b08c1538ee591f169899e452 Mon Sep 17 00:00:00 2001
From: chengke <404835780@qq.com>
Date: Thu, 28 May 2026 16:49:25 +0800
Subject: [PATCH 1/4] refactor: implement smart heading deduplication for PDF
shards and add pre-pass LLM grouping for hierarchical merging.
---
.../formats/markdown/parser.py | 18 +-
.../document_parser/formats/pdf/parser.py | 18 +-
.../formats/pdf/shard_merger.py | 68 +++++++
.../structure/heading_llm_executor.py | 171 ++++++++++++++++--
.../structure/layout_parser.py | 53 ++++++
.../shared/services/ai/prompt_service.py | 63 ++++++-
6 files changed, 361 insertions(+), 30 deletions(-)
diff --git a/apps/worker/app/services/document_parser/formats/markdown/parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py
index 614bab23a..7fdfd4df9 100755
--- a/apps/worker/app/services/document_parser/formats/markdown/parser.py
+++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py
@@ -70,7 +70,13 @@ def is_skip(line):
def heading_md_relocate(md_lines, heading_preds):
- """Relocate markdown headings based on predicted levels (sxjg simplified logic)"""
+ """Relocate markdown headings based on predicted levels (sxjg simplified logic)
+
+ When ``_apply_merge_signals`` has merged continuation rows (level='<')
+ into a preceding heading, the DataFrame's ``heading`` column contains the
+ merged text while ``md_lines`` still has the original truncated text.
+ For positive-level headings, we use the DataFrame's ``heading`` to ensure the merged text is emitted.
+ """
def remove_hash(txt):
return re.sub(r"^\s*(#+)\s*", "", txt)
@@ -85,9 +91,13 @@ def remove_hash(txt):
if pred_level < 0:
line_txt = remove_hash(line_txt)
else:
- # sxjg simplified: remove all #, then add correct number of #
- clean_text = line_txt.lstrip("#").lstrip()
- line_txt = f"{'#' * int(pred_level)} {clean_text}"
+ # Use the DataFrame heading text which may have been updated
+ # by _apply_merge_signals (continuation rows appended).
+ heading_text = str(pred_level_df["heading"].iloc[0]).strip()
+ if not heading_text:
+ # Fallback: strip original line's '#' prefix
+ heading_text = line_txt.lstrip("#").lstrip()
+ line_txt = f"{'#' * int(pred_level)} {heading_text}"
# update lines
md_lines[lid] = line_txt.strip()
diff --git a/apps/worker/app/services/document_parser/formats/pdf/parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py
index c5baf58f3..ec44157bd 100755
--- a/apps/worker/app/services/document_parser/formats/pdf/parser.py
+++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py
@@ -80,7 +80,7 @@ def _parse_oversized_pdf(
eval_md_headings,
merge_html_tables,
)
- from app.services.document_parser.formats.pdf.shard_merger import merge_images
+ from app.services.document_parser.formats.pdf.shard_merger import merge_images, merge_shard_lines
from app.services.document_parser.formats.pdf.shard_splitter import (
bin_pack_shards,
run_doc_agent,
@@ -227,12 +227,16 @@ def _predict_shard_headings(shard_idx: int, shard_out_dir: str) -> ShardHeadingR
shard_heading_results[idx] = future.result()
# 7. Merge: concatenate lines_with_heading (in shard order) + merge images
- all_lines_with_heading: list[str] = []
- total_headings = 0
- for result in shard_heading_results:
- if result is not None:
- all_lines_with_heading.extend(result.lines_with_heading)
- total_headings += result.heading_count
+ all_lines_with_heading: list[str] = merge_shard_lines(
+ [
+ result.lines_with_heading
+ for result in shard_heading_results
+ if result is not None
+ ]
+ )
+ total_headings = sum(
+ 1 for line in all_lines_with_heading if line.startswith("#")
+ )
logger.info(
f"๐ Merged {len(shard_heading_results)} shards: "
diff --git a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py
index 55afa3ca0..0bc46e838 100644
--- a/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py
+++ b/apps/worker/app/services/document_parser/formats/pdf/shard_merger.py
@@ -3,11 +3,79 @@
from __future__ import annotations
import os
+import re
import shutil
from loguru import logger
+def _extract_heading_key(line: str) -> tuple[int, str] | None:
+ """Return (level, text) if *line* is a Markdown heading, else None."""
+ m = re.match(r"^(#+)\s*(.*)", line)
+ if not m:
+ return None
+ return len(m.group(1)), m.group(2).strip()
+
+
+def merge_shard_lines(shard_lines_list: list[list[str]]) -> list[str]:
+ """Concatenate per-shard lines_with_heading in order, removing boundary duplicates.
+
+ When a PDF section-divider page falls at the end of shard N and the same
+ heading opens shard N+1, each shard independently identifies it as a heading,
+ resulting in two consecutive identical headings after naรฏve concatenation.
+
+ Strategy: before appending shard N's lines, if the *last heading* in that
+ shard's output has the same (level, text) as the *first heading* in shard
+ N+1's output, strip the trailing heading (and any non-heading lines that
+ follow it, i.e. the divider page content) from shard N's lines.
+ """
+ if not shard_lines_list:
+ return []
+
+ result: list[str] = []
+ for shard_idx, lines in enumerate(shard_lines_list):
+ if not lines:
+ continue
+
+ # Determine next shard's first heading (if any)
+ next_first_heading: tuple[int, str] | None = None
+ for future_idx in range(shard_idx + 1, len(shard_lines_list)):
+ for next_line in shard_lines_list[future_idx]:
+ key = _extract_heading_key(next_line)
+ if key is not None:
+ next_first_heading = key
+ break
+ if next_first_heading is not None:
+ break
+
+ # Find this shard's last heading and its position
+ lines_to_add = list(lines)
+ if next_first_heading is not None:
+ last_heading_pos: int | None = None
+ last_heading_key: tuple[int, str] | None = None
+ for pos, line in enumerate(lines_to_add):
+ key = _extract_heading_key(line)
+ if key is not None:
+ last_heading_pos = pos
+ last_heading_key = key
+
+ if (
+ last_heading_pos is not None
+ and last_heading_key == next_first_heading
+ ):
+ # Truncate from the last (duplicate) heading onward
+ logger.info(
+ f"๐ shard_{shard_idx}: removing trailing boundary heading "
+ f"'{last_heading_key[1]}' (L{last_heading_key[0]}) duplicated "
+ f"at start of next shard ({len(lines_to_add) - last_heading_pos} lines trimmed)"
+ )
+ lines_to_add = lines_to_add[:last_heading_pos]
+
+ result.extend(lines_to_add)
+
+ return result
+
+
def merge_images(shard_dirs: list[str], target_dir: str) -> None:
"""Copy all images from shard images/ dirs into target_dir/images/."""
target_img_dir = os.path.join(target_dir, "images")
diff --git a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py
index bea5e8540..9db85c0c5 100644
--- a/apps/worker/app/services/document_parser/structure/heading_llm_executor.py
+++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py
@@ -21,18 +21,25 @@
def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame:
"""Collapse consecutive body rows into placeholder rows before LLM chunking.
- The output DataFrame has columns [id, heading, reason]. The ``level``
- column is intentionally NOT forwarded to the LLM โ preliminary estimates
- were found to mislead the model more often than they helped. The naive-
- stage body-text detection (level == -1) is still used here to decide which
- rows become placeholders vs candidates.
+ The output DataFrame has columns [id, heading, note, reason].
+
+ ``note`` is used internally by ``run_merge_pre_pass`` only โ it is
+ NOT forwarded to the main hierarchy LLM:
+ - ``"?"`` marks a heading candidate that is **directly adjacent** to the
+ previous candidate with **no placeholder between them**. This signals
+ to the pre-pass that the pair should be evaluated for possible merging.
+ - ``""`` (empty) for all other rows (normal candidates and placeholders).
+
+ The ``level`` column is intentionally NOT forwarded to the LLM.
"""
if df is None or len(df) == 0:
- return pd.DataFrame(columns=["id", "heading", "reason"])
+ return pd.DataFrame(columns=["id", "heading", "note", "reason"])
rows: list[dict[str, Any]] = []
index = 0
row_count = len(df)
+ prev_was_candidate = False # True when the immediately preceding output row is a candidate
+
while index < row_count:
lvl_raw = df.iloc[index]["level"]
try:
@@ -57,22 +64,28 @@ def compact_for_llm(df: pd.DataFrame) -> pd.DataFrame:
{
"id": f"{start_id}-{end_id}",
"heading": f"[{run_length} BODY LINES]",
+ "note": "",
"reason": PLACEHOLDER_REASON,
}
)
+ prev_was_candidate = False
index = end_index
else:
row = df.iloc[index]
+ # Mark with '?' when directly following another candidate (no placeholder gap)
+ note = "?" if prev_was_candidate else ""
rows.append(
{
"id": int(row["id"]),
"heading": str(row["heading"]),
+ "note": note,
"reason": str(row.get("reason", "") or ""),
}
)
+ prev_was_candidate = True
index += 1
- return pd.DataFrame(rows, columns=["id", "heading", "reason"])
+ return pd.DataFrame(rows, columns=["id", "heading", "note", "reason"])
def split_heading_table(
@@ -105,6 +118,120 @@ def split_heading_table(
return sub_dfs, raw_headings
+def run_merge_pre_pass(
+ compact_df: pd.DataFrame,
+ model_name: str | None = None,
+) -> dict[int, str]:
+ """Focused pre-pass: decide merge/keep for consecutive heading candidate groups.
+
+ Runs BEFORE the main hierarchy LLM call. Scans ``compact_df`` (output of
+ ``compact_for_llm``) for rows whose ``note == "?"``, groups them together
+ with their preceding candidate, and sends all groups in a single focused
+ ``eval-merge-groups`` LLM call.
+
+ Returns a dict ``{row_id: "<"}`` for every row the LLM decides to merge
+ into the previous heading. The caller seeds ``llm_levels`` with these
+ decisions before running the main hierarchy LLM so that merge decisions
+ are final and the main LLM only handles level assignment.
+ """
+ # โโ 1. Collect consecutive groups โโ
+ groups: list[list[dict[str, Any]]] = [] # each element: list of {id, heading}
+ current_group: list[dict[str, Any]] = []
+
+ for _, row in compact_df.iterrows():
+ if row.get("reason") == PLACEHOLDER_REASON:
+ # Body-text placeholder breaks any running group
+ if len(current_group) >= 2:
+ groups.append(current_group)
+ current_group = []
+ continue
+
+ note = str(row.get("note", ""))
+ if note == "?":
+ # Continuation of a consecutive run
+ current_group.append({"id": int(row["id"]), "heading": str(row["heading"])})
+ else:
+ # Start of a new candidate โ flush previous group if large enough
+ if len(current_group) >= 2:
+ groups.append(current_group)
+ current_group = [{"id": int(row["id"]), "heading": str(row["heading"])}]
+
+ if len(current_group) >= 2:
+ groups.append(current_group)
+
+ if not groups:
+ logger.info("merge pre-pass: no consecutive groups found, skipping")
+ return {}
+
+ logger.info(f"merge pre-pass: {len(groups)} consecutive group(s) to evaluate")
+
+ # โโ 2. Format groups for the prompt โโ
+ lines: list[str] = []
+ for g_idx, group in enumerate(groups, start=1):
+ headings_str = " | ".join(f'"{item["heading"]}"' for item in group)
+ lines.append(f"Group {g_idx}: [{headings_str}]")
+ texts = "\n".join(lines)
+
+ # โโ 3. Call LLM directly (bypass df2md โ texts is already formatted) โโ
+ from shared.services.ai.prompt_service import build_prompt
+ from shared.services.ai.openai_compatible_client_sync import get_openai_client
+ from shared.services.ai.response_process_service import eval_response
+
+ try:
+ prompt, temperature, top_p, max_tokens = build_prompt(
+ task="eval-merge-groups",
+ texts=texts,
+ query="",
+ paras={"max_tokens": min(800, len(groups) * 50 + 200)},
+ )
+ messages = [
+ {"role": "system", "content": "you are a document structure expert"},
+ {"role": "user", "content": prompt},
+ ]
+ with stage_timer("heading.merge_pre_pass_llm", group_count=len(groups), model_name=model_name):
+ answer = get_openai_client(model=model_name).chat_completion(
+ messages=messages,
+ model=model_name,
+ max_tokens=max_tokens,
+ temperature=temperature,
+ )
+ result = eval_response(answer)
+ except Exception as exc:
+ logger.warning(f"merge pre-pass LLM call failed: {exc}, skipping pre-pass")
+ return {}
+
+
+ # โโ 4. Parse result โ {id: "<"} โโ
+ merge_ids: dict[int, str] = {}
+ if not isinstance(result, list):
+ logger.warning(f"merge pre-pass: unexpected result type {type(result)}, skipping")
+ return {}
+
+ for item in result:
+ if not isinstance(item, dict):
+ continue
+ g_idx = item.get("group")
+ should_merge = item.get("merge", False)
+ if not should_merge:
+ continue
+ try:
+ g_idx = int(g_idx)
+ except (TypeError, ValueError):
+ continue
+ if g_idx < 1 or g_idx > len(groups):
+ continue
+ group = groups[g_idx - 1]
+ # Mark all rows except the first as "<"
+ for member in group[1:]:
+ merge_ids[member["id"]] = "<"
+ logger.debug(
+ f"merge pre-pass: id={member['id']} '{member['heading'][:50]}' โ '<'"
+ )
+
+ logger.info(f"merge pre-pass: {len(merge_ids)} row(s) flagged for merge")
+ return merge_ids
+
+
def execute_llm_heading_hierarchy(
raw_preds: pd.DataFrame,
prompt_limt: int,
@@ -145,6 +272,15 @@ def execute_llm_heading_hierarchy(
fallback["level"] = -1
return fallback.sort_values("id").reset_index(drop=True)
+ # โโ Merge pre-pass (always on when compact is enabled) โโ
+ # Runs a focused LLM call BEFORE the main hierarchy call to resolve all
+ # consecutive-candidate groups. Main LLM receives clean [id, heading] only.
+ pre_pass_levels: dict[int, str] = {}
+ if compact_enabled:
+ pre_pass_levels = run_merge_pre_pass(
+ preds_for_llm, model_name=model_name
+ )
+
level_dfs, _raw_headings = split_heading_table(
preds_for_llm, threshold=prompt_limt, max_start=max_len, max_end=5
)
@@ -164,7 +300,8 @@ def execute_llm_heading_hierarchy(
model_name=model_name,
):
# โโ Per-chunk independent LLM calls โโ
- llm_levels: dict[int, Any] = {}
+ # Seed with pre-pass merge decisions; main LLM cannot override them
+ llm_levels: dict[int, Any] = dict(pre_pass_levels)
for chunk_idx, chunk_df in enumerate(level_dfs):
# Skip chunks that contain only placeholders
@@ -175,7 +312,8 @@ def execute_llm_heading_hierarchy(
)
continue
- df4llm = chunk_df.drop(columns=["reason", "level"], errors="ignore").copy()
+ # Send only [id, heading] to the main LLM โ no note, no merge hints
+ df4llm = chunk_df.drop(columns=["reason", "level", "note"], errors="ignore").copy()
df4llm["heading"] = df4llm["heading"].apply(clean_md_text_for_llm)
logger.info(
@@ -190,7 +328,11 @@ def execute_llm_heading_hierarchy(
for item in chunk_result:
if isinstance(item, dict) and "id" in item and "level" in item:
try:
- llm_levels[int(item["id"])] = item["level"]
+ row_id = int(item["id"])
+ # Pre-pass merge decisions take priority โ never override
+ if row_id in pre_pass_levels:
+ continue
+ llm_levels[row_id] = item["level"]
except (TypeError, ValueError):
pass
@@ -218,18 +360,21 @@ def execute_llm_heading_hierarchy(
full_preds = raw_preds.copy()[["id", "heading", "level", "reason"]]
- def resolve_level(row_id: Any) -> int:
+ def resolve_level(row_id: Any) -> Any:
try:
int_id = int(row_id)
except (TypeError, ValueError):
return -1
level = llm_levels.get(int_id, -1)
+ # Pre-pass merge decisions arrive as "<"; pass through for _apply_merge_signals
+ if level == "<":
+ return "<"
try:
return int(level)
except (TypeError, ValueError):
return -1
-
- full_preds["level"] = full_preds["id"].map(resolve_level).astype(int)
+
+ full_preds["level"] = full_preds["id"].map(resolve_level)
except Exception as exc:
logger.warning(
diff --git a/apps/worker/app/services/document_parser/structure/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py
index 78c9f56ea..7c2545a6f 100755
--- a/apps/worker/app/services/document_parser/structure/layout_parser.py
+++ b/apps/worker/app/services/document_parser/structure/layout_parser.py
@@ -318,6 +318,56 @@ def _compute_zone_boundaries(toc_hierarchies, coordinate_mode="post_removal"):
return zones
+def _apply_merge_signals(df: pd.DataFrame) -> pd.DataFrame:
+ """Resolve LLM '<' merge signals in-place.
+
+ When the LLM assigns ``level == "<"`` to a row, it means the row's text
+ is a continuation of the *previous* heading and should be appended to it.
+
+ Post-processing:
+ 1. Walk the DataFrame in order.
+ 2. For each ``"<"`` row, append its heading text (space-joined) to the
+ last row whose level was a positive integer (the merge target).
+ 3. Set the merged row's level to -1 so it becomes body text and does not
+ create a spurious heading entry in the tree.
+
+ Rows that carry ``"<"`` but have no preceding valid heading (e.g. the very
+ first row) are simply demoted to -1 without merging.
+ """
+ merge_count = 0
+ last_valid_idx: int | None = None # positional index of last positive-level row
+
+ for pos in range(len(df)):
+ level_val = df.iloc[pos]["level"]
+ if level_val == "<":
+ if last_valid_idx is not None:
+ continuation = str(df.iloc[pos]["heading"]).strip()
+ current_heading = str(df.iloc[last_valid_idx]["heading"]).strip()
+ df.at[df.index[last_valid_idx], "heading"] = (
+ f"{current_heading} {continuation}" if continuation else current_heading
+ )
+ logger.debug(
+ f"๐ Merge signal: row id={df.iloc[pos]['id']} "
+ f"'{continuation[:40]}' โ appended to id={df.iloc[last_valid_idx]['id']}"
+ )
+ else:
+ logger.debug(
+ f"๐ Merge signal at id={df.iloc[pos]['id']} has no preceding heading; demoting"
+ )
+ df.at[df.index[pos], "level"] = -1
+ merge_count += 1
+ else:
+ try:
+ if int(level_val) > 0:
+ last_valid_idx = pos
+ except (TypeError, ValueError):
+ pass
+
+ if merge_count:
+ logger.info(f"๐ Applied {merge_count} '<' merge signal(s)")
+ return df
+
+
def _resolve_first_toc_boundary(toc_hierarchies=None, first_toc_ele_num=None):
"""Resolve the earliest available first-TOC boundary across coordinate sources.
@@ -557,6 +607,9 @@ def _process_single_zone(zone_idx, zone_start, zone_end, zone_toc):
logger.warning("โ ๏ธ No valid headings estimated")
heading_preds = pd.DataFrame()
else:
+ # โโ Apply '<' merge signal before numeric conversion โโ
+ heading_preds = _apply_merge_signals(heading_preds)
+
heading_preds["level"] = (
pd.to_numeric(heading_preds["level"], errors="coerce")
.fillna(-1)
diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py
index e587294e7..20bc6e3fa 100755
--- a/packages/shared-python/shared/services/ai/prompt_service.py
+++ b/packages/shared-python/shared/services/ai/prompt_service.py
@@ -383,7 +383,7 @@ def build_prompt(task, texts, query, **kwargs):
1) HEADING CANDIDATE โ ``id`` is an integer, ``heading`` is the candidate
text. These โ and ONLY these โ are the rows you can evaluate.
- 2) PLACEHOLDER โ ``id`` is ALWAYS a range "start-end" (for a single-line
+ 2) PLACEHOLDER โ ``id`` is ALWAYS a range "start-end" (for a single-line
it is "N-N", e.g. "56-56"); ``heading`` is "[N BODY LINES]"
where N is the number of body lines folded here.
Placeholders are positional markers that tell you how many body lines sit
@@ -432,17 +432,16 @@ def build_prompt(task, texts, query, **kwargs):
A candidate WITHOUT any structural/numbering marker can still be a
heading, but ONLY when ALL of the following hold:
a) The text is short and title-like โ no sentence-ending punctuation.
- b) It is NOT a broken fragment that continues into the next row.
- c) In the input sequence it is IMMEDIATELY followed by a
+ b) In the input sequence it is IMMEDIATELY followed by a
placeholder ``[N BODY LINES]``, or by another candidate with finer granularity.
This is the "section bulk" signal โ the row introduces a body block or a subsection group.
When Rule-2 is satisfied, pick a level consistent with Rule 1.
Rule 3 โ Body text demotion (candidate โ -1):
Demote a candidate to level = -1 when it clearly does NOT serve as a
- section title. Strongest demotion cues are:
- - Two CANDIDATE rows appear adjacent with NO placeholder between them.
- - The text is an isolated broken phrase, fragment, data value, or caption-like snippet (e.g. "Table 3-2", "Figure 4").
+ section title. Strongest demotion cues:
+ - The text is an isolated fragment, data value, or caption-like snippet (e.g. "Table x", "Figure x", "Table/Figure").
+ - The text has sentence-ending punctuation or is clearly prose.
Rule 4 โ Normalise to start at level 1:
The shallowest (the most coarse granularity) heading found MUST be assigned level 1.
@@ -460,6 +459,58 @@ def build_prompt(task, texts, query, **kwargs):
- Do not add any explanations, comments, control characters, or descriptive texts.
"""
+ # ==================== Merge-Group Pre-pass Prompt ====================
+
+ elif task == "eval-merge-groups":
+ # Focused single-question prompt: ONLY decides merge vs. keep for
+ # groups of consecutive heading candidates (no body text between them).
+ # Does NOT assign hierarchy levels โ that is left to the main LLM call.
+ temperature = 0
+ top_p = 0.01
+ max_tokens = kwargs["paras"].get("max_tokens", 800)
+ prompt = f"""
+ You are a PDF heading reconstruction expert.
+
+ A PDF renderer sometimes splits a single long heading title across multiple
+ consecutive lines. You will receive a numbered list of groups. Each group
+ contains 2โ6 consecutive heading candidate lines from a PDF with NO body
+ text between them.
+
+ Your ONLY task: for each group, decide whether the lines should be MERGED
+ into one single heading, or kept as SEPARATE headings in a parent-child
+ relationship.
+
+ **MERGE when ALL hold:**
+ 1. Reading the lines in sequence produces ONE grammatically complete,
+ natural-sounding title โ no missing words, no awkward break.
+ 2. The first line alone is grammatically INCOMPLETE as a standalone title
+ (e.g. ends with a possessive "'s", a preposition "of / for / and",
+ a conjunction, or is otherwise a dangling fragment).
+ 3. No semantic gap: every subsequent line is a direct lexical extension
+ of the first, not a new sub-topic.
+
+ **KEEP SEPARATE when ANY hold:**
+ - The first line is already a complete, self-contained title on its own.
+ - Subsequent lines introduce a different topic or finer sub-topic.
+ - Lines form a clear parent-heading โ child-heading sequence.
+ - **Any subsequent line begins with a numeric or ordinal prefix** such as
+ `01`, `1.`, `(1)`, `โ `, `ไธใ`, `็ฌฌไธ` โ these are numbered sub-items,
+ never continuation fragments of the preceding heading.
+
+ **Generic linguistic signals that indicate MERGE:**
+ - Line ends with a possessive ("Company's", "Board's") โ demands a noun phrase.
+ - Line ends with a preposition ("of", "for", "under", "and") โ phrase is incomplete.
+ - Line ends mid-adjective or mid-noun phrase that continues on the next line.
+
+ Groups to evaluate:
+ {texts}
+
+ Output a JSON array โ one object per group, in the SAME ORDER as the input:
+ [{{"group": 1, "merge": true}}, {{"group": 2, "merge": false}}, ...]
+
+ Output ONLY valid JSON. No markdown fences, no explanations.
+ """
+
# ==================== TOC Heading Evaluation Prompts ====================
elif task == "eval-toc-headings":
From 7d64f847d7bef074fe4be6087a258496261c9cea Mon Sep 17 00:00:00 2001
From: chengke <404835780@qq.com>
Date: Thu, 28 May 2026 23:01:24 +0800
Subject: [PATCH 2/4] refactor: transition to evidence-only retrieval model,
deprecate answer generation
- Updated retrieval architecture to prioritize evidence_text as the primary output, with answer_text now deprecated and always empty.
- Revised documentation and code comments to reflect the new evidence-centric approach.
- Removed legacy components related to answer synthesis and streamlined retrieval processes.
- Added tests for new evidence rendering and discovery selection functionalities.
---
AGENTS.md | 12 +-
apps/api/.env.example | 6 +-
apps/api/app/api/v1/routes/retrieval.py | 21 +-
apps/api/app/mcp/retrieval_server.py | 12 +-
...st_agentic_discovery_selection_contract.py | 96 ++++++
.../test_legacy_evidence_renderer_contract.py | 43 +++
.../tests/contract/test_retrieval_contract.py | 13 +-
apps/worker/.env.example | 6 +-
.../shared-python/shared/core/config/ai.py | 4 -
.../services/retrieval/agentic/__init__.py | 16 +-
.../retrieval/agentic/core/runtime.py | 1 -
.../services/retrieval/agentic/core/types.py | 34 +-
.../retrieval/agentic/discovery/phase.py | 34 --
.../retrieval/agentic/discovery/selection.py | 88 +++---
.../retrieval/agentic/discovery/tools.py | 12 -
.../retrieval/agentic/evidence/builder.py | 58 ----
.../retrieval/agentic/evidence/renderer.py | 12 +-
.../retrieval/agentic/navigation/document.py | 30 +-
.../agentic/navigation/selection_hydration.py | 33 ++
.../retrieval/agentic/navigation/tools.py | 38 ++-
.../retrieval/agentic/orchestrator.py | 222 ++------------
.../services/retrieval/agentic/policy.py | 290 ------------------
.../services/retrieval/agentic/prompts.py | 10 +-
.../services/retrieval/agentic/tools.py | 6 -
.../retrieval/execution/legacy_route.py | 3 +
.../services/retrieval/execution/plan.py | 3 +
.../execution/response_projection.py | 9 +-
.../services/retrieval/execution/routes.py | 6 +-
.../retrieval/hydration/legacy_evidence.py | 55 ++++
.../retrieval/workflow/orchestrator.py | 10 +-
.../services/retrieval/workflow/planner.py | 26 +-
.../retrieval/workflow/runtime_config.py | 2 -
.../retrieval/workflow/step_runner.py | 61 +---
.../retrieval/workflow/synthesizer.py | 141 ---------
.../services/retrieval/workflow/types.py | 18 +-
.../services/retrieval/workflow/wallet.py | 21 +-
36 files changed, 425 insertions(+), 1027 deletions(-)
create mode 100644 apps/api/tests/contract/test_agentic_discovery_selection_contract.py
create mode 100644 apps/api/tests/contract/test_legacy_evidence_renderer_contract.py
delete mode 100644 packages/shared-python/shared/services/retrieval/agentic/policy.py
create mode 100644 packages/shared-python/shared/services/retrieval/hydration/legacy_evidence.py
delete mode 100644 packages/shared-python/shared/services/retrieval/workflow/synthesizer.py
diff --git a/AGENTS.md b/AGENTS.md
index 6582747be..397965b26 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -537,7 +537,7 @@ Core retrieval internals are grouped by ownership:
- `hydration/`: row/path/reference hydration, inline assets, and result assembly.
- `graph/`: document graph publication/query support.
- `stats/`: retrieval hit recording.
-- `workflow/`: query planning, step execution, synthesis, and wallet state.
+- `workflow/`: query planning, retrieve step execution, and wallet state.
- `agentic/core/`: agentic run types, token budgets, runtime config, and traces.
- `agentic/discovery/`: bottom discovery and document selection.
- `agentic/navigation/`: section-tree navigation, selection hydration, and asset tools.
@@ -569,15 +569,15 @@ flowchart LR
The agentic pipeline uses `WorkflowOrchestrator` to handle complex queries via a DAG-based planning and budget-constrained execution engine:
-1. **Planning (`PlannerAgent`)**: The query is analyzed and decomposed into a DAG of steps.
+1. **Planning (`PlannerAgent`)**: The query is analyzed and decomposed into a DAG of retrieval steps.
- Simple queries generate a single `retrieve` step.
- - Complex queries are broken into multiple `retrieve` steps followed by a final `synthesize` step.
-2. **Budget Ledger (`BudgetLedger`)**: A strict token budget mechanism is enforced across the entire DAG execution (e.g., `AGENTIC_MAX_BUDGET=30000`). If the budget is exhausted, the pipeline halts safely and returns the best-effort evidence collected so far.
+ - Complex queries are broken into multiple `retrieve` steps. KNOWHERE does not plan answer synthesis steps.
+2. **Budget Ledger (`BudgetLedger`)**: A strict token budget mechanism is enforced across the entire DAG execution. If the budget is exhausted, the pipeline halts safely and returns the best-effort evidence collected so far.
3. **Execution (`RetrievalAgent`)**: For each `retrieve` step, a multi-phase navigation engine runs:
- **Phase 1 (Discovery)**: 3-channel RRF keyword search and KG document selection.
- **Phase 2 (Navigation)**: Constrained Breadth-First Search (BFS) over the document's section tree. Discovered orphan leaves are merged into the tree to prevent data loss.
- - **Phase 3 (Verdict)**: The LLM evaluates the collected structural outlines + hydrated chunks. Triggers a revision round (max 2) if `NOT_FOUND`.
-4. **Synthesis**: The LLM synthesizes a final `answer_text` and precise citations (`referenced_chunks`) using the unified evidence tree.
+ - **Phase 3 (Evidence Rendering)**: The hydrated document tree is rendered as `evidence_text`.
+4. **Evidence-Only Contract**: Retrieval responses always expose `evidence_text` as the primary output. `answer_text` is retained only as a deprecated empty string. Downstream agents decide whether the evidence is sufficient and synthesize answers outside KNOWHERE.
### Tree Rendering & Hydration
diff --git a/apps/api/.env.example b/apps/api/.env.example
index 609cc7cb6..e83542603 100644
--- a/apps/api/.env.example
+++ b/apps/api/.env.example
@@ -84,8 +84,10 @@ ARK_API_KEY=
# IMAGE_MODEL=qwen3.5-flash
# IMAGE_MODEL_MAX=qwen3.5-flash
-# Optional retrieval overrides have code defaults. Set RETRIEVAL_AGENTIC_ENABLED=false
-# only when you need to fall back to legacy 3-channel RRF mode.
+# Optional retrieval overrides have code defaults. Retrieval is evidence-only:
+# evidence_text is the primary output and answer_text is always empty. Set
+# RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy
+# 3-channel RRF mode.
# File handling defaults
SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md
diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py
index f814ec3ef..e1cd88cb3 100644
--- a/apps/api/app/api/v1/routes/retrieval.py
+++ b/apps/api/app/api/v1/routes/retrieval.py
@@ -17,10 +17,6 @@
router = APIRouter(tags=["Retrieval"])
-def _is_none(value: object) -> bool:
- return value is None
-
-
class ExcludeSection(BaseModel):
document_id: str
section_path: str
@@ -84,12 +80,21 @@ class RetrievalQueryResponse(BaseModel):
namespace: str
query: str
router_used: str
- answer_text: str | None = None
+ evidence_text: str = Field(
+ default="",
+ description="Hierarchical evidence text. Primary output for downstream agents.",
+ )
+ answer_text: str = Field(
+ default="",
+ description=(
+ "DEPRECATED. Always empty; KNOWHERE no longer generates answers. "
+ "Use evidence_text and synthesize answers downstream."
+ ),
+ )
referenced_chunks: list[dict] = Field(default_factory=list)
results: list[dict] = Field(default_factory=list)
- evidence_text: str | None = Field(default=None, exclude_if=_is_none)
- stop_reason: str | None = Field(default=None, exclude_if=_is_none)
- failure_reason: str | None = Field(default=None, exclude_if=_is_none)
+ stop_reason: str | None = None
+ failure_reason: str | None = None
@router.post("/query", response_model=RetrievalQueryResponse)
diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py
index e1bff80e8..e702411de 100644
--- a/apps/api/app/mcp/retrieval_server.py
+++ b/apps/api/app/mcp/retrieval_server.py
@@ -71,10 +71,8 @@ def to_mcp_query_response(response: dict[str, Any]) -> dict[str, Any]:
mcp_response: dict[str, Any] = {
"query": response.get("query"),
"results": results,
+ "evidence_text": response.get("evidence_text") or "",
}
- # Forward evidence_text for agentic mode
- if response.get("evidence_text") is not None:
- mcp_response["evidence_text"] = response["evidence_text"]
return mcp_response
@@ -98,7 +96,8 @@ def create_retrieval_mcp_server(
"knowhere-retrieval",
instructions=(
"Use this server to search published documents. "
- "If you need information before answering, try searching with this tool."
+ "It returns evidence_text and ranked snippets, but never final answers. "
+ "Downstream agents should synthesize from the returned evidence."
),
streamable_http_path=streamable_http_path,
stateless_http=True,
@@ -107,7 +106,10 @@ def create_retrieval_mcp_server(
@server.tool(
name="retrieval.query",
- description="Search published documents and return relevant snippets.",
+ description=(
+ "Search published documents and return relevant snippets plus unified "
+ "evidence_text for downstream answer synthesis."
+ ),
)
async def query_documents(
query: Annotated[str, Field(description="What you want to search for.")],
diff --git a/apps/api/tests/contract/test_agentic_discovery_selection_contract.py b/apps/api/tests/contract/test_agentic_discovery_selection_contract.py
new file mode 100644
index 000000000..12cfdae97
--- /dev/null
+++ b/apps/api/tests/contract/test_agentic_discovery_selection_contract.py
@@ -0,0 +1,96 @@
+from shared.services.retrieval.agentic.core.types import DocTreeNode
+from shared.services.retrieval.agentic.discovery.selection import (
+ _build_discovery_path_selections,
+ _project_discovery_hints,
+)
+
+
+def test_root_discovery_hint_is_projected_for_llm_selection() -> None:
+ hint_lines, hint_by_path = _project_discovery_hints(
+ [
+ {
+ "section_path": "Root",
+ "chunk_id": "chunk_root_relevant",
+ "summary": "document-level market chart",
+ }
+ ],
+ exclude_paths=None,
+ )
+
+ assert hint_lines == [
+ 'โธ path="Root"',
+ " document-level market chart",
+ ]
+ assert hint_by_path["Root"]["chunk_id"] == "chunk_root_relevant"
+
+
+def test_root_discovery_hint_without_llm_selection_does_not_hydrate() -> None:
+ node = DocTreeNode()
+
+ path_selections, chunk_refs = _build_discovery_path_selections(
+ selections=[],
+ hint_by_path={
+ "Root": {
+ "section_path": "Root",
+ "chunk_id": "chunk_root_relevant",
+ }
+ },
+ document_id="doc_root",
+ node=node,
+ )
+
+ assert path_selections == []
+ assert chunk_refs == []
+ assert node.confidence == {}
+
+
+def test_explicit_root_discovery_selection_with_chunk_id_uses_exact_chunk_ref() -> None:
+ node = DocTreeNode()
+
+ path_selections, chunk_refs = _build_discovery_path_selections(
+ selections=[{"path": "Root", "confidence": 0.91}],
+ hint_by_path={
+ "Root": {
+ "section_path": "Root",
+ "chunk_id": "chunk_root_relevant",
+ }
+ },
+ document_id="doc_root",
+ node=node,
+ )
+
+ assert path_selections == []
+ assert chunk_refs == [
+ {
+ "document_id": "doc_root",
+ "chunk_id": "chunk_root_relevant",
+ "section_path": "Root",
+ }
+ ]
+ assert node.confidence["Root"] == 0.91
+
+
+def test_explicit_root_discovery_selection_without_chunk_id_keeps_path_fallback() -> None:
+ node = DocTreeNode()
+
+ path_selections, chunk_refs = _build_discovery_path_selections(
+ selections=[{"path": "Root", "confidence": 0.7}],
+ hint_by_path={
+ "Root": {
+ "section_path": "Root",
+ "chunk_id": "",
+ }
+ },
+ document_id="doc_root",
+ node=node,
+ )
+
+ assert path_selections == [
+ {
+ "path": "Root",
+ "confidence": 0.7,
+ "hydrate_mode": "self_only",
+ }
+ ]
+ assert chunk_refs == []
+ assert node.confidence["Root"] == 0.7
diff --git a/apps/api/tests/contract/test_legacy_evidence_renderer_contract.py b/apps/api/tests/contract/test_legacy_evidence_renderer_contract.py
new file mode 100644
index 000000000..dcf4c086f
--- /dev/null
+++ b/apps/api/tests/contract/test_legacy_evidence_renderer_contract.py
@@ -0,0 +1,43 @@
+from shared.services.retrieval.hydration.legacy_evidence import render_legacy_evidence_text
+
+
+def test_render_legacy_evidence_text_should_group_documents_and_sections() -> None:
+ rows = [
+ {
+ "chunk_id": "c2",
+ "content": "second section content",
+ "sort_order": 2,
+ "source": {
+ "source_file_name": "alpha.pdf",
+ "section_path": "Alpha / Two",
+ },
+ },
+ {
+ "chunk_id": "c1",
+ "content": "first section content\nwith more detail",
+ "sort_order": 1,
+ "source": {
+ "source_file_name": "alpha.pdf",
+ "section_path": "Alpha / One",
+ },
+ },
+ {
+ "chunk_id": "c3",
+ "content": "",
+ "source_file_name": "beta.pdf",
+ "section_path": "Beta / Table",
+ },
+ ]
+
+ evidence_text = render_legacy_evidence_text(rows)
+
+ assert "[Document] alpha.pdf" in evidence_text
+ assert "[Document] beta.pdf" in evidence_text
+ assert "โธ Alpha / One" in evidence_text
+ assert "โธ Alpha / Two" in evidence_text
+ assert " โ first section content" in evidence_text
+ assert " โ with more detail" in evidence_text
+ assert " โ " in evidence_text
+ assert "\u3010\u6587\u6863\u3011" not in evidence_text
+ assert "[\u8868\u683c\u5185\u5bb9]" not in evidence_text
+ assert "[\u56fe\u7247" not in evidence_text
diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py
index 2c9db2ff7..20665a019 100644
--- a/apps/api/tests/contract/test_retrieval_contract.py
+++ b/apps/api/tests/contract/test_retrieval_contract.py
@@ -183,7 +183,7 @@ async def fake_retrieval_run(
captured_requests.append(kwargs)
return AgenticResult(
evidence_text="policy evidence",
- answer_text="policy answer",
+ answer_text="",
referenced_chunks=[
{
"chunk_id": policy_document["chunk_id"],
@@ -350,8 +350,9 @@ async def test_should_return_empty_results_for_an_empty_query(
"namespace": "default",
"query": "",
"router_used": "empty_query_filtered",
+ "evidence_text": "",
+ "answer_text": "",
"results": [],
- "answer_text": None,
"referenced_chunks": [],
}
@@ -741,7 +742,7 @@ async def run_request(
namespace=request.namespace,
query=request.query,
router_used="workflow_single_step",
- answer_text="foreign reference answer",
+ answer_text="",
referenced_chunks=[
{
"chunk_id": foreign_document["chunk_id"],
@@ -820,7 +821,7 @@ async def run_request(
namespace=request.namespace,
query=request.query,
router_used="workflow_single_step",
- answer_text="mismatched section answer",
+ answer_text="",
referenced_chunks=[
{
"chunk_id": visible_chunk["chunk_id"],
@@ -911,7 +912,7 @@ async def fake_retrieval_run(
document = first_document if query == "first shared reference" else second_document
return AgenticResult(
evidence_text=f"evidence for {document['document_id']}",
- answer_text=f"answer for {document['document_id']}",
+ answer_text="",
referenced_chunks=[
{
"chunk_id": shared_chunk_id,
@@ -1022,7 +1023,7 @@ async def fake_retrieval_run(
chunk = first_chunk if query == "first shared section" else second_chunk
return AgenticResult(
evidence_text=f"evidence for {chunk['section_path']}",
- answer_text=f"answer for {chunk['section_path']}",
+ answer_text="",
referenced_chunks=[
{
"chunk_id": shared_chunk_id,
diff --git a/apps/worker/.env.example b/apps/worker/.env.example
index bf5f8251e..4b265d5b1 100644
--- a/apps/worker/.env.example
+++ b/apps/worker/.env.example
@@ -89,8 +89,10 @@ ARK_API_KEY=
# IMAGE_MODEL=qwen3.5-flash
# IMAGE_MODEL_MAX=qwen3.5-flash
-# Optional retrieval overrides have code defaults. Set RETRIEVAL_AGENTIC_ENABLED=false
-# only when you need to fall back to legacy 3-channel RRF mode.
+# Optional retrieval overrides have code defaults. Retrieval is evidence-only:
+# evidence_text is the primary output and answer_text is always empty. Set
+# RETRIEVAL_AGENTIC_ENABLED=false only when you need to fall back to legacy
+# 3-channel RRF mode.
# Required for specific features: billing and analytics
BILLING_ENABLED=false
diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py
index b48567a2d..315511f68 100644
--- a/packages/shared-python/shared/core/config/ai.py
+++ b/packages/shared-python/shared/core/config/ai.py
@@ -61,10 +61,6 @@ class AIConfig(BaseModel):
default=40000,
description="Default token budget issued to each retrieve step.",
)
- RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET: int = Field(
- default=6000,
- description="Default token budget issued to each synthesize step.",
- )
RETRIEVAL_WORKFLOW_PARALLEL_MAX: int = Field(
default=3,
description="Maximum concurrent workflow steps in the same DAG batch.",
diff --git a/packages/shared-python/shared/services/retrieval/agentic/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/__init__.py
index e59b26254..1a8149596 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/__init__.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/__init__.py
@@ -1,15 +1,11 @@
-"""Agentic retrieval orchestration for Knowhere.
+"""Agentic evidence retrieval orchestration for Knowhere.
-Navigate-then-answer loop:
+Flow:
Phase 1: Document selection (discovery + KG LLM select)
- Phase 2: Per-document iterative navigation (navigate_step โ unified action)
- Phase 3: attempt_answer โ DONE (return answer) or NOT_FOUND โ revision
+ Phase 2: Per-document iterative navigation (navigate_step)
+ Phase 3: Render evidence text for downstream agents
Each navigate_step decides action (NAVIGATE/STOP), optional asset tools,
-and section selections in a single LLM call. STOP terminates drill-down.
-After navigation, attempt_answer is called automatically โ its result
-(answer or NOT_FOUND+reason) drives the revision loop.
-
-All tools are thin wrappers around existing retrieval components โ no new
-retrieval algorithms, ranking strategies, or prompts are introduced.
+and section selections in a single LLM call. KNOWHERE does not generate
+final answers; downstream agents decide whether the evidence is sufficient.
"""
diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py
index 7f71cf076..8e7b4f7ac 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py
@@ -17,7 +17,6 @@
def build_config_from_env() -> AgentRunConfig:
return AgentRunConfig(
- max_revisions=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_REVISIONS", "2")),
max_nav_depth=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_DEPTH", "3")),
latency_budget_ms=int(os.environ.get("RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS", "12000")),
token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")),
diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/types.py b/packages/shared-python/shared/services/retrieval/agentic/core/types.py
index 79981b9f6..9547ba584 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/core/types.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/core/types.py
@@ -16,7 +16,6 @@
@dataclass
class AgentRunConfig:
"""Budget and limit configuration for a single agent run."""
- max_revisions: int = 2 # max attempt_answer โ revision cycles
max_nav_depth: int = 3 # max scope_navigate recursion depth
latency_budget_ms: int = 12000
token_budget_total: int = 40000
@@ -80,24 +79,6 @@ def has_content(self) -> bool:
return any(c.has_content() for c in self.children.values())
return False
- def collect_all_paths(self, doc_id: str) -> set[str]:
- """Recursively collect paths of actually-explored sections.
-
- Only ``leaf_content`` and ``children`` represent sections whose
- content was retrieved. ``outline_items`` are structural context
- (titles/summaries shown to the LLM) and must NOT be masked โ
- otherwise revision rounds see 0 candidates and can't re-navigate.
- """
- paths: set[str] = set()
- if self.scope_path:
- paths.add(f'{doc_id}::{self.scope_path}')
- for path in self.leaf_content:
- paths.add(f'{doc_id}::{path}')
- for path, child in self.children.items():
- paths.add(f'{doc_id}::{path}')
- paths.update(child.collect_all_paths(doc_id))
- return paths
-
def flatten_chunk_rows(self) -> list[dict[str, Any]]:
"""Recursively collect all hydrated chunk rows (document order)."""
rows: list[dict[str, Any]] = []
@@ -160,7 +141,7 @@ def collect_referenced_ids(self, *, document_name: str = '') -> list[dict[str, s
return refs
def merge(self, other: 'DocTreeNode') -> None:
- """Additive merge for revision cycles.
+ """Additive merge for navigation results.
Merges outline items, leaf content, children, and confidence from
``other`` into this node. Existing data is preserved; new data is
@@ -198,17 +179,15 @@ class AgenticResult:
- ``evidence_text``: complete hierarchical context for LLM answering
(rendered doc tree with outline + leaf content + inline tables)
- - ``answer_text``: LLM-generated answer to the query based on the
- evidence. Empty string when the evidence was insufficient
- (NOT_FOUND) and max revisions were exhausted.
+ - ``answer_text``: deprecated; always empty because KNOWHERE returns
+ evidence only and downstream agents synthesize answers.
- ``referenced_chunks``: minimal chunk references for hit stats
and frontend display (chunk_id, document_id, chunk_type, etc.)
- ``router_used``: routing path identifier
- ``budget_snapshot``: final budget ledger state at run completion
- - ``stop_reason``: why the run terminated (answer_done / max_revisions /
+ - ``stop_reason``: why the run terminated (evidence_only /
latency_budget / context_budget / no_llm / etc.)
- - ``failure_reason``: semantic reason from the answer attempt when no
- answer could be produced, e.g. evidence was insufficient.
+ - ``failure_reason``: fatal retrieval failure reason, if any.
"""
evidence_text: str
answer_text: str = ''
@@ -242,10 +221,7 @@ class AgentState:
# Phase 2: Per-document navigation results
doc_trees: dict[str, DocTreeNode] = field(default_factory=dict) # doc_id โ DocTreeNode
- # Revision state
- revision_count: int = 0
ever_explored_doc_ids: set[str] = field(default_factory=set)
- seen_section_keys: set[str] = field(default_factory=set) # "{doc_id}::{section_path}"
# Token budget + KG inventory
ledger: BudgetLedger | None = None
diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py
index 5a95061ed..49f932c9c 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py
@@ -125,40 +125,6 @@ async def register_discovery_documents(
state.doc_job_map[did] = job_result_id
-async def select_revision_documents(
- db: AsyncSession,
- *,
- state: AgentState,
- trace: TraceRecorder,
- trace_enabled: bool,
- user_id: str,
- namespace: str,
- query: str,
- exclude_document_ids: list[str],
- bootstrap_llm_fn: LLMFn,
- revision_hint: str,
-) -> str | None:
- try:
- kg_result = await tools.kg_document_select(
- db,
- user_id=user_id,
- namespace=namespace,
- query=query,
- llm_fn=bootstrap_llm_fn,
- exclude_document_ids=list(set(exclude_document_ids)),
- revision_hint=revision_hint,
- budget_snapshot=state.ledger.snapshot() if state.ledger else None,
- )
- except BudgetExceeded:
- logger.info(" agentic: bootstrap budget exhausted during revision doc selection")
- if trace_enabled:
- trace.record_budget_stop("bootstrap_exhausted")
- return "bootstrap_budget"
- state.step_count += 1
- _append_selected_docs(state, kg_result)
- return None
-
-
async def _select_documents(
db: AsyncSession,
*,
diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py
index 55103b94b..40a73f84d 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py
@@ -14,6 +14,7 @@
parse_action_response,
)
from shared.services.retrieval.agentic.navigation.selection_hydration import (
+ hydrate_chunk_refs_into_node,
hydrate_path_selections_into_node,
)
from shared.services.retrieval.agentic.core.types import DocTreeNode
@@ -35,7 +36,6 @@ async def discovery_select_step(
doc_name: str = "",
discovery_hints: list[dict[str, Any]],
exclude_paths: set[str] | None = None,
- revision_hint: str | None = None,
budget_snapshot: dict | None = None,
) -> DocTreeNode:
"""Select and hydrate discovery-found sections after BFS navigation."""
@@ -47,11 +47,11 @@ async def discovery_select_step(
t0 = time.monotonic()
try:
- hint_lines, hint_by_path, root_path_selections = _project_discovery_hints(
+ hint_lines, hint_by_path = _project_discovery_hints(
hints,
exclude_paths=exclude_paths,
)
- if not hint_lines and not root_path_selections:
+ if not hint_lines:
return node
selections: list[dict[str, Any]] = []
@@ -61,7 +61,6 @@ async def discovery_select_step(
doc_name=doc_name,
query=query,
hint_lines=hint_lines,
- revision_hint=revision_hint,
budget_snapshot=budget_snapshot,
)
response = await llm_fn(prompt)
@@ -70,16 +69,23 @@ async def discovery_select_step(
logger.info(
f' discovery_select_step doc="{doc_name}": '
- f"hints={len(hints)} selections={len(selections)} "
- f"root_selections={len(root_path_selections)}"
+ f"hints={len(hints)} selections={len(selections)}"
)
- path_selections = _build_discovery_path_selections(
+ path_selections, chunk_refs = _build_discovery_path_selections(
selections=selections,
hint_by_path=hint_by_path,
- root_path_selections=root_path_selections,
+ document_id=document_id,
node=node,
)
+ await hydrate_chunk_refs_into_node(
+ db,
+ node=node,
+ refs=chunk_refs,
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ )
await hydrate_path_selections_into_node(
db,
node=node,
@@ -107,7 +113,7 @@ def _project_discovery_hints(
hints: list[dict[str, Any]],
*,
exclude_paths: set[str] | None,
-) -> tuple[list[str], dict[str, dict], list[dict[str, Any]]]:
+) -> tuple[list[str], dict[str, dict]]:
exclude_set = {
normalize_section_path(path)
for path in (exclude_paths or set())
@@ -115,7 +121,6 @@ def _project_discovery_hints(
}
hint_lines: list[str] = []
hint_by_path: dict[str, dict] = {}
- root_path_selections: list[dict[str, Any]] = []
for hint in hints:
section_path = normalize_section_path(hint.get("section_path", ""))
if not section_path:
@@ -126,22 +131,12 @@ def _project_discovery_hints(
continue
hint_by_path[section_path] = hint
- if section_path == "Root":
- root_path_selections.append({
- "path": section_path,
- "confidence": float(
- hint.get("discovery_score") or hint.get("score") or 0.7
- ),
- "hydrate_mode": "self_only",
- })
- continue
-
summary = hint.get("summary", "") or ""
hint_lines.append(f'โธ path="{section_path}"')
if summary:
hint_lines.append(f" {summary[:300]}")
- return hint_lines, hint_by_path, root_path_selections
+ return hint_lines, hint_by_path
def _build_discovery_selection_prompt(
@@ -150,25 +145,13 @@ def _build_discovery_selection_prompt(
doc_name: str,
query: str,
hint_lines: list[str],
- revision_hint: str | None,
budget_snapshot: dict | None,
) -> str:
- revision_context = ""
- if revision_hint:
- revision_context = (
- "\nIMPORTANT: This is a REVISION round. "
- "The previous search attempt failed because:\n"
- f'"{revision_hint}"\n'
- "Adjust your selection accordingly. "
- "If no candidate is relevant, return an EMPTY list [].\n"
- )
-
return DISCOVERY_SELECT_PROMPT.format(
doc_name=doc_name or document_id,
budget_block=format_budget_block(budget_snapshot),
items="\n".join(hint_lines),
query=query,
- revision_context=revision_context,
)
@@ -176,31 +159,34 @@ def _build_discovery_path_selections(
*,
selections: list[dict[str, Any]],
hint_by_path: dict[str, dict],
- root_path_selections: list[dict[str, Any]],
+ document_id: str,
node: DocTreeNode,
-) -> list[dict[str, Any]]:
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
valid_selections = [
selection for selection in selections if selection["path"] in hint_by_path
]
- path_selections = list(root_path_selections)
+ path_selections: list[dict[str, Any]] = []
+ chunk_refs: list[dict[str, Any]] = []
for selection in valid_selections:
path = selection["path"]
confidence = selection.get("confidence", 0.7)
node.confidence[path] = confidence
+ hint = hint_by_path[path]
+ if path == "Root":
+ chunk_id = str(hint.get("chunk_id") or "").strip()
+ if chunk_id:
+ chunk_refs.append({
+ "document_id": document_id,
+ "chunk_id": chunk_id,
+ "section_path": path,
+ })
+ continue
+ path_selections.append({
+ "path": path,
+ "confidence": confidence,
+ "hydrate_mode": "self_only",
+ })
+ continue
path_selections.append({"path": path, "confidence": confidence})
- if not path_selections and hint_by_path:
- fallback_path, fallback_hint = next(iter(hint_by_path.items()))
- fallback_confidence = float(
- fallback_hint.get("discovery_score")
- or fallback_hint.get("score")
- or 0.5
- )
- node.confidence[fallback_path] = fallback_confidence
- path_selections.append({
- "path": fallback_path,
- "confidence": fallback_confidence,
- "hydrate_mode": "self_only",
- })
-
- return path_selections
+ return path_selections, chunk_refs
diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py
index e574d1b94..29cae7783 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py
@@ -189,7 +189,6 @@ async def kg_document_select(
query: str,
llm_fn: LLMFn | None,
exclude_document_ids: list[str],
- revision_hint: str | None = None,
**_kwargs: Any,
) -> ToolResult:
"""Select candidate documents from document-level KG."""
@@ -216,20 +215,9 @@ async def kg_document_select(
latency_ms=latency,
)
- revision_context = ""
- if revision_hint:
- revision_context = (
- "\nIMPORTANT: This is a REVISION round. "
- "The previous search attempt failed because:\n"
- f'"{revision_hint}"\n'
- "Adjust your document selection accordingly. "
- "If no document can address this, return an EMPTY array [].\n"
- )
-
file_prompt = FILE_SELECT_PROMPT.format(
overview=overview_text,
query=query,
- revision_context=revision_context,
budget_block=format_budget_block(_kwargs.get("budget_snapshot")),
)
file_response = await llm_fn(file_prompt)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py
index 9e158191a..2b1cc3cf3 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py
@@ -14,44 +14,6 @@
from shared.utils.token_estimate import estimate_tokens
-def with_context_prompt_projection(
- snapshot: dict[str, object],
- *,
- prompt_tokens: int,
-) -> dict[str, object]:
- projected: dict[str, object] = dict(snapshot)
- context_raw = projected.get("context") or {}
- if not isinstance(context_raw, dict):
- return projected
-
- context = dict(context_raw)
- used = int(context.get("used", 0) or 0)
- reserved = int(context.get("reserved", 0) or 0)
- capacity = int(context.get("capacity", 0) or 0)
- projected_used = min(capacity, used + max(int(prompt_tokens), 0))
- projected_remaining = max(capacity - projected_used - reserved, 0)
- context.update(
- {
- "used_projected_before_answer": projected_used,
- "answer_prompt_estimate": max(int(prompt_tokens), 0),
- "remaining": projected_remaining,
- "used_pct": 100
- if capacity <= 0
- else min(100, int(round((projected_used + reserved) * 100 / capacity))),
- }
- )
- if projected_remaining <= 0:
- context["status"] = "EXHAUSTED"
- elif context["used_pct"] >= 80:
- context["status"] = "CRITICAL"
- elif context["used_pct"] >= 50:
- context["status"] = "TIGHT"
- else:
- context["status"] = "HEALTHY"
- projected["context"] = context
- return projected
-
-
def _collect_chunks_by_type(
node: DocTreeNode,
chunk_types: set[str],
@@ -73,10 +35,6 @@ def collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]:
return _collect_chunks_by_type(node, {"image", "table"})
-def collect_image_chunks(node: DocTreeNode) -> list[dict[str, Any]]:
- return _collect_chunks_by_type(node, {"image"})
-
-
def collect_media_chunks_all(
doc_trees: dict[str, DocTreeNode],
) -> list[dict[str, Any]]:
@@ -86,15 +44,6 @@ def collect_media_chunks_all(
return media
-def collect_image_chunks_all(
- doc_trees: dict[str, DocTreeNode],
-) -> list[dict[str, Any]]:
- image_chunks: list[dict[str, Any]] = []
- for tree in doc_trees.values():
- image_chunks.extend(collect_image_chunks(tree))
- return image_chunks
-
-
async def build_asset_url_map(
media_chunks: list[dict[str, Any]],
) -> dict[str, str]:
@@ -104,13 +53,6 @@ async def build_asset_url_map(
)
-async def build_vlm_image_urls(
- doc_trees: dict[str, DocTreeNode],
-) -> list[str]:
- asset_url_map = await build_asset_url_map(collect_image_chunks_all(doc_trees))
- return [url for url in asset_url_map.values() if url]
-
-
def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]:
paths = set(node.leaf_content.keys())
for child in node.children.values():
diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py
index 016e245d4..4b464d004 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py
@@ -17,7 +17,7 @@ def render_unified_doc_tree(
indent = " " * depth
if depth == 0:
- parts.append(f"ใๆๆกฃใ{doc_name}\n")
+ parts.append(f"[Document] {doc_name}\n")
child_prefixes = set(node.children.keys())
@@ -137,7 +137,7 @@ def render_leaf_chunks(
if target_type == "table":
table_html = str(target.get("content", "")).strip()
- content = content.replace(ref_str, f"\n[่กจๆ ผๅ
ๅฎน]\n{table_html}\n")
+ content = content.replace(ref_str, f"\n[Table]\n{table_html}\n")
elif target_type == "image":
file_path = target.get("file_path") or ""
image_description = str(target.get("content", "")).strip()
@@ -146,9 +146,9 @@ def render_leaf_chunks(
asset_url = (asset_lookup or {}).get(target_id, "") if target_id else ""
display_ref = asset_url or file_path
if display_ref:
- content = content.replace(ref_str, f"\n[ๅพ็: {display_ref}]\n{image_description}\n")
+ content = content.replace(ref_str, f"\n[Image: {display_ref}]\n{image_description}\n")
elif image_description:
- content = content.replace(ref_str, f"\n[ๅพ็ๆ่ฟฐ]\n{image_description}\n")
+ content = content.replace(ref_str, f"\n[Image description]\n{image_description}\n")
for line in content.split("\n"):
if line.strip():
@@ -168,14 +168,14 @@ def render_leaf_chunks(
asset_url = (asset_lookup or {}).get(chunk_id, "") if chunk_id else ""
display_ref = asset_url or file_path
if display_ref:
- parts.append(f"{indent}โ [ๅพ็: {display_ref}]")
+ parts.append(f"{indent}โ [Image: {display_ref}]")
if image_description:
for line in image_description.split("\n"):
if line.strip():
parts.append(f"{indent}โ {line}")
elif chunk_type == "table":
table_html = str(chunk.get("content", "")).strip()
- parts.append(f"{indent}โ [่กจๆ ผๅ
ๅฎน]")
+ parts.append(f"{indent}โ [Table]")
if table_html:
for line in table_html.split("\n"):
if line.strip():
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py
index 6720bd28c..6dd1c1e08 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py
@@ -49,7 +49,7 @@ def __init__(
self._llm_fn = llm_fn
self._llm_budget = llm_budget
- async def navigate_selected_documents(self, *, revision_hint: str | None) -> None:
+ async def navigate_selected_documents(self) -> None:
logger.info(
f" agentic: Phase 2 โ navigating {len(self._state.selected_docs)} documents"
)
@@ -57,13 +57,11 @@ async def navigate_selected_documents(self, *, revision_hint: str | None) -> Non
if self._state.elapsed_ms >= self._config.latency_budget_ms:
logger.info(" agentic: latency budget hit during Phase 2, stopping")
break
- await self._navigate_document(doc, revision_hint=revision_hint)
+ await self._navigate_document(doc)
async def _navigate_document(
self,
doc: CandidateDoc,
- *,
- revision_hint: str | None,
) -> None:
job_result_id = self._state.doc_job_map.get(doc.document_id, "")
if not job_result_id:
@@ -82,14 +80,12 @@ async def _navigate_document(
root=root,
doc_name=doc_name,
job_result_id=job_result_id,
- revision_hint=revision_hint,
)
await self._hydrate_discovery_hints(
doc=doc,
root=root,
doc_name=doc_name,
- revision_hint=revision_hint,
)
if not is_discovery_only_doc and doc_pending_assets:
@@ -115,13 +111,8 @@ async def _navigate_bfs(
root: DocTreeNode,
doc_name: str,
job_result_id: str,
- revision_hint: str | None,
) -> list[dict[str, Any]]:
- doc_exclude: set[str] = {
- key.split("::", 1)[1]
- for key in self._state.seen_section_keys
- if key.startswith(f"{doc.document_id}::")
- } if self._state.seen_section_keys else set()
+ doc_exclude: set[str] = set()
pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)]
doc_pending_assets: list[dict[str, Any]] = []
@@ -155,7 +146,6 @@ async def _navigate_bfs(
doc_name=doc_name,
scope_path=scope,
exclude_paths=doc_exclude,
- revision_hint=revision_hint if depth == 0 else None,
budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None,
)
except BudgetExceeded:
@@ -251,7 +241,6 @@ async def _hydrate_discovery_hints(
doc: CandidateDoc,
root: DocTreeNode,
doc_name: str,
- revision_hint: str | None,
) -> None:
doc_hints = self._discovery_by_doc.get(doc.document_id, [])
if not doc_hints or self._llm_fn is None:
@@ -259,10 +248,7 @@ async def _hydrate_discovery_hints(
if self._state.elapsed_ms >= self._config.latency_budget_ms:
return
- discovery_exclude_paths = {
- key.split("::", 1)[1]
- for key in root.collect_all_paths(doc.document_id)
- }
+ discovery_exclude_paths = _collect_leaf_paths(root)
doc_discovery_llm_fn = self._llm_budget.for_discovery(
cast(LLMFn, self._llm_fn),
doc_id=doc.document_id,
@@ -279,7 +265,6 @@ async def _hydrate_discovery_hints(
doc_name=doc_name,
discovery_hints=doc_hints,
exclude_paths=discovery_exclude_paths,
- revision_hint=revision_hint,
budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None,
)
except BudgetExceeded:
@@ -388,6 +373,13 @@ def _merge_step_node(parent_node: DocTreeNode, step_node: DocTreeNode) -> None:
parent_node.confidence = step_node.confidence
+def _collect_leaf_paths(node: DocTreeNode) -> set[str]:
+ paths = set(node.leaf_content.keys())
+ for child in node.children.values():
+ paths.update(_collect_leaf_paths(child))
+ return paths
+
+
def _update_excluded_leaf_paths(
doc_exclude: set[str],
step_node: DocTreeNode,
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py
index e75e32759..45ef16939 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py
@@ -8,6 +8,7 @@
from shared.services.retrieval.agentic.navigation import assets as asset_tools
from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows
from shared.services.retrieval.hydration.path import hydrate_paths_to_rows
+from shared.services.retrieval.hydration.reference import hydrate_referenced_chunk_rows
async def hydrate_path_selections_into_node(
@@ -43,6 +44,38 @@ async def hydrate_path_selections_into_node(
add_chunks_to_node(node, chunks)
+async def hydrate_chunk_refs_into_node(
+ db: AsyncSession,
+ *,
+ node: DocTreeNode,
+ refs: list[dict[str, Any]],
+ user_id: str,
+ namespace: str,
+ document_id: str,
+ job_result_id: str | None = None,
+) -> None:
+ chunks = await hydrate_referenced_chunk_rows(
+ db=db,
+ user_id=user_id,
+ namespace=namespace,
+ refs=refs,
+ )
+ if not chunks:
+ return
+
+ chunks = await _append_connected_asset_targets(db, chunks)
+ resolved_job_result_id = job_result_id or _find_job_result_id(chunks)
+ if resolved_job_result_id:
+ await _attach_root_asset_owners(
+ db,
+ document_id=document_id,
+ job_result_id=resolved_job_result_id,
+ chunks=chunks,
+ )
+
+ add_chunks_to_node(node, chunks)
+
+
async def _append_connected_asset_targets(
db: AsyncSession, chunks: list[dict[str, Any]]
) -> list[dict[str, Any]]:
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
index 79753a2f3..fe9b0b870 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
@@ -42,7 +42,6 @@ async def navigate_step(
doc_name: str = "",
scope_path: str | list[str] | None = None,
exclude_paths: set[str] | None = None,
- revision_hint: str | None = None,
budget_snapshot: dict | None = None,
) -> tuple[str, list[str], DocTreeNode, list[dict]]:
"""Navigate one document scope and hydrate selected sections."""
@@ -69,6 +68,9 @@ async def navigate_step(
selectable = {
item["path"]: item for item in items if item.get("selectable", False)
}
+ visible_items = {
+ item["path"]: item for item in items if item.get("show_summary", True)
+ }
total_images, total_tables = await count_assets_under_scope(
db,
document_id=document_id,
@@ -86,7 +88,6 @@ async def navigate_step(
budget_snapshot=budget_snapshot,
items_text=items_text,
tools_block=tools_block,
- revision_hint=revision_hint,
)
response = await llm_fn(prompt)
@@ -106,18 +107,19 @@ async def navigate_step(
node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None)
node.outline_items = [item for item in items if item.get("show_summary", True)]
- valid_selections = [
+ raw_valid_selections = [
selection
for selection in selections
- if selection["path"] in selectable and selection["path"] not in scope_path_set
+ if selection["path"] in visible_items and selection["path"] not in scope_path_set
]
+ valid_selections = _dedupe_selected_ancestors(raw_valid_selections)
pending: list[dict] = []
path_selections: list[dict[str, Any]] = []
for selection in valid_selections:
path = selection["path"]
confidence = selection.get("confidence", 0.7)
- item = selectable[path]
+ item = visible_items[path]
node.confidence[path] = confidence
if item.get("is_leaf"):
@@ -160,7 +162,6 @@ def _build_navigation_prompt(
budget_snapshot: dict | None,
items_text: str,
tools_block: str,
- revision_hint: str | None,
) -> str:
if not scope_paths:
scope_header = "Current scope: root (document top level)"
@@ -169,7 +170,7 @@ def _build_navigation_prompt(
else:
scope_header = f"Current scope: navigating into {len(scope_paths)} sections"
- prompt = ACTION_PROMPT.format(
+ return ACTION_PROMPT.format(
doc_name=doc_name or document_id,
doc_id=document_id,
scope_header=scope_header,
@@ -178,9 +179,20 @@ def _build_navigation_prompt(
query=query,
tools_block=tools_block,
)
- if revision_hint:
- prompt += (
- "\n\nIMPORTANT: Previous round feedback: "
- f'"{revision_hint}". Adjust your selections accordingly.'
- )
- return prompt
+
+
+def _dedupe_selected_ancestors(selections: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ """Keep coarser selected ancestors when both parent and child paths are selected."""
+ selected_paths = [str(selection.get("path") or "") for selection in selections]
+ kept: list[dict[str, Any]] = []
+ for selection in selections:
+ path = str(selection.get("path") or "")
+ if not path:
+ continue
+ if any(
+ other != path and path.startswith(other + " / ")
+ for other in selected_paths
+ ):
+ continue
+ kept.append(selection)
+ return kept
diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
index 0e1cc51b2..d6e3ac2ba 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
@@ -1,40 +1,35 @@
-"""Retrieval Agent orchestrator โ navigate-then-answer loop.
+"""Retrieval Agent orchestrator โ evidence-only navigation loop.
Flow:
Phase 1: Document selection (bottom_discovery + kg_document_select)
Phase 2: Per-document navigation (iterative BFS via navigate_step)
- Phase 3: Render evidence โ attempt_answer
- โ DONE (has answer) โ return answer + evidence
- โ NOT_FOUND + reason โ revision_hint โ re-select docs + re-navigate
- (exclude seen paths) โ re-attempt
- โ max_revisions โ return best available
+ Phase 3: Render evidence text for downstream agents
The orchestrator drives navigation via an iterative BFS queue per document,
calling navigate_step at each level. Each navigate_step is a single LLM call
that decides action (NAVIGATE/STOP), asset tools (FIND_IMAGES/FIND_TABLES),
and section selections. STOP terminates the drill-down for that scope.
-After navigation completes, attempt_answer is called automatically.
+
+KNOWHERE does not generate final answers. Downstream agents decide whether the
+returned evidence is sufficient for their task and may call retrieval again.
"""
from __future__ import annotations
import os
-from typing import Any, cast
+from typing import Any
from loguru import logger
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetLedger
+from shared.services.retrieval.agentic.core.budget import BudgetLedger
from shared.services.retrieval.agentic.discovery.phase import (
register_discovery_documents,
run_initial_discovery,
- select_revision_documents,
)
from shared.services.retrieval.agentic.navigation.document import DocumentNavigationRunner
from shared.services.retrieval.agentic.evidence.builder import (
- build_vlm_image_urls as _build_vlm_image_urls,
render_evidence as _render_evidence,
trim_evidence_to_budget as _trim_evidence_to_budget,
- with_context_prompt_projection as _with_context_prompt_projection,
)
from shared.services.retrieval.agentic.core.runtime import (
AgentLlmBudget,
@@ -46,13 +41,12 @@
AgentRunConfig,
AgentState,
AgenticResult,
- ToolResult,
)
from shared.services.retrieval.llm_adapter import LLMFn
class RetrievalAgent:
- """Agentic retrieval orchestrator โ navigate-then-answer loop.
+ """Agentic retrieval orchestrator โ navigate and return evidence.
Usage::
@@ -60,8 +54,8 @@ class RetrievalAgent:
result = await agent.run(
db, user_id=..., namespace=..., query=..., llm_fn=..., ...
)
- # result.evidence_text โ hierarchical context
- # result.answer_text โ LLM answer (empty if NOT_FOUND after max revisions)
+ # result.evidence_text โ hierarchical context for downstream agents
+ # result.answer_text โ deprecated, always empty
# result.referenced_chunks โ chunk IDs for hit stats / frontend
The agent requires a valid ``llm_fn`` for LLM-driven navigation.
@@ -92,19 +86,10 @@ async def run(
) -> AgenticResult:
"""Run the agentic retrieval pipeline.
- Returns an ``AgenticResult`` containing the rendered evidence
- text, LLM answer, and referenced chunk IDs. Never raises โ
- errors are captured in trace and the best available result
- is returned.
+ Returns an ``AgenticResult`` containing rendered evidence text
+ and referenced chunk IDs. Never raises โ errors are captured in
+ trace and the best available evidence is returned.
"""
- from shared.services.retrieval.agentic.policy import (
- attempt_answer,
- estimate_attempt_answer_prompt_tokens,
- )
- from shared.services.retrieval.llm_adapter import create_retrieval_vlm_fn
-
- vlm_fn = create_retrieval_vlm_fn()
-
config = config or _build_config_from_env()
exclude_document_ids = exclude_document_ids or []
exclude_sections = exclude_sections or []
@@ -153,11 +138,9 @@ async def run(
logger.warning('agentic: no llm_fn provided โ running discovery-only mode')
bootstrap_llm_fn: LLMFn | None = None
- context_llm_fn: LLMFn | None = None
llm_budget = AgentLlmBudget(state)
if llm_fn is not None:
bootstrap_llm_fn = llm_budget.for_pool(llm_fn, pool='bootstrap')
- context_llm_fn = llm_budget.for_pool(llm_fn, pool='context')
discovery_rows = await run_initial_discovery(
db,
@@ -228,20 +211,14 @@ async def run(
for doc in state.selected_docs
})
- # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
- # Phase 2 + 3 Loop: Navigate โ Render โ Attempt Answer โ (Revise)
- # โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
- answer_text = ''
+ # Phase 2 + 3: navigate once, then render evidence for downstream agents.
evidence_text = ''
- revision_hint: str | None = None
- stop_reason = 'max_revisions'
+ stop_reason = 'evidence_only'
failure_reason = ''
- for round_idx in range(config.max_revisions + 1):
- if state.elapsed_ms >= config.latency_budget_ms:
- stop_reason = 'latency_budget'
- break
-
+ if state.elapsed_ms >= config.latency_budget_ms:
+ stop_reason = 'latency_budget'
+ else:
navigation_runner = DocumentNavigationRunner(
db=db,
state=state,
@@ -255,167 +232,17 @@ async def run(
llm_fn=llm_fn,
llm_budget=llm_budget,
)
- await navigation_runner.navigate_selected_documents(revision_hint=revision_hint)
-
- # โโ Phase 3: Render evidence + attempt_answer โโโโโโโโโโโโโโโโ
- budget_snapshot_before_answer = state.ledger.snapshot() if state.ledger else None
- context_remaining = (
- state.ledger.remaining('context') if state.ledger else config.token_budget_total
- )
- answer_prompt_overhead = estimate_attempt_answer_prompt_tokens(
- query=query,
- evidence_text='',
- state=state,
- config=config,
- budget_snapshot=budget_snapshot_before_answer,
- )
+ await navigation_runner.navigate_selected_documents()
+ context_remaining = state.ledger.remaining('context') if state.ledger else config.token_budget_total
evidence_text = await _trim_evidence_to_budget(
db,
doc_trees=state.doc_trees,
doc_id_to_name=state.doc_id_to_name,
- context_remaining=max(context_remaining - answer_prompt_overhead, 0),
+ context_remaining=context_remaining,
user_id=user_id,
namespace=namespace,
ledger=state.ledger,
)
- answer_prompt_tokens = estimate_attempt_answer_prompt_tokens(
- query=query,
- evidence_text=evidence_text,
- state=state,
- config=config,
- budget_snapshot=budget_snapshot_before_answer,
- )
- budget_snapshot_for_answer = (
- _with_context_prompt_projection(
- state.ledger.snapshot(),
- prompt_tokens=answer_prompt_tokens,
- )
- if state.ledger else None
- )
- if budget_snapshot_for_answer is not None:
- answer_prompt_tokens = estimate_attempt_answer_prompt_tokens(
- query=query,
- evidence_text=evidence_text,
- state=state,
- config=config,
- budget_snapshot=budget_snapshot_for_answer,
- )
- budget_snapshot_for_answer = _with_context_prompt_projection(
- state.ledger.snapshot(),
- prompt_tokens=answer_prompt_tokens,
- )
- context_budget = (budget_snapshot_for_answer.get('context') or {})
- logger.info(
- ' agentic: answer context projection '
- f'prompt_tokens={answer_prompt_tokens} '
- f'remaining={context_budget.get("remaining")}/'
- f'{context_budget.get("capacity")} '
- f'status={context_budget.get("status")}'
- )
-
- if context_llm_fn is None:
- stop_reason = 'no_llm'
- break
-
- # Collect image URLs from evidence for VLM switch
- evidence_image_urls: list[str] = []
- if vlm_fn:
- evidence_image_urls = await _build_vlm_image_urls(state.doc_trees)
-
- async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
- return await llm_budget.call(cast(LLMFn, _vlm_fn), prompt, pool='context')
-
- # Auto-trigger attempt_answer (VLM if images present)
- try:
- status, answer_text, reason = await attempt_answer(
- context_llm_fn,
- query=query,
- evidence_text=evidence_text,
- state=state,
- config=config,
- vlm_fn=vlm_context_call if vlm_fn else None,
- image_urls=evidence_image_urls or None,
- budget_snapshot=budget_snapshot_for_answer,
- )
- except BudgetExceeded:
- logger.info(' agentic: context budget exhausted before attempt_answer')
- if trace_enabled:
- trace.record_budget_stop('context_exhausted')
- status, answer_text, reason = 'NOT_FOUND', '', 'context budget exhausted'
- stop_reason = 'context_budget'
- break
- state.step_count += 1
-
- if trace_enabled:
- trace.record_step(
- 'attempt_answer', ToolResult(
- status=status,
- payload={
- 'answer_length': len(answer_text),
- 'reason': reason,
- },
- ),
- decision_reason=f'phase_3_answer_r{round_idx}',
- )
-
- logger.info(
- f' agentic: attempt_answer status={status} '
- f'round={round_idx}/{config.max_revisions} '
- f'answer_len={len(answer_text)} reason="{reason}"'
- )
-
- if status == 'DONE':
- stop_reason = 'answer_done'
- failure_reason = ''
- break
-
- # โโ NOT_FOUND: prepare revision โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
- failure_reason = reason
- if round_idx >= config.max_revisions:
- stop_reason = 'max_revisions'
- break
-
- state.revision_count += 1
- revision_hint = reason
- logger.info(f' agentic: starting revision {state.revision_count}, hint="{reason}"')
-
- # Record all explored paths for masking (deepest-first removal)
- for doc_id, doc_tree in state.doc_trees.items():
- state.seen_section_keys.update(doc_tree.collect_all_paths(doc_id))
-
- # Clear doc selection for re-exploration (preserve doc_trees for merge)
- state.selected_docs.clear()
-
- # Re-run KG select with revision hint
- if bootstrap_llm_fn is None:
- stop_reason = 'no_llm'
- break
- revision_stop_reason = await select_revision_documents(
- db,
- state=state,
- trace=trace,
- trace_enabled=trace_enabled,
- user_id=user_id,
- namespace=namespace,
- query=query,
- exclude_document_ids=exclude_document_ids,
- bootstrap_llm_fn=bootstrap_llm_fn,
- revision_hint=revision_hint,
- )
- if revision_stop_reason is not None:
- stop_reason = revision_stop_reason
- break
-
- if not state.selected_docs:
- logger.info(' agentic: revision found no new docs โ stopping')
- stop_reason = 'no_new_docs'
- break
-
- if state.ledger is not None:
- await state.ledger.allocate_doc_caps({
- doc.document_id: chunks_count_by_doc.get(doc.document_id, 1)
- for doc in state.selected_docs
- })
# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
# Final Assembly
@@ -436,7 +263,7 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
seen_ref_ids.add(cid)
all_refs.append(ref)
- # Re-render final evidence (may have been updated in last revision)
+ # Re-render when navigation only produced structural context.
if not evidence_text or evidence_text == '(no evidence collected)':
evidence_text = await _render_evidence(
db,
@@ -445,7 +272,7 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
result = AgenticResult(
evidence_text=evidence_text,
- answer_text=answer_text,
+ answer_text='',
referenced_chunks=all_refs,
router_used=router_used,
budget_snapshot=state.ledger.snapshot() if state.ledger else None,
@@ -456,9 +283,8 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
logger.info(
f'agentic retrieval DONE: {len(all_refs)} referenced chunks, '
f'evidence_text={len(evidence_text)} chars, '
- f'answer_text={len(answer_text)} chars, '
f'router={router_used}, steps={state.step_count}, '
- f'stop_reason={stop_reason}, revisions={state.revision_count}, '
+ f'stop_reason={stop_reason}, '
f'{state.elapsed_ms}ms'
)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py
deleted file mode 100644
index 6c201b074..000000000
--- a/packages/shared-python/shared/services/retrieval/agentic/policy.py
+++ /dev/null
@@ -1,290 +0,0 @@
-"""LLM answer-attempt tool for agentic retrieval.
-
-Provides ``attempt_answer()`` โ a single LLM call that tries to answer
-the user's query using the collected evidence.
-
-Returns one of two outcomes:
- - answer_text (non-empty) โ the evidence was sufficient, answer is ready
- - NOT_FOUND + reason โ the evidence was insufficient, triggers a revision
-"""
-from __future__ import annotations
-
-import json
-import os
-import re
-from ipaddress import ip_address
-from typing import Any
-from urllib.parse import urlparse
-
-from loguru import logger
-
-from shared.services.retrieval.agentic.core.budget import BudgetExceeded
-from shared.services.retrieval.agentic.core.types import AgentRunConfig, AgentState
-from shared.services.retrieval.llm_adapter import LLMFn
-from shared.utils.token_estimate import estimate_tokens
-
-
-def _parse_answer_response(text: str) -> dict[str, Any] | None:
- """Extract a JSON answer object from LLM response text."""
- text = text.strip()
- parsed = _load_json_object(text)
- if parsed is not None:
- return parsed
- match = re.search(r'\{.*\}', text, re.DOTALL)
- if match:
- return _load_json_object(match.group())
- return None
-
-
-def _load_json_object(raw_value: str) -> dict[str, Any] | None:
- try:
- parsed = json.loads(raw_value)
- except (json.JSONDecodeError, ValueError):
- return None
- return parsed if isinstance(parsed, dict) else None
-
-
-def _looks_like_json_wrapper(text: str) -> bool:
- """Detect malformed JSON-ish answer wrappers without exposing them."""
- stripped = text.strip()
- if not stripped:
- return False
- if stripped.startswith('{') or stripped.endswith('}'):
- return True
- return bool(re.search(r'"(?:status|answer|reason)"\s*:', stripped))
-
-
-def _is_external_http_url(url: str) -> bool:
- parsed = urlparse(str(url))
- if parsed.scheme not in {'http', 'https'} or not parsed.hostname:
- return False
- host = parsed.hostname.strip().lower()
- if host in {'localhost', 'ip6-localhost', 'ip6-loopback'} or host.endswith('.local'):
- return False
- try:
- addr = ip_address(host)
- except ValueError:
- return True
- return not (addr.is_private or addr.is_loopback or addr.is_link_local)
-
-
-def _budget_line_parts(budget_snapshot: dict | None, pool_name: str) -> dict[str, Any]:
- pool = ((budget_snapshot or {}).get(pool_name) or {})
- capacity = pool.get('capacity', 'unknown')
- remaining = pool.get('remaining', 'unknown')
- used_pct = pool.get('used_pct', 'unknown')
- remaining_pct: int | str = 'unknown'
- try:
- capacity_int = int(capacity)
- remaining_int = int(remaining)
- remaining_pct = 0 if capacity_int <= 0 else max(
- 0,
- min(100, round(remaining_int * 100 / capacity_int)),
- )
- except (TypeError, ValueError):
- remaining_pct = 'unknown'
- return {
- 'status': pool.get('status', 'HEALTHY'),
- 'remaining': remaining,
- 'capacity': capacity,
- 'used_pct': used_pct,
- 'remaining_pct': remaining_pct,
- }
-
-
-async def attempt_answer(
- llm_fn: LLMFn,
- *,
- query: str,
- evidence_text: str,
- state: AgentState,
- config: AgentRunConfig,
- vlm_fn: LLMFn | None = None,
- image_urls: list[str] | None = None,
- budget_snapshot: dict | None = None,
-) -> tuple[str, str, str]:
- """Attempt to answer the query using collected evidence.
-
- When ``vlm_fn`` is provided and ``image_urls`` is non-empty, the
- answer is generated using the VLM in multimodal message format
- (text + image_url parts) so the model can actually *see* chart
- images rather than only reading their text descriptions.
-
- Falls back to ``llm_fn`` (text-only) when VLM is unavailable or
- when no images are present.
-
- Returns (status, answer_text, reason) where:
- - status='DONE', answer_text=, reason=''
- โ evidence was sufficient, answer is ready
- - status='NOT_FOUND', answer_text='', reason=
- โ evidence was insufficient, reason is used as revision_hint
- """
- prompt_text = build_attempt_answer_prompt(
- query=query,
- evidence_text=evidence_text,
- state=state,
- config=config,
- budget_snapshot=budget_snapshot,
- )
-
- verbose = os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', '') == 'true'
-
- # Decide whether to use VLM (multimodal) or text-only LLM
- effective_fn = llm_fn
- effective_input: Any = prompt_text
-
- usable_image_urls = [url for url in image_urls or [] if _is_external_http_url(url)]
- if image_urls and len(usable_image_urls) != len(image_urls):
- logger.info(
- f' [attempt_answer] skipped {len(image_urls) - len(usable_image_urls)} '
- 'non-public image URLs for VLM'
- )
-
- if vlm_fn and usable_image_urls:
- # Build multimodal message: text evidence + image_url parts
- content_parts: list[dict[str, Any]] = [
- {'type': 'text', 'text': prompt_text},
- ]
- for url in usable_image_urls[:20]: # Cap at 20 images to avoid token overflow
- content_parts.append({
- 'type': 'image_url',
- 'image_url': {'url': url},
- })
- effective_input = [{'role': 'user', 'content': content_parts}]
- effective_fn = vlm_fn
- logger.info(
- f' [attempt_answer] using VLM with {len(usable_image_urls)} image URLs '
- f'(capped at {min(len(usable_image_urls), 20)})'
- )
-
- if verbose:
- logger.info(
- f'[attempt_answer PROMPT]\n'
- f'{prompt_text}'
- )
-
- try:
- raw_response = await effective_fn(effective_input)
- except BudgetExceeded:
- raise
- except Exception as exc:
- if effective_fn is llm_fn:
- raise
- logger.warning(f' [attempt_answer] VLM failed, falling back to text LLM: {exc}')
- raw_response = await llm_fn(prompt_text)
- logger.info(f' [attempt_answer] raw={repr(raw_response[:300])}')
-
- if verbose:
- logger.info(
- f'[attempt_answer RESPONSE]\n'
- f'{raw_response}'
- )
-
- if not raw_response.strip() and effective_fn is not llm_fn:
- return 'NOT_FOUND', '', 'VLM returned empty response for multimodal evidence'
-
- parsed = _parse_answer_response(raw_response)
- if not parsed:
- if _looks_like_json_wrapper(raw_response):
- return 'NOT_FOUND', '', 'attempt_answer returned malformed JSON'
- # Keep plain-text fallback for providers that ignore JSON mode entirely.
- return 'DONE', raw_response.strip(), 'parse_error โ treating plain text response as answer'
-
- status = str(parsed.get('status', 'DONE')).strip().upper()
- answer = str(parsed.get('answer', '')).strip()
- reason = str(parsed.get('reason', '')).strip()
-
- if status == 'NOT_FOUND':
- return 'NOT_FOUND', '', reason or 'LLM returned NOT_FOUND without reason'
-
- # Any status other than NOT_FOUND โ treat as DONE
- if not answer:
- answer = reason or '(empty answer)'
- return 'DONE', answer, ''
-
-
-def build_attempt_answer_prompt(
- *,
- query: str,
- evidence_text: str,
- state: AgentState,
- config: AgentRunConfig,
- budget_snapshot: dict | None = None,
-) -> str:
- """Build the final answer prompt so trimming can estimate it beforehand."""
- planning = _budget_line_parts(budget_snapshot, 'planning')
- context = _budget_line_parts(budget_snapshot, 'context')
- return _ATTEMPT_ANSWER_PROMPT.format(
- query=query,
- evidence_context=evidence_text,
- revision_count=state.revision_count,
- max_revisions=config.max_revisions,
- planning_status=planning['status'],
- planning_remaining=planning['remaining'],
- planning_capacity=planning['capacity'],
- planning_used_pct=planning['used_pct'],
- planning_remaining_pct=planning['remaining_pct'],
- context_status=context['status'],
- context_remaining=context['remaining'],
- context_capacity=context['capacity'],
- context_used_pct=context['used_pct'],
- context_remaining_pct=context['remaining_pct'],
- )
-
-
-def estimate_attempt_answer_prompt_tokens(
- *,
- query: str,
- evidence_text: str,
- state: AgentState,
- config: AgentRunConfig,
- budget_snapshot: dict | None = None,
-) -> int:
- """Estimate the exact prompt shape that will be charged to context budget."""
- return estimate_tokens(build_attempt_answer_prompt(
- query=query,
- evidence_text=evidence_text,
- state=state,
- config=config,
- budget_snapshot=budget_snapshot,
- ))
-
-
-_ATTEMPT_ANSWER_PROMPT = """\
-You are a knowledge retrieval assistant. Answer the user's query based
-STRICTLY on the provided evidence. Do NOT use any external knowledge.
-
-QUERY: "{query}"
-
-EVIDENCE CONTEXT:
-The following evidence is organized by document in a unified hierarchy.
-Each document shows its structural outline (section titles + summaries)
-with retrieved content (โ lines) inline under the relevant sections.
-
-{evidence_context}
-
-REVISION: {revision_count} of {max_revisions} revisions used.
-Planning budget: {planning_status} ({planning_used_pct}% used, {planning_remaining_pct}% remaining, {planning_remaining}/{planning_capacity} remaining).
-Context budget: {context_status} ({context_used_pct}% used, {context_remaining_pct}% remaining, {context_remaining}/{context_capacity} remaining); the evidence may have been trimmed.
-
-INSTRUCTIONS:
-1. If the evidence contains enough information to answer the query,
- compose a clear and comprehensive answer. Return:
- {{"status": "DONE", "answer": ""}}
-
-2. If the evidence does NOT contain sufficient information to answer
- the query, return NOT_FOUND with a specific reason explaining what
- information is missing. This reason will be used to guide the next
- search round. Return:
- {{"status": "NOT_FOUND", "reason": ""}}
-
-IMPORTANT:
-- Base your judgment ONLY on the actual retrieved content (โ lines),
- not just section titles or summaries.
-- Be specific in your NOT_FOUND reason โ mention exactly what data,
- section, or detail you expected but didn't find.
-- When the evidence partially covers the query, still return DONE with
- the available information and note any gaps in your answer.
-
-Return ONLY a JSON object, no other text.
-"""
diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py
index e48e82901..1002a68b1 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/prompts.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py
@@ -18,7 +18,6 @@
=== End Overview ===
User query: {query}
-{revision_context}
Based on the query, select documents that may contain relevant information.
If NO document in the corpus is relevant to the query, return an EMPTY array [].
Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"]
@@ -41,7 +40,6 @@
=== End Discovery Candidates ===
User query: {query}
-{revision_context}
Select section paths whose content is needed to answer the query.
If none are relevant, return an EMPTY list [].
@@ -59,8 +57,8 @@
{budget_block}
{scope_header}
Below is the document's section tree.
-Sections tagged [SELECT] are within the current scope and may be selected.
-Other sections are shown as structural context only (not selectable).
+Sections tagged [SELECT] are the recommended selection granularity for this scope.
+Other visible sections are structural context and may be selected when you need to drill into that broader scope.
Nodes marked [Leaf] have no further sub-sections.
=== Section Tree ===
@@ -75,7 +73,7 @@
NAVIGATE โ Drill into selected sections for detailed content.
Consider this when the query targets specific topics and you need deeper text evidence.
- Select one or more [SELECT] sections.
+ Prefer one or more [SELECT] sections, or choose a broader visible section when needed.
STOP โ Current scope evidence is sufficient. No further drill-down.
Consider this when:
@@ -86,7 +84,7 @@
{tools_block}
When action is NAVIGATE, provide selections:
-- You may ONLY select sections marked with [SELECT].
+- Select visible section paths from the tree above; prefer [SELECT] paths when they fit.
When action is STOP, selections must be empty.
diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py
index bce60a71b..8daad41ba 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py
@@ -60,7 +60,6 @@ async def kg_document_select(
query: str,
llm_fn: LLMFn | None,
exclude_document_ids: list[str],
- revision_hint: str | None = None,
**kwargs: Any,
) -> ToolResult:
return await discovery_tools.kg_document_select(
@@ -70,7 +69,6 @@ async def kg_document_select(
query=query,
llm_fn=llm_fn,
exclude_document_ids=exclude_document_ids,
- revision_hint=revision_hint,
**kwargs,
)
@@ -104,7 +102,6 @@ async def navigate_step(
doc_name: str = "",
scope_path: str | list[str] | None = None,
exclude_paths: set[str] | None = None,
- revision_hint: str | None = None,
budget_snapshot: dict | None = None,
) -> tuple[str, list[str], DocTreeNode, list[dict]]:
return await navigation_tools.navigate_step(
@@ -118,7 +115,6 @@ async def navigate_step(
doc_name=doc_name,
scope_path=scope_path,
exclude_paths=exclude_paths,
- revision_hint=revision_hint,
budget_snapshot=budget_snapshot,
)
@@ -134,7 +130,6 @@ async def discovery_select_step(
doc_name: str = "",
discovery_hints: list[dict[str, Any]],
exclude_paths: set[str] | None = None,
- revision_hint: str | None = None,
budget_snapshot: dict | None = None,
) -> DocTreeNode:
return await discovery_selection.discovery_select_step(
@@ -147,6 +142,5 @@ async def discovery_select_step(
doc_name=doc_name,
discovery_hints=discovery_hints,
exclude_paths=exclude_paths,
- revision_hint=revision_hint,
budget_snapshot=budget_snapshot,
)
diff --git a/packages/shared-python/shared/services/retrieval/execution/legacy_route.py b/packages/shared-python/shared/services/retrieval/execution/legacy_route.py
index 6b1976f53..b0d5dd9e1 100644
--- a/packages/shared-python/shared/services/retrieval/execution/legacy_route.py
+++ b/packages/shared-python/shared/services/retrieval/execution/legacy_route.py
@@ -10,6 +10,7 @@
from shared.services.retrieval.graph.query_service import GraphQueryService
from shared.services.retrieval.search.ranking import rank_retrieval_candidates
from shared.services.retrieval.execution.response_projection import attach_citation
+from shared.services.retrieval.hydration.legacy_evidence import render_legacy_evidence_text
from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results
from shared.services.retrieval.execution.route_types import (
RetrievalRouteContext,
@@ -77,6 +78,8 @@ async def run_legacy_retrieval_route(
"namespace": context.namespace,
"query": context.query,
"router_used": router_used,
+ "evidence_text": render_legacy_evidence_text(results),
+ "answer_text": "",
"results": results,
}
return RetrievalRouteOutcome(
diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py
index 89fc85048..2dda78192 100644
--- a/packages/shared-python/shared/services/retrieval/execution/plan.py
+++ b/packages/shared-python/shared/services/retrieval/execution/plan.py
@@ -86,6 +86,9 @@ async def execute(self) -> dict[str, Any]:
"namespace": request.namespace,
"query": request.query,
"router_used": "empty_query_filtered",
+ "evidence_text": "",
+ "answer_text": "",
+ "referenced_chunks": [],
"results": [],
}
diff --git a/packages/shared-python/shared/services/retrieval/execution/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py
index 90e811972..9c50a7e97 100644
--- a/packages/shared-python/shared/services/retrieval/execution/response_projection.py
+++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py
@@ -35,15 +35,12 @@ async def project_public_retrieval_response(response: dict[str, Any]) -> dict[st
'namespace': response.get('namespace'),
'query': response.get('query'),
'router_used': response.get('router_used'),
+ 'evidence_text': response.get('evidence_text') or '',
+ 'answer_text': '',
+ 'referenced_chunks': response.get('referenced_chunks') or [],
'results': [],
}
- if response.get('answer_text') is not None:
- public_response['answer_text'] = response['answer_text']
- if response.get('referenced_chunks') is not None:
- public_response['referenced_chunks'] = response['referenced_chunks']
- if response.get('evidence_text') is not None:
- public_response['evidence_text'] = response['evidence_text']
if response.get('stop_reason') is not None:
public_response['stop_reason'] = response['stop_reason']
if response.get('failure_reason') is not None:
diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py
index efb24f240..36e227712 100644
--- a/packages/shared-python/shared/services/retrieval/execution/routes.py
+++ b/packages/shared-python/shared/services/retrieval/execution/routes.py
@@ -10,6 +10,7 @@
from shared.services.retrieval.execution.response_projection import (
attach_citation,
)
+from shared.services.retrieval.hydration.legacy_evidence import render_legacy_evidence_text
from shared.services.retrieval.execution.route_types import (
RetrievalRouteContext,
RetrievalRouteOutcome,
@@ -83,6 +84,8 @@ async def _try_run_small_corpus_route(
"namespace": context.namespace,
"query": context.query,
"router_used": "small_corpus_all",
+ "evidence_text": render_legacy_evidence_text(results),
+ "answer_text": "",
"results": results,
}
return RetrievalRouteOutcome(
@@ -126,6 +129,7 @@ async def _run_agentic_route(
allowed_chunk_types=context.allowed_chunk_types,
)
response = workflow_result.to_api_response()
+ response["answer_text"] = ""
response["referenced_chunks"] = resolved_references.refs
response["results"] = [attach_citation(row) for row in assembled_workflow_rows]
@@ -148,7 +152,7 @@ async def _run_agentic_route(
response["failure_reason"] = last_retrieve.failure_reason
completion_detail = (
- f"chunks | answer={len(workflow_result.answer_text)} chars | "
+ f"chunks | evidence={len(response.get('evidence_text') or '')} chars | "
f"router={workflow_result.router_used}"
)
return RetrievalRouteOutcome(
diff --git a/packages/shared-python/shared/services/retrieval/hydration/legacy_evidence.py b/packages/shared-python/shared/services/retrieval/hydration/legacy_evidence.py
new file mode 100644
index 000000000..63b4ea358
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/hydration/legacy_evidence.py
@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any
+
+
+def render_legacy_evidence_text(rows: list[dict[str, Any]]) -> str:
+ """Render assembled retrieval rows into evidence-only context."""
+ grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
+ for row in rows:
+ doc_name = _source_value(row, "source_file_name") or "Unknown document"
+ grouped_rows[doc_name].append(row)
+
+ parts: list[str] = []
+ for doc_name in sorted(grouped_rows):
+ parts.append(f"[Document] {doc_name}")
+ last_section = object()
+ for row in sorted(grouped_rows[doc_name], key=_row_sort_key):
+ section_path = _source_value(row, "section_path") or doc_name
+ if section_path != last_section:
+ parts.append(f"โธ {section_path}")
+ last_section = section_path
+ _append_content_lines(parts, row.get("content"))
+
+ return "\n".join(parts)
+
+
+def _source_value(row: dict[str, Any], key: str) -> str:
+ source = row.get("source")
+ if isinstance(source, dict):
+ value = source.get(key)
+ if value:
+ return str(value)
+ value = row.get(key)
+ return str(value) if value else ""
+
+
+def _row_sort_key(row: dict[str, Any]) -> tuple[str, int, str]:
+ section_path = _source_value(row, "section_path")
+ try:
+ sort_order = int(row.get("sort_order") or 0)
+ except (TypeError, ValueError):
+ sort_order = 0
+ chunk_id = str(row.get("chunk_id") or "")
+ return section_path, sort_order, chunk_id
+
+
+def _append_content_lines(parts: list[str], content: object) -> None:
+ text = str(content or "").strip()
+ if not text:
+ return
+ for line in text.splitlines():
+ stripped = line.strip()
+ if stripped:
+ parts.append(f" โ {stripped}")
diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
index 234927d04..619404181 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
@@ -21,7 +21,6 @@
from shared.services.retrieval.workflow.run_request import WorkflowRunRequest
from shared.services.retrieval.workflow.runtime_config import WorkflowRuntimeConfig
from shared.services.retrieval.workflow.step_runner import WorkflowStepRunner
-from shared.services.retrieval.workflow.synthesizer import compose_final_answer
from shared.services.retrieval.workflow.types import StepResult, WorkflowResult
from shared.services.retrieval.workflow.wallet import BudgetWallet
@@ -140,7 +139,6 @@ async def run_request(
wallet = BudgetWallet(
total=config.wallet_total_budget,
per_retrieve_step_default=config.per_retrieve_step_budget,
- per_synthesize_step_default=config.per_synthesize_step_budget,
)
ledgers = await wallet.allocate(plan)
results_by_id: dict[str, StepResult] = {}
@@ -167,8 +165,8 @@ async def run_request(
for step in batch:
await wallet.reclaim(step.id, ledgers[step.id])
- answer_text = compose_final_answer(plan, results_by_id)
ordered_results = [results_by_id[step.id] for step in plan.steps if step.id in results_by_id]
+ evidence_chars = sum(len(step_result.evidence_text or "") for step_result in ordered_results)
reference_projection = WorkflowReferenceProjection()
referenced_chunks = reference_projection.dedupe(
ref for step_result in ordered_results for ref in step_result.referenced_chunks
@@ -176,17 +174,17 @@ async def run_request(
api_results = reference_projection.to_api_results(referenced_chunks)
elapsed_ms = int((time.monotonic() - t0) * 1000)
logger.info(
- 'workflow retrieval DONE: steps={} refs={} answer_chars={} elapsed={}ms',
+ 'workflow retrieval DONE: steps={} refs={} evidence_chars={} elapsed={}ms',
len(ordered_results),
len(referenced_chunks),
- len(answer_text),
+ evidence_chars,
elapsed_ms,
)
return WorkflowResult(
namespace=request.namespace,
query=request.query,
router_used='workflow_decomposed' if len(plan.steps) > 1 else 'workflow_single_step',
- answer_text=answer_text,
+ answer_text="",
plan=plan,
steps=ordered_results,
referenced_chunks=referenced_chunks,
diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py
index 82b574677..ae413f658 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/planner.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py
@@ -10,7 +10,7 @@
from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetLedger
from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage
-from shared.services.retrieval.workflow.types import FinalStrategy, OutputRole, PlannedStep, QueryPlan, StepKind
+from shared.services.retrieval.workflow.types import OutputRole, PlannedStep, QueryPlan, StepKind
from shared.utils.token_estimate import estimate_tokens
@@ -27,7 +27,6 @@
}
],
"final_strategy": "concat_final_parts",
- "final_template": "",
}
_PLANNER_PROMPT = """\
@@ -48,10 +47,10 @@
- max_steps = {max_steps}
- Each retrieve step costs ~{per_step_budget} tokens; do NOT plan more
retrieve steps than the wallet can afford.
- - synthesize steps must have non-empty depends_on.
- - final_strategy must be one of: concat_final_parts, last_synthesize, template.
- - step_kind must be retrieve or synthesize.
- - output_role must be final_part, intermediate, or consumed_by_synthesis.
+ - final_strategy must be concat_final_parts.
+ - step_kind must be retrieve.
+ - output_role must be final_part or intermediate.
+ - KNOWHERE returns evidence only; do not plan answer synthesis steps.
Return ONLY a JSON object matching this schema (think first, then answer):
{schema}
@@ -191,13 +190,11 @@ def _parse_plan_response(text: str, *, original_query: str, max_steps: int) -> Q
)
)
- final_strategy = _coerce_final_strategy(data.get("final_strategy"))
return QueryPlan(
original_query=original_query,
steps=steps,
- final_strategy=final_strategy,
+ final_strategy="concat_final_parts",
reasoning_summary=str(data.get("reasoning_summary") or "").strip(),
- final_template=str(data.get("final_template") or "").strip() or None,
)
@@ -220,20 +217,13 @@ def _extract_json_object(text: str) -> dict[str, Any]:
def _coerce_step_kind(value: Any) -> StepKind:
raw = str(value or "retrieve").strip().lower()
- if raw not in {"retrieve", "synthesize"}:
+ if raw != "retrieve":
raise ValueError(f"unsupported step_kind: {value}")
return raw # type: ignore[return-value]
def _coerce_output_role(value: Any) -> OutputRole:
raw = str(value or "final_part").strip().lower()
- if raw not in {"final_part", "intermediate", "consumed_by_synthesis"}:
+ if raw not in {"final_part", "intermediate"}:
raise ValueError(f"unsupported output_role: {value}")
return raw # type: ignore[return-value]
-
-
-def _coerce_final_strategy(value: Any) -> FinalStrategy:
- raw = str(value or "concat_final_parts").strip().lower()
- if raw not in {"concat_final_parts", "last_synthesize", "template"}:
- raise ValueError(f"unsupported final_strategy: {value}")
- return raw # type: ignore[return-value]
diff --git a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py
index bcb6ab262..28630dff0 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/runtime_config.py
@@ -10,7 +10,6 @@ class WorkflowRuntimeConfig:
planner_budget: int = 4000
wallet_total_budget: int = 200000
per_retrieve_step_budget: int = 40000
- per_synthesize_step_budget: int = 6000
max_steps: int = 5
parallel_max: int = 3
@@ -20,7 +19,6 @@ def from_env(cls) -> "WorkflowRuntimeConfig":
planner_budget=_env_int("RETRIEVAL_PLANNER_THINKING_BUDGET", 4000),
wallet_total_budget=_env_int("RETRIEVAL_WALLET_TOTAL_BUDGET", 200000),
per_retrieve_step_budget=_env_int("RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET", 40000),
- per_synthesize_step_budget=_env_int("RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET", 6000),
max_steps=_env_int("RETRIEVAL_DECOMPOSITION_MAX_STEPS", 5),
parallel_max=_env_int("RETRIEVAL_WORKFLOW_PARALLEL_MAX", 3),
)
diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
index 18a2e7059..c14328aaf 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
@@ -12,9 +12,7 @@
from shared.services.retrieval.agentic.orchestrator import RetrievalAgent
from shared.services.retrieval.agentic.core.types import AgenticResult
from shared.services.retrieval.llm_adapter import LLMFn
-from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection
from shared.services.retrieval.workflow.run_request import WorkflowStepRequest
-from shared.services.retrieval.workflow.synthesizer import synthesize_step
from shared.services.retrieval.workflow.types import PlannedStep, StepResult, StepStatus
DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]]
@@ -32,7 +30,6 @@ def __init__(
self._db_factory = db_factory
self._parent_run_id = parent_run_id
self._agent_factory = agent_factory or RetrievalAgent
- self._references = WorkflowReferenceProjection()
async def run_step(
self,
@@ -45,9 +42,6 @@ async def run_step(
llm_fn: LLMFn | None,
) -> None:
async with semaphore:
- if step.step_kind == "synthesize":
- await self._run_synthesize_step(step, ledger, results_by_id, llm_fn)
- return
await self._run_retrieve_step(
step=step,
ledger=ledger,
@@ -96,64 +90,15 @@ async def _run_retrieve_step(
depends_on=step.depends_on,
output_role=step.output_role,
status="error",
- error=str(exc),
- budget_snapshot=ledger.snapshot(),
- )
-
- async def _run_synthesize_step(
- self,
- step: PlannedStep,
- ledger: BudgetLedger,
- results_by_id: dict[str, StepResult],
- llm_fn: LLMFn | None,
- ) -> None:
- if llm_fn is None:
- results_by_id[step.id] = StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status="skipped",
answer_text="",
- error="llm unavailable for synthesis",
- budget_snapshot=ledger.snapshot(),
- )
- return
- prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id}
- try:
- answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger)
- refs = self._references.dedupe(ref for result in prior.values() for ref in result.referenced_chunks)
- results_by_id[step.id] = StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status="done",
- answer_text=answer,
- referenced_chunks=refs,
- budget_snapshot=ledger.snapshot(),
- )
- except Exception as exc:
- results_by_id[step.id] = StepResult(
- step_id=step.id,
- sub_query=step.sub_query,
- step_kind=step.step_kind,
- depends_on=step.depends_on,
- output_role=step.output_role,
- status="budget_stop" if "budget" in str(exc).lower() else "error",
- answer_text="(budget exhausted)" if "budget" in str(exc).lower() else "",
error=str(exc),
budget_snapshot=ledger.snapshot(),
)
def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult:
- if result.answer_text:
- status: StepStatus = "done"
- elif result.failure_reason:
- status = "not_found"
+ if result.failure_reason:
+ status: StepStatus = "not_found"
elif "budget" in (result.stop_reason or ""):
status = "budget_stop"
else:
@@ -165,7 +110,7 @@ def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepR
depends_on=step.depends_on,
output_role=step.output_role,
status=status,
- answer_text=result.answer_text,
+ answer_text="",
evidence_text=result.evidence_text,
referenced_chunks=result.referenced_chunks,
budget_snapshot=result.budget_snapshot,
diff --git a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py
deleted file mode 100644
index a81bb517d..000000000
--- a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py
+++ /dev/null
@@ -1,141 +0,0 @@
-"""Synthesis helpers for decomposed retrieval workflows."""
-from __future__ import annotations
-
-import re
-
-from shared.services.retrieval.agentic.core.budget import BudgetLedger
-from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage
-from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult
-from shared.utils.token_estimate import estimate_tokens
-
-
-_SYNTHESIZE_PROMPT = """\
-You are composing an intermediate or final answer for a retrieval workflow.
-Use ONLY the prior step outputs below. Do not use external knowledge.
-
-Current step id: {step_id}
-Current step task: {sub_query}
-
-Prior step outputs:
-{prior_outputs}
-
-Return a concise answer that directly satisfies the current step task.
-If the prior outputs are insufficient, explain the missing information.
-"""
-
-
-async def synthesize_step(
- step: PlannedStep,
- *,
- prior_results: dict[str, StepResult],
- llm_fn: LLMFn,
- ledger: BudgetLedger | None,
-) -> str:
- prior_outputs = _format_prior_outputs(step.depends_on, prior_results)
- prompt = _SYNTHESIZE_PROMPT.format(
- step_id=step.id,
- sub_query=step.sub_query,
- prior_outputs=prior_outputs,
- )
- if ledger is None:
- return (await llm_fn(prompt)).strip()
-
- est = estimate_tokens(prompt)
- reserved = await ledger.try_reserve("context", est)
- if not reserved:
- raise RuntimeError("synthesis context budget exhausted")
- try:
- response = await llm_fn(prompt)
- except Exception:
- await ledger.refund("context", est=est)
- raise
- usage = current_llm_usage.get() or {}
- actual = int(usage.get("prompt_tokens") or est)
- await ledger.commit("context", actual=actual, est=est)
- return response.strip()
-
-
-def compose_final_answer(plan: QueryPlan, results: dict[str, StepResult]) -> str:
- """Compose workflow final answer according to planner strategy."""
- if plan.final_strategy == "last_synthesize":
- for step in reversed(plan.steps):
- result = results.get(step.id)
- if result and step.step_kind == "synthesize" and result.status == "done":
- return result.answer_text
- return _concat_final_parts(plan, results)
-
- if plan.final_strategy == "template" and plan.final_template:
- return _render_template(plan.final_template, results).strip()
-
- return _concat_final_parts(plan, results)
-
-
-def _concat_final_parts(plan: QueryPlan, results: dict[str, StepResult]) -> str:
- parts: list[str] = []
- for step in plan.steps:
- if step.output_role != "final_part":
- continue
- result = results.get(step.id)
- if not result or result.status not in ("done", "budget_stop") or not result.answer_text:
- continue
- parts.append(result.answer_text.strip())
- if parts:
- return "\n\n".join(parts)
-
- fallback_parts = [
- result.answer_text.strip()
- for step in plan.steps
- if (result := results.get(step.id)) and result.answer_text.strip()
- ]
- if fallback_parts:
- return "\n\n".join(fallback_parts)
-
- missing_reasons = [
- result.failure_reason.strip()
- for step in plan.steps
- if (result := results.get(step.id))
- and result.status == "not_found"
- and result.failure_reason.strip()
- ]
- if missing_reasons:
- return "ๆช่ฝๅบไบๅฝๅ็ฅ่ฏๅบ่ฏๆฎๅ็ญ่ฏฅ้ฎ้ข๏ผ" + "๏ผ".join(dict.fromkeys(missing_reasons))
-
- if any((result := results.get(step.id)) and result.status == "budget_stop" for step in plan.steps):
- return "Unable to return a valid answer because the retrieval budget was exhausted."
-
- return ""
-
-
-def _format_prior_outputs(depends_on: list[str], prior_results: dict[str, StepResult]) -> str:
- lines: list[str] = []
- for step_id in depends_on:
- result = prior_results.get(step_id)
- if not result:
- lines.append(f"## {step_id}\n(status: missing)\n")
- continue
- lines.append(
- "\n".join(
- [
- f"## {step_id}",
- f"Sub-query: {result.sub_query}",
- f"Status: {result.status}",
- "Answer:",
- result.answer_text or "(empty)",
- "",
- ]
- )
- )
- return "\n".join(lines) if lines else "(no prior outputs)"
-
-
-def _render_template(template: str, results: dict[str, StepResult]) -> str:
- def _replace(match: re.Match[str]) -> str:
- step_id = match.group(1)
- field = match.group(2)
- result = results.get(step_id)
- if not result:
- return ""
- attr = "answer_text" if field == "answer" else field
- return str(getattr(result, attr, ""))
-
- return re.sub(r"\{\{\s*steps\.([A-Za-z0-9_-]+)\.(answer_text|answer|evidence_text|status)\s*\}\}", _replace, template)
diff --git a/packages/shared-python/shared/services/retrieval/workflow/types.py b/packages/shared-python/shared/services/retrieval/workflow/types.py
index b9147856f..0141d05c1 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/types.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/types.py
@@ -5,9 +5,9 @@
from typing import Any, Literal
-StepKind = Literal["retrieve", "synthesize"]
-OutputRole = Literal["final_part", "intermediate", "consumed_by_synthesis"]
-FinalStrategy = Literal["concat_final_parts", "last_synthesize", "template"]
+StepKind = Literal["retrieve"]
+OutputRole = Literal["final_part", "intermediate"]
+FinalStrategy = Literal["concat_final_parts"]
StepStatus = Literal["done", "skipped", "error", "budget_stop", "not_found"]
@@ -49,7 +49,6 @@ class QueryPlan:
steps: list[PlannedStep]
final_strategy: FinalStrategy = "concat_final_parts"
reasoning_summary: str = ""
- final_template: str | None = None
planner_status: str = "planned"
planner_error: str | None = None
@@ -89,9 +88,8 @@ def from_dict(data: dict[str, Any], *, original_query: str | None = None) -> "Qu
plan = QueryPlan(
original_query=str(data.get("original_query") or original_query or ""),
steps=steps,
- final_strategy=data.get("final_strategy", "concat_final_parts"),
+ final_strategy="concat_final_parts",
reasoning_summary=str(data.get("reasoning_summary") or ""),
- final_template=data.get("final_template"),
planner_status=str(data.get("planner_status") or "cached"),
planner_error=data.get("planner_error"),
)
@@ -106,8 +104,6 @@ def to_dict(self) -> dict[str, Any]:
"final_strategy": self.final_strategy,
"planner_status": self.planner_status,
}
- if self.final_template:
- data["final_template"] = self.final_template
if self.planner_error:
data["planner_error"] = self.planner_error
return data
@@ -127,15 +123,13 @@ def validate(self) -> None:
raise ValueError("query plan step id cannot be empty")
if not step.sub_query.strip():
raise ValueError(f"query plan step {step.id} sub_query cannot be empty")
- if step.step_kind not in ("retrieve", "synthesize"):
+ if step.step_kind != "retrieve":
raise ValueError(f"unsupported step_kind: {step.step_kind}")
- if step.output_role not in ("final_part", "intermediate", "consumed_by_synthesis"):
+ if step.output_role not in ("final_part", "intermediate"):
raise ValueError(f"unsupported output_role: {step.output_role}")
missing = [dep for dep in step.depends_on if dep not in id_set]
if missing:
raise ValueError(f"step {step.id} depends on unknown steps: {missing}")
- if step.step_kind == "synthesize" and not step.depends_on:
- raise ValueError(f"synthesize step {step.id} must depend on prior steps")
self.topological_batches()
def topological_batches(self) -> list[list[PlannedStep]]:
diff --git a/packages/shared-python/shared/services/retrieval/workflow/wallet.py b/packages/shared-python/shared/services/retrieval/workflow/wallet.py
index 624f03f7c..da9d67eac 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/wallet.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/wallet.py
@@ -10,7 +10,6 @@
_RETRIEVE_FLOOR = 4000
-_SYNTHESIZE_FLOOR = 1500
def _env_float(name: str, default: float) -> float:
@@ -33,7 +32,6 @@ class BudgetWallet:
total: int
per_retrieve_step_default: int
- per_synthesize_step_default: int
# Read from env for consistency with _build_config_from_env()
planning_ratio: float = field(
default_factory=lambda: _env_float('RETRIEVAL_AGENTIC_PLANNING_RATIO', 0.5)
@@ -66,8 +64,7 @@ async def allocate(self, plan: QueryPlan) -> dict[str, BudgetLedger]:
scale = max(self.total, 1) / requested_total
allocations = {}
for step in plan.steps:
- floor = _RETRIEVE_FLOOR if step.step_kind == "retrieve" else _SYNTHESIZE_FLOOR
- allocations[step.id] = max(floor, int(requested[step.id] * scale))
+ allocations[step.id] = max(_RETRIEVE_FLOOR, int(requested[step.id] * scale))
scaled_total = sum(allocations.values())
if scaled_total > self.total:
@@ -79,9 +76,7 @@ async def allocate(self, plan: QueryPlan) -> dict[str, BudgetLedger]:
):
if excess <= 0:
break
- step = next(s for s in plan.steps if s.id == step_id)
- floor = _RETRIEVE_FLOOR if step.step_kind == "retrieve" else _SYNTHESIZE_FLOOR
- reducible = max(amount - floor, 0)
+ reducible = max(amount - _RETRIEVE_FLOOR, 0)
delta = min(reducible, excess)
allocations[step_id] = amount - delta
excess -= delta
@@ -118,19 +113,11 @@ def snapshot(self) -> dict[str, object]:
}
def _requested_for_step(self, step: PlannedStep) -> int:
- if step.step_kind == "synthesize":
- return max(self.per_synthesize_step_default, _SYNTHESIZE_FLOOR)
+ del step
return max(self.per_retrieve_step_default, _RETRIEVE_FLOOR)
def _new_ledger(self, step: PlannedStep, total: int) -> BudgetLedger:
- if step.step_kind == "synthesize":
- # Put almost all tokens into context for pure synthesis calls.
- return BudgetLedger(
- total=max(total, 1),
- planning_ratio=0.0,
- bootstrap=0,
- per_doc_min_share=0,
- )
+ del step
return BudgetLedger(
total=max(total, 1),
planning_ratio=self.planning_ratio,
From ee7259b6e1f6d4560f38494df6b47a2d01af7225 Mon Sep 17 00:00:00 2001
From: chengke <404835780@qq.com>
Date: Fri, 29 May 2026 11:05:10 +0800
Subject: [PATCH 3/4] feat: implement navigation decision tracing and refactor
navigation steps to return structured NavigateStepResult
---
AGENTS.md | 3 +-
apps/api/app/api/v1/routes/retrieval.py | 9 +++
apps/api/app/mcp/retrieval_server.py | 5 ++
.../services/retrieval/agentic/core/types.py | 21 ++++++
.../retrieval/agentic/navigation/document.py | 70 +++++++++++++++----
.../agentic/navigation/section_tree.py | 4 ++
.../retrieval/agentic/navigation/tools.py | 19 +++--
.../retrieval/agentic/orchestrator.py | 30 +++++++-
.../services/retrieval/agentic/prompts.py | 19 ++---
.../execution/response_projection.py | 2 +
.../services/retrieval/execution/routes.py | 8 +++
.../retrieval/workflow/step_runner.py | 1 +
.../services/retrieval/workflow/types.py | 2 +
13 files changed, 161 insertions(+), 32 deletions(-)
diff --git a/AGENTS.md b/AGENTS.md
index 397965b26..a625cdf56 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -682,8 +682,7 @@ cd apps/worker && uv run worker.py # Celery worker
| Script | Purpose |
|:---|:---|
-| `debug_parse.py` | End-to-end parsing with `MockRedis`, `LOCAL_DEBUG=1` |
-| `debug_hierarchy_llm.py` | Test heading recognition LLM calls |
+| `debug_parse.py` | Unified parsing debug: all formats, `--stop-at profile/hierarchy/full`, `--run-db` |
| `debug_agentic_e2e.py` | End-to-end agentic retrieval test |
| `debug_profiler.py` | Document profiler testing |
| `debug_toc_detection.py` | TOC detection and hierarchy building |
diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py
index e1cd88cb3..9f8efabb0 100644
--- a/apps/api/app/api/v1/routes/retrieval.py
+++ b/apps/api/app/api/v1/routes/retrieval.py
@@ -95,6 +95,15 @@ class RetrievalQueryResponse(BaseModel):
results: list[dict] = Field(default_factory=list)
stop_reason: str | None = None
failure_reason: str | None = None
+ decision_trace: list[dict] | None = Field(
+ default=None,
+ description=(
+ "Per-step navigation decisions from agentic retrieval. "
+ "Each entry has phase, document, action, reason, stop_type, "
+ "selected_paths, and hydrated_count. Use this to understand "
+ "why KNOWHERE stopped or made specific navigation choices."
+ ),
+ )
@router.post("/query", response_model=RetrievalQueryResponse)
diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py
index e702411de..2f5846fa0 100644
--- a/apps/api/app/mcp/retrieval_server.py
+++ b/apps/api/app/mcp/retrieval_server.py
@@ -74,6 +74,11 @@ def to_mcp_query_response(response: dict[str, Any]) -> dict[str, Any]:
"evidence_text": response.get("evidence_text") or "",
}
+ if response.get("stop_reason") is not None:
+ mcp_response["stop_reason"] = response["stop_reason"]
+ if response.get("decision_trace") is not None:
+ mcp_response["decision_trace"] = response["decision_trace"]
+
return mcp_response
diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/types.py b/packages/shared-python/shared/services/retrieval/agentic/core/types.py
index 9547ba584..7eb2dd0e5 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/core/types.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/core/types.py
@@ -163,6 +163,24 @@ def merge(self, other: 'DocTreeNode') -> None:
self.reparent_leaf_content()
+@dataclass
+class NavigateStepResult:
+ """Return type for navigate_step โ typed replacement for raw tuple."""
+ action: str # "NAVIGATE" or "STOP"
+ tools: list[str] = field(default_factory=list)
+ node: DocTreeNode = field(default_factory=DocTreeNode)
+ pending: list[dict[str, Any]] = field(default_factory=list)
+ reason: str = ""
+ stop_type: str = "" # only for STOP: sufficient_outline | no_relevant_child | ...
+
+ @staticmethod
+ def stop(scope_path: str | None = None) -> 'NavigateStepResult':
+ return NavigateStepResult(
+ action="STOP",
+ node=DocTreeNode.empty(scope_path),
+ )
+
+
@dataclass
class CandidateDoc:
"""A document selected by kg_document_select."""
@@ -188,6 +206,8 @@ class AgenticResult:
- ``stop_reason``: why the run terminated (evidence_only /
latency_budget / context_budget / no_llm / etc.)
- ``failure_reason``: fatal retrieval failure reason, if any.
+ - ``decision_trace``: per-step navigation decisions with reasons,
+ exposed to downstream agents for stop/retry/modify-query decisions.
"""
evidence_text: str
answer_text: str = ''
@@ -196,6 +216,7 @@ class AgenticResult:
budget_snapshot: dict[str, Any] | None = None
stop_reason: str = ''
failure_reason: str = ''
+ decision_trace: list[dict[str, Any]] = field(default_factory=list)
@dataclass
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py
index 6dd1c1e08..f7e711fa5 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py
@@ -16,6 +16,7 @@
AgentState,
CandidateDoc,
DocTreeNode,
+ NavigateStepResult,
ToolResult,
)
from shared.services.retrieval.llm_adapter import LLMFn
@@ -48,6 +49,11 @@ def __init__(
self._discovery_by_doc = discovery_by_doc
self._llm_fn = llm_fn
self._llm_budget = llm_budget
+ self._decision_steps: list[dict[str, Any]] = []
+
+ @property
+ def decision_steps(self) -> list[dict[str, Any]]:
+ return list(self._decision_steps)
async def navigate_selected_documents(self) -> None:
logger.info(
@@ -135,7 +141,7 @@ async def _navigate_bfs(
depth=depth,
)
try:
- action, asset_tools, step_node, drill_paths = await tools.navigate_step(
+ nav_result = await tools.navigate_step(
self._db,
document_id=doc.document_id,
job_result_id=job_result_id,
@@ -158,27 +164,24 @@ async def _navigate_bfs(
await self._collect_assets(
doc=doc,
scope=scope,
- step_node=step_node,
- asset_tools=asset_tools,
+ step_node=nav_result.node,
+ asset_tools=nav_result.tools,
pending_assets=doc_pending_assets,
round_scope="nav",
)
- _merge_step_node(parent_node, step_node)
- _update_excluded_leaf_paths(doc_exclude, step_node, drill_paths)
- _queue_drill_paths(pending, parent_node, drill_paths, depth)
+ _merge_step_node(parent_node, nav_result.node)
+ _update_excluded_leaf_paths(doc_exclude, nav_result.node, nav_result.pending)
+ _queue_drill_paths(pending, parent_node, nav_result.pending, depth)
parent_node.reparent_leaf_content()
self._record_navigation_step(
doc=doc,
scope=scope,
depth=depth,
- action=action,
- asset_tools=asset_tools,
- step_node=step_node,
- drill_paths=drill_paths,
+ nav_result=nav_result,
)
if self._state.ledger is not None:
self._state.ledger.mark_explored(
- chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()),
+ chunks=sum(len(chunks) for chunks in nav_result.node.leaf_content.values()),
)
return doc_pending_assets
@@ -287,6 +290,16 @@ async def _hydrate_discovery_hints(
),
decision_reason=f"discovery_{doc.source_file_name}",
)
+ self._decision_steps.append({
+ "phase": "discovery_select",
+ "document": doc_name,
+ "document_id": doc.document_id,
+ "action": "select" if discovery_node.has_content() else "skip",
+ "reason": "",
+ "candidate_count": len(doc_hints),
+ "hydrated_count": len(discovery_node.leaf_content),
+ "selected_paths": list(discovery_node.leaf_content.keys()),
+ })
root.merge(discovery_node)
if self._state.ledger is not None:
self._state.ledger.mark_explored(
@@ -332,11 +345,17 @@ def _record_navigation_step(
doc: CandidateDoc,
scope: str | list[str] | None,
depth: int,
- action: str,
- asset_tools: list[str],
- step_node: DocTreeNode,
- drill_paths: list[dict[str, Any]],
+ nav_result: NavigateStepResult,
) -> None:
+ action = nav_result.action
+ step_node = nav_result.node
+ asset_tools = nav_result.tools
+ drill_paths = nav_result.pending
+ reason = nav_result.reason
+ stop_type = nav_result.stop_type
+ selected_paths = list(step_node.confidence.keys())
+ hydrated_paths = list(step_node.leaf_content.keys())
+
if self._trace_enabled:
self._trace.record_step(
"navigate_step",
@@ -347,19 +366,40 @@ def _record_navigation_step(
"scope": scope if isinstance(scope, str) else (scope or "root"),
"depth": depth,
"action": action,
+ "reason": reason,
+ "stop_type": stop_type,
"asset_tools": asset_tools,
+ "selected_paths": selected_paths,
+ "hydrated_paths": hydrated_paths,
"outline_count": len(step_node.outline_items),
"leaf_count": len(step_node.leaf_content),
+ "hydrated_count": sum(len(c) for c in step_node.leaf_content.values()),
"pending_drills": len(drill_paths),
},
),
decision_reason=f"nav_d{depth}_{doc.source_file_name}",
)
+
+ doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "")
+ self._decision_steps.append({
+ "phase": "navigate",
+ "document": doc_name,
+ "document_id": doc.document_id,
+ "action": action,
+ "reason": reason,
+ "stop_type": stop_type,
+ "depth": depth,
+ "selected_paths": selected_paths,
+ "hydrated_paths": hydrated_paths,
+ "hydrated_count": sum(len(c) for c in step_node.leaf_content.values()),
+ })
+
scope_log = scope if isinstance(scope, str) else (", ".join(scope) if scope else "root")
logger.info(
f" agentic step {self._state.step_count}: navigate_step "
f'doc="{doc.source_file_name}" scope={scope_log} '
f"depth={depth} action={action} tools={asset_tools} "
+ f"reason=\"{reason[:80]}\" stop_type={stop_type} "
f"outline={len(step_node.outline_items)} "
f"leaves={len(step_node.leaf_content)} "
f"drills={len(drill_paths)}"
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py
index 7125fdf4f..46a3ecc9d 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py
@@ -56,6 +56,10 @@ async def load_child_sections(
continue
normalized_path = normalize_section_path(path)
parts = split_section_path(normalized_path)
+ # Treat Root as a virtual L1 node so it appears in the
+ # navigation tree and the LLM can decide whether to drill in.
+ if normalized_path == "Root" and not parts:
+ parts = ["Root"]
all_sections[normalized_path] = {
"title": title or parts[-1] if parts else normalized_path,
"summary": summary or "",
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
index fe9b0b870..44f007779 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
@@ -26,7 +26,7 @@
from shared.services.retrieval.agentic.navigation.selection_hydration import (
hydrate_path_selections_into_node,
)
-from shared.services.retrieval.agentic.core.types import DocTreeNode
+from shared.services.retrieval.agentic.core.types import DocTreeNode, NavigateStepResult
from shared.services.retrieval.llm_adapter import LLMFn
@@ -43,7 +43,7 @@ async def navigate_step(
scope_path: str | list[str] | None = None,
exclude_paths: set[str] | None = None,
budget_snapshot: dict | None = None,
-) -> tuple[str, list[str], DocTreeNode, list[dict]]:
+) -> NavigateStepResult:
"""Navigate one document scope and hydrate selected sections."""
scope_paths = (
scope_path if isinstance(scope_path, list)
@@ -63,7 +63,7 @@ async def navigate_step(
exclude_paths=exclude_paths,
)
if not items:
- return "STOP", [], empty, []
+ return NavigateStepResult.stop(scope_paths[0] if scope_paths else None)
selectable = {
item["path"]: item for item in items if item.get("selectable", False)
@@ -95,6 +95,8 @@ async def navigate_step(
action = parsed["action"]
selected_tools = parsed["tools"]
selections = parsed["selections"]
+ reason = parsed.get("reason", "")
+ stop_type = parsed.get("stop_type", "")
scope_label = ", ".join(scope_paths) if scope_paths else "root"
logger.info(
@@ -146,13 +148,20 @@ async def navigate_step(
job_result_id=job_result_id,
)
- return action, selected_tools, node, pending
+ return NavigateStepResult(
+ action=action,
+ tools=selected_tools,
+ node=node,
+ pending=pending,
+ reason=reason,
+ stop_type=stop_type,
+ )
except BudgetExceeded:
raise
except Exception as exc:
logger.error(f" navigate_step failed for doc={document_id}: {exc}")
- return "STOP", [], empty, []
+ return NavigateStepResult.stop(scope_paths[0] if scope_paths else None)
def _build_navigation_prompt(
*,
document_id: str,
diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
index d6e3ac2ba..43548d5a5 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
@@ -196,8 +196,13 @@ async def run(
discovery_by_doc: dict[str, list[dict[str, Any]]] = {}
for row in discovery_rows:
doc_id = row.get('document_id', '')
- if doc_id:
- discovery_by_doc.setdefault(doc_id, []).append(row)
+ # Root chunks are navigable via the section tree; exclude them
+ # from discovery hints where the bare "Root" label gives the LLM
+ # no actionable information.
+ section_path = str(row.get('section_path', '') or '').strip()
+ if not doc_id or section_path == 'Root':
+ continue
+ discovery_by_doc.setdefault(doc_id, []).append(row)
await register_discovery_documents(
db,
@@ -215,6 +220,25 @@ async def run(
evidence_text = ''
stop_reason = 'evidence_only'
failure_reason = ''
+ decision_trace: list[dict[str, Any]] = []
+
+ # Record KG document selection as the first decision trace entry
+ if state.selected_docs:
+ decision_trace.append({
+ 'phase': 'kg_select',
+ 'action': 'select',
+ 'reason': f'{len(state.selected_docs)} document(s) selected for navigation',
+ 'documents': [
+ {
+ 'document': doc.source_file_name,
+ 'document_id': doc.document_id,
+ 'confidence': doc.confidence,
+ 'reason': doc.reason,
+ 'source': doc.source,
+ }
+ for doc in state.selected_docs
+ ],
+ })
if state.elapsed_ms >= config.latency_budget_ms:
stop_reason = 'latency_budget'
@@ -233,6 +257,7 @@ async def run(
llm_budget=llm_budget,
)
await navigation_runner.navigate_selected_documents()
+ decision_trace.extend(navigation_runner.decision_steps)
context_remaining = state.ledger.remaining('context') if state.ledger else config.token_budget_total
evidence_text = await _trim_evidence_to_budget(
db,
@@ -278,6 +303,7 @@ async def run(
budget_snapshot=state.ledger.snapshot() if state.ledger else None,
stop_reason=stop_reason,
failure_reason=failure_reason,
+ decision_trace=decision_trace,
)
logger.info(
diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py
index 1002a68b1..ca18d8048 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/prompts.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py
@@ -88,10 +88,13 @@
When action is STOP, selections must be empty.
+Always include a "reason" field (1-2 sentences) explaining your choice.
+When action is STOP, also include "stop_type" from: sufficient_outline, no_relevant_child, evidence_sufficient, budget_conserve.
+
Return ONLY a JSON object:
-{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}}
+{{"action": "NAVIGATE", "reason": "...", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}}
or
-{{"action": "STOP", "tools": [...], "selections": []}}
+{{"action": "STOP", "reason": "...", "stop_type": "...", "tools": [...], "selections": []}}
Do not include any explanation.
"""
@@ -100,7 +103,7 @@ def parse_action_response(text: str) -> dict:
"""Parse the unified navigation response from an LLM."""
text = text.strip()
asset_tools = {"FIND_IMAGES", "FIND_TABLES"}
- default = {"action": "NAVIGATE", "tools": [], "selections": []}
+ default = {"action": "NAVIGATE", "tools": [], "selections": [], "reason": "", "stop_type": ""}
def extract(data: dict) -> dict:
action = str(data.get("action", "NAVIGATE")).strip().upper()
@@ -117,8 +120,11 @@ def extract(data: dict) -> dict:
else:
tools = []
+ reason = str(data.get("reason") or "").strip()[:500]
+ stop_type = str(data.get("stop_type") or "").strip()[:50] if action == "STOP" else ""
+
if action == "STOP":
- return {"action": action, "tools": tools, "selections": []}
+ return {"action": action, "tools": tools, "selections": [], "reason": reason, "stop_type": stop_type}
selections_val = data.get("selections") or []
selections = []
@@ -131,7 +137,7 @@ def extract(data: dict) -> dict:
"confidence": confidence or 0.7,
})
- return {"action": action, "tools": tools, "selections": selections}
+ return {"action": action, "tools": tools, "selections": selections, "reason": reason, "stop_type": ""}
data = _parse_json_object(text)
if data is not None:
@@ -156,13 +162,10 @@ def format_budget_block(snapshot: dict | None) -> str:
if not snapshot:
return ""
planning = snapshot.get("planning") or {}
- context = snapshot.get("context") or {}
return (
"=== Resource Status ===\n"
f"Planning Budget: {planning.get('status', 'HEALTHY')} "
f"({planning.get('used_pct', 0)}% used)\n"
- f"Context Budget: {context.get('status', 'HEALTHY')} "
- f"({context.get('used_pct', 0)}% used)\n"
f"KG Coverage: {snapshot.get('explored_chunks', 0)}/"
f"{snapshot.get('total_chunks', 0)} chunks explored\n"
f"Docs Explored: {snapshot.get('explored_docs', 0)}/"
diff --git a/packages/shared-python/shared/services/retrieval/execution/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py
index 9c50a7e97..80cf15aa7 100644
--- a/packages/shared-python/shared/services/retrieval/execution/response_projection.py
+++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py
@@ -45,6 +45,8 @@ async def project_public_retrieval_response(response: dict[str, Any]) -> dict[st
public_response['stop_reason'] = response['stop_reason']
if response.get('failure_reason') is not None:
public_response['failure_reason'] = response['failure_reason']
+ if response.get('decision_trace') is not None:
+ public_response['decision_trace'] = response['decision_trace']
projected_rows = await enrich_rows_with_retrieval_asset_urls(
response.get('results', []),
diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py
index 36e227712..2bb24958e 100644
--- a/packages/shared-python/shared/services/retrieval/execution/routes.py
+++ b/packages/shared-python/shared/services/retrieval/execution/routes.py
@@ -151,6 +151,14 @@ async def _run_agentic_route(
if last_retrieve.failure_reason:
response["failure_reason"] = last_retrieve.failure_reason
+ # Merge decision traces from all retrieve steps
+ all_decision_trace: list[dict] = []
+ for step in workflow_result.steps:
+ if step.decision_trace:
+ all_decision_trace.extend(step.decision_trace)
+ if all_decision_trace:
+ response["decision_trace"] = all_decision_trace
+
completion_detail = (
f"chunks | evidence={len(response.get('evidence_text') or '')} chars | "
f"router={workflow_result.router_used}"
diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
index c14328aaf..49f80f384 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py
@@ -117,4 +117,5 @@ def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepR
router_used=result.router_used,
stop_reason=result.stop_reason,
failure_reason=result.failure_reason,
+ decision_trace=result.decision_trace,
)
diff --git a/packages/shared-python/shared/services/retrieval/workflow/types.py b/packages/shared-python/shared/services/retrieval/workflow/types.py
index 0141d05c1..23366c242 100644
--- a/packages/shared-python/shared/services/retrieval/workflow/types.py
+++ b/packages/shared-python/shared/services/retrieval/workflow/types.py
@@ -170,6 +170,7 @@ class StepResult:
router_used: str = ""
stop_reason: str = ""
failure_reason: str = ""
+ decision_trace: list[dict[str, Any]] = field(default_factory=list)
error: str | None = None
def to_api_dict(self) -> dict[str, Any]:
@@ -188,6 +189,7 @@ def to_api_dict(self) -> dict[str, Any]:
"router_used": self.router_used,
"stop_reason": self.stop_reason,
"failure_reason": self.failure_reason,
+ "decision_trace": self.decision_trace if self.decision_trace else None,
"error": self.error,
}
From ab8824581672eb9aaee2e5fcb1e9e6ca323f6797 Mon Sep 17 00:00:00 2001
From: chengke <404835780@qq.com>
Date: Fri, 29 May 2026 11:08:20 +0800
Subject: [PATCH 4/4] fix: remove unused variable and import (lint F841/F401)
---
.../shared/services/retrieval/agentic/navigation/tools.py | 1 -
1 file changed, 1 deletion(-)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
index 44f007779..c577d6699 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py
@@ -52,7 +52,6 @@ async def navigate_step(
)
scope_path_set = set(scope_paths)
- empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None)
try:
items = await load_child_sections(