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/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 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/ai/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py index 8bf155f21..c751e3253 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,13 @@ def build_mock_chat_completion_response( model_name, task_name, ) + # 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": + return _build_discovery_select_mock_response(prompt_text) return _build_mock_response(task_name) @@ -68,6 +77,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 +161,136 @@ def _detect_mock_task(prompt_text: str) -> str: return "default" +def _extract_first_section_path(prompt_text: str) -> str | None: + """Pull the first path value from a COLLECTOR_PROMPT section tree block. + + 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 ===", + prompt_text, + re.DOTALL | re.IGNORECASE, + ) + if not tree_match: + return None + tree_block = tree_match.group(1) + # 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. + + 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) + 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 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, + re.DOTALL | re.IGNORECASE, + ) + if not candidates_match: + return None + block = candidates_match.group(1) + # 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 + + +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]}', 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..3f1563b44 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py @@ -17,8 +17,8 @@ def build_config_from_env() -> AgentRunConfig: return AgentRunConfig( - 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")), + 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", "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")), @@ -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/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/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..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,15 +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 - 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 orphan heading if the child has content. + # Prevents empty orphan nodes from polluting evidence_text. if child_text.strip(): + parts.append(f"{indent}▸ {title}") parts.append(child_text) return "\n".join(parts) @@ -180,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 f7e711fa5..167615c7c 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,34 @@ 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, + collected_paths=collected_paths, ) + # 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 +143,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 +191,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 +204,162 @@ 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()) + + # 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, + 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() + + # 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()) + 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 +387,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, }, @@ -244,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: @@ -251,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, @@ -343,41 +508,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 +545,176 @@ 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: +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 - 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)) + + # 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_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..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 [] @@ -110,7 +112,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 +177,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 +217,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..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 ( @@ -288,12 +287,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, ) 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'},