diff --git a/AGENTS.md b/AGENTS.md
index 1bec3165d..a47ed0712 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -44,19 +44,20 @@ knowhereapi-main/
│ ├── web/ # Frontend (separate repo: knowhere-dashboard)
│ └── docs/ # Internal documentation
├── packages/
-│ ├── shared-python/shared/ # Shared library (pip: knowhere-shared)
-│ │ ├── models/database/ # SQLAlchemy ORM models
-│ │ ├── models/schemas/ # Pydantic request/response schemas
-│ │ ├── services/retrieval/ # Core retrieval engine
-│ │ ├── services/chunks/ # DataFrame → ChunkPayload conversion
-│ │ ├── services/ai/ # LLM prompt service & AI client
-│ │ └── utils/ # Text, file, and chunk utilities
-│ ├── sdk-python/ # Public Python SDK
-│ ├── sdk-typescript/ # Public Node.js SDK
-│ └── openapi-specs/ # OpenAPI spec definitions
+│ └── shared-python/shared/ # Shared library (pip: knowhere-shared)
+│ ├── models/database/ # SQLAlchemy ORM models
+│ ├── models/schemas/ # Pydantic request/response schemas
+│ ├── services/retrieval/ # Core retrieval engine
+│ ├── services/chunks/ # DataFrame → ChunkPayload conversion
+│ ├── services/ai/ # LLM prompt service & AI client
+│ └── utils/ # Text, file, and chunk utilities
└── deploy/ # Docker Compose & deployment scripts
```
+> **SDKs live in standalone repos:**
+> - Python SDK → [`Ontos-AI/knowhere-python-sdk`](https://github.com/Ontos-AI/knowhere-python-sdk)
+> - Node SDK → [`Ontos-AI/knowhere-node-sdk`](https://github.com/Ontos-AI/knowhere-node-sdk)
+
---
## End-to-End Pipeline Overview
@@ -99,7 +100,7 @@ flowchart TB
subgraph RETRIEVE["⑤ Retrieval (shared)"]
Query["GET /v1/retrieval/query"] --> Pipeline["run_retrieval_query"]
Pipeline --> Channels["3-Channel BM25 (path/content/term)"]
- Pipeline --> Agentic["RetrievalAgent.run (LLM-driven)"]
+ Pipeline --> Agentic["WorkflowOrchestrator (Planner + DAG)"]
Channels --> RRF["RRF Fusion"]
Agentic --> Hydrate["hydrate_paths_to_rows"]
RRF --> Rank["_rank_candidates_by_path"]
@@ -528,7 +529,7 @@ debug CSVs (`preds_*.csv`) are saved alongside for troubleshooting.
### Two Retrieval Modes
-The system supports two modes controlled by `RETRIEVAL_AGENTIC_ENABLED`:
+The system supports two modes, controlled globally by `RETRIEVAL_AGENTIC_ENABLED` and locally via the per-request `use_agentic` toggle.
#### Legacy Mode (3-Channel RRF)
@@ -546,33 +547,21 @@ flowchart LR
```
**Channel weights** (default): path=1.0, content=2.0, term=1.5
-
**RRF formula**: `score = weight / (k + rank + 1)` per channel, summed across channels.
-#### Agentic Mode (LLM-driven Navigation)
-
-#### Agentic Mode (LLM-driven Navigation)
-
-The agentic pipeline uses a deterministic multi-phase orchestration engine:
-
-**Phase 1: Discovery + Document Selection**
-- **Bottom Discovery**: Always runs first. Executes a 3-channel RRF keyword search across the entire Knowledge Base, returning top high-relevance chunks and their parent documents (`discovery_auto`).
-- **KG Document Select**: The LLM analyzes the KB-wide overview (from `knowledge_graph.json`) and selects highly relevant documents.
-- *Merge Strategy*: Documents found by Bottom Discovery but omitted by the LLM are automatically appended to the selected documents list to ensure no blind spots.
+#### Agentic Mode (Workflow Orchestrator)
-**Phase 2: Per-Document Navigation & Discovery Merging**
-For each selected document, the agent performs a constrained Breadth-First Search (BFS):
-1. **Scope Navigation**: The document's section tree is dynamically rendered to the LLM.
- - *Path-Based Hierarchy*: Child nodes are strictly filtered using structural path prefixes (e.g., `child_path.startswith(parent_path + ' / ')`) to maintain structural integrity and eliminate L2 duplicate rendering.
- - *Visual Constraints*: Actionable drill-down paths are explicitly prefixed with `[SELECT]` tags. The LLM system prompt tightly constrains the model to only pick paths with this tag, preventing redundant re-selection of the current scope.
-2. **Discovery Select**: The LLM reviews the specific paths flagged by Phase 1's Bottom Discovery for the current document. Selected discovery paths are hydrated into leaf chunks (with `job_result_id` dynamically extracted from the chunks) and merged directly into the BFS document tree.
- - *Reparenting*: The `DocTreeNode.merge()` process reparents these discovered leaf chunks into the closest matching navigated child node.
- - *Orphan Leaves*: Discovered chunks whose paths are not explicitly covered by the BFS `outline_items` are rendered cleanly as `[Leaf]` items (orphans) beneath their appropriate parent, ensuring no relevant data is lost even if the BFS did not explicitly drill into that path.
+The agentic pipeline uses `WorkflowOrchestrator` to handle complex queries via a DAG-based planning and budget-constrained execution engine:
-**Phase 3: Verdict & Revision**
-The combined document tree (BFS Navigation + Discovery) is rendered as unified evidence. The tree naturally displays structural context (outlines) alongside hydrated chunk rows (for selected leaf paths). The LLM attempts to answer the user's query:
-- `DONE`: Evidence is sufficient (or partially covers the query), exit and return final results.
-- `NOT_FOUND`: Evidence lacks sufficient information. Discard current evidence and trigger another revision round with a generated hint (max 2 rounds).
+1. **Planning (`PlannerAgent`)**: The query is analyzed and decomposed into a DAG of steps.
+ - Simple queries generate a single `retrieve` step.
+ - Complex queries are broken into multiple `retrieve` steps followed by a final `synthesize` step.
+2. **Budget Ledger (`BudgetLedger`)**: A strict token budget mechanism is enforced across the entire DAG execution (e.g., `AGENTIC_MAX_BUDGET=30000`). If the budget is exhausted, the pipeline halts safely and returns the best-effort evidence collected so far.
+3. **Execution (`RetrievalAgent`)**: For each `retrieve` step, a multi-phase navigation engine runs:
+ - **Phase 1 (Discovery)**: 3-channel RRF keyword search and KG document selection.
+ - **Phase 2 (Navigation)**: Constrained Breadth-First Search (BFS) over the document's section tree. Discovered orphan leaves are merged into the tree to prevent data loss.
+ - **Phase 3 (Verdict)**: The LLM evaluates the collected structural outlines + hydrated chunks. Triggers a revision round (max 2) if `NOT_FOUND`.
+4. **Synthesis**: The LLM synthesizes a final `answer_text` and precise citations (`referenced_chunks`) using the unified evidence tree.
### Tree Rendering & Hydration
diff --git a/apps/api/.env.example b/apps/api/.env.example
index 54c9a9709..7f51ce287 100644
--- a/apps/api/.env.example
+++ b/apps/api/.env.example
@@ -85,6 +85,18 @@ NORMOL_MODEL=deepseek-chat
HIERARCHY_LLM_MODEL=qwen3.6-flash
IMAGE_MODEL=qwen3.5-flash
IMAGE_MODEL_MAX=qwen3.5-flash
+RETRIEVAL_DECOMPOSITION_ENABLED=false
+RETRIEVAL_PLANNER_MODEL=
+RETRIEVAL_PLANNER_THINKING_BUDGET=4000
+RETRIEVAL_DECOMPOSITION_MAX_STEPS=5
+RETRIEVAL_WALLET_TOTAL_BUDGET=200000
+RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET=40000
+RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET=6000
+RETRIEVAL_WORKFLOW_PARALLEL_MAX=3
+
+# Agentic retrieval (LLM-driven hierarchical navigation).
+# Set to false to fall back to legacy 3-channel RRF mode.
+RETRIEVAL_AGENTIC_ENABLED=true
# File handling defaults
SUPPORTED_EXTENSIONS=.doc,.docx,.pdf,.txt,.xls,.xlsx,.csv,.pptx,.jpg,.jpeg,.png,.md
@@ -120,4 +132,4 @@ ILOVEAPI_TIMEOUT=120
# Legacy parser compatibility fields.
ALL_DF_COLS=content,path,type,length,keywords,summary,know_id,tokens,connectto,addtime,page_nums
-SPLIT_CHAR=-->
+SPLIT_CHAR=/
diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py
index e0106f0bc..1ef332e4b 100644
--- a/apps/api/app/api/v1/routes/retrieval.py
+++ b/apps/api/app/api/v1/routes/retrieval.py
@@ -52,6 +52,10 @@ class RetrievalQueryRequest(BaseModel):
internal_recall_k: int | None = Field(
None, ge=1, description="Override per-channel recall count"
)
+ use_agentic: bool | None = Field(
+ None,
+ description="Per-request agentic mode toggle. true=force agentic, false=force legacy, null=use server default.",
+ )
@field_validator("channels")
@classmethod
@@ -63,7 +67,16 @@ def validate_channels(cls, v: list[str]) -> list[str]:
return v
-@router.post("/query")
+class RetrievalQueryResponse(BaseModel):
+ namespace: str
+ query: str
+ router_used: str
+ answer_text: str | None = None
+ referenced_chunks: list[dict] = Field(default_factory=list)
+ results: list[dict] = Field(default_factory=list)
+
+
+@router.post("/query", response_model=RetrievalQueryResponse)
async def query_retrieval(
payload: RetrievalQueryRequest,
current_user: CurrentUser = Depends(with_current_user),
@@ -85,4 +98,5 @@ async def query_retrieval(
rerank=payload.rerank,
threshold=payload.threshold,
internal_recall_k=payload.internal_recall_k,
+ use_agentic=payload.use_agentic,
)
diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py
index abca22c68..ea21e7de6 100644
--- a/apps/api/tests/contract/test_demo_documents_contract.py
+++ b/apps/api/tests/contract/test_demo_documents_contract.py
@@ -140,6 +140,7 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge(
"shared.services.storage.result_storage.get_result_storage",
lambda: fake_result_storage,
)
+ monkeypatch.setenv("RETRIEVAL_AGENTIC_ENABLED", "false")
async with developer_api_client_factory() as api_client:
empty_cached_response = await api_client.post(
diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py
index 2deca35fe..a95ed4d72 100644
--- a/apps/api/tests/contract/test_retrieval_contract.py
+++ b/apps/api/tests/contract/test_retrieval_contract.py
@@ -182,6 +182,8 @@ async def test_should_return_empty_results_for_an_empty_query(
"query": "",
"router_used": "empty_query_filtered",
"results": [],
+ "answer_text": None,
+ "referenced_chunks": [],
}
diff --git a/apps/worker/.env.example b/apps/worker/.env.example
index 49bb7aae4..a57b458fc 100644
--- a/apps/worker/.env.example
+++ b/apps/worker/.env.example
@@ -85,6 +85,18 @@ NORMOL_MODEL=deepseek-chat
HIERARCHY_LLM_MODEL=deepseek-chat
IMAGE_MODEL=qwen3.5-flash
IMAGE_MODEL_MAX=qwen3.5-flash
+RETRIEVAL_DECOMPOSITION_ENABLED=false
+RETRIEVAL_PLANNER_MODEL=
+RETRIEVAL_PLANNER_THINKING_BUDGET=4000
+RETRIEVAL_DECOMPOSITION_MAX_STEPS=5
+RETRIEVAL_WALLET_TOTAL_BUDGET=200000
+RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET=40000
+RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET=6000
+RETRIEVAL_WORKFLOW_PARALLEL_MAX=3
+
+# Agentic retrieval (LLM-driven hierarchical navigation).
+# Set to false to fall back to legacy 3-channel RRF mode.
+RETRIEVAL_AGENTIC_ENABLED=true
# Required for specific features: billing and analytics
BILLING_ENABLED=false
@@ -110,5 +122,5 @@ MAX_FILE_SIZE=104857600
# Legacy parser compatibility fields.
ALL_DF_COLS=content,path,type,length,keywords,summary,know_id,tokens,connectto,addtime,page_nums
-SPLIT_CHAR=-->
+SPLIT_CHAR=/
diff --git a/docs/agentic-rag-audit-20260513.md b/docs/agentic-rag-audit-20260513.md
new file mode 100644
index 000000000..fe09703ce
--- /dev/null
+++ b/docs/agentic-rag-audit-20260513.md
@@ -0,0 +1,353 @@
+# Agentic RAG 流程审计更新
+
+日期：2026-05-13
+
+范围：基于 `AGENTS.md`、`.agent/skills/agentic_debug_patterns/SKILL.md`、既有 trace 目录 `/Users/wuchengke/Desktop/agentic_e2e_traces/20260513_183340`，以及额外运行的典型用例结果。
+
+额外 trace 输出目录：
+
+- `/Users/wuchengke/Desktop/agentic_e2e_traces/20260513_190749_extra`
+- `/Users/wuchengke/Desktop/agentic_e2e_traces/20260513_210646_extra_batch`
+
+## 结论摘要
+
+这套 agentic RAG 的总体方向是成立的：外层 workflow 可以把复杂问题拆成多个 retrieve/synthesize step 并并发执行；内层 retrieval agent 能基于 KG 选文档、树导航、发现补充路径，并把 connected image/table 内嵌回证据树。尤其是全局图片/图表类问题，已有路径可以通过 `connect_to` 找回资源所属文本 section。
+
+但从 harness 工程师和真实用户视角看，目前仍有几个核心逻辑风险：
+
+1. 图表资源的“证据渲染归属”和“返回引用归属”不一致。渲染树里通常能用 `connect_to` 找到底层 owner section，但 `referenced_chunks`/citation 仍可能显示物理路径 `Root`，这会直接破坏用户理解图表出处。
+2. discovery merge 过于积极。即使 BFS 已经 `STOP`，后置 discovery 仍会把深层或邻近年份路径并入证据，导致 outline 类问题和窄 section 问题出现噪声。
+3. 预算和状态分类混淆。无证据问题会被包装成 `budget_stop`，掩盖真实原因；同时 bootstrap/revision 小预算耗尽时，整体 wallet 仍可能很充足，用户看到的失败原因不准确。
+4. 多文件/多 step 并发在外层有效，但单 step 内的多 doc 导航仍偏串行；更重要的是 `discovery_auto` 会把弱相关文档强行并入，容易在跨年份、跨主题问题上污染预算和证据。
+5. 回答 JSON 解析不够稳健。`attempt_answer` 返回含换行的 JSON-like 文本时会解析失败，单步用户可能看到 JSON wrapper。
+6. trace DB schema 与 ORM 不一致，导致 agentic trace 入库失败，削弱 harness 可观测性。
+
+## 本次补跑用例
+
+### T1_Outline_Extra
+
+Query：
+
+> 民生证券这份利率专题研报的整体结构是什么？包含哪些主要章节？
+
+结果：
+
+- Router：`workflow_single_step`
+- LLM calls：4
+- refs：13
+- elapsed：约 13s
+- action：`kg_document_select -> navigate -> discovery_select -> attempt_answer`
+
+观察：
+
+- `navigate` 在 root 层正确选择 `STOP`，这对“整体结构/主要章节”类问题是合理的。
+- 但后续 `discovery_select` 又选入了深层路径，如 `2 阶段性调整.../2.1.1 基本面企稳` 和 `5、2024：“资产荒”的极致演绎`。
+- 最终 evidence 约 7395 chars，answer 只有约 149 chars，说明证据明显过量。
+- `referenced_chunks` 里部分 image/table 的 section 显示为 `Root`，但 evidence tree 实际把它们挂在更具体 leaf section 下。
+
+判断：
+
+这是 discovery merge 策略的问题。对 root outline 查询，BFS 已经完成任务后，不应默认再并入深层 discovery 结果。否则用户问“目录结构”，结果引用中会混入某些深层图表，影响可信度。
+
+### T2_Deep_Section_Extra
+
+Query：
+
+> 2016年债市走牛的几个阶段中，机构行为是如何推动行情演绎的？有哪些相关图表说明？
+
+结果：
+
+- Router：`workflow_single_step`
+- LLM calls：4
+- refs：51
+- elapsed：约 29s
+- evidence：约 22960 chars
+- wallet context：`TIGHT`
+
+观察：
+
+- `navigate` 选择了 `NAVIGATE`，并带 `FIND_IMAGES`、`FIND_TABLES`，方向正确。
+- 但 root scope 的 asset tool 拉入了过多全局资源；同时 discovery 又选中父级 `1、2016：机构行为助推行情演绎`，导致 hydration 范围扩大。
+- evidence 里实际有图2、图3、图4、图5等图题和图片描述，但模型回答中仍说“未提供图表具体标题/编号”。
+- refs 达到 51，包含不少 2018、2019、2023、2024 等非目标年份资源。
+- 2016 相关图片在引用元数据中仍有 `section=Root` 的情况，虽然它们通过 `connect_to` 在 evidence 中被放回了具体 section。
+
+判断：
+
+这是窄 section + 图表问题的典型失败形态：导航方向正确，但工具作用域过宽、discovery 过宽、证据渲染噪声大，导致模型虽然拿到了图表，却没有稳定提取图题和归属。
+
+### T3_Compare_Extra
+
+Query：
+
+> 对比2024年和2025年AI安全市场规模，并结合证据给出变化原因。
+
+结果：
+
+- Router：`workflow_decomposed`
+- LLM calls：27
+- refs：1
+- elapsed：约 38s
+- plan：s1 查 2024 市场规模，s2 查 2025 市场规模，s3 查变化原因，s4 synthesize
+
+观察：
+
+- 外层 workflow 确认可以并发执行多个 retrieve step，三个 retrieve step 的 KG select 和 navigate 调用是交错发生的。
+- 三个 retrieve step 最终都进入 revision，然后以 `budget_stop` 结束。
+- 总体 wallet 仍有大量剩余，但 bootstrap/revision 局部预算先被耗尽，最终对用户呈现为“预算停止”。
+- 实际语义更接近：KB 中缺少可支撑 2024/2025 AI 安全市场规模对比的证据。
+- `discovery_auto` 因年份词匹配，把债券研报等弱相关文档带入候选，造成预算消耗和路径污染。
+
+判断：
+
+这是预算状态和无证据状态混淆。对用户来说，“知识库没有足够证据”和“预算不够”是两类完全不同的反馈；当前状态分类会误导用户，也会误导 harness 判断。
+
+## 核心问题清单
+
+### 1. 图表资源归属在 citation 层丢失
+
+涉及核心设计：
+
+- 每个独立图表/图片/表格都应通过 `connect_to` 找到底层 section 归属。
+- 物理资源 chunk 的 `path` 可能是 `images/...` 或 `tables/...`，甚至 DB section 可能挂在 `Root`。
+- 逻辑归属应以 text chunk 的 `metadata.connect_to[].target` 为准，`target` 指向 image/table chunk_id。
+
+当前表现：
+
+- evidence tree 渲染阶段多数情况下能用 owner path 把资源挂回 leaf section。
+- 但最终 `referenced_chunks`/citation 仍可能使用资源 chunk 自身的 `section_path`，因此显示 `Root`。
+
+影响：
+
+- 用户看到图表出处为 `Root`，无法判断它属于哪个章节。
+- 对图表比较、章节归因、报告复核非常不友好。
+- 这和“每个独立图表都有 `connect_to` 找到一个底层 section 归属”的设计要求冲突。
+
+建议：
+
+- citation/ref 组装时优先使用 `owner_section_path`，只有不存在时才回退到物理 `section_path`。
+- 返回结构中建议同时保留：
+ - `owner_section_path`：逻辑归属，用于用户展示和排序。
+ - `physical_section_path`：数据库/资源物理挂载位置，用于调试。
+ - `connect_to_source_chunk_id`：是哪一个 text chunk 证明了该资源归属。
+- 对 image/table 引用增加断言：若存在 `connect_to` owner，则展示 section 不应为 `Root`。
+
+### 2. discovery merge 对 STOP 和 outline 查询缺少门控
+
+当前表现：
+
+- T1 root outline 查询已经由 `navigate` 正确 `STOP`。
+- 后续 `discovery_select` 仍并入深层路径和资源。
+
+影响：
+
+- 简单结构问题证据膨胀。
+- 引用混入深层内容，用户会怀疑答案是不是依据了错误章节。
+- 预算被无谓消耗。
+
+建议：
+
+- 对 outline/structure/catalogue 类意图设置 discovery gate：
+ - 若 root STOP 且问题不要求“细节/图表/数据”，跳过 discovery hydration。
+ - 或只允许 discovery 返回 top-level structural sections，不 hydrate leaf content/assets。
+- `discovery_select` 的 prompt 应明确区分：
+ - structure query：只补结构遗漏。
+ - evidence query：可补 leaf 内容。
+ - asset query：可补 image/table。
+
+### 3. 图表工具作用域过宽
+
+当前表现：
+
+- T2 中 `NAVIGATE + FIND_IMAGES/FIND_TABLES` 方向正确，但 root 或父级 scope asset extraction 拉入大量非目标年份图表。
+- 后续 trimming 虽然会删一部分，但已经消耗 context 和模型注意力。
+
+影响：
+
+- 窄问题变成大范围 evidence dump。
+- 模型可能拿到正确图题却没有稳定使用，反而回答“没有具体标题”。
+- refs 过多，前端引用列表不可读。
+
+建议：
+
+- 当 action 为 `NAVIGATE` 且有 selected leaf paths 时，asset tools 默认只对 selected paths 或其 owner-linked assets 生效。
+- 只有 action 为 root `STOP` 且 query 明确要求“列出全部图表/图片/表格”时，才允许文档级全量 asset pull。
+- 对 `FIND_IMAGES/FIND_TABLES` 的输出增加 owner filter：资源必须能通过 `connect_to` 归属到当前 selected subtree。
+
+### 4. 预算分配与状态管理需要区分技术预算和语义失败
+
+当前表现：
+
+- T3 三个 retrieve step 最终都是 `budget_stop`。
+- 但总 wallet 明显还有剩余，真正失败原因是没有足够证据。
+- bootstrap/revision 局部预算耗尽被升级成 step 级 budget stop。
+
+影响：
+
+- 用户会以为“系统钱/上下文不够”，而不是“知识库无证据”。
+- harness 也难以判断是预算策略问题、检索召回问题还是 KB 数据缺失。
+
+建议：
+
+- step status 拆分：
+ - `not_found_no_evidence`
+ - `not_found_low_confidence`
+ - `budget_exhausted_bootstrap`
+ - `budget_exhausted_context`
+ - `budget_exhausted_total`
+- synthesize 时保留每个 retrieve step 的 semantic reason，不要只看 stop_reason 字符串。
+- revision loop 中，如果第一轮和第二轮文档选择高度重复且 verdict 是“KB 缺证据”，应提前停止，避免继续烧 bootstrap。
+- `BudgetWallet` 的 reclaimed budget 如果暂不重分配，snapshot 文案应避免暗示这些预算已重新可用。
+
+### 5. Planner 缺少 KB inventory，导致 plan reasoning 误报
+
+当前表现：
+
+- T4 中 planner reasoning 出现 “knowledge base is empty”。
+- 实际 trace 中 KB 并不为空。
+
+判断：
+
+`QueryPlanner.plan()` 支持 `kb_total_docs/kb_total_chunks` 参数，但 workflow 调用路径没有传入真实 inventory，默认值为 0。
+
+影响：
+
+- plan reasoning 不可信。
+- 对调试和用户解释都很危险。
+
+建议：
+
+- `_load_or_plan()` 前读取当前 namespace 的 KB inventory，并传给 planner。
+- workflow plan cache key 应包含 KB version 或文档集合 fingerprint，否则 KB 更新后可能复用旧 plan。
+
+### 6. `attempt_answer` JSON 解析不稳健
+
+当前表现：
+
+- T4 中 `attempt_answer` 返回 JSON-like 内容，但因 raw newline 或不合规转义导致 parse 失败。
+- parse 失败后逻辑把原始字符串当作 DONE answer。
+
+影响：
+
+- 单步用户可能看到 `{"status":"DONE","answer":...}` wrapper。
+- synth step 可能能“洗掉”问题，但 single-step 场景会暴露。
+
+建议：
+
+- 增加 tolerant JSON repair，只修复回答字段中的裸换行/控制字符。
+- 如果解析失败且文本明显以 JSON object 开头，不应直接 `DONE raw`，而应降级重试或抽取 `answer` 字段。
+
+### 7. trace DB schema 与 ORM 不一致
+
+当前表现：
+
+- `retrieval_runs.parent_run_id/workflow_step_id/workflow_plan` 在 ORM 中存在。
+- alembic migration 中未创建这些列。
+- trace create_run 报 `UndefinedColumnError`。
+
+影响：
+
+- DB trace 不可用。
+- harness 只能依赖 Markdown trace，无法做结构化聚合和回归分析。
+
+建议：
+
+- 补 migration。
+- 增加一个轻量 schema contract test，覆盖 `RetrievalTraceRecorder.create_run()`。
+
+### 8. 多文件并发导航的现状
+
+已确认：
+
+- 外层 workflow retrieve steps 使用 topological batch 并发执行。
+- T3 中多个 retrieve step 的 KG select/navigate 调用交错，说明并发有效。
+
+风险：
+
+- 单个 retrieve step 内 selected docs 仍偏串行。
+- `discovery_auto` 追加的弱相关文档没有足够 domain guard，T3 因年份匹配引入了债券研报。
+
+建议：
+
+- 对 `discovery_auto` 文档追加设置最低 domain relevance：
+ - 文档 title/summary/keywords 至少命中主题实体。
+ - 或要求 bottom chunk 与 query 的非时间词、非通用词有足够 overlap。
+- 单 step 多 doc 可考虑并发，但要先修好 doc relevance guard，否则并发只会更快地放大噪声。
+
+## 遗留与冗余代码观察
+
+### Legacy retrieval 路径仍和 agentic 路径混杂
+
+`run_retrieval_query()` 中同时存在 agentic workflow 和 legacy 3-channel RRF 排序/graph routing。若 agentic 已是主路径，建议把 legacy 路径隔离为明确 fallback，避免后续改动时误改两套逻辑。
+
+### 旧 graph/discovery helper 有疑似未使用分支
+
+`agentic/orchestrator.py` 附近存在 `_grep_discover_document_ids`、`_expand_by_edges` 等老式发现逻辑痕迹。若主流程已经切到 bottom discovery + KG select，应确认这些 helper 是否仍被调用；未调用则标记删除或迁移到测试辅助。
+
+### path dedup 当前依赖“一叶一文本 chunk”隐含前提
+
+当前 `_hydrate_paths_to_rows` 用 path-level `seen_paths` 是安全的，因为解析模型近似保持“一 leaf section 一个 text chunk”。但如果未来 parser 把一个 leaf section 拆成多个 text chunks，path-level dedup 会丢内容。
+
+建议：
+
+- 在注释和测试中写明该前提。
+- 或把 dedup key 改成 `(document_id, section_path, chunk_id)`，再在 render 层控制同 section 合并。
+
+## 建议优先级
+
+P0：
+
+1. 修复 image/table citation 归属：优先展示 `connect_to` owner section，不再把有 owner 的图表显示成 `Root`。
+2. 修复 trace DB migration，恢复 harness 结构化观测。
+3. 修复 `attempt_answer` JSON parse fallback，避免把 wrapper 暴露给用户。
+
+P1：
+
+1. 对 root STOP/outline query 增加 discovery gate。
+2. 收紧 asset tool 作用域：`NAVIGATE + selected paths` 时只找 selected subtree 的 connected assets。
+3. 拆分 `budget_stop` 与 `not_found` 状态，synthesize 阶段保留真实失败原因。
+
+P2：
+
+1. 给 planner 传真实 KB inventory，并把 KB fingerprint 纳入 plan cache key。
+2. 给 `discovery_auto` 增加 domain relevance guard。
+3. 清理 legacy helper 和未使用 discovery/graph 分支。
+4. 为 path dedup 增加未来多 chunk leaf 的保护测试。
+
+## 建议回归用例
+
+1. Outline STOP 不应 hydrate 深层 leaf：
+ - Query：`民生证券这份利率专题研报的整体结构是什么？包含哪些主要章节？`
+ - 断言：refs 中不应出现大量 image/table；深层 section 不应被 discovery 自动并入。
+
+2. 2016 section 图表归属：
+ - Query：`2016年债市走牛的几个阶段中，机构行为是如何推动行情演绎的？有哪些相关图表说明？`
+ - 断言：所有相关 image/table citation 的展示 section 应为 2016 底层 section，而不是 `Root`。
+
+3. 全量图表查询：
+ - Query：`列出AI安全大模型报告中所有的图表和图片，并简要描述每张图的内容。`
+ - 断言：允许 root/global asset pull，但每个独立图表仍应有 owner section；确实无底层 owner 的封面/前言图要显式标记为 document-level。
+
+4. KB 无证据查询：
+ - Query：`对比2024年和2025年AI安全市场规模，并结合证据给出变化原因。`
+ - 断言：返回状态应是 no evidence / insufficient evidence，而不是 generic `budget_stop`。
+
+5. Planner inventory：
+ - 构造非空 KB。
+ - 断言 planner reasoning 不得出现 “knowledge base is empty”。
+
+## 代码落点索引
+
+- Workflow orchestration：`packages/shared-python/shared/services/retrieval/workflow/orchestrator.py`
+- Planner：`packages/shared-python/shared/services/retrieval/workflow/planner.py`
+- Workflow budget wallet：`packages/shared-python/shared/services/retrieval/workflow/wallet.py`
+- Inner agent orchestrator：`packages/shared-python/shared/services/retrieval/agentic/orchestrator.py`
+- Inner agent tools：`packages/shared-python/shared/services/retrieval/agentic/tools.py`
+- Answer policy / JSON parse：`packages/shared-python/shared/services/retrieval/agentic/policy.py`
+- Navigation tree render：`packages/shared-python/shared/services/retrieval/agent_navigate.py`
+- Retrieval entry / hydration：`packages/shared-python/shared/services/retrieval/app_service.py`
+- Retrieval trace：`packages/shared-python/shared/services/retrieval/agentic/trace.py`
+- ORM retrieval tables：`packages/shared-python/shared/models/database/document.py`
+- Migration：`apps/api/alembic/versions/e5f6a7b8c9d0_add_agentic_retrieval_tables.py`
+- Debug harness：`apps/worker/debug_agentic_e2e.py`
+
diff --git a/docs/external-services.md b/docs/external-services.md
deleted file mode 100644
index 32ecb405f..000000000
--- a/docs/external-services.md
+++ /dev/null
@@ -1,51 +0,0 @@
-# External Service Dependencies
-
-Knowhere API is the backend for . Public product
-documentation can also link to when deeper setup
-references are helpful.
-
-## Required For Local Startup
-
-Needs these dependencies to run the backend
-surface locally:
-
-- PostgreSQL for the application database
-- Redis for Celery and short-lived state
-- S3-compatible storage for uploads and result assets
-- one OpenAI-compatible LLM provider key for retrieval and parsing flows
-
-The repo-managed `deploy/local-dev` stack provides PostgreSQL, Redis, and
-LocalStack so the default `env.example` files can use a coherent local baseline.
-
-## Required Only For Specific Features
-
-- MinerU:
- required only if you want MinerU-backed document parsing flows
-- iLoveAPI:
- required only for conversion paths such as PPTX-to-PDF
-- QStash:
- required only if you want queued outbound webhook delivery
-- Stripe:
- required only if you enable billing and checkout flows
-- OAuth provider credentials:
- required only if you run dashboard-linked auth flows
-
-## Optional Observability And Analytics
-
-- Logfire for distributed tracing export
-
-These integrations are intentionally optional. Leaving them empty should not
-block a local backend bootstrap.
-
-## Minimum Viable Local Configuration
-
-The smallest supported local setup is:
-
-1. copy `apps/api/env.example` and `apps/worker/env.example`
-2. keep the default local PostgreSQL, Redis, and LocalStack values
-3. add one real LLM provider key such as `DS_KEY`
-4. run `deploy/local-dev/start-dev.sh`
-5. start the API and worker with `uv run`
-
-That path is the baseline public developer workflow. Additional providers should
-only be configured when you need the matching feature set.
diff --git a/docs/testing-guidance.md b/docs/testing-guidance.md
deleted file mode 100644
index 8d6df6b49..000000000
--- a/docs/testing-guidance.md
+++ /dev/null
@@ -1,82 +0,0 @@
-# Testing Guidance
-
-## Goal
-
-Keep the main test suites focused on the project surface rather than internal implementation details.
-
-The important things to verify are:
-
-- HTTP contract: request shape, response shape, status codes, and headers
-- Observable side effects: database writes, database updates, Redis state, and queued work
-- Runtime guarantees: auth behavior, rate limiting, conflict handling, and validation handling
-- Migration and persistence correctness: schema, constraints, and SQL-backed behavior
-
-The important things to avoid are:
-
-- Internal function boundary assertions
-- Repository or service call-sequence assertions
-- Mock-heavy tests that change the behavior under test in a material way
-
-## Test Taxonomy
-
-- `apps/api/tests/contract`
- API endpoint and surface specifications
-- `apps/api/tests/support`
- app bootstrap, test environment, database reset, Redis reset, and seed helpers
-- `apps/api/tests/migrations`
- Alembic and schema guarantees
-- `apps/worker/tests/contract`
- worker entrypoint, queued-work boundary, and durable side-effect specifications
-- `apps/api/tests` and `apps/worker/tests`
- narrow app-level component tests only when they protect pure logic or deterministic edge-case parsing
-
-Do not add a standalone shared-package test tree for behavior that belongs to the API or worker surface.
-
-## Contract Test Rules
-
-- A contract test must call the project surface, not an internal helper.
-- A contract test must assert an externally visible result or durable side effect.
-- API contract tests should use the real FastAPI lifespan.
-- Contract tests should use real PostgreSQL where SQL behavior depends on it.
-- API and worker contract tests use `fakeredis` for Redis behavior while keeping the same Redis service interfaces.
-- Mock only hard-to-control external boundaries such as third-party HTTP, storage providers, time, or filesystem edges.
-
-## Naming Rules
-
-- File names should follow the surface area being specified.
-- Contract test functions should prefer `test_should_`.
-- Test names should describe user-visible behavior, not implementation details.
-
-## Fixture Boundaries
-
-- `apps/api/tests/support` owns API bootstrap, environment setup, lifespan control, and seed data.
-- API contract tests should not override core dependencies such as auth, database access, or rate limiting for the behavior under test.
-- Worker contract tests own worker task entrypoints, queued-work boundaries, and durable task outcomes.
-
-## Coverage Expectations
-
-- API contract coverage should track every mounted router group in `apps/api/app/api/v1/api_v1.py`.
-- Worker contract coverage should track every registered Celery task in `apps/worker/app/core/tasks`.
-- When a stronger contract test replaces an old mock-heavy test, remove the weaker test or reduce it to a narrow component test.
-
-## Local Environment
-
-- API and worker contract tests require PostgreSQL server binaries and contrib extensions for `pytest-postgresql`.
-- API and worker contract tests do not require a running local PostgreSQL or Redis service.
-- API and worker contract tests use isolated `pytest-postgresql` processes and `fakeredis`.
-
-## Commands
-
-- `uv run python apps/api/scripts/ensure_test_environment.py --install`
-- `uv run python apps/api/scripts/ensure_test_environment.py`
-- `uv run pytest apps/api/tests/contract -q`
-- `uv run pytest apps/api/tests/migrations -q`
-- `uv run pytest apps/api/tests -q`
-- `uv run pytest apps/worker/tests/contract -q`
-- `uv run pytest apps/api/tests apps/worker/tests/contract -q`
-
-## Failure Triage
-
-- Re-run the smallest affected suite first.
-- Re-run from a clean `pytest-postgresql` process when the failure depends on database side effects.
-- Prefer fixing the harness or production behavior over adding more mocks.
diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py
index 2b25b5571..2164df3d2 100644
--- a/packages/shared-python/shared/core/config/ai.py
+++ b/packages/shared-python/shared/core/config/ai.py
@@ -35,6 +35,38 @@ class AIConfig(BaseModel):
default="qwen3.5-flash",
description="Higher-capability image model for OCR and ask-image Q&A",
)
+ RETRIEVAL_DECOMPOSITION_ENABLED: bool = Field(
+ default=False,
+ description="Enable query-decomposition workflow before agentic retrieval.",
+ )
+ RETRIEVAL_PLANNER_MODEL: str = Field(
+ default="",
+ description="Reasoning-capable model used by the workflow query planner.",
+ )
+ RETRIEVAL_PLANNER_THINKING_BUDGET: int = Field(
+ default=4000,
+ description="Token budget for the query planner thinking call.",
+ )
+ RETRIEVAL_DECOMPOSITION_MAX_STEPS: int = Field(
+ default=5,
+ description="Maximum number of planned workflow steps.",
+ )
+ RETRIEVAL_WALLET_TOTAL_BUDGET: int = Field(
+ default=200000,
+ description="Total workflow token wallet for decomposed retrieval.",
+ )
+ RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET: int = Field(
+ default=40000,
+ description="Default token budget issued to each retrieve step.",
+ )
+ RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET: int = Field(
+ default=6000,
+ description="Default token budget issued to each synthesize step.",
+ )
+ RETRIEVAL_WORKFLOW_PARALLEL_MAX: int = Field(
+ default=3,
+ description="Maximum concurrent workflow steps in the same DAG batch.",
+ )
# Runtime LLM controls.
LLM_MOCK_ENABLED: bool = Field(
diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py
index 8b5ba0a23..41bed45bb 100644
--- a/packages/shared-python/shared/models/database/document.py
+++ b/packages/shared-python/shared/models/database/document.py
@@ -380,6 +380,9 @@ class RetrievalRun(Base):
result_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
final_doc_ids: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)
result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)
+ parent_run_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True, index=True)
+ workflow_step_id: Mapped[Optional[str]] = mapped_column(String(64), nullable=True)
+ workflow_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True)
latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
error: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py
index a25d45cdc..832e81eb6 100644
--- a/packages/shared-python/shared/services/retrieval/__init__.py
+++ b/packages/shared-python/shared/services/retrieval/__init__.py
@@ -8,10 +8,11 @@
)
from .graph_service import DocumentGraphService, GraphQueryService, GraphScope
from .hit_stats_service import record_retrieval_hits
-from .llm_adapter import create_retrieval_llm_fn
+from .llm_adapter import create_retrieval_llm_fn, create_retrieval_planner_fn
__all__ = [
"create_retrieval_llm_fn",
+ "create_retrieval_planner_fn",
"run_retrieval_query",
"merge_channels_rrf",
"DocumentGraphService",
diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py
index 3f7a2d316..d5b1bed78 100644
--- a/packages/shared-python/shared/services/retrieval/agent_navigate.py
+++ b/packages/shared-python/shared/services/retrieval/agent_navigate.py
@@ -38,9 +38,35 @@
"""
-_SCOPE_NAV_PROMPT = """\
+
+_DISCOVERY_SELECT_PROMPT = """\
You are a document navigation assistant.
+Document: "{doc_name}"
+
+{budget_block}
+After navigating the document's section tree, the following section paths
+were additionally discovered via keyword and semantic search.
+They may contain relevant evidence not found through hierarchical navigation.
+
+=== Discovery Candidates ===
+{items}
+=== End Discovery Candidates ===
+
+User query: {query}
+{revision_context}
+Select section paths whose content is needed to answer the query.
+If none are relevant, return an EMPTY list [].
+
+Return ONLY a JSON object:
+{{"selections": [{{"path": "...", "confidence": }}, ...]}}
+Do not include any explanation.
+"""
+
+
+_ACTION_PROMPT = """\
+You are a document navigation agent.
+
Document: "{doc_name}" (id: {doc_id})
{budget_block}
@@ -56,44 +82,106 @@
User query: {query}
-Select sections to drill into for more detailed content.
-- You may ONLY select sections marked with [SELECT]. Do NOT select any other sections.
-- Select sections whose content is needed to answer the query.
-- If the titles and summaries already visible are sufficient (e.g. the query asks for an outline or overview), return an EMPTY list [].
-- When budget is TIGHT, prefer fewer high-confidence selections over broad exploration.
-- When budget is CRITICAL, be very selective — only pick paths with strong relevance. Return [] if current evidence already suffices.
+=== Available Actions ===
+
+Choose ONE action:
+
+NAVIGATE — Drill into selected sections for detailed content.
+ Consider this when the query targets specific topics and you need deeper text evidence.
+ Select one or more [SELECT] sections.
+
+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.
+
+{tools_block}
+
+When action is NAVIGATE, provide selections:
+- You may ONLY select sections marked with [SELECT].
+
+When action is STOP, selections must be empty.
Return ONLY a JSON object:
-{{"selections": [{{"path": "...", "confidence": }}, ...]}}
+{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}}
+or
+{{"action": "STOP", "tools": [...], "selections": []}}
Do not include any explanation.
"""
-_DISCOVERY_SELECT_PROMPT = """\
-You are a document navigation assistant.
+def _parse_action_response(text: str) -> dict:
+ """Parse the unified action response from LLM.
-Document: "{doc_name}"
+ Returns dict with keys:
+ action: 'NAVIGATE' | 'STOP'
+ tools: list[str] (subset of FIND_IMAGES, FIND_TABLES)
+ selections: list[dict] (each has 'path' and optional 'confidence')
-{budget_block}
-After navigating the document's section tree, the following section paths
-were additionally discovered via keyword and semantic search.
-They may contain relevant evidence not found through hierarchical navigation.
+ When action is STOP, selections are forced to empty.
+ """
+ import json as _json
+ import re as _re
-=== Discovery Candidates ===
-{items}
-=== End Discovery Candidates ===
+ text = text.strip()
+ _ASSET_TOOLS = {'FIND_IMAGES', 'FIND_TABLES'}
+ default = {'action': 'NAVIGATE', 'tools': [], 'selections': []}
-User query: {query}
-{revision_context}
-Select section paths whose content is needed to answer the query.
-If none are relevant, return an EMPTY list [].
-When budget is TIGHT, prefer fewer high-confidence candidates.
-When budget is CRITICAL, be very selective — only pick paths with strong relevance. Return [] if evidence suffices.
+ def _extract(data: dict) -> dict:
+ action = str(data.get('action', 'NAVIGATE')).strip().upper()
+ if action not in ('NAVIGATE', 'STOP'):
+ action = 'NAVIGATE'
-Return ONLY a JSON object:
-{{"selections": [{{"path": "...", "confidence": }}, ...]}}
-Do not include any explanation.
-"""
+ tools_val = data.get('tools') or []
+ if isinstance(tools_val, list):
+ tools = [str(t).strip().upper() for t in tools_val if str(t).strip().upper() in _ASSET_TOOLS]
+ else:
+ tools = []
+
+ # STOP → no selections allowed
+ if action == 'STOP':
+ return {'action': action, 'tools': tools, 'selections': []}
+
+ selections_val = data.get('selections') or []
+ selections = []
+ if isinstance(selections_val, list):
+ for s in selections_val:
+ if isinstance(s, dict) and s.get('path'):
+ conf = _normalize_confidence(s.get('confidence', 0.7))
+ selections.append({'path': str(s['path']), 'confidence': conf or 0.7})
+
+ return {'action': action, 'tools': tools, 'selections': selections}
+
+ # Try JSON parse
+ try:
+ data = _json.loads(text)
+ if isinstance(data, dict):
+ return _extract(data)
+ except (ValueError, _json.JSONDecodeError):
+ pass
+
+ # Try extracting JSON from markdown fences
+ fence_match = _re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, _re.DOTALL)
+ if fence_match:
+ try:
+ data = _json.loads(fence_match.group(1).strip())
+ if isinstance(data, dict):
+ return _extract(data)
+ except (ValueError, _json.JSONDecodeError):
+ pass
+
+ # Try finding a JSON object anywhere
+ brace_match = _re.search(r'\{.*\}', text, _re.DOTALL)
+ if brace_match:
+ try:
+ data = _json.loads(brace_match.group())
+ if isinstance(data, dict):
+ return _extract(data)
+ except (ValueError, _json.JSONDecodeError):
+ pass
+
+ return default
def _format_budget_block(snapshot: dict | None) -> str:
@@ -287,7 +375,7 @@ def _render_item(item: dict, include_summary: bool) -> str:
indent = " " * (level - 1)
prefix = '▸' if level == 1 else '└'
level_tag = f'[L{level}]'
- select_tag = '[SELECT] ' if show else ''
+ select_tag = '[SELECT] ' if item.get('selectable', False) else ''
lines: list[str] = []
lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}')
@@ -459,7 +547,7 @@ async def _load_child_sections(
db: AsyncSession,
document_id: str,
job_result_id: str,
- scope_path: str | None = None,
+ scope_path: str | list[str] | None = None,
exclude_paths: set[str] | None = None,
) -> list[dict]:
"""Load the Continuous Context Tree for *scope_path*.
@@ -468,17 +556,16 @@ async def _load_child_sections(
{path, title, summary, chunk_count, image_count, table_count,
level, show_summary, is_leaf}
- The tree contains three categories of nodes:
- 1. Ancestors of scope_path + their siblings → show_summary=False (title only)
- 2. Children of scope_path (2 depth bands) → show_summary=True (with summary)
- 3. Everything else → pruned (not returned)
-
- When scope_path is None (root), all items are category 2.
+ scope_path can be:
+ - None: root scope, all items are selectable (2 depth bands).
+ - str: single scope, descendants are selectable.
+ - list[str]: multi-scope, descendants of ALL paths are selectable
+ simultaneously — used when the LLM selected multiple drill-down
+ paths in the previous step.
- level: absolute depth in the document (1-based)
- show_summary: controls whether _format_items_for_llm renders summary
- - exclude_paths: paths already seen in prior revision rounds;
- any path matching (exact or subtree) is skipped from category 2
+ - exclude_paths: paths already hydrated; skipped from selectable items
"""
# ── Fetch all sections for this document revision ────────────────────
stmt = (
@@ -497,12 +584,20 @@ async def _load_child_sections(
if not section_rows:
return []
- scope = normalize_section_path(scope_path) if scope_path else ''
- scope_parts = split_section_path(scope)
- scope_depth = len(scope_parts)
+ # ── Normalize scope(s) ───────────────────────────────────────────────
+ # Multi-scope: list of paths to expand simultaneously
+ if isinstance(scope_path, list):
+ scope_list = [normalize_section_path(p) for p in scope_path]
+ elif scope_path:
+ scope_list = [normalize_section_path(scope_path)]
+ else:
+ scope_list = [] # root
+
+ # For logging, derive representative scope info
+ scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0
logger.debug(
- f' _load_child_sections: scope={scope!r} scope_parts={scope_parts} '
+ f' _load_child_sections: scopes={scope_list or ["root"]} '
f'scope_depth={scope_depth} exclude_paths={_excl if (_excl := exclude_paths or set()) else "none"} '
f'total_sections={len(section_rows)}'
)
@@ -524,134 +619,105 @@ async def _load_child_sections(
}
# ── Build the set of ancestor prefixes for pruning ────────────────────
- # e.g. scope = "A / B / K" → ancestor_prefixes = {"A", "A / B", "A / B / K"}
+ # For multi-scope, union all ancestor prefixes from all scope paths
ancestor_prefixes: set[str] = set()
- for i in range(1, scope_depth + 1):
- ancestor_prefixes.add(' / '.join(scope_parts[:i]))
+ for sp in scope_list:
+ sp_parts = split_section_path(sp)
+ for i in range(1, len(sp_parts) + 1):
+ ancestor_prefixes.add(' / '.join(sp_parts[:i]))
# ── Classify each section ────────────────────────────────────────────
_excl = exclude_paths or set()
items_by_path: dict[str, dict] = {}
- scope_child_depths: set[int] = set()
+ # Per-scope depth bands: track child depths separately per scope
+ per_scope_child_depths: dict[str, set[int]] = {sp: set() for sp in scope_list} if scope_list else {}
+ root_child_depths: set[int] = set() # used when scope_list is empty (root)
+
+ def _make_item(path: str, meta: dict, show_summary: bool) -> dict:
+ return {
+ 'path': path,
+ 'title': meta['title'],
+ 'summary': meta['summary'],
+ 'level': meta['depth'],
+ 'sort_order': meta['sort_order'],
+ 'chunk_count': 0,
+ 'image_count': 0,
+ 'table_count': 0,
+ 'section_id': meta['section_id'],
+ 'show_summary': show_summary,
+ }
+
+ def _is_excluded(path: str) -> bool:
+ return bool(_excl and any(
+ path == ep or path.startswith(ep + ' / ') for ep in _excl
+ ))
for path, meta in all_sections.items():
parts = meta['parts']
depth = meta['depth']
- if scope_depth == 0:
+ if not scope_list:
# Root scope: everything is a potential child
- if depth < 1:
+ if depth < 1 or _is_excluded(path):
continue
- # Skip excluded paths
- if _excl and any(
- path == ep or path.startswith(ep + ' / ')
- for ep in _excl
- ):
- continue
- scope_child_depths.add(depth)
- items_by_path[path] = {
- 'path': path,
- 'title': meta['title'],
- 'summary': meta['summary'],
- 'level': depth,
- 'sort_order': meta['sort_order'],
- 'chunk_count': 0,
- 'image_count': 0,
- 'table_count': 0,
- 'section_id': meta['section_id'],
- 'show_summary': True, # will be refined after depth band selection
- }
+ root_child_depths.add(depth)
+ items_by_path[path] = _make_item(path, meta, show_summary=True)
continue
- # --- Non-root scope ---
+ # --- Non-root scope(s) ---
+ # Check if this path is a descendant of ANY scope in scope_list
+ matched_scope: str | None = None
+ for sp in scope_list:
+ sp_parts = split_section_path(sp)
+ sp_depth = len(sp_parts)
+ if depth > sp_depth and parts[:sp_depth] == sp_parts:
+ matched_scope = sp
+ break
+
+ if matched_scope:
+ # Category 2: descendant of a scope path → selectable
+ if _is_excluded(path):
+ continue
+ per_scope_child_depths[matched_scope].add(depth)
+ items_by_path[path] = _make_item(path, meta, show_summary=True)
+ continue
- # Category 1: Ancestors and their siblings (structural context)
- # A node is an ancestor/sibling if its depth <= scope_depth AND
- # its parent prefix matches the scope's ancestry chain.
- if depth <= scope_depth:
- # Check: is this node in the ancestry chain or a sibling of one?
+ # Category 1: structural context (ancestors of scope paths only)
+ # Only show nodes that are on the ancestor chain of a scope path.
+ # Non-scope siblings (e.g. 法律声明, 前言 when navigating into
+ # chapters 一~六) are pruned to reduce token waste and prevent
+ # summary overflow in _format_items_for_llm.
+ max_scope_depth = max(len(split_section_path(sp)) for sp in scope_list)
+ if depth <= max_scope_depth:
if depth == 1:
- # All L1 nodes are either the ancestor or its siblings
- items_by_path[path] = {
- 'path': path,
- 'title': meta['title'],
- 'summary': meta['summary'],
- 'level': depth,
- 'sort_order': meta['sort_order'],
- 'chunk_count': 0,
- 'image_count': 0,
- 'table_count': 0,
- 'section_id': meta['section_id'],
- 'show_summary': False,
- }
- elif depth <= scope_depth:
- # For deeper ancestors/siblings: their parent must be in the
- # ancestor chain. e.g. "A / C" is a sibling of "A / B" only
- # if "A" is an ancestor of scope.
+ if path in ancestor_prefixes:
+ items_by_path.setdefault(path, _make_item(path, meta, show_summary=False))
+ else:
parent_prefix = ' / '.join(parts[:-1])
if parent_prefix in ancestor_prefixes:
- items_by_path[path] = {
- 'path': path,
- 'title': meta['title'],
- 'summary': meta['summary'],
- 'level': depth,
- 'sort_order': meta['sort_order'],
- 'chunk_count': 0,
- 'image_count': 0,
- 'table_count': 0,
- 'section_id': meta['section_id'],
- 'show_summary': False,
- }
+ items_by_path.setdefault(path, _make_item(path, meta, show_summary=False))
continue
- # Category 2: Descendants of scope_path (children to explore)
- # depth > scope_depth is guaranteed by the continue at line above
- is_descendant = parts[:scope_depth] == scope_parts
- if is_descendant:
- # Skip excluded paths
- is_excluded = _excl and any(
- path == ep or path.startswith(ep + ' / ')
- for ep in _excl
- )
- if is_excluded:
- logger.debug(f' _load_child_sections: EXCLUDED descendant path={path!r}')
- continue
- scope_child_depths.add(depth)
- items_by_path[path] = {
- 'path': path,
- 'title': meta['title'],
- 'summary': meta['summary'],
- 'level': depth,
- 'sort_order': meta['sort_order'],
- 'chunk_count': 0,
- 'image_count': 0,
- 'table_count': 0,
- 'section_id': meta['section_id'],
- 'show_summary': True,
- }
- continue
- else:
- logger.debug(
- f' _load_child_sections: NOT descendant path={path!r} '
- f'parts[:scope_depth]={parts[:scope_depth]} != scope_parts={scope_parts}'
- )
-
- # Category 3: Everything else → pruned (not added)
+ # Category 3: pruned
if not items_by_path:
return []
- # ── Limit children to 2 depth bands (relative to scope) ─────────────
- if scope_child_depths:
- if scope_depth == 0:
- allowed_depths = sorted(scope_child_depths)[:2]
- else:
- allowed_depths = sorted(scope_child_depths)[:2]
- allowed_set = set(allowed_depths)
- to_remove = []
- for path, item in items_by_path.items():
- if item['show_summary'] and item['level'] not in allowed_set:
- to_remove.append(path)
+ # ── Limit children to 2 depth bands (relative to each scope) ────────
+ allowed_set: set[int] = set()
+ if scope_list:
+ for sp, depths in per_scope_child_depths.items():
+ if depths:
+ allowed_set.update(sorted(depths)[:2])
+ else:
+ if root_child_depths:
+ allowed_set.update(sorted(root_child_depths)[:2])
+
+ 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]
@@ -806,48 +872,30 @@ async def _load_child_sections(
)
item['is_leaf'] = not has_descendants
+ # ── Assign selectability ──────────────────────────────────────────────
+ # Rule: in the 2-band window, only the DEEPER band is selectable.
+ # Leaf nodes at the shallower band are still selectable (no children
+ # to drill into). Structural context (show_summary=False) is never
+ # selectable.
+ if allowed_set:
+ shallowest_band = min(allowed_set)
+ for item in sorted_items:
+ if not item.get('show_summary', True):
+ # Structural context → never selectable
+ item['selectable'] = False
+ elif item['level'] == shallowest_band and not item.get('is_leaf', False):
+ # Shallowest band, non-leaf → grouping header, not selectable
+ item['selectable'] = False
+ else:
+ item['selectable'] = True
+ else:
+ for item in sorted_items:
+ item['selectable'] = item.get('show_summary', True)
+
return sorted_items
# ---------------------------------------------------------------------------
-# LLM response parser (for scope_navigate)
-# ---------------------------------------------------------------------------
-
-def _parse_scope_nav_response(text: str) -> list[dict[str, Any]]:
- """Parse selections JSON from scope navigation LLM response.
-
- Returns list of {"path": str, "confidence": float}.
- """
- text = text.strip()
- # Try direct parse
- try:
- data = json.loads(text)
- except (json.JSONDecodeError, ValueError):
- # Extract JSON object from markdown wrapper
- match = re.search(r'\{.*\}', text, re.DOTALL)
- if not match:
- return []
- try:
- data = json.loads(match.group())
- except (json.JSONDecodeError, ValueError):
- return []
-
- if not isinstance(data, dict):
- return []
-
- selections: list[dict[str, Any]] = []
- for item in (data.get('selections') or []):
- if not isinstance(item, dict):
- continue
- path = str(item.get('path') or '').strip()
- if not path:
- continue
- confidence = _normalize_confidence(item.get('confidence'))
- if confidence is None:
- confidence = 0.7
- selections.append({'path': path, 'confidence': confidence})
-
- return selections
# ---------------------------------------------------------------------------
diff --git a/packages/shared-python/shared/services/retrieval/agentic/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/__init__.py
index e4b5a0695..e59b26254 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/__init__.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/__init__.py
@@ -2,10 +2,11 @@
Navigate-then-answer loop:
Phase 1: Document selection (discovery + KG LLM select)
- Phase 2: Per-document iterative navigation (scope_navigate_step)
+ Phase 2: Per-document iterative navigation (navigate_step — unified action)
Phase 3: attempt_answer → DONE (return answer) or NOT_FOUND → revision
-Navigation auto-terminates when the LLM returns empty selections.
+Each navigate_step decides action (NAVIGATE/STOP), optional asset tools,
+and section selections in a single LLM call. STOP terminates drill-down.
After navigation, attempt_answer is called automatically — its result
(answer or NOT_FOUND+reason) drives the revision loop.
diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
index b8b766327..77bf80bca 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
@@ -2,7 +2,7 @@
Flow:
Phase 1: Document selection (bottom_discovery + kg_document_select)
- Phase 2: Per-document navigation (iterative BFS scope_navigate_step)
+ Phase 2: Per-document navigation (iterative BFS via navigate_step)
Phase 3: Render evidence → attempt_answer
→ DONE (has answer) → return answer + evidence
→ NOT_FOUND + reason → revision_hint → re-select docs + re-navigate
@@ -10,9 +10,10 @@
→ max_revisions → return best available
The orchestrator drives navigation via an iterative BFS queue per document,
-calling scope_navigate_step at each level. Navigation auto-terminates when
-the LLM returns empty selections. After navigation completes, attempt_answer
-is called automatically — no separate verdict step needed.
+calling navigate_step at each level. Each navigate_step is a single LLM call
+that decides action (NAVIGATE/STOP), asset tools (FIND_IMAGES/FIND_TABLES),
+and section selections. STOP terminates the drill-down for that scope.
+After navigation completes, attempt_answer is called automatically.
"""
from __future__ import annotations
@@ -100,6 +101,117 @@ async def _build_asset_url_map(
return url_map
+def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]:
+ """Recursively collect all leaf_content keys across the entire tree."""
+ paths = set(node.leaf_content.keys())
+ for child in node.children.values():
+ paths.update(_collect_all_leaf_paths(child))
+ return paths
+
+
+def _collect_visible_paths(node: DocTreeNode) -> set[str]:
+ """Collect all outline_items paths across the entire tree.
+
+ These are sections that are "visible" in the rendered tree (shown to the
+ LLM during navigation) even if no chunks have been hydrated into them yet.
+ Used as fallback targets for asset reconciliation.
+ """
+ paths = {item['path'] for item in node.outline_items if item.get('path')}
+ for child in node.children.values():
+ paths.update(_collect_visible_paths(child))
+ return paths
+
+
+def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None:
+ """Walk up a section path to find the closest ancestor in target_paths.
+
+ Example: path="kb/file/Ch1/S1.1/S1.1.1", target_paths={"kb/file/Ch1/S1.1"}
+ → returns "kb/file/Ch1/S1.1"
+
+ Uses the ' / ' separator convention from the section path format.
+ """
+ parts = path.split(' / ')
+ # Walk from most specific to least specific (skip the full path itself)
+ for i in range(len(parts) - 1, 0, -1):
+ ancestor = ' / '.join(parts[:i])
+ if ancestor in target_paths:
+ return ancestor
+ return None
+
+
+def _reconcile_deferred_assets(
+ tree: DocTreeNode,
+ pending_assets: list[dict],
+) -> None:
+ """Place collected assets into the tree based on final navigated paths.
+
+ Called ONCE after the entire BFS + discovery merge completes for a
+ document. Asset placement uses a two-tier strategy:
+
+ 1. **Exact match**: If the asset's ``owner_section_path`` matches a
+ leaf_content key, place directly (existing behavior).
+ 2. **Closest visible ancestor**: If exact match fails, walk up the
+ owner_section_path hierarchy to find the nearest ancestor that
+ appears in either leaf_content or outline_items. This handles
+ the case where the LLM stopped navigation early (e.g. at root)
+ but still requested images/tables — assets at L3 get attributed
+ to the visible L2 section on their path.
+ """
+ final_paths = _collect_all_leaf_paths(tree)
+ visible_paths = _collect_visible_paths(tree)
+ all_target_paths = final_paths | visible_paths
+
+ if not all_target_paths:
+ return
+
+ # Collect existing chunk_ids to avoid duplicates
+ existing_ids = {
+ str(row.get('chunk_id') or '')
+ for row in tree.flatten_chunk_rows()
+ if row.get('chunk_id')
+ }
+
+ placed = 0
+ ancestor_placed = 0
+ for asset in pending_assets:
+ chunk_id = str(asset.get('chunk_id') or '')
+ if chunk_id and chunk_id in existing_ids:
+ continue # already in tree via hydrate_connected_target_rows
+
+ owner_path = (
+ asset.get('owner_section_path')
+ or asset.get('section_path')
+ )
+ if not owner_path:
+ continue
+
+ # Tier 1: exact match in leaf_content or visible outline
+ target_path = owner_path if owner_path in all_target_paths else None
+
+ # Tier 2: closest visible ancestor fallback
+ if target_path is None:
+ target_path = _find_closest_ancestor(owner_path, all_target_paths)
+ if target_path:
+ ancestor_placed += 1
+
+ if target_path is None:
+ continue # no visible ancestor → discard
+
+ # Place into root; reparent_leaf_content will move to correct child
+ tree.add_leaf_chunks(target_path, [asset])
+ if chunk_id:
+ existing_ids.add(chunk_id)
+ placed += 1
+
+ if placed:
+ tree.reparent_leaf_content()
+ logger.info(
+ f' deferred asset reconcile: {placed}/{len(pending_assets)} '
+ f'assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths '
+ f'(ancestor_fallback={ancestor_placed})'
+ )
+
+
def _build_config_from_env() -> AgentRunConfig:
"""Read agent config from environment, with sensible defaults."""
return AgentRunConfig(
@@ -412,6 +524,9 @@ async def run(
channels: list[str] | None = None,
channel_weights: dict[str, float] | None = None,
config: AgentRunConfig | None = None,
+ ledger: BudgetLedger | None = None,
+ parent_run_id: str | None = None,
+ workflow_step_id: str | None = None,
) -> AgenticResult:
"""Run the agentic retrieval pipeline.
@@ -431,7 +546,7 @@ async def run(
exclude_sections = exclude_sections or []
state = AgentState()
- state.ledger = BudgetLedger(
+ state.ledger = ledger or BudgetLedger(
total=config.token_budget_total,
planning_ratio=config.planning_ratio,
bootstrap=config.bootstrap_budget,
@@ -455,6 +570,8 @@ async def run(
'exclude_sections': exclude_sections,
'signal_paths': signal_paths,
},
+ parent_run_id=parent_run_id,
+ workflow_step_id=workflow_step_id,
)
trace_enabled = os.environ.get('RETRIEVAL_AGENTIC_TRACE_ENABLED', 'true') == 'true'
@@ -711,9 +828,11 @@ async def _context_llm_call(prompt):
if key.startswith(f'{doc.document_id}::')
} if state.seen_section_keys else set()
- # BFS queue: (scope_path, parent_node, depth)
+ # BFS queue: (scope_path(s), parent_node, depth)
+ # scope can be: None (root), str, or list[str] (multi-scope)
root = DocTreeNode(scope_path=None)
- pending: list[tuple[str | None, DocTreeNode, int]] = [(None, root, 0)]
+ pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)]
+ doc_pending_assets: list[dict] = [] # deferred asset reconcile
while pending:
if state.elapsed_ms >= config.latency_budget_ms:
@@ -736,14 +855,16 @@ async def _context_llm_call(prompt):
depth=depth,
)
- # ★ Step 1: Tool selection (agent decides which asset tools)
+ # ★ Unified navigate step (supports multi-scope batching)
try:
- tool_choices = await tools.tool_select_step(
+ action, asset_tools, step_node, drill_paths = await tools.navigate_step(
db,
document_id=doc.document_id,
job_result_id=job_result_id,
query=query,
llm_fn=doc_llm_fn,
+ user_id=user_id,
+ namespace=namespace,
doc_name=doc_name,
scope_path=scope,
exclude_paths=doc_exclude,
@@ -751,37 +872,18 @@ async def _context_llm_call(prompt):
budget_snapshot=state.ledger.snapshot() if state.ledger else None,
)
except BudgetExceeded:
- logger.info(' agentic: planning budget exhausted during tool selection')
+ logger.info(' agentic: planning budget exhausted during navigation')
if trace_enabled:
trace.record_budget_stop('planning_exhausted')
break
state.step_count += 1
- if trace_enabled:
- trace.record_step(
- 'tool_select_step', ToolResult(
- status='selected',
- payload={
- 'document_id': doc.document_id,
- 'scope': scope or 'root',
- 'depth': depth,
- 'tool_choice': tool_choices or ['NAVIGATE'],
- },
- ),
- decision_reason=f'tool_r{round_idx}_d{depth}_{doc.source_file_name}',
- )
-
- logger.info(
- f' agentic step {state.step_count}: tool_select_step '
- f'doc="{doc.source_file_name}" scope={scope or "root"} '
- f'depth={depth} tools={tool_choices or ["NAVIGATE"]}'
- )
-
- pending_scope_assets: list[dict] = []
- for asset_tool in tool_choices:
+ # ★ Asset collection (deferred reconcile) — runs if LLM selected tools
+ # scope is passed directly: None (root), str, or list[str] (multi-scope).
+ # asset_filter_step handles all forms natively.
+ for asset_tool in asset_tools:
if asset_tool not in ('FIND_IMAGES', 'FIND_TABLES'):
continue
- # ★ Asset collection (programmatic extraction)
asset_type = 'image' if asset_tool == 'FIND_IMAGES' else 'table'
asset_chunks = await tools.asset_filter_step(
db,
@@ -791,15 +893,16 @@ async def _context_llm_call(prompt):
asset_type=asset_type,
)
if asset_chunks:
- pending_scope_assets.extend(asset_chunks)
+ doc_pending_assets.extend(asset_chunks)
+ scope_display = scope if isinstance(scope, list) else (scope or 'root')
if trace_enabled:
trace.record_step(
'asset_filter_step', ToolResult(
status='filtered' if asset_chunks else 'empty',
payload={
'document_id': doc.document_id,
- 'scope': scope or 'root',
+ 'scope': scope_display,
'asset_type': asset_type,
'chunks_found': len(asset_chunks) if asset_chunks else 0,
},
@@ -809,33 +912,9 @@ async def _context_llm_call(prompt):
logger.info(
f' agentic step {state.step_count}: asset_filter_step '
- f'doc="{doc.source_file_name}" scope={scope or "root"} '
+ f'doc="{doc.source_file_name}" scope={scope_display} '
f'type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}'
)
- # ★ Fallthrough: always proceed to NAVIGATE.
-
- # ★ Step 2: NAVIGATE (existing scope_navigate_step)
- try:
- step_node, drill_paths = await tools.scope_navigate_step(
- db,
- document_id=doc.document_id,
- job_result_id=job_result_id,
- query=query,
- llm_fn=doc_llm_fn,
- user_id=user_id,
- namespace=namespace,
- doc_name=doc_name,
- scope_path=scope,
- exclude_paths=doc_exclude,
- revision_hint=revision_hint if depth == 0 else None,
- budget_snapshot=state.ledger.snapshot() if state.ledger else None,
- )
- except BudgetExceeded:
- logger.info(' agentic: planning budget exhausted during navigation')
- if trace_enabled:
- trace.record_budget_stop('planning_exhausted')
- break
- state.step_count += 1
# Merge step result into parent node
parent_node.outline_items = step_node.outline_items
@@ -843,66 +922,34 @@ async def _context_llm_call(prompt):
parent_node.add_leaf_chunks(leaf_path, chunks)
parent_node.confidence = step_node.confidence
- # ★ Step 2.5: Reconcile pending assets with navigated leaf content
- if pending_scope_assets:
- # Collect all chunk_ids already present in any leaf_content
- existing_ids = {
- str(row.get('chunk_id') or '')
- for row in parent_node.flatten_chunk_rows()
- if row.get('chunk_id')
- }
- # Filter out assets already inlined
- supplementary = [
- a for a in pending_scope_assets
- if str(a.get('chunk_id') or '') not in existing_ids
- ]
- if supplementary:
- # Only place assets into sections already in the
- # navigated tree. If the owner path isn't part of
- # the tree, fall back to the current scope.
- navigated_paths = set(parent_node.leaf_content.keys()) | set(parent_node.children.keys())
- for asset in supplementary:
- owner_path = (
- asset.get('owner_section_path')
- or asset.get('section_path')
- or scope
- )
- if owner_path and owner_path in navigated_paths:
- parent_node.add_leaf_chunks(str(owner_path), [asset])
- elif scope:
- parent_node.add_leaf_chunks(str(scope), [asset])
-
# Accumulate hydrated leaf paths into doc_exclude
- # so subsequent drill-downs don't re-show them as [SELECT].
- # Skip paths that are pending drill-down (non-leaf hybrid nodes
- # hydrated via self_only) — their children must remain selectable.
drill_path_set = {sel['path'] for sel in drill_paths}
for leaf_path in step_node.leaf_content:
if leaf_path not in drill_path_set:
doc_exclude.add(leaf_path)
- # Queue non-leaf selections for further drill-down
- for sel in drill_paths:
- child = DocTreeNode(scope_path=sel['path'])
- parent_node.children[sel['path']] = child
- pending.append((sel['path'], child, depth + 1))
-
- # Re-parent leaf paths that belong to a child's subtree.
- # When the LLM selects both a parent section (non-leaf)
- # and one of its children (leaf) at the same depth, the
- # leaf chunks are stored on the parent node. Move them
- # into the child node so render_unified_doc_tree nests
- # them correctly instead of rendering orphans at root.
+ # Queue non-leaf selections as a SINGLE batched item
+ # (all drill paths expand simultaneously in the next call)
+ if drill_paths:
+ for sel in drill_paths:
+ child = DocTreeNode(scope_path=sel['path'])
+ parent_node.children[sel['path']] = child
+ batch_scope = [sel['path'] for sel in drill_paths]
+ pending.append((batch_scope, parent_node, depth + 1))
+
+ # Re-parent leaf paths that belong to a child's subtree
parent_node.reparent_leaf_content()
if trace_enabled:
trace.record_step(
- 'scope_navigate_step', ToolResult(
- status='navigated' if step_node.has_content() else 'empty',
+ 'navigate_step', ToolResult(
+ status=f'{action.lower()}' + (' (content)' if step_node.has_content() else ''),
payload={
'document_id': doc.document_id,
- 'scope': scope or 'root',
+ 'scope': scope if isinstance(scope, str) else (scope or 'root'),
'depth': depth,
+ 'action': action,
+ 'asset_tools': asset_tools,
'outline_count': len(step_node.outline_items),
'leaf_count': len(step_node.leaf_content),
'pending_drills': len(drill_paths),
@@ -911,10 +958,12 @@ async def _context_llm_call(prompt):
decision_reason=f'nav_r{round_idx}_d{depth}_{doc.source_file_name}',
)
+ scope_log = scope if isinstance(scope, str) else (', '.join(scope) if scope else 'root')
logger.info(
- f' agentic step {state.step_count}: scope_navigate_step '
- f'doc="{doc.source_file_name}" scope={scope or "root"} '
- f'depth={depth} outline={len(step_node.outline_items)} '
+ f' agentic step {state.step_count}: navigate_step '
+ f'doc="{doc.source_file_name}" scope={scope_log} '
+ f'depth={depth} action={action} tools={asset_tools} '
+ f'outline={len(step_node.outline_items)} '
f'leaves={len(step_node.leaf_content)} '
f'drills={len(drill_paths)}'
)
@@ -975,6 +1024,38 @@ async def _context_llm_call(prompt):
chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()),
)
+ # ── Deferred asset reconcile ──────────────────────────────
+ # Assets were collected across all BFS depths but NOT placed
+ # into the tree yet. Now that the final navigated paths are
+ # known (BFS + discovery), filter and place only those assets
+ # whose owner path matches a navigated leaf.
+ if not is_b_class and doc_pending_assets:
+ # Inject doc file name as a visible root-level path, but
+ # ONLY when BFS stopped at root (no children = STOP action).
+ if doc_name and not root.children and not any(
+ item.get('path') == doc_name for item in root.outline_items
+ ):
+ root.outline_items.insert(0, {'path': doc_name, 'level': 0})
+ _reconcile_deferred_assets(root, doc_pending_assets)
+ if trace_enabled:
+ trace.record_step(
+ 'deferred_asset_reconcile', ToolResult(
+ status='reconciled',
+ payload={
+ 'document_id': doc.document_id,
+ 'pending_count': len(doc_pending_assets),
+ 'placed_count': sum(
+ 1 for a in doc_pending_assets
+ if str(a.get('chunk_id') or '') in {
+ str(r.get('chunk_id') or '')
+ for r in root.flatten_chunk_rows()
+ }
+ ),
+ },
+ ),
+ decision_reason=f'deferred_reconcile_r{round_idx}_{doc.source_file_name}',
+ )
+
# Merge or store doc tree
if doc.document_id in state.doc_trees:
state.doc_trees[doc.document_id].merge(root)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py
index d11e4da7b..4774d2491 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py
@@ -22,8 +22,8 @@
_format_items_for_llm,
_load_child_sections,
_parse_json_array,
- _parse_scope_nav_response,
- _SCOPE_NAV_PROMPT,
+ _parse_action_response,
+ _ACTION_PROMPT,
_DISCOVERY_SELECT_PROMPT,
_FILE_SELECT_PROMPT,
_format_budget_block,
@@ -362,184 +362,6 @@ async def kg_document_select(
return ToolResult(status='error', error=str(e), latency_ms=latency)
-# ---------------------------------------------------------------------------
-# Tool: tool_select_step (lightweight LLM router)
-# ---------------------------------------------------------------------------
-
-_TOOL_SELECT_PROMPT = """\
-You are a document navigation agent.
-
-Document: "{doc_name}"
-
-{budget_block}
-{scope_header}
-Below is a summary of the current scope's sections:
-
-{tree_summary}
-
-User query: {query}
-
-=== Available Actions ===
-
-NAVIGATE (always performed)
- Drill into specific sections to explore detailed content.
- This action always runs — you do not need to select it.
-
-FIND_IMAGES (optional, additive)
- Also extract image/chart/diagram assets under this scope.
- Select this when the query asks about images, charts, figures, or visual content.
-
-FIND_TABLES (optional, additive)
- Also extract table/data assets under this scope.
- Select this when the query asks about tables, tabular data, or structured data.
-
-You may select ZERO, ONE, or BOTH optional actions.
-Navigation always happens regardless of your selection.
-
-Return ONLY a JSON object:
-{{"tools": []}} — navigate only, no extra assets
-{{"tools": ["FIND_IMAGES"]}} — navigate + extract images
-{{"tools": ["FIND_TABLES"]}} — navigate + extract tables
-{{"tools": ["FIND_IMAGES", "FIND_TABLES"]}} — navigate + extract both
-When budget is TIGHT, prefer fewer extra actions.
-When budget is CRITICAL, return empty tools unless assets directly answer the query.
-Do not include any explanation.
-"""
-
-
-def _parse_tool_choice(text: str) -> list[str]:
- """Parse tool choices from LLM response.
-
- Returns a list of selected tools (subset of FIND_IMAGES, FIND_TABLES).
- NAVIGATE is always implicit — an empty list means "navigate only".
- """
- import json as _json
- import re as _re
-
- text = text.strip()
- _ASSET_TOOLS = {'FIND_IMAGES', 'FIND_TABLES'}
-
- def _extract_from_data(data: dict) -> list[str]:
- # New format: {"tools": [...]}
- tools_val = data.get('tools')
- if isinstance(tools_val, list):
- return [str(t).strip().upper() for t in tools_val if str(t).strip().upper() in _ASSET_TOOLS]
- # Legacy format: {"tool": "..."}
- tool_val = str(data.get('tool', '')).strip().upper()
- if tool_val in _ASSET_TOOLS:
- return [tool_val]
- if tool_val == 'NAVIGATE':
- return []
- return []
-
- # Try JSON parse
- try:
- data = _json.loads(text)
- if isinstance(data, dict):
- return _extract_from_data(data)
- except (ValueError, _json.JSONDecodeError):
- pass
-
- # Accept a JSON object wrapped in markdown
- match = _re.search(r'\{.*?\}', text, _re.DOTALL)
- if match:
- try:
- data = _json.loads(match.group())
- if isinstance(data, dict):
- return _extract_from_data(data)
- except (ValueError, _json.JSONDecodeError):
- pass
-
- # Fallback: scan for tool names in raw text
- upper = text.upper()
- result: list[str] = []
- if 'FIND_IMAGES' in upper:
- result.append('FIND_IMAGES')
- if 'FIND_TABLES' in upper:
- result.append('FIND_TABLES')
- return result
-
-
-async def tool_select_step(
- db: AsyncSession,
- *,
- document_id: str,
- job_result_id: str,
- query: str,
- llm_fn: LLMFn,
- doc_name: str = '',
- scope_path: str | None = None,
- exclude_paths: set[str] | None = None,
- revision_hint: str | None = None,
- budget_snapshot: dict | None = None,
-) -> list[str]:
- """Route to the appropriate tools for the current scope.
-
- Returns a list of asset tools to run (FIND_IMAGES, FIND_TABLES).
- NAVIGATE always runs implicitly after any asset extraction.
-
- Optimization: if the current scope has no image or table chunks,
- skips the LLM call and returns an empty list (navigate only).
- """
- items = await _load_child_sections(
- db, document_id, job_result_id, scope_path,
- exclude_paths=exclude_paths,
- )
- if not items:
- return []
-
- # Build lightweight tree summary (titles + counts only, no summaries)
- summary_lines = []
- for item in items:
- if not item.get('show_summary'):
- continue
- title = item.get('title', '')
- img = item.get('image_count', 0)
- tbl = item.get('table_count', 0)
- txt = item.get('chunk_count', 0)
- counts = f'text={txt}'
- if img > 0:
- counts += f' image={img}'
- if tbl > 0:
- counts += f' table={tbl}'
- summary_lines.append(f'- {title} [{counts}]')
-
- tree_summary = '\n'.join(summary_lines) or '(empty)'
-
- # Check if scope has ANY images or tables — skip prompt if none
- total_images = sum(i.get('image_count', 0) for i in items)
- total_tables = sum(i.get('table_count', 0) for i in items)
- if total_images == 0 and total_tables == 0:
- return [] # no assets → skip tool selection, navigate only
-
- scope_header = (
- f'Current scope: "{scope_path}"' if scope_path
- else 'Current scope: root (document top level)'
- )
- prompt = _TOOL_SELECT_PROMPT.format(
- doc_name=doc_name or document_id,
- scope_header=scope_header,
- budget_block=_format_budget_block(budget_snapshot),
- tree_summary=tree_summary,
- query=query,
- )
- if revision_hint:
- prompt += (
- f'\n\nIMPORTANT: This is a REVISION round. '
- f'The previous search attempt failed because:\n'
- f'"{revision_hint}"\n'
- f'Adjust your tool selection accordingly.'
- )
- response = await llm_fn(prompt)
-
- # Parse tool choices
- asset_tools = _parse_tool_choice(response)
- logger.info(
- f' tool_select_step scope={scope_path or "root"}: '
- f'tools={asset_tools or ["NAVIGATE"]} images={total_images} tables={total_tables}'
- )
- return asset_tools
-
# ---------------------------------------------------------------------------
# Tool: asset_filter_step (programmatic asset extraction)
@@ -550,7 +372,7 @@ async def asset_filter_step(
*,
document_id: str,
job_result_id: str,
- scope_path: str | None,
+ scope_path: str | list[str] | None,
asset_type: str, # 'image' | 'table'
) -> list[dict[str, Any]]:
"""Extract assets from all descendants under scope_path.
@@ -561,22 +383,36 @@ async def asset_filter_step(
Also collects standalone asset chunks (image/table) that exist directly
under the scope but are not referenced via connect_to.
+
+ scope_path can be:
+ - None: root scope (entire document)
+ - str: single scope path
+ - list[str]: multiple scope paths (queried simultaneously)
"""
from shared.models.database.document import DocumentChunk, DocumentSection
t0 = time.monotonic()
try:
- # 1. Find all section_ids under scope_path
+ # 1. Find all section_ids under scope_path(s)
+ # Normalize scope to list for uniform handling
+ scope_list = (
+ scope_path if isinstance(scope_path, list)
+ else [scope_path] if scope_path
+ else []
+ )
+
section_stmt = (
select(DocumentSection.section_id, DocumentSection.section_path)
.where(DocumentSection.document_id == document_id)
.where(DocumentSection.job_result_id == job_result_id)
)
- if scope_path:
- section_stmt = section_stmt.where(
- (DocumentSection.section_path == scope_path) |
- (DocumentSection.section_path.like(f'{scope_path} / %'))
- )
+ if scope_list:
+ from sqlalchemy import or_
+ scope_filters = []
+ for sp in scope_list:
+ scope_filters.append(DocumentSection.section_path == sp)
+ scope_filters.append(DocumentSection.section_path.like(f'{sp} / %'))
+ section_stmt = section_stmt.where(or_(*scope_filters))
section_result = await db.execute(section_stmt)
section_rows = section_result.all()
section_ids = {row[0] for row in section_rows}
@@ -635,6 +471,20 @@ async def asset_filter_step(
]
owner_by_target_id = _build_connected_owner_map(text_row_dicts)
+ # Replace synthetic "Root" owner with the document's source_file_name.
+ # Root is a hybrid node whose real path is the file name (e.g.
+ # "32_安全大模型技术与市场研究报告_1.docx"); the DB stores the
+ # synthetic label "Root" which cannot match any outline node.
+ if any(v == 'Root' for v in owner_by_target_id.values()):
+ doc_stmt = select(Document.source_file_name).where(
+ Document.document_id == document_id
+ )
+ doc_file_name = (await db.execute(doc_stmt)).scalar() or ''
+ if doc_file_name:
+ for tid in list(owner_by_target_id):
+ if owner_by_target_id[tid] == 'Root':
+ owner_by_target_id[tid] = doc_file_name
+
# Collect connected target IDs for batch-loading
connected_target_ids: set[str] = set(owner_by_target_id.keys())
@@ -689,8 +539,9 @@ async def asset_filter_step(
# Root / top-level aggregation sections
if not owner_section_path:
own_section_path = section_path_by_id.get(row[4])
- if own_section_path and ' / ' not in own_section_path:
- # Reject document-root level sections as fallback owners
+ if own_section_path and own_section_path == 'Root':
+ # Reject only the synthetic Root aggregation label;
+ # legitimate L1 sections (e.g. "前言") are valid owners.
logger.warning(
f' asset_filter_step: rejecting root-level owner fallback '
f'chunk_id={chunk_id} section_path={own_section_path}'
@@ -731,13 +582,11 @@ async def asset_filter_step(
logger.error(f' asset_filter_step failed: {e}')
return []
-
# ---------------------------------------------------------------------------
-# Tool: scope_navigate_step (single-step navigation)
+# Tool: navigate_step (unified action — merges tool_select + scope_navigate)
# ---------------------------------------------------------------------------
-
-async def scope_navigate_step(
+async def navigate_step(
db: AsyncSession,
*,
document_id: str,
@@ -747,73 +596,149 @@ async def scope_navigate_step(
user_id: str,
namespace: str,
doc_name: str = '',
- scope_path: str | None = None,
+ scope_path: str | list[str] | None = None,
exclude_paths: set[str] | None = None,
revision_hint: str | None = None,
budget_snapshot: dict | None = None,
-) -> tuple[DocTreeNode, list[dict]]:
- """Single navigation step — one LLM call, no recursion.
+) -> tuple[str, list[str], DocTreeNode, list[dict]]:
+ """Unified navigation step — one LLM call for action + tools + selections.
+
+ scope_path can be:
+ - None: root scope
+ - str: single scope to drill into
+ - list[str]: multiple scopes to expand simultaneously
Returns:
- - node: DocTreeNode with outline_items (current scope local items only)
- and leaf_content (hydrated leaf selections)
- - pending: list of {path, confidence, mode} for non-leaf selections
- (orchestrator queues these for further drill-down)
+ - action: 'STOP' | 'NAVIGATE'
+ - asset_tools: list of asset tools to run (FIND_IMAGES, FIND_TABLES)
+ - node: DocTreeNode with outline_items and leaf_content
+ - pending: list of {path, confidence} for non-leaf drill-downs (empty when STOP)
"""
from shared.services.retrieval.app_service import _hydrate_paths_to_rows
- empty = DocTreeNode.empty(scope_path)
+ # Normalize scope for internal use
+ scope_paths: list[str] = (
+ scope_path if isinstance(scope_path, list)
+ else [scope_path] if scope_path
+ else []
+ )
+ # Set of scope path strings (for filtering selections)
+ scope_path_set = set(scope_paths)
+
+ empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None)
try:
- # 1. Load continuous context tree
+ # 1. Load continuous context tree (supports multi-scope)
items = await _load_child_sections(
db, document_id, job_result_id, scope_path,
exclude_paths=exclude_paths,
)
if not items:
- return empty, []
+ return 'STOP', [], empty, []
- # 2. Build selectable index (only current-scope items with summary)
- selectable = {item['path']: item for item in items if item.get('show_summary', True)}
+ # 2. Build selectable index
+ selectable = {item['path']: item for item in items if item.get('selectable', False)}
- # 3. Format full tree and call LLM
- text, overflowed = _format_items_for_llm(items)
- scope_header = (
- f'Current scope: navigating into "{scope_path}"'
- if scope_path else
- 'Current scope: root (document top level)'
+ # 3. Count ALL image/table chunks under the scope subtree(s)
+ from shared.models.database.document import DocumentChunk, DocumentSection
+ from sqlalchemy import func as sa_func
+
+ scope_section_stmt = (
+ select(DocumentSection.section_id)
+ .where(DocumentSection.document_id == document_id)
+ .where(DocumentSection.job_result_id == job_result_id)
)
- prompt = _SCOPE_NAV_PROMPT.format(
+ if scope_paths:
+ from sqlalchemy import or_
+ scope_filters = []
+ for sp in scope_paths:
+ scope_filters.append(DocumentSection.section_path == sp)
+ scope_filters.append(DocumentSection.section_path.like(f'{sp} / %'))
+ scope_section_stmt = scope_section_stmt.where(or_(*scope_filters))
+ scope_section_ids = await db.execute(scope_section_stmt)
+ all_section_ids = [r[0] for r in scope_section_ids.all()]
+
+ total_images = 0
+ total_tables = 0
+ if all_section_ids:
+ count_stmt = (
+ select(
+ DocumentChunk.chunk_type,
+ sa_func.count(DocumentChunk.id),
+ )
+ .where(DocumentChunk.document_id == document_id)
+ .where(DocumentChunk.job_result_id == job_result_id)
+ .where(DocumentChunk.section_id.in_(all_section_ids))
+ .where(DocumentChunk.chunk_type.in_(['image', 'table']))
+ .group_by(DocumentChunk.chunk_type)
+ )
+ count_result = await db.execute(count_stmt)
+ for chunk_type, cnt in count_result.all():
+ if chunk_type == 'image':
+ total_images = cnt
+ elif chunk_type == 'table':
+ total_tables = cnt
+
+ tools_block = ''
+ if total_images > 0 or total_tables > 0:
+ tools_lines = ['\nOptional asset tools (usable with NAVIGATE or STOP):\n']
+ if total_images > 0:
+ tools_lines.append(
+ f' FIND_IMAGES — Extract all image/chart assets under this scope ({total_images} available).\n'
+ )
+ if total_tables > 0:
+ tools_lines.append(
+ f' FIND_TABLES — Extract all table/data assets under this scope ({total_tables} available).\n'
+ )
+ tools_block = ''.join(tools_lines)
+
+ # 4. Format tree and build prompt
+ text, overflowed = _format_items_for_llm(items)
+ 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'
+ prompt = _ACTION_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=text,
query=query,
+ tools_block=tools_block,
)
if revision_hint:
prompt += (
f'\n\nIMPORTANT: Previous round feedback: '
- f'"{revision_hint}". Select specific sections this time.'
+ f'"{revision_hint}". Adjust your selections accordingly.'
)
+
+ # 5. Single LLM call
response = await llm_fn(prompt)
- selections = _parse_scope_nav_response(response)
+ parsed = _parse_action_response(response)
+ action = parsed['action']
+ asset_tools = parsed['tools']
+ selections = parsed['selections']
+ scope_label = ', '.join(scope_paths) if scope_paths else 'root'
logger.info(
- f' scope_navigate_step scope={scope_path or "root"}: '
- f'selections={len(selections)}, selectable={len(selectable)}, '
+ f' navigate_step scope={scope_label}: '
+ f'action={action} tools={asset_tools} '
+ f'selections={len(selections)} selectable={len(selectable)} '
f'overflowed={overflowed}'
)
- # 4. Build node with LOCAL items only (no ancestors/siblings)
- node = DocTreeNode(scope_path=scope_path)
+ # 6. Build node with LOCAL items only
+ node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None)
local_items = [item for item in items if item.get('show_summary', True)]
node.outline_items = local_items
- # 5. Dispatch selections (guard: never re-select scope_path itself)
+ # 7. Dispatch selections (only present when action == NAVIGATE)
valid_selections = [
s for s in selections
- if s['path'] in selectable and s['path'] != scope_path
+ if s['path'] in selectable and s['path'] not in scope_path_set
]
pending: list[dict] = []
@@ -827,8 +752,8 @@ async def scope_navigate_step(
if item.get('is_leaf'):
path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'})
else:
- pending.append(sel)
- # ★ NEW: Also hydrate this node's OWN direct chunks (not descendants)
+ # Non-leaf → will be batched into a single next call
+ pending.append({'path': path, 'confidence': conf})
path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'self_only'})
if path_selections:
@@ -847,17 +772,12 @@ async def scope_navigate_step(
exclude_sections=[],
)
if connected:
- # Resolve owner_section_path for connected assets:
- # map target_chunk_id → the section_path of the text chunk
- # that references it via connect_to.
_owner_map = _build_connected_owner_map(chunks)
for c in connected:
if not c.get('owner_section_path'):
c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or ''))
chunks = chunks + connected
- # Resolve Root-stranded assets to their true owner sections
- # via document-wide connect_to lookup
_root_map = await _resolve_root_asset_owners(
db,
document_id=document_id,
@@ -867,24 +787,23 @@ async def scope_navigate_step(
if _root_map:
for c in chunks:
if c.get('owner_section_path'):
- continue # already resolved by batch-level owner map
+ continue
cid = str(c.get('chunk_id') or '')
if cid in _root_map:
c['owner_section_path'] = _root_map[cid]
for chunk in chunks:
- # Distribute chunk to its real path or fallback to the selection path
real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path')
if real_path:
node.add_leaf_chunks(str(real_path), [chunk])
- return node, pending
+ return action, asset_tools, node, pending
except BudgetExceeded:
raise
except Exception as e:
- logger.error(f' scope_navigate_step failed for doc={document_id}: {e}')
- return empty, []
+ logger.error(f' navigate_step failed for doc={document_id}: {e}')
+ return 'STOP', [], empty, []
# ---------------------------------------------------------------------------
@@ -927,13 +846,15 @@ async def discovery_select_step(
t0 = time.monotonic()
try:
- # 1. Format hints for LLM
+ # 1. Format hints for LLM (deduplicate by section_path)
hint_lines: list[str] = []
hint_by_path: dict[str, dict] = {}
for h in hints:
sp = h.get('section_path', '')
if not sp or sp == 'Root':
continue
+ if sp in hint_by_path:
+ continue # skip duplicate section_path
title = sp.rsplit(' / ', 1)[-1] if ' / ' in sp else sp
summary = h.get('summary', '') or ''
hint_lines.append(f'▸ path="{sp}" {title} [Leaf]')
@@ -965,7 +886,9 @@ async def discovery_select_step(
revision_context=revision_context,
)
response = await llm_fn(prompt)
- selections = _parse_scope_nav_response(response)
+ # Parse {"selections": [...]} response — reuse action parser's extraction
+ parsed = _parse_action_response(response)
+ selections = parsed.get('selections', [])
logger.info(
f' discovery_select_step doc="{doc_name}": '
diff --git a/packages/shared-python/shared/services/retrieval/agentic/trace.py b/packages/shared-python/shared/services/retrieval/agentic/trace.py
index 8a58a243c..9c48e7a30 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/trace.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/trace.py
@@ -50,6 +50,9 @@ def __init__(
top_k: int = 10,
data_type: int = 1,
filters: dict[str, Any] | None = None,
+ parent_run_id: str | None = None,
+ workflow_step_id: str | None = None,
+ workflow_plan: dict[str, Any] | None = None,
) -> None:
self._db = db
self._run_id = f'aret_{uuid4().hex[:12]}'
@@ -60,6 +63,9 @@ def __init__(
self._top_k = top_k
self._data_type = data_type
self._filters = filters or {}
+ self._parent_run_id = parent_run_id
+ self._workflow_step_id = workflow_step_id
+ self._workflow_plan = workflow_plan
self._steps: list[dict[str, Any]] = []
self._start_time = time.monotonic()
self._created = False
@@ -86,6 +92,9 @@ async def create_run(self) -> None:
agentic_enabled=True,
cache_hit=False,
result_count=0,
+ parent_run_id=self._parent_run_id,
+ workflow_step_id=self._workflow_step_id,
+ workflow_plan=self._workflow_plan,
latency_ms=0,
created_at=_now_utc(),
)
@@ -179,6 +188,10 @@ async def complete(
}
if budget_snapshot is not None:
provenance['budget_snapshot'] = budget_snapshot
+ if self._parent_run_id:
+ provenance['parent_run_id'] = self._parent_run_id
+ if self._workflow_step_id:
+ provenance['workflow_step_id'] = self._workflow_step_id
stmt = (
update(RetrievalRun)
diff --git a/packages/shared-python/shared/services/retrieval/agentic/types.py b/packages/shared-python/shared/services/retrieval/agentic/types.py
index 498bf08b5..a448a5f6b 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/types.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/types.py
@@ -44,7 +44,7 @@ class ToolResult:
class DocTreeNode:
"""Unified navigation result tree for one document.
- Produced by ``scope_navigate_step``. Captures the full
+ Produced by ``navigate_step``. Captures the full
navigation outcome for rendering as a single hierarchy:
- ``outline_items``: section tree items at this scope level
@@ -218,7 +218,7 @@ class AgentState:
"""Mutable state carried through the 2-phase orchestrator.
Phase 1: Document selection (discovery + KG)
- Phase 2: Per-document navigation (scope_navigate_step per doc)
+ Phase 2: Per-document navigation (navigate_step per doc)
Phase 3: Assembly + final verdict
"""
# Timing
diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py
index c066a1c16..96c125e49 100644
--- a/packages/shared-python/shared/services/retrieval/app_service.py
+++ b/packages/shared-python/shared/services/retrieval/app_service.py
@@ -432,8 +432,6 @@ async def _to_public_response(response: dict[str, Any]) -> dict[str, Any]:
}
# Forward agentic evidence fields when present
- if response.get('evidence_text') is not None:
- public_response['evidence_text'] = response['evidence_text']
if response.get('answer_text') is not None:
public_response['answer_text'] = response['answer_text']
if response.get('referenced_chunks') is not None:
@@ -982,6 +980,7 @@ async def run_retrieval_query(
rerank: bool = False,
threshold: float = 0.0,
internal_recall_k: int | None = None,
+ use_agentic: bool | None = None,
) -> dict[str, Any]:
"""Checkerboard retrieval: 3 independent channels -> RRF -> agent/graph union -> assembly."""
t_start = time.monotonic()
@@ -1015,6 +1014,8 @@ async def run_retrieval_query(
rerank=rerank,
threshold=threshold,
internal_recall_k=internal_recall_k,
+ # Always True: agentic mode now always routes through workflow
+ decomposition_enabled=True,
)
cache_version: int | None = None
@@ -1096,22 +1097,25 @@ async def run_retrieval_query(
logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms')
return await _to_public_response(response)
- # ══ Route: agentic vs legacy ══
- _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'false') == 'true'
+ # ══ Route: agentic (unified workflow) vs legacy ══
+ if use_agentic is not None:
+ _agentic_enabled = use_agentic
+ else:
+ _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'true') == 'true'
if _agentic_enabled:
- # ── AGENTIC path (all errors self-contained, no fallback to legacy) ──
- from shared.services.retrieval.agentic.orchestrator import RetrievalAgent
- from shared.services.retrieval.llm_adapter import create_retrieval_llm_fn as _create_llm
-
- llm_fn = _create_llm()
- agent = RetrievalAgent()
- agentic_result = await agent.run(
+ # ── Unified agentic path via WorkflowOrchestrator ──
+ # Simple queries: planner returns a single-step plan (no decomposition).
+ # Complex queries: planner returns a multi-step plan with synthesize.
+ # Both go through the same code path.
+ from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator
+
+ workflow = WorkflowOrchestrator()
+ workflow_result = await workflow.run(
db,
user_id=user_id,
namespace=namespace,
query=query,
top_k=top_k,
- llm_fn=llm_fn,
exclude_document_ids=exclude_document_ids,
exclude_sections=exclude_sections,
data_type=data_type,
@@ -1120,11 +1124,10 @@ async def run_retrieval_query(
channels=channels,
channel_weights=channel_weights,
)
- router_used = agentic_result.router_used
- # Generate asset URLs for media chunks in referenced_chunks
+ # Enrich referenced_chunks with asset URLs (images/tables)
enriched_refs: list[dict[str, Any]] = []
- for ref in agentic_result.referenced_chunks:
+ for ref in workflow_result.referenced_chunks:
enriched = dict(ref)
chunk_type = _normalize_chunk_type(ref.get('chunk_type'))
artifact_ref = ref.get('file_path', '')
@@ -1140,30 +1143,9 @@ async def run_retrieval_query(
logger.warning(f'Failed to generate agentic asset URL (ignored): {e}')
enriched_refs.append(enriched)
- # Build backward-compatible results[] from referenced_chunks
- # (minimal: chunk_id + document_id + chunk_type + section_path)
- results = [
- {
- 'chunk_id': ref.get('chunk_id'),
- 'document_id': ref.get('document_id'),
- 'chunk_type': ref.get('chunk_type'),
- 'source': {
- 'document_id': ref.get('document_id'),
- 'section_path': ref.get('section_path'),
- },
- }
- for ref in enriched_refs
- ]
-
- response = {
- "namespace": namespace,
- "query": query,
- "router_used": router_used,
- "results": results,
- "evidence_text": agentic_result.evidence_text,
- "answer_text": agentic_result.answer_text,
- "referenced_chunks": enriched_refs,
- }
+ response = workflow_result.to_api_response()
+ # Override referenced_chunks with enriched versions
+ response['referenced_chunks'] = enriched_refs
if cache_version is not None:
try:
@@ -1180,7 +1162,7 @@ async def run_retrieval_query(
try:
schedule_retrieval_hit_stats_update(
user_id=user_id, namespace=namespace,
- results=agentic_result.referenced_chunks,
+ results=enriched_refs,
)
except Exception as e:
logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}")
@@ -1190,14 +1172,15 @@ async def run_retrieval_query(
f'\n{"█" * 70}\n'
f' ✅ AGENTIC RETRIEVAL COMPLETE: '
f'{len(enriched_refs)} chunks | '
- f'evidence={len(agentic_result.evidence_text)} chars | '
- f'answer={len(agentic_result.answer_text)} chars | '
- f'router={router_used} | {elapsed_total}ms\n'
+ f'answer={len(workflow_result.answer_text)} chars | '
+ f'router={workflow_result.router_used} | {elapsed_total}ms\n'
f'{"█" * 70}'
)
return await _to_public_response(response)
+
else:
+
# ── LEGACY path (existing code, unchanged) ──
# ── Channel execution ──
diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py
index b7c2bd88b..bb165eb41 100644
--- a/packages/shared-python/shared/services/retrieval/cache_service.py
+++ b/packages/shared-python/shared/services/retrieval/cache_service.py
@@ -8,6 +8,7 @@
from shared.services.redis import RedisServiceFactory
_RETRIEVAL_CACHE_TTL_SECONDS = 300
+_WORKFLOW_PLAN_CACHE_TTL_SECONDS = 600
_VERSION_FALLBACK = 0
@@ -42,6 +43,7 @@ def _cache_shape_digest(
rerank: bool = False,
threshold: float = 0.0,
internal_recall_k: int | None = None,
+ decomposition_enabled: bool | None = None,
) -> str:
normalized_excludes = sorted(exclude_document_ids)
normalized_sections = _normalize_exclude_sections(exclude_sections)
@@ -55,6 +57,7 @@ def _cache_shape_digest(
str(rerank),
str(threshold),
str(internal_recall_k),
+ str(decomposition_enabled),
]
)
payload = f"{query}|{top_k}|{'|'.join(normalized_excludes)}|{'|'.join(normalized_sections)}|{extra}"
@@ -178,3 +181,37 @@ async def set_cached_retrieval_query_result(
response,
ex=_RETRIEVAL_CACHE_TTL_SECONDS,
)
+
+
+def _workflow_plan_cache_key(*, user_id: str, namespace: str, query: str) -> str:
+ digest = hashlib.sha256(query.encode("utf-8")).hexdigest()
+ return f"retrieval:workflow:plan:{user_id}:{namespace}:{digest}"
+
+
+async def get_cached_workflow_plan(
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+) -> dict[str, Any] | None:
+ redis_service = RedisServiceFactory.get_service()
+ cached = await redis_service.get(
+ _workflow_plan_cache_key(user_id=user_id, namespace=namespace, query=query),
+ default=None,
+ )
+ return cached if isinstance(cached, dict) else None
+
+
+async def set_cached_workflow_plan(
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ plan: dict[str, Any],
+) -> None:
+ redis_service = RedisServiceFactory.get_service()
+ await redis_service.set(
+ _workflow_plan_cache_key(user_id=user_id, namespace=namespace, query=query),
+ plan,
+ ex=_WORKFLOW_PLAN_CACHE_TTL_SECONDS,
+ )
diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py
index c6566d5f3..c5e592b5e 100644
--- a/packages/shared-python/shared/services/retrieval/llm_adapter.py
+++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py
@@ -54,6 +54,21 @@ def _resolve_default_model() -> str:
return getattr(settings, 'NORMOL_MODEL', None) or 'deepseek-chat'
+def _resolve_planner_model(*, thinking: bool) -> str:
+ configured = getattr(settings, 'RETRIEVAL_PLANNER_MODEL', '') or ''
+ if configured:
+ return configured
+ if getattr(settings, 'DS_KEY', ''):
+ return 'deepseek-reasoner' if thinking else 'deepseek-chat'
+ if getattr(settings, 'ALI_API_KEYS', ''):
+ return 'qwq-32b-preview' if thinking else 'qwen-plus'
+ if getattr(settings, 'GLM_API_KEY', ''):
+ return 'glm-4-plus' if thinking else 'glm-4-flash'
+ if getattr(settings, 'GPT_API_KEY', ''):
+ return 'o3-mini' if thinking else 'gpt-4o-mini'
+ return getattr(settings, 'NORMOL_MODEL', None) or 'deepseek-chat'
+
+
def create_retrieval_llm_fn(
*,
model: str | None = None,
@@ -89,6 +104,37 @@ async def llm_fn(prompt: LLMFnInput) -> str:
return llm_fn
+def create_retrieval_planner_fn(
+ *,
+ thinking: bool = True,
+ model: str | None = None,
+ max_tokens: int = 8192,
+) -> LLMFn | None:
+ """Create a reasoning-capable LLM callable for query planning."""
+ if not _has_llm_credentials():
+ logger.debug('retrieval: no LLM credentials configured, workflow planner disabled')
+ return None
+
+ effective_model = model or _resolve_planner_model(thinking=thinking)
+
+ async def llm_fn(prompt: LLMFnInput) -> str:
+ from shared.utils.OpenAICompatibleClientSync import get_openai_client
+
+ client = get_openai_client(model=effective_model)
+ current_llm_usage.set(None)
+ result, usage = await asyncio.to_thread(
+ client.chat_completion_with_usage,
+ cast(Any, prompt),
+ model=effective_model,
+ temperature=0.0,
+ max_tokens=max_tokens,
+ )
+ current_llm_usage.set(usage)
+ return result
+
+ return llm_fn
+
+
def create_retrieval_vlm_fn(
*,
model: str | None = None,
diff --git a/packages/shared-python/shared/services/retrieval/workflow/__init__.py b/packages/shared-python/shared/services/retrieval/workflow/__init__.py
new file mode 100644
index 000000000..67f896d3d
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/__init__.py
@@ -0,0 +1,13 @@
+"""Query-decomposition retrieval workflow."""
+from .types import PlannedStep, QueryPlan, StepResult, WorkflowResult
+from .wallet import BudgetWallet
+from .orchestrator import WorkflowOrchestrator
+
+__all__ = [
+ "BudgetWallet",
+ "PlannedStep",
+ "QueryPlan",
+ "StepResult",
+ "WorkflowResult",
+ "WorkflowOrchestrator",
+]
diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
new file mode 100644
index 000000000..9c5a20aff
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py
@@ -0,0 +1,382 @@
+"""Workflow orchestrator for decomposed retrieval queries."""
+from __future__ import annotations
+
+import asyncio
+import os
+import time
+from typing import Any
+from uuid import uuid4
+
+from loguru import logger
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from shared.core.database import get_db_context
+from shared.services.retrieval.agentic.budget import BudgetLedger
+from shared.services.retrieval.agentic.orchestrator import RetrievalAgent
+from shared.services.retrieval.agentic.types import AgenticResult
+from shared.services.retrieval.cache_service import (
+ get_cached_workflow_plan,
+ set_cached_workflow_plan,
+)
+from shared.services.retrieval.llm_adapter import (
+ create_retrieval_llm_fn,
+ create_retrieval_planner_fn,
+)
+from shared.services.retrieval.workflow.planner import QueryPlanner
+from shared.services.retrieval.workflow.synthesizer import compose_final_answer, synthesize_step
+from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult, WorkflowResult
+from shared.services.retrieval.workflow.wallet import BudgetWallet
+
+
+class WorkflowOrchestrator:
+ """Plan and execute a query workflow DAG."""
+
+ def __init__(self) -> None:
+ self.parent_run_id = f'wret_{uuid4().hex[:12]}'
+
+ async def run(
+ self,
+ db: AsyncSession,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int = 1,
+ signal_paths: list[str] | None = None,
+ filter_mode: str = 'delete',
+ channels: list[str] | None = None,
+ channel_weights: dict[str, float] | None = None,
+ llm_fn=None,
+ ) -> WorkflowResult:
+ t0 = time.monotonic()
+ llm_fn = llm_fn or create_retrieval_llm_fn()
+ planner_llm = create_retrieval_planner_fn(thinking=True)
+ planner_budget = _env_int('RETRIEVAL_PLANNER_THINKING_BUDGET', 4000)
+ wallet_total = _env_int('RETRIEVAL_WALLET_TOTAL_BUDGET', 200000)
+ per_retrieve = _env_int('RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET', 40000)
+ per_synthesize = _env_int('RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET', 6000)
+ max_steps = _env_int('RETRIEVAL_DECOMPOSITION_MAX_STEPS', 5)
+
+ planner_ledger = BudgetLedger(
+ total=planner_budget,
+ planning_ratio=0.0,
+ bootstrap=planner_budget,
+ per_doc_min_share=0,
+ )
+ plan = await self._load_or_plan(
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ planner_llm=planner_llm,
+ planner_ledger=planner_ledger,
+ max_steps=max_steps,
+ wallet_total=wallet_total,
+ per_retrieve=per_retrieve,
+ )
+
+ wallet = BudgetWallet(
+ total=wallet_total,
+ per_retrieve_step_default=per_retrieve,
+ per_synthesize_step_default=per_synthesize,
+ )
+ ledgers = await wallet.allocate(plan)
+ results_by_id: dict[str, StepResult] = {}
+ sem = asyncio.Semaphore(_env_int('RETRIEVAL_WORKFLOW_PARALLEL_MAX', 3))
+
+ for batch in plan.topological_batches():
+ await asyncio.gather(
+ *[
+ self._run_step(
+ db,
+ step=step,
+ ledger=ledgers[step.id],
+ results_by_id=results_by_id,
+ semaphore=sem,
+ user_id=user_id,
+ namespace=namespace,
+ top_k=step.top_k or top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ data_type=step.data_type or data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ llm_fn=llm_fn,
+ )
+ for step in batch
+ ]
+ )
+ for step in batch:
+ await wallet.reclaim(step.id, ledgers[step.id])
+
+ answer_text = compose_final_answer(plan, results_by_id)
+ ordered_results = [results_by_id[step.id] for step in plan.steps if step.id in results_by_id]
+ referenced_chunks = _dedupe_references(
+ ref for step_result in ordered_results for ref in step_result.referenced_chunks
+ )
+ api_results = _references_to_results(referenced_chunks)
+ elapsed_ms = int((time.monotonic() - t0) * 1000)
+ logger.info(
+ 'workflow retrieval DONE: steps={} refs={} answer_chars={} elapsed={}ms',
+ len(ordered_results),
+ len(referenced_chunks),
+ len(answer_text),
+ elapsed_ms,
+ )
+ return WorkflowResult(
+ namespace=namespace,
+ query=query,
+ router_used='workflow_decomposed' if len(plan.steps) > 1 else 'workflow_single_step',
+ answer_text=answer_text,
+ plan=plan,
+ steps=ordered_results,
+ referenced_chunks=referenced_chunks,
+ results=api_results,
+ final_strategy_used=plan.final_strategy,
+ wallet_snapshot=wallet.snapshot(),
+ planner_snapshot=planner_ledger.snapshot(),
+ parent_run_id=self.parent_run_id,
+ )
+
+ async def _load_or_plan(
+ self,
+ *,
+ user_id: str,
+ namespace: str,
+ query: str,
+ planner_llm,
+ planner_ledger: BudgetLedger,
+ max_steps: int,
+ wallet_total: int,
+ per_retrieve: int,
+ ) -> QueryPlan:
+ try:
+ cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query)
+ if cached:
+ return QueryPlan.from_dict(cached, original_query=query)
+ except Exception as exc:
+ logger.warning(f'workflow plan cache read failed (ignored): {exc}')
+
+ planner = QueryPlanner(
+ llm_fn=planner_llm,
+ planner_ledger=planner_ledger,
+ max_steps=max_steps,
+ total_budget=wallet_total,
+ per_step_budget=per_retrieve,
+ )
+ plan = await planner.plan(query=query)
+ try:
+ await set_cached_workflow_plan(
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ plan=plan.to_dict(),
+ )
+ except Exception as exc:
+ logger.warning(f'workflow plan cache write failed (ignored): {exc}')
+ return plan
+
+ async def _run_step(
+ self,
+ db: AsyncSession,
+ *,
+ step: PlannedStep,
+ ledger: BudgetLedger,
+ results_by_id: dict[str, StepResult],
+ semaphore: asyncio.Semaphore,
+ user_id: str,
+ namespace: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int,
+ signal_paths: list[str] | None,
+ filter_mode: str,
+ channels: list[str] | None,
+ channel_weights: dict[str, float] | None,
+ llm_fn,
+ ) -> None:
+ async with semaphore:
+ if step.step_kind == 'synthesize':
+ await self._run_synthesize_step(step, ledger, results_by_id, llm_fn)
+ return
+ await self._run_retrieve_step(
+ db,
+ step=step,
+ ledger=ledger,
+ results_by_id=results_by_id,
+ user_id=user_id,
+ namespace=namespace,
+ top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ data_type=data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ llm_fn=llm_fn,
+ )
+
+ async def _run_retrieve_step(
+ self,
+ db: AsyncSession,
+ *,
+ step: PlannedStep,
+ ledger: BudgetLedger,
+ results_by_id: dict[str, StepResult],
+ user_id: str,
+ namespace: str,
+ top_k: int,
+ exclude_document_ids: list[str],
+ exclude_sections: list[dict[str, str]],
+ data_type: int,
+ signal_paths: list[str] | None,
+ filter_mode: str,
+ channels: list[str] | None,
+ channel_weights: dict[str, float] | None,
+ llm_fn,
+ ) -> None:
+ try:
+ # AsyncSession is not safe for concurrent use. Workflow steps may
+ # run in the same topological batch, so each retrieve step opens an
+ # isolated session and leaves the parent session untouched.
+ async with get_db_context() as step_db:
+ agentic_result = await RetrievalAgent().run(
+ step_db,
+ user_id=user_id,
+ namespace=namespace,
+ query=step.sub_query,
+ top_k=top_k,
+ llm_fn=llm_fn,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ data_type=data_type,
+ signal_paths=signal_paths,
+ filter_mode=filter_mode,
+ channels=channels,
+ channel_weights=channel_weights,
+ ledger=ledger,
+ parent_run_id=self.parent_run_id,
+ workflow_step_id=step.id,
+ )
+ results_by_id[step.id] = _step_result_from_agentic(step, agentic_result)
+ except Exception as exc:
+ logger.exception(f'workflow retrieve step failed: step_id={step.id}')
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status='error',
+ error=str(exc),
+ budget_snapshot=ledger.snapshot(),
+ )
+
+ async def _run_synthesize_step(
+ self,
+ step: PlannedStep,
+ ledger: BudgetLedger,
+ results_by_id: dict[str, StepResult],
+ llm_fn,
+ ) -> None:
+ if llm_fn is None:
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status='skipped',
+ answer_text='',
+ error='llm unavailable for synthesis',
+ budget_snapshot=ledger.snapshot(),
+ )
+ return
+ prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id}
+ try:
+ answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger)
+ refs = _dedupe_references(
+ ref for result in prior.values() for ref in result.referenced_chunks
+ )
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status='done',
+ answer_text=answer,
+ referenced_chunks=refs,
+ budget_snapshot=ledger.snapshot(),
+ )
+ except Exception as exc:
+ results_by_id[step.id] = StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status='budget_stop' if 'budget' in str(exc).lower() else 'error',
+ answer_text='(budget exhausted)' if 'budget' in str(exc).lower() else '',
+ error=str(exc),
+ budget_snapshot=ledger.snapshot(),
+ )
+
+
+def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult:
+ status = 'budget_stop' if 'budget' in (result.stop_reason or '') else 'done'
+ return StepResult(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ step_kind=step.step_kind,
+ depends_on=step.depends_on,
+ output_role=step.output_role,
+ status=status, # type: ignore[arg-type]
+ answer_text=result.answer_text,
+ evidence_text=result.evidence_text,
+ referenced_chunks=result.referenced_chunks,
+ budget_snapshot=result.budget_snapshot,
+ router_used=result.router_used,
+ stop_reason=result.stop_reason,
+ )
+
+
+def _dedupe_references(refs) -> list[dict[str, Any]]:
+ seen: set[str] = set()
+ out: list[dict[str, Any]] = []
+ for ref in refs:
+ chunk_id = str(ref.get('chunk_id') or '')
+ key = chunk_id or str(ref)
+ if key in seen:
+ continue
+ seen.add(key)
+ out.append(dict(ref))
+ return out
+
+
+def _references_to_results(refs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ return [
+ {
+ 'chunk_id': ref.get('chunk_id'),
+ 'document_id': ref.get('document_id'),
+ 'chunk_type': ref.get('chunk_type'),
+ 'source': {
+ 'document_id': ref.get('document_id'),
+ 'section_path': ref.get('section_path'),
+ },
+ }
+ for ref in refs
+ ]
+
+
+def _env_int(name: str, default: int) -> int:
+ try:
+ return int(os.environ.get(name, str(default)))
+ except (TypeError, ValueError):
+ return default
diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py
new file mode 100644
index 000000000..e34713db6
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py
@@ -0,0 +1,239 @@
+"""Query planner for decomposed retrieval workflows."""
+from __future__ import annotations
+
+import json
+import re
+import time
+from typing import Any
+
+from loguru import logger
+
+from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger
+from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage
+from shared.services.retrieval.workflow.types import FinalStrategy, OutputRole, PlannedStep, QueryPlan, StepKind
+from shared.utils.token_estimate import estimate_tokens
+
+
+_PLAN_SCHEMA = {
+ "reasoning_summary": "",
+ "steps": [
+ {
+ "id": "s1",
+ "sub_query": "",
+ "step_kind": "retrieve",
+ "depends_on": [],
+ "output_role": "final_part",
+ "top_k": 10,
+ }
+ ],
+ "final_strategy": "concat_final_parts",
+ "final_template": "",
+}
+
+_PLANNER_PROMPT = """\
+You are a retrieval workflow planner. Think step by step before answering.
+
+User query: {query}
+Knowledge base inventory: {kb_total_docs} docs / {kb_total_chunks} chunks.
+Wallet status: total_budget={total_budget} tokens (planner_used={planner_used}).
+
+Decide whether the query needs decomposition into multiple sub-queries.
+Most queries are single-step (return a 1-step plan with the original query).
+Only decompose when:
+ - The query asks for a comparison across distinct entities/time periods
+ - The query asks for a derived computation that requires multiple facts
+ - The query bundles 2+ independent informational asks
+
+Hard constraints:
+ - max_steps = {max_steps}
+ - Each retrieve step costs ~{per_step_budget} tokens; do NOT plan more
+ retrieve steps than the wallet can afford.
+ - synthesize steps must have non-empty depends_on.
+ - final_strategy must be one of: concat_final_parts, last_synthesize, template.
+ - step_kind must be retrieve or synthesize.
+ - output_role must be final_part, intermediate, or consumed_by_synthesis.
+
+Return ONLY a JSON object matching this schema (think first, then answer):
+{schema}
+"""
+
+
+class QueryPlanner:
+ """LLM-backed planner with strict fallback to a single retrieve step."""
+
+ def __init__(
+ self,
+ *,
+ llm_fn: LLMFn | None,
+ planner_ledger: BudgetLedger | None,
+ max_steps: int,
+ total_budget: int,
+ per_step_budget: int,
+ ) -> None:
+ self._llm_fn = llm_fn
+ self._ledger = planner_ledger
+ self._max_steps = max(max_steps, 1)
+ self._total_budget = max(total_budget, 1)
+ self._per_step_budget = max(per_step_budget, 1)
+
+ async def plan(
+ self,
+ *,
+ query: str,
+ kb_total_docs: int = 0,
+ kb_total_chunks: int = 0,
+ ) -> QueryPlan:
+ if self._llm_fn is None:
+ return QueryPlan.single_step(query, reason="planner_llm_unavailable")
+
+ prompt = _PLANNER_PROMPT.format(
+ query=query,
+ kb_total_docs=kb_total_docs,
+ kb_total_chunks=kb_total_chunks,
+ total_budget=self._total_budget,
+ planner_used=self._planner_used(),
+ max_steps=self._max_steps,
+ per_step_budget=self._per_step_budget,
+ schema=json.dumps(_PLAN_SCHEMA, ensure_ascii=False, indent=2),
+ )
+
+ try:
+ raw = await self._call_llm_with_budget(prompt)
+ plan = _parse_plan_response(
+ raw,
+ original_query=query,
+ max_steps=self._max_steps,
+ )
+ plan.validate()
+ return plan
+ except Exception as exc:
+ logger.warning(f"workflow planner failed, falling back to single step: {exc}")
+ plan = QueryPlan.single_step(query, reason="planner_fallback_single_step")
+ plan.planner_status = "fallback"
+ plan.planner_error = str(exc)
+ return plan
+
+ def _planner_used(self) -> int:
+ if self._ledger is None:
+ return 0
+ snapshot = self._ledger.snapshot()
+ used = 0
+ for pool in ("bootstrap", "planning", "context"):
+ pool_state = snapshot.get(pool) or {}
+ if isinstance(pool_state, dict):
+ used += int(pool_state.get("used") or 0)
+ return used
+
+ async def _call_llm_with_budget(self, prompt: str) -> str:
+ if self._ledger is None:
+ return await self._llm_fn(prompt) # type: ignore[misc]
+
+ est = estimate_tokens(prompt)
+ reserved = await self._ledger.try_reserve("bootstrap", est)
+ if not reserved:
+ raise BudgetExceeded("planner bootstrap budget exhausted")
+
+ t0 = time.monotonic()
+ try:
+ response = await self._llm_fn(prompt) # type: ignore[misc]
+ except Exception:
+ await self._ledger.refund("bootstrap", est=est)
+ raise
+ usage = current_llm_usage.get() or {}
+ actual = int(usage.get("prompt_tokens") or est)
+ await self._ledger.commit("bootstrap", actual=actual, est=est)
+ logger.info(
+ "workflow planner llm call: est_tokens={} actual_tokens={} latency={}ms",
+ est,
+ actual,
+ int((time.monotonic() - t0) * 1000),
+ )
+ return response
+
+
+def _parse_plan_response(text: str, *, original_query: str, max_steps: int) -> QueryPlan:
+ data = _extract_json_object(text)
+ if not isinstance(data, dict):
+ raise ValueError("planner response is not a JSON object")
+
+ steps_data = data.get("steps")
+ if not isinstance(steps_data, list) or not steps_data:
+ raise ValueError("planner response must include non-empty steps[]")
+ if len(steps_data) > max_steps:
+ raise ValueError(f"planner returned {len(steps_data)} steps, max is {max_steps}")
+
+ steps: list[PlannedStep] = []
+ for index, item in enumerate(steps_data, start=1):
+ if not isinstance(item, dict):
+ raise ValueError("planner step must be an object")
+ step_id = str(item.get("id") or f"s{index}").strip()
+ step_kind = _coerce_step_kind(item.get("step_kind"))
+ output_role = _coerce_output_role(item.get("output_role"))
+ depends_on_raw = item.get("depends_on") or []
+ if not isinstance(depends_on_raw, list):
+ raise ValueError(f"step {step_id} depends_on must be a list")
+ top_k = item.get("top_k")
+ data_type = item.get("data_type")
+ steps.append(
+ PlannedStep(
+ id=step_id,
+ sub_query=str(item.get("sub_query") or original_query).strip(),
+ step_kind=step_kind,
+ depends_on=[str(dep).strip() for dep in depends_on_raw if str(dep).strip()],
+ output_role=output_role,
+ top_k=int(top_k) if top_k is not None else None,
+ data_type=int(data_type) if data_type is not None else None,
+ metadata={
+ key: value
+ for key, value in item.items()
+ if key not in {"id", "sub_query", "step_kind", "depends_on", "output_role", "top_k", "data_type"}
+ },
+ )
+ )
+
+ final_strategy = _coerce_final_strategy(data.get("final_strategy"))
+ return QueryPlan(
+ original_query=original_query,
+ steps=steps,
+ final_strategy=final_strategy,
+ reasoning_summary=str(data.get("reasoning_summary") or "").strip(),
+ final_template=str(data.get("final_template") or "").strip() or None,
+ )
+
+
+def _extract_json_object(text: str) -> dict[str, Any]:
+ text = text.strip()
+ try:
+ parsed = json.loads(text)
+ if isinstance(parsed, dict):
+ return parsed
+ except (ValueError, json.JSONDecodeError):
+ pass
+ match = re.search(r"\{.*\}", text, re.DOTALL)
+ if not match:
+ raise ValueError("no JSON object found")
+ parsed = json.loads(match.group())
+ if not isinstance(parsed, dict):
+ raise ValueError("extracted JSON is not an object")
+ return parsed
+
+
+def _coerce_step_kind(value: Any) -> StepKind:
+ raw = str(value or "retrieve").strip().lower()
+ if raw not in {"retrieve", "synthesize"}:
+ raise ValueError(f"unsupported step_kind: {value}")
+ return raw # type: ignore[return-value]
+
+
+def _coerce_output_role(value: Any) -> OutputRole:
+ raw = str(value or "final_part").strip().lower()
+ if raw not in {"final_part", "intermediate", "consumed_by_synthesis"}:
+ raise ValueError(f"unsupported output_role: {value}")
+ return raw # type: ignore[return-value]
+
+
+def _coerce_final_strategy(value: Any) -> FinalStrategy:
+ raw = str(value or "concat_final_parts").strip().lower()
+ if raw not in {"concat_final_parts", "last_synthesize", "template"}:
+ raise ValueError(f"unsupported final_strategy: {value}")
+ return raw # type: ignore[return-value]
diff --git a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py
new file mode 100644
index 000000000..0e825eed9
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py
@@ -0,0 +1,125 @@
+"""Synthesis helpers for decomposed retrieval workflows."""
+from __future__ import annotations
+
+import re
+
+from shared.services.retrieval.agentic.budget import BudgetLedger
+from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage
+from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult
+from shared.utils.token_estimate import estimate_tokens
+
+
+_SYNTHESIZE_PROMPT = """\
+You are composing an intermediate or final answer for a retrieval workflow.
+Use ONLY the prior step outputs below. Do not use external knowledge.
+
+Current step id: {step_id}
+Current step task: {sub_query}
+
+Prior step outputs:
+{prior_outputs}
+
+Return a concise answer that directly satisfies the current step task.
+If the prior outputs are insufficient, explain the missing information.
+"""
+
+
+async def synthesize_step(
+ step: PlannedStep,
+ *,
+ prior_results: dict[str, StepResult],
+ llm_fn: LLMFn,
+ ledger: BudgetLedger | None,
+) -> str:
+ prior_outputs = _format_prior_outputs(step.depends_on, prior_results)
+ prompt = _SYNTHESIZE_PROMPT.format(
+ step_id=step.id,
+ sub_query=step.sub_query,
+ prior_outputs=prior_outputs,
+ )
+ if ledger is None:
+ return (await llm_fn(prompt)).strip()
+
+ est = estimate_tokens(prompt)
+ reserved = await ledger.try_reserve("context", est)
+ if not reserved:
+ raise RuntimeError("synthesis context budget exhausted")
+ try:
+ response = await llm_fn(prompt)
+ except Exception:
+ await ledger.refund("context", est=est)
+ raise
+ usage = current_llm_usage.get() or {}
+ actual = int(usage.get("prompt_tokens") or est)
+ await ledger.commit("context", actual=actual, est=est)
+ return response.strip()
+
+
+def compose_final_answer(plan: QueryPlan, results: dict[str, StepResult]) -> str:
+ """Compose workflow final answer according to planner strategy."""
+ if plan.final_strategy == "last_synthesize":
+ for step in reversed(plan.steps):
+ result = results.get(step.id)
+ if result and step.step_kind == "synthesize" and result.status == "done":
+ return result.answer_text
+ return _concat_final_parts(plan, results)
+
+ if plan.final_strategy == "template" and plan.final_template:
+ return _render_template(plan.final_template, results).strip()
+
+ return _concat_final_parts(plan, results)
+
+
+def _concat_final_parts(plan: QueryPlan, results: dict[str, StepResult]) -> str:
+ parts: list[str] = []
+ for step in plan.steps:
+ if step.output_role != "final_part":
+ continue
+ result = results.get(step.id)
+ if not result or result.status not in ("done", "budget_stop") or not result.answer_text:
+ continue
+ parts.append(result.answer_text.strip())
+ if parts:
+ return "\n\n".join(parts)
+
+ fallback_parts = [
+ result.answer_text.strip()
+ for step in plan.steps
+ if (result := results.get(step.id)) and result.answer_text.strip()
+ ]
+ return "\n\n".join(fallback_parts)
+
+
+def _format_prior_outputs(depends_on: list[str], prior_results: dict[str, StepResult]) -> str:
+ lines: list[str] = []
+ for step_id in depends_on:
+ result = prior_results.get(step_id)
+ if not result:
+ lines.append(f"## {step_id}\n(status: missing)\n")
+ continue
+ lines.append(
+ "\n".join(
+ [
+ f"## {step_id}",
+ f"Sub-query: {result.sub_query}",
+ f"Status: {result.status}",
+ "Answer:",
+ result.answer_text or "(empty)",
+ "",
+ ]
+ )
+ )
+ return "\n".join(lines) if lines else "(no prior outputs)"
+
+
+def _render_template(template: str, results: dict[str, StepResult]) -> str:
+ def _replace(match: re.Match[str]) -> str:
+ step_id = match.group(1)
+ field = match.group(2)
+ result = results.get(step_id)
+ if not result:
+ return ""
+ attr = "answer_text" if field == "answer" else field
+ return str(getattr(result, attr, ""))
+
+ return re.sub(r"\{\{\s*steps\.([A-Za-z0-9_-]+)\.(answer_text|answer|evidence_text|status)\s*\}\}", _replace, template)
diff --git a/packages/shared-python/shared/services/retrieval/workflow/types.py b/packages/shared-python/shared/services/retrieval/workflow/types.py
new file mode 100644
index 000000000..e196fd656
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/types.py
@@ -0,0 +1,230 @@
+"""Types for query-decomposition retrieval workflows."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+StepKind = Literal["retrieve", "synthesize"]
+OutputRole = Literal["final_part", "intermediate", "consumed_by_synthesis"]
+FinalStrategy = Literal["concat_final_parts", "last_synthesize", "template"]
+StepStatus = Literal["done", "skipped", "error", "budget_stop"]
+
+
+@dataclass
+class PlannedStep:
+ """A single node in the query workflow DAG."""
+
+ id: str
+ sub_query: str
+ step_kind: StepKind = "retrieve"
+ depends_on: list[str] = field(default_factory=list)
+ output_role: OutputRole = "final_part"
+ top_k: int | None = None
+ data_type: int | None = None
+ metadata: dict[str, Any] = field(default_factory=dict)
+
+ def to_dict(self) -> dict[str, Any]:
+ data: dict[str, Any] = {
+ "id": self.id,
+ "sub_query": self.sub_query,
+ "step_kind": self.step_kind,
+ "depends_on": list(self.depends_on),
+ "output_role": self.output_role,
+ }
+ if self.top_k is not None:
+ data["top_k"] = self.top_k
+ if self.data_type is not None:
+ data["data_type"] = self.data_type
+ if self.metadata:
+ data["metadata"] = dict(self.metadata)
+ return data
+
+
+@dataclass
+class QueryPlan:
+ """Planner output consumed by ``WorkflowOrchestrator``."""
+
+ original_query: str
+ steps: list[PlannedStep]
+ final_strategy: FinalStrategy = "concat_final_parts"
+ reasoning_summary: str = ""
+ final_template: str | None = None
+ planner_status: str = "planned"
+ planner_error: str | None = None
+
+ @staticmethod
+ def single_step(query: str, *, reason: str = "single_step") -> "QueryPlan":
+ return QueryPlan(
+ original_query=query,
+ steps=[
+ PlannedStep(
+ id="s1",
+ sub_query=query,
+ step_kind="retrieve",
+ depends_on=[],
+ output_role="final_part",
+ )
+ ],
+ final_strategy="concat_final_parts",
+ reasoning_summary=reason,
+ )
+
+ @staticmethod
+ def from_dict(data: dict[str, Any], *, original_query: str | None = None) -> "QueryPlan":
+ steps = [
+ PlannedStep(
+ id=str(item.get("id") or f"s{idx}"),
+ sub_query=str(item.get("sub_query") or original_query or ""),
+ step_kind=item.get("step_kind", "retrieve"),
+ depends_on=[str(dep) for dep in item.get("depends_on") or []],
+ output_role=item.get("output_role", "final_part"),
+ top_k=int(item["top_k"]) if item.get("top_k") is not None else None,
+ data_type=int(item["data_type"]) if item.get("data_type") is not None else None,
+ metadata=dict(item.get("metadata") or {}),
+ )
+ for idx, item in enumerate(data.get("steps") or [], start=1)
+ if isinstance(item, dict)
+ ]
+ plan = QueryPlan(
+ original_query=str(data.get("original_query") or original_query or ""),
+ steps=steps,
+ final_strategy=data.get("final_strategy", "concat_final_parts"),
+ reasoning_summary=str(data.get("reasoning_summary") or ""),
+ final_template=data.get("final_template"),
+ planner_status=str(data.get("planner_status") or "cached"),
+ planner_error=data.get("planner_error"),
+ )
+ plan.validate()
+ return plan
+
+ def to_dict(self) -> dict[str, Any]:
+ data: dict[str, Any] = {
+ "original_query": self.original_query,
+ "reasoning_summary": self.reasoning_summary,
+ "steps": [step.to_dict() for step in self.steps],
+ "final_strategy": self.final_strategy,
+ "planner_status": self.planner_status,
+ }
+ if self.final_template:
+ data["final_template"] = self.final_template
+ if self.planner_error:
+ data["planner_error"] = self.planner_error
+ return data
+
+ def step_by_id(self) -> dict[str, PlannedStep]:
+ return {step.id: step for step in self.steps}
+
+ def validate(self) -> None:
+ if not self.steps:
+ raise ValueError("query plan must contain at least one step")
+ ids = [step.id for step in self.steps]
+ if len(ids) != len(set(ids)):
+ raise ValueError("query plan step ids must be unique")
+ id_set = set(ids)
+ for step in self.steps:
+ if not step.id.strip():
+ raise ValueError("query plan step id cannot be empty")
+ if not step.sub_query.strip():
+ raise ValueError(f"query plan step {step.id} sub_query cannot be empty")
+ if step.step_kind not in ("retrieve", "synthesize"):
+ raise ValueError(f"unsupported step_kind: {step.step_kind}")
+ if step.output_role not in ("final_part", "intermediate", "consumed_by_synthesis"):
+ raise ValueError(f"unsupported output_role: {step.output_role}")
+ missing = [dep for dep in step.depends_on if dep not in id_set]
+ if missing:
+ raise ValueError(f"step {step.id} depends on unknown steps: {missing}")
+ if step.step_kind == "synthesize" and not step.depends_on:
+ raise ValueError(f"synthesize step {step.id} must depend on prior steps")
+ self.topological_batches()
+
+ def topological_batches(self) -> list[list[PlannedStep]]:
+ """Return executable batches; steps in a batch have no mutual dependency."""
+ remaining = {step.id: step for step in self.steps}
+ completed: set[str] = set()
+ batches: list[list[PlannedStep]] = []
+ while remaining:
+ ready = [
+ step
+ for step in self.steps
+ if step.id in remaining and set(step.depends_on).issubset(completed)
+ ]
+ if not ready:
+ raise ValueError("query plan contains a dependency cycle")
+ batches.append(ready)
+ for step in ready:
+ remaining.pop(step.id, None)
+ completed.add(step.id)
+ return batches
+
+
+@dataclass
+class StepResult:
+ """Execution result for one planned step."""
+
+ step_id: str
+ sub_query: str
+ step_kind: StepKind
+ depends_on: list[str]
+ output_role: OutputRole
+ status: StepStatus = "done"
+ answer_text: str = ""
+ evidence_text: str | None = None
+ referenced_chunks: list[dict[str, Any]] = field(default_factory=list)
+ budget_snapshot: dict[str, Any] | None = None
+ child_run_id: str | None = None
+ router_used: str = ""
+ stop_reason: str = ""
+ error: str | None = None
+
+ def to_api_dict(self) -> dict[str, Any]:
+ return {
+ "step_id": self.step_id,
+ "sub_query": self.sub_query,
+ "step_kind": self.step_kind,
+ "depends_on": list(self.depends_on),
+ "output_role": self.output_role,
+ "status": self.status,
+ "answer_text": self.answer_text,
+ "evidence_text": self.evidence_text,
+ "referenced_chunks": list(self.referenced_chunks),
+ "budget_snapshot": self.budget_snapshot,
+ "child_run_id": self.child_run_id,
+ "router_used": self.router_used,
+ "stop_reason": self.stop_reason,
+ "error": self.error,
+ }
+
+
+@dataclass
+class WorkflowResult:
+ """Top-level response from query workflow execution."""
+
+ namespace: str
+ query: str
+ router_used: str
+ answer_text: str
+ plan: QueryPlan | None = None
+ steps: list[StepResult] = field(default_factory=list)
+ referenced_chunks: list[dict[str, Any]] = field(default_factory=list)
+ results: list[dict[str, Any]] = field(default_factory=list)
+ final_strategy_used: FinalStrategy | None = None
+ wallet_snapshot: dict[str, Any] | None = None
+ planner_snapshot: dict[str, Any] | None = None
+ parent_run_id: str | None = None
+
+ def to_api_response(self) -> dict[str, Any]:
+ return {
+ "namespace": self.namespace,
+ "query": self.query,
+ "router_used": self.router_used,
+ "answer_text": self.answer_text,
+ "referenced_chunks": self.referenced_chunks,
+ "results": self.results,
+ "plan": self.plan.to_dict() if self.plan else None,
+ "steps": [step.to_api_dict() for step in self.steps] if self.steps else None,
+ "final_strategy_used": self.final_strategy_used,
+ "wallet_snapshot": self.wallet_snapshot,
+ "planner_snapshot": self.planner_snapshot,
+ "parent_run_id": self.parent_run_id,
+ }
diff --git a/packages/shared-python/shared/services/retrieval/workflow/wallet.py b/packages/shared-python/shared/services/retrieval/workflow/wallet.py
new file mode 100644
index 000000000..0a540a120
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/workflow/wallet.py
@@ -0,0 +1,150 @@
+"""Top-level wallet for decomposed retrieval workflows."""
+from __future__ import annotations
+
+import asyncio
+import os
+from dataclasses import dataclass, field
+
+from shared.services.retrieval.agentic.budget import BudgetLedger
+from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan
+
+
+_RETRIEVE_FLOOR = 4000
+_SYNTHESIZE_FLOOR = 1500
+
+
+def _env_float(name: str, default: float) -> float:
+ try:
+ return float(os.environ.get(name, str(default)))
+ except (TypeError, ValueError):
+ return default
+
+
+def _env_int(name: str, default: int) -> int:
+ try:
+ return int(os.environ.get(name, str(default)))
+ except (TypeError, ValueError):
+ return default
+
+
+@dataclass
+class BudgetWallet:
+ """Issue per-step ``BudgetLedger`` instances under a workflow hard cap."""
+
+ total: int
+ per_retrieve_step_default: int
+ per_synthesize_step_default: int
+ # Read from env for consistency with _build_config_from_env()
+ planning_ratio: float = field(
+ default_factory=lambda: _env_float('RETRIEVAL_AGENTIC_PLANNING_RATIO', 0.5)
+ )
+ bootstrap_budget: int = field(
+ default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET', 2000)
+ )
+ per_doc_min_share: int = field(
+ default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE', 1500)
+ )
+ _allocations: dict[str, int] = field(default_factory=dict, init=False)
+ _reclaimed: dict[str, int] = field(default_factory=dict, init=False)
+ _ledgers: dict[str, BudgetLedger] = field(default_factory=dict, init=False)
+ _lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)
+
+ async def allocate(self, plan: QueryPlan) -> dict[str, BudgetLedger]:
+ """Pre-allocate one ledger per step, scaling down when needed."""
+ async with self._lock:
+ requested = {
+ step.id: self._requested_for_step(step)
+ for step in plan.steps
+ }
+ requested_total = sum(requested.values())
+ if requested_total <= 0:
+ requested_total = 1
+
+ if requested_total <= self.total:
+ allocations = requested
+ else:
+ scale = max(self.total, 1) / requested_total
+ allocations = {}
+ for step in plan.steps:
+ floor = _RETRIEVE_FLOOR if step.step_kind == "retrieve" else _SYNTHESIZE_FLOOR
+ allocations[step.id] = max(floor, int(requested[step.id] * scale))
+
+ scaled_total = sum(allocations.values())
+ if scaled_total > self.total:
+ # Respect the hard cap by reducing largest allocations first,
+ # never dropping below the per-kind floors.
+ excess = scaled_total - self.total
+ for step_id, amount in sorted(
+ allocations.items(), key=lambda item: item[1], reverse=True
+ ):
+ if excess <= 0:
+ break
+ step = next(s for s in plan.steps if s.id == step_id)
+ floor = _RETRIEVE_FLOOR if step.step_kind == "retrieve" else _SYNTHESIZE_FLOOR
+ reducible = max(amount - floor, 0)
+ delta = min(reducible, excess)
+ allocations[step_id] = amount - delta
+ excess -= delta
+
+ self._allocations = allocations
+ self._ledgers = {
+ step.id: self._new_ledger(step, allocations[step.id])
+ for step in plan.steps
+ }
+ return dict(self._ledgers)
+
+ async def reclaim(self, step_id: str, ledger: BudgetLedger) -> None:
+ """Record unused capacity after a step completes."""
+ async with self._lock:
+ allocated = self._allocations.get(step_id, 0)
+ used = self._ledger_used(ledger)
+ self._reclaimed[step_id] = max(allocated - used, 0)
+
+ def total_used(self) -> int:
+ return sum(self._ledger_used(ledger) for ledger in self._ledgers.values())
+
+ def snapshot(self) -> dict[str, object]:
+ return {
+ "total": self.total,
+ "allocated": sum(self._allocations.values()),
+ "used": self.total_used(),
+ "remaining": max(self.total - self.total_used(), 0),
+ "allocations": dict(self._allocations),
+ "reclaimed": dict(self._reclaimed),
+ "steps": {
+ step_id: ledger.snapshot()
+ for step_id, ledger in self._ledgers.items()
+ },
+ }
+
+ def _requested_for_step(self, step: PlannedStep) -> int:
+ if step.step_kind == "synthesize":
+ return max(self.per_synthesize_step_default, _SYNTHESIZE_FLOOR)
+ return max(self.per_retrieve_step_default, _RETRIEVE_FLOOR)
+
+ def _new_ledger(self, step: PlannedStep, total: int) -> BudgetLedger:
+ if step.step_kind == "synthesize":
+ # Put almost all tokens into context for pure synthesis calls.
+ return BudgetLedger(
+ total=max(total, 1),
+ planning_ratio=0.0,
+ bootstrap=0,
+ per_doc_min_share=0,
+ )
+ return BudgetLedger(
+ total=max(total, 1),
+ planning_ratio=self.planning_ratio,
+ bootstrap=min(self.bootstrap_budget, max(total, 1)),
+ per_doc_min_share=self.per_doc_min_share,
+ )
+
+ @staticmethod
+ def _ledger_used(ledger: BudgetLedger) -> int:
+ snapshot = ledger.snapshot()
+ total = 0
+ for pool in ("bootstrap", "planning", "context"):
+ pool_state = snapshot.get(pool) or {}
+ if isinstance(pool_state, dict):
+ total += int(pool_state.get("used") or 0)
+ total += int(pool_state.get("reserved") or 0)
+ return total