From 0b97f4c3ca144cfdab74e2047477f00cccdd3d28 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 29 May 2026 17:20:36 +0800 Subject: [PATCH 1/6] refactor: transition agentic navigation to collector-agent architecture with decoupled step-based collection and batch hydration --- apps/api/app/api/v1/routes/retrieval.py | 4 +- .../retrieval/agentic/core/runtime.py | 6 +- .../services/retrieval/agentic/core/types.py | 50 ++- .../retrieval/agentic/evidence/builder.py | 5 +- .../retrieval/agentic/evidence/renderer.py | 4 +- .../retrieval/agentic/navigation/document.py | 292 ++++++++++++------ .../agentic/navigation/section_counts.py | 12 +- .../navigation/section_prompt_projection.py | 107 ++++++- .../agentic/navigation/section_tree.py | 36 +-- .../retrieval/agentic/navigation/tools.py | 255 ++++++++------- .../retrieval/agentic/orchestrator.py | 8 +- .../services/retrieval/agentic/prompts.py | 150 ++++++--- .../services/retrieval/agentic/tools.py | 8 +- 13 files changed, 639 insertions(+), 298 deletions(-) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 9f8efabb0..508f9bf8e 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -99,8 +99,8 @@ class RetrievalQueryResponse(BaseModel): 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 " + "Each entry has phase, document, action, reason, collected_paths, " + "and drill_into. Use this to understand " "why KNOWHERE stopped or made specific navigation choices." ), ) 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 8e7b4f7ac..594184e92 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,7 @@ def build_config_from_env() -> AgentRunConfig: return AgentRunConfig( - max_nav_depth=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_DEPTH", "3")), + max_nav_steps=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_STEPS", "6")), 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")), planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")), @@ -104,7 +104,7 @@ def for_document( llm_fn: LLMFn, *, doc_id: str, - depth: int, + step: int = 0, ) -> LLMFn: async def _call(prompt: Any) -> str: return await self.call( @@ -112,7 +112,7 @@ async def _call(prompt: Any) -> str: prompt, pool="planning", doc_id=doc_id, - priority="low" if depth >= 2 else "normal", + priority="low" if step >= 4 else "normal", ) return _call 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 7eb2dd0e5..651ce014b 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,7 @@ @dataclass class AgentRunConfig: """Budget and limit configuration for a single agent run.""" - max_nav_depth: int = 3 # max scope_navigate recursion depth + max_nav_steps: int = 6 # max navigation steps per document (no depth limit) latency_budget_ms: int = 12000 token_budget_total: int = 40000 planning_ratio: float = 0.5 @@ -79,6 +79,17 @@ def has_content(self) -> bool: return any(c.has_content() for c in self.children.values()) return False + def has_leaf_content(self) -> bool: + """Check if this tree has any actual hydrated chunk content (not just outline). + + Unlike ``has_content()`` which returns True for outline-only trees, + this method only returns True when real text/table/image chunks have + been hydrated into leaf_content. + """ + if self.leaf_content: + return True + return any(c.has_leaf_content() for c in self.children.values()) + def flatten_chunk_rows(self) -> list[dict[str, Any]]: """Recursively collect all hydrated chunk rows (document order).""" rows: list[dict[str, Any]] = [] @@ -108,10 +119,16 @@ def add_leaf_chunks(self, path: str, chunks: list[dict[str, Any]]) -> None: existing.append(chunk) def reparent_leaf_content(self) -> None: - """Move descendant leaf paths into matching child nodes.""" + """Move descendant leaf paths into matching child nodes. + + Only moves true descendants (prefix match). Content whose path + exactly equals a child key stays here — the renderer handles the + case where a path is both a child and a leaf (section with own + content *and* sub-sections). + """ for child_path, child in list(self.children.items()): for leaf_path in list(self.leaf_content.keys()): - if leaf_path == child_path or leaf_path.startswith(child_path + ' / '): + if leaf_path.startswith(child_path + ' / '): child.add_leaf_chunks(leaf_path, self.leaf_content.pop(leaf_path)) child.reparent_leaf_content() @@ -165,13 +182,32 @@ def merge(self, other: 'DocTreeNode') -> None: @dataclass class NavigateStepResult: - """Return type for navigate_step — typed replacement for raw tuple.""" - action: str # "NAVIGATE" or "STOP" + """Return type for navigate_step — Collector Agent model. + + Each step returns: + - ``collect``: paths to add to the evidence collection (full hydration) + - ``drill``: paths to explore deeper in subsequent steps + - ``action``: navigation direction — DRILL/BACK/STOP + - ``tools``: optional asset tools (FIND_IMAGES/FIND_TABLES) + - ``node``: outline tree node for rendering context + - ``reason``: LLM reasoning for trace + """ + action: str = "STOP" # DRILL | BACK | STOP + collect: list[dict[str, Any]] = field(default_factory=list) + drill: list[dict[str, Any]] = field(default_factory=list) 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 | ... + + @property + def drill_into(self) -> str | None: + """Single drill target path, or None.""" + return self.drill[0]["path"] if self.drill else None + + @property + def is_terminal(self) -> bool: + """True when navigation should stop (STOP or empty collect+drill).""" + return self.action == "STOP" or (not self.collect and not self.drill) @staticmethod def stop(scope_path: str | None = None) -> 'NavigateStepResult': 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 2b1cc3cf3..ee0e68e7f 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py @@ -140,7 +140,8 @@ async def render_evidence( evidence_parts: list[str] = [] for doc_id, doc_tree in doc_trees.items(): - if doc_tree.has_content(): + # Only render if there is actual evidence (hydrated chunks or outline content). + if doc_tree.has_leaf_content() or doc_tree.has_content(): doc_name = doc_id_to_name.get(doc_id, doc_id) rendered = render_unified_doc_tree( doc_tree, @@ -150,7 +151,7 @@ async def render_evidence( if rendered.strip(): evidence_parts.append(rendered) - return "\n\n".join(evidence_parts) if evidence_parts else "(no evidence collected)" + return "\n\n".join(evidence_parts) if evidence_parts else "" def _iter_leaf_content(node: DocTreeNode): 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 4b464d004..9e974b78a 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py @@ -88,9 +88,11 @@ def min_sort(path: str) -> float: elif render_type == "orphan_child": path = cast(str, data) title = path.rsplit(" / ", 1)[-1] if " / " in path else path - parts.append(f"{indent}▸ {title} [DrillDown]") child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) + # Only render the [DrillDown] heading if the child has content. + # Prevents empty orphan nodes from polluting evidence_text. if child_text.strip(): + parts.append(f"{indent}▸ {title} [DrillDown]") parts.append(child_text) return "\n".join(parts) 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 f7e711fa5..cfece903f 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py @@ -1,4 +1,20 @@ -"""Per-document navigation for agentic retrieval.""" +"""Per-document navigation for agentic retrieval. + +Collector Agent architecture +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +The navigation loop uses a Collector Agent model where each step +independently produces two decisions: + +1. **collect**: paths to add to the evidence collection +2. **action**: navigation direction (DRILL/BACK/STOP) + +The ``collected_paths`` list accumulates across all steps. After +navigation completes (or is interrupted), a single batch hydration +pass loads content for all collected paths. + +Asset collection (images/tables) still runs during navigation so LLM +tool requests are honoured, but assets are reconciled after hydration. +""" from __future__ import annotations from typing import Any, cast @@ -19,6 +35,9 @@ NavigateStepResult, ToolResult, ) +from shared.services.retrieval.agentic.navigation.selection_hydration import ( + hydrate_path_selections_into_node, +) from shared.services.retrieval.llm_adapter import LLMFn @@ -80,20 +99,33 @@ async def _navigate_document( root = DocTreeNode(scope_path=None) doc_pending_assets: list[dict[str, Any]] = [] + # Phase 2A: Collector Agent navigation (summary-only, no content hydration) + collected_paths: list[dict[str, Any]] = [] if not is_discovery_only_doc: - doc_pending_assets = await self._navigate_bfs( + doc_pending_assets, collected_paths = await self._navigate_collector( doc=doc, root=root, doc_name=doc_name, job_result_id=job_result_id, ) + # Phase 2B: Discovery hints (independent hydration path) await self._hydrate_discovery_hints( doc=doc, root=root, doc_name=doc_name, ) + # Phase 2C: Batch hydrate all collected paths + if not is_discovery_only_doc and collected_paths: + await self._hydrate_collected( + doc=doc, + root=root, + job_result_id=job_result_id, + collected_paths=collected_paths, + ) + + # Phase 2D: Reconcile assets into hydrated tree if not is_discovery_only_doc and doc_pending_assets: self._reconcile_pending_assets( doc=doc, @@ -110,35 +142,43 @@ async def _navigate_document( if self._state.ledger is not None: self._state.ledger.mark_explored(docs=1) - async def _navigate_bfs( + async def _navigate_collector( self, *, doc: CandidateDoc, root: DocTreeNode, doc_name: str, job_result_id: str, - ) -> list[dict[str, Any]]: + ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Collector Agent navigation loop. + + Returns (doc_pending_assets, collected_paths). + """ doc_exclude: set[str] = set() - pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)] + nav_trace: list[dict[str, Any]] = [] + collected_paths: list[dict[str, Any]] = [] doc_pending_assets: list[dict[str, Any]] = [] - while pending: + # Scope stack for BACK support: each entry is a scope path (None = root) + scope_stack: list[str | None] = [None] + step_count = 0 + + while step_count < self._config.max_nav_steps: if self._state.elapsed_ms >= self._config.latency_budget_ms: break - - scope, parent_node, depth = pending.pop(0) - if depth >= self._config.max_nav_depth: - continue if self._llm_fn is None: break if self._state.ledger and self._state.ledger.status("planning") in ("CRITICAL", "EXHAUSTED"): - logger.info(" agentic: planning budget critical, ending BFS for current doc") + logger.info(" agentic: planning budget critical, ending navigation for current doc") break + current_scope = scope_stack[-1] + step_count += 1 + doc_llm_fn = self._llm_budget.for_document( cast(LLMFn, self._llm_fn), doc_id=doc.document_id, - depth=depth, + step=step_count, ) try: nav_result = await tools.navigate_step( @@ -150,9 +190,11 @@ async def _navigate_bfs( user_id=self._user_id, namespace=self._namespace, doc_name=doc_name, - scope_path=scope, + scope_path=current_scope, exclude_paths=doc_exclude, budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, + nav_trace=nav_trace if nav_trace else None, + collected_paths=collected_paths, ) except BudgetExceeded: logger.info(" agentic: planning budget exhausted during navigation") @@ -161,43 +203,125 @@ async def _navigate_bfs( break self._state.step_count += 1 + # Asset collection runs during navigation (images/tables). await self._collect_assets( doc=doc, - scope=scope, + scope=current_scope, step_node=nav_result.node, asset_tools=nav_result.tools, pending_assets=doc_pending_assets, round_scope="nav", ) - _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() + + # Merge outline + confidence into root tree + _merge_step_node(root, nav_result.node) + + # ── Process COLLECT ────────────────────────────────────────── + collected_in_step: list[str] = [] + for coll_item in nav_result.collect: + path = coll_item["path"] + coll_item["collected_at_step"] = step_count + coll_item["scope_context"] = current_scope or "root" + collected_paths.append(coll_item) + collected_in_step.append(path) + # Collected paths should be excluded from future navigation + doc_exclude.add(path) + + # ── Build trace entry ──────────────────────────────────────── + trace_entry: dict[str, Any] = { + "step": step_count, + "scope": current_scope or "root", + "action": nav_result.action, + "drill_into": nav_result.drill_into, + "collected": collected_in_step, + "reason": nav_result.reason, + } + nav_trace.append(trace_entry) + + # ── Record decision step ───────────────────────────────────── self._record_navigation_step( doc=doc, - scope=scope, - depth=depth, + scope=current_scope, + step_num=step_count, nav_result=nav_result, + collected_in_step=collected_in_step, ) - if self._state.ledger is not None: - self._state.ledger.mark_explored( - chunks=sum(len(chunks) for chunks in nav_result.node.leaf_content.values()), - ) - return doc_pending_assets + # ── Process navigation action ──────────────────────────────── + if nav_result.action == "DRILL" and nav_result.drill_into: + drill_path = nav_result.drill_into + # Create child node in tree for the drill target + target_parent = _find_target_node(root, drill_path) + target_parent.children.setdefault(drill_path, DocTreeNode(scope_path=drill_path)) + scope_stack.append(drill_path) + + elif nav_result.action == "BACK": + if len(scope_stack) > 1: + scope_stack.pop() + else: + # Already at root, treat as STOP + logger.info(" agentic: BACK at root scope, treating as STOP") + break + + elif nav_result.action == "STOP" or nav_result.is_terminal: + break + + return doc_pending_assets, collected_paths + + async def _hydrate_collected( + self, + *, + doc: CandidateDoc, + root: DocTreeNode, + job_result_id: str, + collected_paths: list[dict[str, Any]], + ) -> None: + """Batch-hydrate all collected paths after navigation completes.""" + if not collected_paths: + return + + # Deduplicate: keep highest confidence per path + deduped: dict[str, dict[str, Any]] = {} + for item in collected_paths: + path = item["path"] + if path not in deduped or item.get("confidence", 0) > deduped[path].get("confidence", 0): + deduped[path] = item + unique_selections = list(deduped.values()) + + await hydrate_path_selections_into_node( + self._db, + node=root, + path_selections=unique_selections, + user_id=self._user_id, + namespace=self._namespace, + document_id=doc.document_id, + job_result_id=job_result_id, + ) + + # Single reparent pass — tree structure is final. + root.reparent_leaf_content() + + # Accurate budget accounting — count only actually-hydrated chunks. + if self._state.ledger is not None: + total_chunks = len(root.flatten_chunk_rows()) + self._state.ledger.mark_explored(chunks=total_chunks) + logger.info( + f" agentic: hydrate_collected doc={doc.document_id} " + f"collected={len(unique_selections)} hydrated_chunks={total_chunks}" + ) async def _collect_assets( self, *, doc: CandidateDoc, - scope: str | list[str] | None, + scope: str | None, step_node: DocTreeNode, asset_tools: list[str], pending_assets: list[dict[str, Any]], round_scope: str, ) -> None: selected_asset_scopes = list(step_node.confidence.keys()) - asset_scope = selected_asset_scopes or scope + asset_scope: str | list[str] | None = selected_asset_scopes or scope for asset_tool in asset_tools: if asset_tool not in ("FIND_IMAGES", "FIND_TABLES"): continue @@ -225,7 +349,7 @@ async def _collect_assets( payload={ "document_id": doc.document_id, "scope": scope_display, - "navigation_scope": scope if isinstance(scope, str) else (scope or "root"), + "navigation_scope": scope or "root", "asset_type": asset_type, "chunks_found": len(asset_chunks) if asset_chunks else 0, }, @@ -343,41 +467,34 @@ def _record_navigation_step( self, *, doc: CandidateDoc, - scope: str | list[str] | None, - depth: int, + scope: str | None, + step_num: int, nav_result: NavigateStepResult, + collected_in_step: list[str], ) -> 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()) + drill_into = nav_result.drill_into if self._trace_enabled: self._trace.record_step( "navigate_step", ToolResult( - status=f"{action.lower()}" + (" (content)" if step_node.has_content() else ""), + status=f"{action.lower()}", payload={ "document_id": doc.document_id, - "scope": scope if isinstance(scope, str) else (scope or "root"), - "depth": depth, + "scope": scope or "root", + "step": step_num, "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), + "drill_into": drill_into, + "collected_count": len(collected_in_step), + "collected_paths": collected_in_step, + "asset_tools": nav_result.tools, + "outline_count": len(nav_result.node.outline_items), }, ), - decision_reason=f"nav_d{depth}_{doc.source_file_name}", + decision_reason=f"nav_s{step_num}_{doc.source_file_name}", ) doc_name = doc.source_file_name or self._state.doc_id_to_name.get(doc.document_id, "") @@ -387,60 +504,55 @@ def _record_navigation_step( "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()), + "step": step_num, + "drill_into": drill_into, + "collected_paths": collected_in_step, + "collected_count": len(collected_in_step), }) - scope_log = scope if isinstance(scope, str) else (", ".join(scope) if scope else "root") + scope_log = scope or "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)}" + f"step={step_num} action={action} tools={nav_result.tools} " + f'reason="{reason[:80]}" ' + f"collected={len(collected_in_step)} " + f"drill_into={drill_into} " + f"outline={len(nav_result.node.outline_items)}" ) -def _merge_step_node(parent_node: DocTreeNode, step_node: DocTreeNode) -> None: - parent_node.outline_items = step_node.outline_items - for leaf_path, chunks in step_node.leaf_content.items(): - parent_node.add_leaf_chunks(leaf_path, chunks) - parent_node.confidence = step_node.confidence +def _find_target_node(node: DocTreeNode, path: str) -> DocTreeNode: + """Walk the tree to find the deepest existing node that owns *path*. + + Only recurse when *path* is a true descendant of a child (prefix match). + An exact match means the item belongs to the section itself, which is + managed by the *parent* node — the renderer already handles the case + where a path appears in both ``children`` and ``leaf_content``. + """ + for child_path, child in node.children.items(): + if path.startswith(child_path + " / "): + return _find_target_node(child, path) + return node + + +def _merge_step_node(root: DocTreeNode, step_node: DocTreeNode) -> None: + """Route outline items and confidence from *step_node* to correct tree positions.""" + for item in step_node.outline_items: + path = item.get("path", "") + target = _find_target_node(root, path) + existing = {i.get("path") for i in target.outline_items} + if path not in existing: + target.outline_items.append(item) + + for path, conf in step_node.confidence.items(): + target = _find_target_node(root, path) + target.confidence[path] = max(target.confidence.get(path, 0), conf) def _collect_leaf_paths(node: DocTreeNode) -> set[str]: + """Collect all paths that have been hydrated (leaf_content).""" 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, - drill_paths: list[dict[str, Any]], -) -> None: - drill_path_set = {str(selection["path"]) for selection in drill_paths} - for leaf_path in step_node.leaf_content: - if leaf_path not in drill_path_set: - doc_exclude.add(leaf_path) - - -def _queue_drill_paths( - pending: list[tuple[str | list[str] | None, DocTreeNode, int]], - parent_node: DocTreeNode, - drill_paths: list[dict[str, Any]], - depth: int, -) -> None: - if not drill_paths: - return - for selection in drill_paths: - child = DocTreeNode(scope_path=selection["path"]) - parent_node.children[selection["path"]] = child - batch_scope = [selection["path"] for selection in drill_paths] - pending.append((batch_scope, parent_node, depth + 1)) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py index 74ba465d3..64e2499fd 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py @@ -33,7 +33,7 @@ async def attach_section_counts( ) sid_to_path = {meta["section_id"]: path for path, meta in all_sections.items()} - for section_id, (text_count, image_count, table_count) in section_id_counts.items(): + for section_id, (text_count, image_count, table_count, total_chars) in section_id_counts.items(): chunk_path = sid_to_path.get(section_id, "") if not chunk_path: continue @@ -45,6 +45,7 @@ async def attach_section_counts( item["chunk_count"] += text_count item["image_count"] += image_count item["table_count"] += table_count + item["total_chars"] += total_chars await _attach_connected_asset_counts( db, @@ -61,7 +62,7 @@ async def _load_direct_chunk_counts( document_id: str, job_result_id: str, all_section_ids: list[str], -) -> dict[str, tuple[int, int, int]]: +) -> dict[str, tuple[int, int, int, int]]: chunk_stmt = ( select( DocumentChunk.section_id, @@ -80,6 +81,9 @@ async def _load_direct_chunk_counts( (DocumentChunk.chunk_type == "table", literal_column("1")), ) ).label("table_count"), + func.coalesce( + func.sum(func.length(DocumentChunk.content)), 0 + ).label("total_chars"), ) .where(DocumentChunk.document_id == document_id) .where(DocumentChunk.job_result_id == job_result_id) @@ -88,8 +92,8 @@ async def _load_direct_chunk_counts( ) chunk_rows = (await db.execute(chunk_stmt)).all() return { - section_id: (int(text_count), int(image_count), int(table_count)) - for section_id, text_count, image_count, table_count in chunk_rows + section_id: (int(text_count), int(image_count), int(table_count), int(total_chars)) + for section_id, text_count, image_count, table_count, total_chars in chunk_rows } diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py index b7a251805..2b4c04fd3 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py @@ -1,34 +1,45 @@ -"""Prompt projection for agentic section navigation.""" +"""Prompt projection for agentic section navigation (Collector Agent model).""" from __future__ import annotations +from typing import Any + from shared.utils.text_utils import truncate_content_preview def format_items_for_llm( items: list[dict], max_chars: int = 20000, + collected_paths: set[str] | None = None, ) -> tuple[str, bool]: - """Format section items with hierarchy, selectability, counts, and summaries.""" + """Format section items with hierarchy, token estimates, and collection marks.""" if not items: return "(no items available)", False - full_text = "\n".join(_render_item(item, include_summary=True) for item in items) + coll = collected_paths or set() + full_text = "\n".join(_render_item(item, include_summary=True, collected=coll) for item in items) if len(full_text) <= max_chars: return full_text, False - slim_text = "\n".join(_render_item(item, include_summary=False) for item in items) + slim_text = "\n".join(_render_item(item, include_summary=False, collected=coll) for item in items) return slim_text[:max_chars], True -def _render_item(item: dict, include_summary: bool) -> str: +def _render_item(item: dict, include_summary: bool, collected: set[str]) -> str: level = item.get("level", 1) show_summary = item.get("show_summary", True) is_leaf = item.get("is_leaf", False) - leaf_tag = " [Leaf]" if is_leaf else "" path = item.get("path", "") summary = item.get("summary") or "" + # Check if this path (or an ancestor) is already collected + is_collected = _is_path_collected(path, collected) + collected_tag = "[✓] " if is_collected else "" + + leaf_tag = " [Leaf]" if is_leaf else "" + + # Counts and token estimate counts_str = "" + token_str = "" if show_summary: count_parts: list[str] = [] chunk_count = item.get("chunk_count", 0) @@ -42,13 +53,22 @@ def _render_item(item: dict, include_summary: bool) -> str: count_parts.append(f"table={table_count}") counts_str = f' [{" ".join(count_parts)}]' if count_parts else "" + total_chars = item.get("total_chars", 0) + if total_chars > 0: + # Approximate tokens: Chinese ~2 chars/token, English ~4 chars/token + # Use conservative 2 chars/token for mixed content + tokens = total_chars / 2 + if tokens >= 1000: + token_str = f" ~{tokens / 1000:.1f}k tokens" + else: + token_str = f" ~{int(tokens)} tokens" + indent = " " * (level - 1) prefix = "▸" if level == 1 else "└" level_tag = f"[L{level}]" - select_tag = "[SELECT] " if item.get("selectable", False) else "" lines = [ - f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}' + f'{indent}{prefix} {collected_tag}{level_tag} path="{path}"{counts_str}{token_str}{leaf_tag}' ] if include_summary and show_summary and summary: @@ -57,3 +77,74 @@ def _render_item(item: dict, include_summary: bool) -> str: lines.append(f"{sub_indent}{clipped}") return "\n".join(lines) + + +def _is_path_collected(path: str, collected: set[str]) -> bool: + """Check if path itself or any ancestor is in the collected set.""" + if path in collected: + return True + for coll_path in collected: + if path.startswith(coll_path + " / "): + return True + return False + + +def format_collection_status( + collected_paths: list[dict[str, Any]], +) -> str: + """Render the collection status block for the navigation prompt.""" + if not collected_paths: + return "" + + lines = [f"=== Collection Status ({len(collected_paths)} items) ==="] + for item in collected_paths: + path = item.get("path", "") + conf = item.get("confidence", 0) + step = item.get("collected_at_step", "?") + outline = item.get("outline", False) + mode_tag = " [outline]" if outline else "" + lines.append(f'✓ "{path}" (step {step}, conf={conf:.1f}{mode_tag})') + lines.append("=== End Collection ===") + return "\n".join(lines) + + +def format_nav_trace( + nav_trace: list[dict[str, Any]], + collected_paths: list[dict[str, Any]], +) -> str: + """Render the unified navigation trace block (includes scope, actions, and collection).""" + if not nav_trace and not collected_paths: + return "" + + lines = ["=== Navigation Trace ==="] + for entry in nav_trace: + step = entry.get("step", "?") + scope = entry.get("scope", "root") + action = entry.get("action", "?") + reason = entry.get("reason", "") + + action_display = action + drill_into = entry.get("drill_into") + if action == "DRILL" and drill_into: + action_display = f'DRILL "{drill_into}"' + + lines.append(f"Step {step}: scope={scope} → {action_display}") + + # Show what was collected in this step + step_collected = entry.get("collected", []) + if step_collected: + paths_display = ", ".join(f'"{c}"' for c in step_collected) + lines.append(f" collected: {paths_display}") + + if reason: + lines.append(f" reason: {reason}") + lines.append("") + + # Append current collection summary + if collected_paths: + total = len(collected_paths) + lines.append(f"[Current] collection: {total} items") + lines.append("Do NOT re-collect paths marked [✓] below.") + + lines.append("=== End Trace ===") + return "\n".join(lines) 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 46a3ecc9d..47c38ffe3 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 @@ -110,7 +110,14 @@ async def load_child_sections( item.pop("sort_order", None) item.pop("section_id", None) - _mark_leaf_and_selectable(sorted_items, all_section_paths=set(all_sections.keys()), allowed_set=allowed_set) + # Mark leaf status (no descendants in full section list) + all_section_paths = set(all_sections.keys()) + for item in sorted_items: + item_path = item["path"] + item["is_leaf"] = not any( + path != item_path and path.startswith(item_path + " / ") + for path in all_section_paths + ) return sorted_items @@ -168,6 +175,7 @@ def _make_item(path: str, meta: dict, show_summary: bool) -> dict: "chunk_count": 0, "image_count": 0, "table_count": 0, + "total_chars": 0, "section_id": meta["section_id"], "show_summary": show_summary, } @@ -207,29 +215,3 @@ def _resolve_allowed_depths(items_by_path: dict[str, dict], scope_list: list[str return allowed_set -def _mark_leaf_and_selectable( - sorted_items: list[dict], - *, - all_section_paths: set[str], - allowed_set: set[int], -) -> None: - for item in sorted_items: - item_path = item["path"] - has_descendants = any( - path != item_path and path.startswith(item_path + " / ") - for path in all_section_paths - ) - item["is_leaf"] = not has_descendants - - if allowed_set: - shallowest_band = min(allowed_set) - for item in sorted_items: - if not item.get("show_summary", True): - item["selectable"] = False - elif item["level"] == shallowest_band and not item.get("is_leaf", False): - item["selectable"] = False - else: - item["selectable"] = True - else: - for item in sorted_items: - item["selectable"] = item.get("show_summary", True) 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 c577d6699..65ab2fae1 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py @@ -1,8 +1,16 @@ -"""Agentic retrieval navigation tools. +"""Agentic retrieval navigation tools — Collector Agent model. -This Module owns document-scope navigation and post-navigation discovery -selection. It keeps the LLM prompt, section traversal, hydration, and asset -owner reconciliation local to the navigation seam. +Collector Agent architecture +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Each ``navigate_step`` returns two independent decisions: + +- **collect**: paths the agent adds to its evidence collection. + Collected paths are hydrated with full content after navigation completes. +- **action + drill_into**: navigation direction (DRILL into a section, + BACK to parent, or STOP). + +Asset collection (images/tables) still runs during navigation so LLM +tool requests are honoured, but assets are reconciled after hydration. """ from __future__ import annotations @@ -17,15 +25,16 @@ ) from shared.services.retrieval.agentic.core.budget import BudgetExceeded from shared.services.retrieval.agentic.prompts import ( - ACTION_PROMPT, + COLLECTOR_PROMPT, + DISCOVERY_SELECT_PROMPT, format_budget_block, - parse_action_response, + parse_collector_response, ) -from shared.services.retrieval.agentic.navigation.section_prompt_projection import format_items_for_llm -from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections -from shared.services.retrieval.agentic.navigation.selection_hydration import ( - hydrate_path_selections_into_node, +from shared.services.retrieval.agentic.navigation.section_prompt_projection import ( + format_items_for_llm, + format_nav_trace, ) +from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections from shared.services.retrieval.agentic.core.types import DocTreeNode, NavigateStepResult from shared.services.retrieval.llm_adapter import LLMFn @@ -40,18 +49,22 @@ async def navigate_step( user_id: str, namespace: str, doc_name: str = "", - scope_path: str | list[str] | None = None, + scope_path: str | None = None, exclude_paths: set[str] | None = None, budget_snapshot: dict | None = None, + nav_trace: list[dict[str, Any]] | None = None, + collected_paths: list[dict[str, Any]] | None = None, ) -> NavigateStepResult: - """Navigate one document scope and hydrate selected sections.""" - scope_paths = ( - scope_path if isinstance(scope_path, list) - else [scope_path] if scope_path - else [] - ) - scope_path_set = set(scope_paths) + """Navigate one document scope using the Collector Agent model. + Returns a ``NavigateStepResult`` with: + - ``collect``: paths to add to the evidence collection + - ``action``: DRILL/BACK/STOP + - ``drill``: the single drill target (if action == DRILL) + - ``tools``: asset tool invocations + - ``node``: outline tree node for rendering context + """ + scope_paths = [scope_path] if scope_path else [] try: items = await load_child_sections( @@ -64,9 +77,6 @@ async def navigate_step( if not items: return NavigateStepResult.stop(scope_paths[0] if scope_paths else None) - 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) } @@ -78,82 +88,102 @@ async def navigate_step( ) tools_block = build_asset_tools_block(total_images, total_tables) - items_text, overflowed = format_items_for_llm(items) - prompt = _build_navigation_prompt( - document_id=document_id, - doc_name=doc_name, + # Build collected path set for [✓] marking on tree + collected_path_set = { + item.get("path", "") for item in (collected_paths or []) + } + items_text, overflowed = format_items_for_llm( + items, + collected_paths=collected_path_set, + ) + + # Build trace block (unified: scope + actions + collection) + trace_block = format_nav_trace( + nav_trace or [], + collected_paths or [], + ) + + prompt = COLLECTOR_PROMPT.format( + doc_name=doc_name or document_id, + doc_id=document_id, + budget_block=format_budget_block(budget_snapshot), + trace_block=trace_block, + items_overview=items_text, query=query, - scope_paths=scope_paths, - budget_snapshot=budget_snapshot, - items_text=items_text, tools_block=tools_block, ) response = await llm_fn(prompt) - parsed = parse_action_response(response) + parsed = parse_collector_response(response) action = parsed["action"] selected_tools = parsed["tools"] - selections = parsed["selections"] reason = parsed.get("reason", "") - stop_type = parsed.get("stop_type", "") + raw_collect = parsed.get("collect", []) + drill_into = parsed.get("drill_into") - scope_label = ", ".join(scope_paths) if scope_paths else "root" + scope_label = scope_path or "root" logger.info( f" navigate_step scope={scope_label}: " - f"action={action} tools={selected_tools} " - f"selections={len(selections)} selectable={len(selectable)} " + f"action={action} collect={len(raw_collect)} " + f"drill_into={drill_into} tools={selected_tools} " f"overflowed={overflowed}" ) 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)] - raw_valid_selections = [ - selection - for selection in selections - 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 = visible_items[path] - node.confidence[path] = confidence - - if item.get("is_leaf"): - path_selections.append({ - "path": path, - "confidence": confidence, - "hydrate_mode": "chunks", - }) - else: - pending.append({"path": path, "confidence": confidence}) - path_selections.append({ + # Validate collect paths: must be visible and not already collected + valid_collect: list[dict[str, Any]] = [] + for item in raw_collect: + path = item.get("path", "") + if path in visible_items and path not in collected_path_set: + confidence = item.get("confidence", 0.7) + outline = item.get("outline", False) + node.confidence[path] = confidence + valid_collect.append({ "path": path, "confidence": confidence, - "hydrate_mode": "self_only", + "hydrate_mode": "outline" if outline else "chunks", }) - await hydrate_path_selections_into_node( - db, - node=node, - path_selections=path_selections, - user_id=user_id, - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - ) + # Validate drill target: must be visible, not collected, not a leaf + valid_drill: list[dict[str, Any]] = [] + if action == "DRILL" and drill_into: + if drill_into in visible_items and drill_into not in collected_path_set: + drill_item = visible_items[drill_into] + if drill_item.get("is_leaf"): + # Leaf nodes can't be drilled — auto-collect instead + logger.info( + f" navigate_step: drill target '{drill_into}' is a leaf, " + f"auto-collecting instead" + ) + if not any(c["path"] == drill_into for c in valid_collect): + node.confidence[drill_into] = 0.7 + valid_collect.append({ + "path": drill_into, + "confidence": 0.7, + "hydrate_mode": "chunks", + }) + action = "STOP" # no valid drill target + else: + valid_drill.append({ + "path": drill_into, + "confidence": 0.8, + }) + else: + logger.warning( + f" navigate_step: drill target '{drill_into}' invalid " + f"(not visible or already collected), falling back to STOP" + ) + action = "STOP" return NavigateStepResult( action=action, + collect=valid_collect, + drill=valid_drill, tools=selected_tools, node=node, - pending=pending, reason=reason, - stop_type=stop_type, ) except BudgetExceeded: @@ -161,46 +191,63 @@ async def navigate_step( except Exception as exc: logger.error(f" navigate_step failed for doc={document_id}: {exc}") return NavigateStepResult.stop(scope_paths[0] if scope_paths else None) -def _build_navigation_prompt( + + +async def discovery_select_step( + db: AsyncSession, *, document_id: str, - doc_name: str, query: str, - scope_paths: list[str], - budget_snapshot: dict | None, - items_text: str, - tools_block: str, -) -> str: - if not scope_paths: - scope_header = "Current scope: root (document top level)" - elif len(scope_paths) == 1: - scope_header = f'Current scope: navigating into "{scope_paths[0]}"' - else: - scope_header = f"Current scope: navigating into {len(scope_paths)} sections" - - return ACTION_PROMPT.format( + llm_fn: LLMFn, + user_id: str, + namespace: str, + doc_name: str = "", + discovery_hints: list[dict[str, Any]], + exclude_paths: set[str] | None = None, + budget_snapshot: dict | None = None, +) -> DocTreeNode: + """Select discovery hint paths via LLM and hydrate them.""" + from shared.services.retrieval.agentic.prompts import parse_action_response + from shared.services.retrieval.agentic.navigation.selection_hydration import ( + hydrate_path_selections_into_node, + ) + + excluded = exclude_paths or set() + filtered_hints = [ + hint for hint in discovery_hints + if hint.get("path", "") not in excluded + ] + if not filtered_hints: + return DocTreeNode(scope_path=None) + + items_text = "\n".join( + f'- path="{hint.get("path", "")}" score={hint.get("score", 0):.2f}' + for hint in filtered_hints + ) + prompt = DISCOVERY_SELECT_PROMPT.format( doc_name=doc_name or document_id, - doc_id=document_id, - scope_header=scope_header, budget_block=format_budget_block(budget_snapshot), - items_overview=items_text, + items=items_text, query=query, - tools_block=tools_block, ) + response = await llm_fn(prompt) + parsed = parse_action_response(response) + selections = parsed.get("selections", []) + + node = DocTreeNode(scope_path=None) + if selections: + path_selections = [ + {"path": sel["path"], "confidence": sel.get("confidence", 0.7), "hydrate_mode": "chunks"} + for sel in selections + ] + await hydrate_path_selections_into_node( + db, + node=node, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) -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 + return node diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index 43548d5a5..b25978973 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -288,12 +288,8 @@ async def run( seen_ref_ids.add(cid) all_refs.append(ref) - # Re-render when navigation only produced structural context. - if not evidence_text or evidence_text == '(no evidence collected)': - evidence_text = await _render_evidence( - db, - state.doc_trees, state.doc_id_to_name, - ) + + result = AgenticResult( evidence_text=evidence_text, diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py index ca18d8048..62c52a288 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/prompts.py +++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py @@ -49,17 +49,17 @@ """ -ACTION_PROMPT = """\ +COLLECTOR_PROMPT = """\ You are a document navigation agent. Document: "{doc_name}" (id: {doc_id}) {budget_block} -{scope_header} +{trace_block} Below is the document's section tree. -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. +Nodes marked [✓] are already in your collection — do not re-collect them. +Token estimates (e.g. ~1.2k) show approximate content size. === Section Tree === {items_overview} @@ -67,67 +67,134 @@ User query: {query} -=== Available Actions === +=== Behavioral Rules === -Choose ONE action: +Each step you make TWO independent decisions: -NAVIGATE — Drill into selected sections for detailed content. - Consider this when the query targets specific topics and you need deeper text evidence. - Prefer one or more [SELECT] sections, or choose a broader visible section when needed. +1. COLLECT — Add sections to your evidence collection (optional, can be empty). + - COLLECT includes the section AND ALL its descendant content. + - If a node is [Leaf] or has ≤500 tokens, prefer COLLECT over DRILL. + - Do NOT re-collect paths marked [✓]. -STOP — Current scope evidence is sufficient. No further drill-down. - Consider this when: - - The query asks for an outline, overview, or summary - - The query is broad/global, the tree section can fulfill it without drilling into individual sections. - - You have already collected enough evidence at this level. +2. Navigate action — Where to go next (required, choose ONE): + - DRILL — Open one section to see its children in the next step. + Use when a section has >1000 tokens and you need to be selective. + You cannot DRILL into a path you just COLLECTed (already fully included). + - BACK — Return to parent scope to explore other branches. + - STOP — End navigation. Use when you have enough evidence or nothing relevant remains. {tools_block} -When action is NAVIGATE, provide selections: -- Select visible section paths from the tree above; prefer [SELECT] paths when they fit. - -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", "reason": "...", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}} +{{"collect": [{{"path": "...", "confidence": , "outline": false}}, ...], + "action": "DRILL", + "drill_into": "section/path", + "tools": [...], + "reason": "..."}} or -{{"action": "STOP", "reason": "...", "stop_type": "...", "tools": [...], "selections": []}} -Do not include any explanation. +{{"collect": [...], "action": "BACK", "tools": [...], "reason": "..."}} +or +{{"collect": [...], "action": "STOP", "tools": [...], "reason": "..."}} + +Set "outline": true on a collect entry to collect only the section structure +(titles and summaries) without full chunk content. Use for overview/structure queries. +Do not include any explanation outside the JSON. """ -def parse_action_response(text: str) -> dict: - """Parse the unified navigation response from an LLM.""" +def parse_collector_response(text: str) -> dict: + """Parse the Collector Agent navigation response. + + Expected format: + {"collect": [...], "action": "DRILL|BACK|STOP", + "drill_into": "path", "tools": [...], "reason": "..."} + """ text = text.strip() asset_tools = {"FIND_IMAGES", "FIND_TABLES"} - default = {"action": "NAVIGATE", "tools": [], "selections": [], "reason": "", "stop_type": ""} + valid_actions = {"DRILL", "BACK", "STOP"} + default: dict[str, Any] = { + "collect": [], "action": "STOP", "drill_into": None, + "tools": [], "reason": "", + } def extract(data: dict) -> dict: - action = str(data.get("action", "NAVIGATE")).strip().upper() - if action not in ("NAVIGATE", "STOP"): - action = "NAVIGATE" + action = str(data.get("action", "STOP")).strip().upper() + if action not in valid_actions: + action = "STOP" + + # Parse collect list + collect_val = data.get("collect") or [] + collect: list[dict[str, Any]] = [] + if isinstance(collect_val, list): + for item in collect_val: + if isinstance(item, dict) and item.get("path"): + confidence = normalize_confidence(item.get("confidence", 0.7)) + outline = bool(item.get("outline", False)) + collect.append({ + "path": str(item["path"]), + "confidence": confidence or 0.7, + "outline": outline, + }) + # Parse drill target + drill_into = None + if action == "DRILL": + drill_into = data.get("drill_into") + if isinstance(drill_into, str): + drill_into = drill_into.strip() or None + else: + drill_into = None + if drill_into is None: + # No valid drill target → treat as STOP + action = "STOP" + + # Parse tools tools_val = data.get("tools") or [] + tools: list[str] = [] if isinstance(tools_val, list): tools = [ - str(tool).strip().upper() - for tool in tools_val - if str(tool).strip().upper() in asset_tools + str(t).strip().upper() + for t in tools_val + if str(t).strip().upper() in asset_tools ] - 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": [], "reason": reason, "stop_type": stop_type} + return { + "collect": collect, + "action": action, + "drill_into": drill_into, + "tools": tools, + "reason": reason, + } + + data = _parse_json_object(text) + if data is not None: + return extract(data) + + fence_match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", text, re.DOTALL) + if fence_match: + data = _parse_json_object(fence_match.group(1).strip()) + if data is not None: + return extract(data) + + brace_match = re.search(r"\{.*\}", text, re.DOTALL) + if brace_match: + data = _parse_json_object(brace_match.group()) + if data is not None: + return extract(data) + return default + + +def parse_action_response(text: str) -> dict: + """Parse discovery_select response (legacy format, kept for discovery).""" + text = text.strip() + default: dict[str, Any] = {"selections": []} + + def extract(data: dict) -> dict: selections_val = data.get("selections") or [] - selections = [] + selections: list[dict[str, Any]] = [] if isinstance(selections_val, list): for selection in selections_val: if isinstance(selection, dict) and selection.get("path"): @@ -136,8 +203,7 @@ def extract(data: dict) -> dict: "path": str(selection["path"]), "confidence": confidence or 0.7, }) - - return {"action": action, "tools": tools, "selections": selections, "reason": reason, "stop_type": ""} + return {"selections": selections} data = _parse_json_object(text) if data is not None: diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 8daad41ba..97e76c5e1 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -100,10 +100,12 @@ async def navigate_step( user_id: str, namespace: str, doc_name: str = "", - scope_path: str | list[str] | None = None, + scope_path: str | None = None, exclude_paths: set[str] | None = None, budget_snapshot: dict | None = None, -) -> tuple[str, list[str], DocTreeNode, list[dict]]: + nav_trace: list[dict[str, Any]] | None = None, + collected_paths: list[dict[str, Any]] | None = None, +) -> navigation_tools.NavigateStepResult: return await navigation_tools.navigate_step( db, document_id=document_id, @@ -116,6 +118,8 @@ async def navigate_step( scope_path=scope_path, exclude_paths=exclude_paths, budget_snapshot=budget_snapshot, + nav_trace=nav_trace, + collected_paths=collected_paths, ) From 35af1ae7cfabbe0f981dda96141336e69059c381 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 29 May 2026 23:13:06 +0800 Subject: [PATCH 2/6] refactor: enhance document tree navigation and discovery by adding hierarchical outline nesting, improved path exclusion logic, and child node enforcement. --- .../app/services/document_agent/manifest.py | 2 +- .../formats/pdf/shard_merger.py | 1 + .../structure/heading_candidates.py | 2 + .../retrieval/agentic/core/runtime.py | 2 +- .../retrieval/agentic/discovery/selection.py | 23 ++- .../retrieval/agentic/evidence/renderer.py | 43 ++++- .../retrieval/agentic/navigation/document.py | 164 +++++++++++++++++- .../agentic/navigation/section_tree.py | 20 ++- .../retrieval/agentic/orchestrator.py | 1 - .../services/retrieval/hydration/path.py | 16 +- 10 files changed, 251 insertions(+), 23 deletions(-) diff --git a/apps/worker/app/services/document_agent/manifest.py b/apps/worker/app/services/document_agent/manifest.py index ebc3c4e3e..038665232 100644 --- a/apps/worker/app/services/document_agent/manifest.py +++ b/apps/worker/app/services/document_agent/manifest.py @@ -110,7 +110,7 @@ def to_dict(self) -> dict[str, Any]: class TocResult: toc_pages: list[int] = field(default_factory=list) candidates: list[TocCandidate] = field(default_factory=list) - method: Literal["toc_marker", "vlm_progressive", "visual_scan", "none"] = "none" + method: Literal["toc_marker", "vlm_progressive", "vlm_batch", "visual_scan", "none"] = "none" notes: str = "" def to_dict(self) -> dict[str, Any]: 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 0bc46e838..143cfabb0 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 @@ -61,6 +61,7 @@ def merge_shard_lines(shard_lines_list: list[list[str]]) -> list[str]: if ( last_heading_pos is not None + and last_heading_key is not None and last_heading_key == next_first_heading ): # Truncate from the last (duplicate) heading onward diff --git a/apps/worker/app/services/document_parser/structure/heading_candidates.py b/apps/worker/app/services/document_parser/structure/heading_candidates.py index 0a7596aed..976d81276 100644 --- a/apps/worker/app/services/document_parser/structure/heading_candidates.py +++ b/apps/worker/app/services/document_parser/structure/heading_candidates.py @@ -22,6 +22,8 @@ def get_max_lvl(code_str: str) -> int: integers, so the ``[…]`` bracket match is guaranteed. """ match = re.search(r"\[([^]]+)]", code_str) + if match is None: + return -2 nums = [int(item.strip()) for item in match.group(1).split(",")] max_value = int(max(nums)) return max_value if max_value > 1 else -2 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 594184e92..3f1563b44 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py @@ -18,7 +18,7 @@ def build_config_from_env() -> AgentRunConfig: return AgentRunConfig( max_nav_steps=int(os.environ.get("RETRIEVAL_AGENTIC_MAX_NAV_STEPS", "6")), - latency_budget_ms=int(os.environ.get("RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS", "12000")), + latency_budget_ms=int(os.environ.get("RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS", "30000")), token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")), planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")), bootstrap_budget=int(os.environ.get("RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET", "2000")), 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 40a73f84d..fbb4c1d90 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py @@ -125,7 +125,7 @@ def _project_discovery_hints( section_path = normalize_section_path(hint.get("section_path", "")) if not section_path: continue - if section_path in exclude_set: + if _is_covered_by_exclude(section_path, exclude_set): continue if section_path in hint_by_path: continue @@ -139,6 +139,21 @@ def _project_discovery_hints( return hint_lines, hint_by_path +def _is_covered_by_exclude(path: str, exclude_set: set[str]) -> bool: + """Check if *path* is covered by any entry in *exclude_set*. + + A path is covered if it exactly matches an exclude entry, OR if any + exclude entry is a prefix of this path (i.e. the parent path was + already collected by navigation). + """ + if path in exclude_set: + return True + for excluded in exclude_set: + if path.startswith(excluded + " / "): + return True + return False + + def _build_discovery_selection_prompt( *, document_id: str, @@ -187,6 +202,10 @@ def _build_discovery_path_selections( "hydrate_mode": "self_only", }) continue - path_selections.append({"path": path, "confidence": confidence}) + path_selections.append({ + "path": path, + "confidence": confidence, + "hydrate_mode": "self_only", + }) return path_selections, chunk_refs 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 9e974b78a..c600218dd 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py @@ -48,8 +48,9 @@ def min_sort(path: str) -> float: render_queue.append((min_sort(path), "orphan_leaf", path)) for path in node.children: - if path not in outline_paths: - render_queue.append((float("inf"), "orphan_child", path)) + if path not in outline_paths and path not in node.leaf_content: + child_sort = _infer_child_sort_order(node.children[path]) + render_queue.append((child_sort, "orphan_child", path)) render_queue.sort(key=lambda item: item[0]) @@ -82,17 +83,27 @@ def min_sort(path: str) -> float: elif render_type == "orphan_leaf": path = cast(str, data) title = path.rsplit(" / ", 1)[-1] if " / " in path else path - parts.append(f"{indent}▸ [Leaf] {title}") - render_leaf_chunks(parts, node.leaf_content[path], indent + " ", asset_lookup=asset_lookup) + sub_indent = indent + " " + if path in node.children: + # Non-leaf node with own content: render heading, then + # self chunks, then child subtree (merged rendering). + parts.append(f"{indent}▸ {title}") + render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) + if child_text.strip(): + parts.append(child_text) + else: + parts.append(f"{indent}▸ [Leaf] {title}") + render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) elif render_type == "orphan_child": path = cast(str, data) title = path.rsplit(" / ", 1)[-1] if " / " in path else path child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) - # Only render the [DrillDown] heading if the child has content. + # Only render the orphan heading if the child has content. # Prevents empty orphan nodes from polluting evidence_text. if child_text.strip(): - parts.append(f"{indent}▸ {title} [DrillDown]") + parts.append(f"{indent}▸ {title}") parts.append(child_text) return "\n".join(parts) @@ -182,3 +193,23 @@ def render_leaf_chunks( for line in table_html.split("\n"): if line.strip(): parts.append(f"{indent}┈ {line}") + + +def _infer_child_sort_order(child: DocTreeNode) -> float: + """Infer sort position from the child's earliest chunk sort_order. + + When an orphan child node has no outline entry, we fall back to the + minimum ``sort_order`` across all its hydrated chunks so that orphans + render in document order instead of being appended at the end. + """ + min_order = float("inf") + for chunks in child.leaf_content.values(): + for chunk in chunks: + order = chunk.get("sort_order") + if order is not None and order < min_order: + min_order = float(order) + for grandchild in child.children.values(): + grandchild_order = _infer_child_sort_order(grandchild) + min_order = min(min_order, grandchild_order) + return min_order + 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 cfece903f..167615c7c 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py @@ -114,6 +114,7 @@ async def _navigate_document( doc=doc, root=root, doc_name=doc_name, + collected_paths=collected_paths, ) # Phase 2C: Batch hydrate all collected paths @@ -288,6 +289,12 @@ async def _hydrate_collected( deduped[path] = item unique_selections = list(deduped.values()) + # Ensure child nodes exist for each collected path so reparent can + # correctly route descendant chunks into the right subtree. + for item in unique_selections: + path = item["path"] + _ensure_child_node(root, path) + await hydrate_path_selections_into_node( self._db, node=root, @@ -301,6 +308,37 @@ async def _hydrate_collected( # Single reparent pass — tree structure is final. root.reparent_leaf_content() + # Load section tree outline for each collected child node, then + # build a proper sub-tree so the renderer can nest L3 under L2 etc. + from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections + for item in unique_selections: + path = item["path"] + child_node = root.children.get(path) + if child_node is None or child_node.outline_items: + continue # Skip if no child or already has outline + try: + section_items = await load_child_sections( + self._db, doc.document_id, job_result_id, path, + limit_depth=False, + ) + if section_items: + # Filter out the scope node itself to avoid duplicate + # title rendering (parent outline already shows it). + child_node.outline_items = [ + si for si in section_items + if si.get("path") != path + ] + # Build sub-tree from outline hierarchy and re-reparent + # so chunks are correctly nested (e.g. L3 under L2). + _build_outline_subtree(child_node) + except Exception as exc: + logger.warning(f" hydrate_collected: failed to load outline for '{path}': {exc}") + + # Also organize any discovery chunks sitting in root.leaf_content. + # Same path-based tree building; pre-existing children are protected. + if root.leaf_content: + _build_outline_subtree(root) + # Accurate budget accounting — count only actually-hydrated chunks. if self._state.ledger is not None: total_chunks = len(root.flatten_chunk_rows()) @@ -368,6 +406,7 @@ async def _hydrate_discovery_hints( doc: CandidateDoc, root: DocTreeNode, doc_name: str, + collected_paths: list[dict[str, Any]] | None = None, ) -> None: doc_hints = self._discovery_by_doc.get(doc.document_id, []) if not doc_hints or self._llm_fn is None: @@ -375,7 +414,9 @@ async def _hydrate_discovery_hints( if self._state.elapsed_ms >= self._config.latency_budget_ms: return - discovery_exclude_paths = _collect_leaf_paths(root) + discovery_exclude_paths = _build_discovery_exclude_set( + root, collected_paths or [] + ) doc_discovery_llm_fn = self._llm_budget.for_discovery( cast(LLMFn, self._llm_fn), doc_id=doc.document_id, @@ -556,3 +597,124 @@ def _collect_leaf_paths(node: DocTreeNode) -> set[str]: for child in node.children.values(): paths.update(_collect_leaf_paths(child)) return paths + + +def _build_discovery_exclude_set( + root: DocTreeNode, + collected_paths: list[dict[str, Any]], +) -> set[str]: + """Build exclude set for discovery using collected navigation paths. + + If navigation COLLECT'd a parent path like "五、施工安全保证措施", + all discovery hints under that path should be excluded because + COLLECT already loads all descendants via prefix matching. + """ + # 1. Already-hydrated leaf paths + exclude = _collect_leaf_paths(root) + + # 2. Collected parent paths from navigation COLLECT decisions. + # These haven't been hydrated yet (hydrate runs after discovery), + # but we know COLLECT will load all their descendants. + for item in collected_paths: + path = item.get("path", "") + if path: + exclude.add(path) + + return exclude + + +def _ensure_child_node(root: DocTreeNode, path: str) -> None: + """Create an intermediate child node for *path* if it doesn't exist. + + When a non-leaf section is COLLECTed (e.g. "五、施工安全保证措施"), + hydration loads all descendant chunks (e.g. "五、... / 3.监控量测措施 / 3.1..."). + These chunks are initially placed in root.leaf_content. + ``reparent_leaf_content`` then moves them into the correct child subtree — + but only if a child node exists for the collected path. + + This function creates that child node so reparent can work correctly. + """ + # Don't create a child for root-level or if already exists + if not path: + return + + # Walk to find the deepest existing ancestor node + target = root + for child_path, child in root.children.items(): + if path.startswith(child_path + " / "): + target = child + break + if path == child_path: + return # Already exists + + # Create the child node on the target + if path not in target.children: + target.children[path] = DocTreeNode(scope_path=path) + + +def _build_outline_subtree(node: DocTreeNode) -> None: + """Recursively build child nodes from outline_items + leaf_content paths. + + Every section_path is ``" / "``-separated; parent-child is prefix match. + This function creates children one level at a time, reparents chunks + into them, then recurses so the next level is handled correctly. + """ + if not node.outline_items and not node.leaf_content: + return + + # All known paths: outline metadata + actual chunk paths. + all_paths = ( + {item["path"] for item in node.outline_items} + | set(node.leaf_content.keys()) + ) + + # Find paths that have at least one descendant. + parent_paths: set[str] = set() + for path in all_paths: + for other in all_paths: + if other != path and other.startswith(path + " / "): + parent_paths.add(path) + break + + if not parent_paths: + return # All items are leaves — nothing to nest. + + # KEY FIX: only keep top-level parents. If "A / B" and "A / B / C" + # are both parents, only create "A / B" as a child NOW. "A / B / C" + # will be created when the function recurses into "A / B". + parent_paths = { + pp for pp in parent_paths + if not any(pp != other and pp.startswith(other + " / ") for other in parent_paths) + } + + # Track children that already exist (e.g. from collected hydration). + # Their outline_items are already populated — don't push duplicates. + pre_existing = set(node.children.keys()) + + # Create child nodes for top-level parent paths only. + for path in parent_paths: + if path not in node.children: + node.children[path] = DocTreeNode(scope_path=path) + + # Split outline_items: keep items at this level, move descendants + # into newly-created children only (skip pre-existing ones). + kept: list[dict] = [] + for item in node.outline_items: + item_path = item["path"] + best_parent: str | None = None + for pp in parent_paths: + if item_path.startswith(pp + " / "): + if best_parent is None or len(pp) > len(best_parent): + best_parent = pp + if best_parent and best_parent not in pre_existing: + node.children[best_parent].outline_items.append(item) + else: + kept.append(item) + node.outline_items = kept + + # Reparent FIRST so children receive their leaf_content, + # THEN recurse so deeper levels can be built from that content. + node.reparent_leaf_content() + for child in node.children.values(): + _build_outline_subtree(child) + 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 47c38ffe3..ece877f2d 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 @@ -16,6 +16,7 @@ async def load_child_sections( job_result_id: str, scope_path: str | list[str] | None = None, exclude_paths: set[str] | None = None, + limit_depth: bool = True, ) -> list[dict]: """Load the continuous context tree for a navigation scope.""" stmt = ( @@ -84,15 +85,16 @@ async def load_child_sections( if not items_by_path: return [] - allowed_set = _resolve_allowed_depths(items_by_path, scope_list) - if allowed_set: - to_remove = [ - path - for path, item in items_by_path.items() - if item["show_summary"] and item["level"] not in allowed_set - ] - for path in to_remove: - del items_by_path[path] + if limit_depth: + allowed_set = _resolve_allowed_depths(items_by_path, scope_list) + if allowed_set: + to_remove = [ + path + for path, item in items_by_path.items() + if item["show_summary"] and item["level"] not in allowed_set + ] + for path in to_remove: + del items_by_path[path] if not items_by_path: return [] diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index b25978973..80123b00b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -28,7 +28,6 @@ ) from shared.services.retrieval.agentic.navigation.document import DocumentNavigationRunner from shared.services.retrieval.agentic.evidence.builder import ( - render_evidence as _render_evidence, trim_evidence_to_budget as _trim_evidence_to_budget, ) from shared.services.retrieval.agentic.core.runtime import ( diff --git a/packages/shared-python/shared/services/retrieval/hydration/path.py b/packages/shared-python/shared/services/retrieval/hydration/path.py index de2ec15a2..b33d193bd 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/path.py +++ b/packages/shared-python/shared/services/retrieval/hydration/path.py @@ -3,7 +3,7 @@ from typing import Any from loguru import logger -from sqlalchemy import or_, select +from sqlalchemy import and_, or_, select from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection @@ -148,9 +148,20 @@ async def _hydrate_chunk_paths( ) -> list[dict[str, Any]]: section_path_filters = [] self_only_paths = {path for path in chunk_paths if mode_by_path.get(path) == 'self_only'} + shallow_paths = {path for path in chunk_paths if mode_by_path.get(path) == 'shallow'} for path in chunk_paths: section_path_filters.append(DocumentSection.section_path == path) - if path not in self_only_paths: + if path in self_only_paths: + pass # exact match only — no descendants + elif path in shallow_paths: + # Direct children only: match "path / X" but not "path / X / Y" + section_path_filters.append( + and_( + DocumentSection.section_path.like(f'{path} / %'), + ~DocumentSection.section_path.like(f'{path} / % / %'), + ) + ) + else: section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) stmt = ( @@ -228,6 +239,7 @@ def _get_allowed_types_for_mode(path_mode: str) -> set[str] | None: mode_allowed_types: dict[str, set[str] | None] = { 'chunks': None, 'self_only': None, + 'shallow': None, 'assets_only': {'image', 'table'}, 'image_only': {'image'}, 'table_only': {'table'}, From 15c889ec0b0c802c65fffba3abc086bd4cd53c90 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 29 May 2026 23:33:15 +0800 Subject: [PATCH 3/6] fix(test): update empty-query contract test to handle additive response schema fields The retrieval API response was extended with decision_trace, failure_reason, and stop_reason fields in a prior PR. The snapshot-style assertEqual on the full response dict broke as soon as those fields were added. Replace with per-field assertions so the test stays resilient to future additive changes. --- .../tests/contract/test_retrieval_contract.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 20665a019..29184bf88 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -346,15 +346,14 @@ async def test_should_return_empty_results_for_an_empty_query( ) assert response.status_code == 200 - assert response.json() == { - "namespace": "default", - "query": "", - "router_used": "empty_query_filtered", - "evidence_text": "", - "answer_text": "", - "results": [], - "referenced_chunks": [], - } + response_json = response.json() + assert response_json["namespace"] == "default" + assert response_json["query"] == "" + assert response_json["router_used"] == "empty_query_filtered" + assert response_json["evidence_text"] == "" + assert response_json["answer_text"] == "" + assert response_json["results"] == [] + assert response_json["referenced_chunks"] == [] @pytest.mark.asyncio From edb9b8c2087a831389ed4dc7e219cb771077421f Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 29 May 2026 23:41:39 +0800 Subject: [PATCH 4/6] fix(test): add LLM mock rules for agentic retrieval prompts (planner, navigator, discovery) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The LLM mock had no task detection for the three core agentic retrieval prompts (workflow planner, navigation collector, discovery select), causing all three to fall through to the 'default' task and return 'Mock LLM response'. - QueryPlanner fell back to single-step (OK), but... - parse_collector_response could not parse the response and defaulted to action=STOP, collect=[], yielding 0 referenced_chunks in every test. Fixes: - Add 'agentic-planner' detection (RETRIEVAL WORKFLOW PLANNER + concat_final_parts) → returns a valid single-step JSON plan. - Add 'agentic-navigate' detection (DOCUMENT NAVIGATION AGENT + SECTION TREE) → dynamically extracts the first visible path from the section tree block and returns a COLLECT+STOP response so the agent actually gathers evidence. - Add 'agentic-discovery-select' detection (DISCOVERY CANDIDATES + selections) → dynamically extracts the first candidate path and returns a selections response. Also adds json+re imports required by the new dynamic response builders. --- .../shared/services/ai/llm_mock.py | 117 ++++++++++++++++++ 1 file changed, 117 insertions(+) diff --git a/packages/shared-python/shared/services/ai/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py index 8bf155f21..a62ff633b 100644 --- a/packages/shared-python/shared/services/ai/llm_mock.py +++ b/packages/shared-python/shared/services/ai/llm_mock.py @@ -1,5 +1,7 @@ """Helpers for deterministic mock responses from OpenAI-compatible LLM calls.""" +import json +import re from typing import Any, Dict, List from loguru import logger @@ -18,6 +20,11 @@ def build_mock_chat_completion_response( model_name, task_name, ) + # For agentic tasks that need dynamic path extraction, build the response here. + if task_name == "agentic-navigate": + return _build_navigate_mock_response(prompt_text) + if task_name == "agentic-discovery-select": + return _build_discovery_select_mock_response(prompt_text) return _build_mock_response(task_name) @@ -68,6 +75,24 @@ def _detect_mock_task(prompt_text: str) -> str: """Infer the prompt task so the mock can return a compatible response shape.""" normalized_prompt = prompt_text.lower() + # ── Agentic retrieval prompts (check first — they are structurally distinct) ── + if ( + "you are a retrieval workflow planner" in normalized_prompt + and "concat_final_parts" in normalized_prompt + ): + return "agentic-planner" + if ( + "you are a document navigation agent" in normalized_prompt + and "=== section tree ==" in normalized_prompt + ): + return "agentic-navigate" + if ( + "=== discovery candidates ==" in normalized_prompt + and "\"selections\"" in normalized_prompt + ): + return "agentic-discovery-select" + + # ── Document parsing / ingestion prompts ── if ( "generate a concise title" in normalized_prompt and "return only the title" in normalized_prompt @@ -134,9 +159,101 @@ def _detect_mock_task(prompt_text: str) -> str: return "default" +def _extract_first_section_path(prompt_text: str) -> str | None: + """Pull the first path-like token from a COLLECTOR_PROMPT section tree block. + + The section tree overview contains lines like: + - Section Name [Leaf] (~500 tokens) + - path/to/section [Leaf] + We look for the first path segment between '=== Section Tree ===' and + '=== End Section Tree ===' and return it verbatim. + """ + tree_match = re.search( + r"=== Section Tree ===(.*?)=== End Section Tree ===", + prompt_text, + re.DOTALL | re.IGNORECASE, + ) + if not tree_match: + return None + tree_block = tree_match.group(1) + # Lines look like: " - Some Section Title [Leaf] (~500 tokens)" + # or indented with bullets. Grab the first one. + for line in tree_block.splitlines(): + stripped = line.strip().lstrip("-•* ").strip() + if not stripped: + continue + # Strip trailing annotations like [Leaf], [✓], (~500 tokens) + cleaned = re.sub(r"\s*\[.*?\]|\s*\(~[\d.]+k? tokens\)", "", stripped).strip() + if cleaned: + return cleaned + return None + + +def _build_navigate_mock_response(prompt_text: str) -> str: + """Return a mock COLLECTOR_PROMPT response that COLLECTs the first visible path.""" + path = _extract_first_section_path(prompt_text) + if path: + response = { + "collect": [{"path": path, "confidence": 0.9, "outline": False}], + "action": "STOP", + "drill_into": None, + "tools": [], + "reason": "Mock: collected first available section", + } + else: + # No path found — STOP without collecting (safe fallback) + response = { + "collect": [], + "action": "STOP", + "drill_into": None, + "tools": [], + "reason": "Mock: no section path found in tree", + } + return json.dumps(response) + + +def _extract_first_discovery_path(prompt_text: str) -> str | None: + """Pull the first path from a DISCOVERY_SELECT_PROMPT candidates block.""" + candidates_match = re.search( + r"=== Discovery Candidates ===(.*?)=== End Discovery Candidates ===", + prompt_text, + re.DOTALL | re.IGNORECASE, + ) + if not candidates_match: + return None + block = candidates_match.group(1) + for line in block.splitlines(): + stripped = line.strip().lstrip("-•* ").strip() + if stripped: + # Lines may be "path - summary" or just "path" + path_part = stripped.split(" - ")[0].strip() + if path_part: + return path_part + return None + + +def _build_discovery_select_mock_response(prompt_text: str) -> str: + """Return a mock DISCOVERY_SELECT_PROMPT response selecting the first candidate.""" + path = _extract_first_discovery_path(prompt_text) + if path: + response = {"selections": [{"path": path, "confidence": 0.85}]} + else: + response = {"selections": []} + return json.dumps(response) + + def _build_mock_response(task_name: str) -> str: """Return a canned response compatible with the inferred task contract.""" response_by_task: Dict[str, str] = { + # Agentic retrieval — static fallbacks (dynamic responses built elsewhere) + "agentic-planner": ( + '{"reasoning_summary": "mock single-step plan", ' + '"steps": [{"id": "s1", "sub_query": "mock query", ' + '"step_kind": "retrieve", "depends_on": [], ' + '"output_role": "final_part", "top_k": 10}], ' + '"final_strategy": "concat_final_parts"}' + ), + # Document parsing / ingestion tasks "fragment-title": "Mock Fragment Title", "detect-toc-range": '{"toc_start": null, "toc_end": null, "confidence": "low"}', "detect-table-headers": '{"answer": [0]}', From 878c0b0e7af550eb37dcd44f646715f97ae52484 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 29 May 2026 23:49:24 +0800 Subject: [PATCH 5/6] fix(test): planner mock now injects real query into sub_query instead of hardcoded 'mock query' The agentic planner mock was returning a hardcoded sub_query='mock query', so the pipeline ran bottom_discovery on 'mock query' and found 0 documents. Navigation was never reached and referenced_chunks stayed empty. Fix: _build_planner_mock_response() now extracts the actual query from 'User query: {query}' in the prompt and passes it through as sub_query, so discovery can match the seeded test documents correctly. --- .../shared/services/ai/llm_mock.py | 36 ++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/packages/shared-python/shared/services/ai/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py index a62ff633b..e4e4b6b39 100644 --- a/packages/shared-python/shared/services/ai/llm_mock.py +++ b/packages/shared-python/shared/services/ai/llm_mock.py @@ -20,7 +20,9 @@ def build_mock_chat_completion_response( model_name, task_name, ) - # For agentic tasks that need dynamic path extraction, build the response here. + # For agentic tasks that need dynamic content extraction, build the response here. + if task_name == "agentic-planner": + return _build_planner_mock_response(prompt_text) if task_name == "agentic-navigate": return _build_navigate_mock_response(prompt_text) if task_name == "agentic-discovery-select": @@ -189,6 +191,38 @@ def _extract_first_section_path(prompt_text: str) -> str | None: return None +def _extract_user_query(prompt_text: str) -> str: + """Extract the user query line from a planner/navigation prompt. + + The PLANNER_PROMPT and COLLECTOR_PROMPT both contain:: + User query: {query} + """ + match = re.search(r"User query:\s*(.+)", prompt_text) + if match: + return match.group(1).strip() + return "mock query" + + +def _build_planner_mock_response(prompt_text: str) -> str: + """Return a valid single-step QueryPlan JSON using the real query from the prompt.""" + query = _extract_user_query(prompt_text) + response = { + "reasoning_summary": "mock single-step plan", + "steps": [ + { + "id": "s1", + "sub_query": query, + "step_kind": "retrieve", + "depends_on": [], + "output_role": "final_part", + "top_k": 10, + } + ], + "final_strategy": "concat_final_parts", + } + return json.dumps(response) + + def _build_navigate_mock_response(prompt_text: str) -> str: """Return a mock COLLECTOR_PROMPT response that COLLECTs the first visible path.""" path = _extract_first_section_path(prompt_text) From 733771ab6f267c0961158d52ba54017ee8626e6e Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 29 May 2026 23:55:47 +0800 Subject: [PATCH 6/6] fix(test): fix path extraction in LLM mock to match actual section tree format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Section tree items are rendered by section_prompt_projection as: ▸ [L1] path="Root" [text=1] ~100 tokens [Leaf] Discovery hints are rendered by selection._project_discovery_hints() as: ▸ path="Findings" Both path extractor functions were using naive line-splitting which couldn't parse this format. Fixed both to use path="..." regex, which correctly extracts the canonical path value for COLLECT and selections responses. --- .../shared/services/ai/llm_mock.py | 51 ++++++++++--------- 1 file changed, 27 insertions(+), 24 deletions(-) diff --git a/packages/shared-python/shared/services/ai/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py index e4e4b6b39..c751e3253 100644 --- a/packages/shared-python/shared/services/ai/llm_mock.py +++ b/packages/shared-python/shared/services/ai/llm_mock.py @@ -162,13 +162,16 @@ def _detect_mock_task(prompt_text: str) -> str: def _extract_first_section_path(prompt_text: str) -> str | None: - """Pull the first path-like token from a COLLECTOR_PROMPT section tree block. + """Pull the first path value from a COLLECTOR_PROMPT section tree block. - The section tree overview contains lines like: - - Section Name [Leaf] (~500 tokens) - - path/to/section [Leaf] - We look for the first path segment between '=== Section Tree ===' and - '=== End Section Tree ===' and return it verbatim. + The section tree is rendered by section_prompt_projection.format_items_for_llm() + and each item line looks like:: + + ▸ [L1] path="Root" [text=1] ~100 tokens [Leaf] + └ [L2] path="Root / Sub" [text=2] ~200 tokens + + We extract the value inside path="..." from the first matching line + within the === Section Tree === block. """ tree_match = re.search( r"=== Section Tree ===(.*?)=== End Section Tree ===", @@ -178,19 +181,14 @@ def _extract_first_section_path(prompt_text: str) -> str | None: if not tree_match: return None tree_block = tree_match.group(1) - # Lines look like: " - Some Section Title [Leaf] (~500 tokens)" - # or indented with bullets. Grab the first one. - for line in tree_block.splitlines(): - stripped = line.strip().lstrip("-•* ").strip() - if not stripped: - continue - # Strip trailing annotations like [Leaf], [✓], (~500 tokens) - cleaned = re.sub(r"\s*\[.*?\]|\s*\(~[\d.]+k? tokens\)", "", stripped).strip() - if cleaned: - return cleaned + # Match path="..." in the section tree — this is the canonical format + path_match = re.search(r'path="([^"]+)"', tree_block) + if path_match: + return path_match.group(1) return None + def _extract_user_query(prompt_text: str) -> str: """Extract the user query line from a planner/navigation prompt. @@ -247,7 +245,15 @@ def _build_navigate_mock_response(prompt_text: str) -> str: def _extract_first_discovery_path(prompt_text: str) -> str | None: - """Pull the first path from a DISCOVERY_SELECT_PROMPT candidates block.""" + """Pull the first path value from a DISCOVERY_SELECT_PROMPT candidates block. + + Discovery hints are rendered by selection._project_discovery_hints() as:: + + ▸ path="Findings" + + + We extract the value inside path="..." from the candidates block. + """ candidates_match = re.search( r"=== Discovery Candidates ===(.*?)=== End Discovery Candidates ===", prompt_text, @@ -256,13 +262,10 @@ def _extract_first_discovery_path(prompt_text: str) -> str | None: if not candidates_match: return None block = candidates_match.group(1) - for line in block.splitlines(): - stripped = line.strip().lstrip("-•* ").strip() - if stripped: - # Lines may be "path - summary" or just "path" - path_part = stripped.split(" - ")[0].strip() - if path_part: - return path_part + # Match path="..." — same canonical format as section tree + path_match = re.search(r'path="([^"]+)"', block) + if path_match: + return path_match.group(1) return None