diff --git a/AGENTS.md b/AGENTS.md
index 6127f2d11..48a62326b 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -22,3 +22,680 @@ feat/alice/add-document-preview
fix/bob/chunk-position-range
refactor/chris/extract-chunk-converter
```
+
+---
+
+## Project Structure
+
+```text
+knowhereapi-main/
+├── apps/
+│ ├── api/ # FastAPI REST API (port 5005)
+│ │ ├── app/
+│ │ │ ├── api/v1/routes/ # Endpoint handlers
+│ │ │ ├── services/ # Business logic (auth, knowledge, billing)
+│ │ │ └── repositories/ # Data access layer
+│ │ └── main.py # Entrypoint, runs migrations on start
+│ ├── worker/ # Celery worker for async document processing
+│ │ ├── app/
+│ │ │ ├── services/document_parser/ # All parser modules
+│ │ │ └── services/workload/ # Celery task handlers
+│ │ └── worker.py # Celery entrypoint
+│ ├── 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
+└── deploy/ # Docker Compose & deployment scripts
+```
+
+---
+
+## End-to-End Pipeline Overview
+
+```mermaid
+flowchart TB
+ subgraph INGEST["① Document Ingestion (API)"]
+ Upload["POST /v1/documents"] --> Job["Create Job + S3 Upload"]
+ Job --> Queue["Celery Task Queue"]
+ end
+
+ subgraph PARSE["② Document Parsing (Worker)"]
+ Queue --> Router["parse_service.checkerboard_inject_parse"]
+ Router --> Profiler["doc_profiler.profile_document"]
+ Profiler --> PDF["pdf_parser → MinerU"]
+ Profiler --> DOCX["doc_parser.parse_docx"]
+ Profiler --> PPTX["pptx_parser → iLoveAPI → PDF"]
+ Profiler --> XLSX["table_parser.parse_xlsx"]
+ Profiler --> MD["md_parser.parse_md"]
+ Profiler --> IMG["image_parser.parse_image"]
+ PDF --> DF["pd.DataFrame (ALL_DF_COLS)"]
+ DOCX --> DF
+ PPTX --> DF
+ XLSX --> DF
+ MD --> DF
+ IMG --> DF
+ end
+
+ subgraph CONVERT["③ Chunk Conversion"]
+ DF --> Converter["dataframe_chunk_converter.dataframe_to_chunks"]
+ Converter --> Chunks["list[ChunkPayload]"]
+ end
+
+ subgraph PUBLISH["④ Publication (shared)"]
+ Chunks --> Dedup["RetrievalPublicationService._dedup_chunks_by_content"]
+ Dedup --> DocState["publish_document_state → Documents/Sections/Chunks"]
+ DocState --> Graph["publish_document_graph → GraphNodes/GraphEdges"]
+ end
+
+ 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)"]
+ Channels --> RRF["RRF Fusion"]
+ Agentic --> Hydrate["hydrate_paths_to_rows"]
+ RRF --> Rank["_rank_candidates_by_path"]
+ Hydrate --> Rank
+ Rank --> Assemble["assemble_retrieval_results"]
+ Assemble --> Results["Cited Evidence Results"]
+ end
+```
+
+---
+
+## Stage ②: Document Parsing Pipeline
+
+### Entry Point
+
+`apps/worker/app/services/document_parser/parse_service.py` →
+`checkerboard_inject_parse()`
+
+This is the universal entry for all file types. It:
+
+1. **Profiles** the document via `doc_profiler.profile_document()` to detect
+ file type, page count, and special categories (e.g. `atlas`).
+2. **Routes** to the appropriate parser based on file extension.
+3. **Post-processes**: cleans up unreferenced images, compresses PNG→JPG.
+4. Returns `(output_dir, parsed_df)` — the parsed DataFrame.
+
+### Parser Routing Table
+
+| Extension | Parser Module | Strategy |
+|:---|:---|:---|
+| `.pdf` | `pdf_parser.parse_pdfs` | MinerU API → `md_parser` → `layout_parser.pred_titles` |
+| `.docx` | `doc_parser.parse_docx` + `convert_doc2dics` | OXML iteration → heading detection → hierarchical tree |
+| `.doc` | `legacy_converter.doc_to_docx` → `.docx` pipeline | LibreOffice headless conversion first |
+| `.pptx` | `pptx_parser.parse_pptx` | iLoveAPI PPTX→PDF → MinerU pipeline |
+| `.xlsx` | `table_parser.parse_xlsx` | Sheet-by-sheet HTML table extraction |
+| `.xls` | `legacy_converter.xls_to_xlsx` → `.xlsx` pipeline | LibreOffice conversion first |
+| `.md` | `md_parser.parse_md` | Markdown heading parsing + LLM summaries |
+| `.txt` | `txt_parser.parse_texts` → `md_parser` | Read lines then route to MD parser |
+| `.png/.jpg` | `image_parser.parse_image` | VLM image description + OCR |
+| `.fragment` | `fragment_parser.parse_fragment` | Raw text fragment ingestion |
+
+### Heading Detection: `layout_parser.pred_titles()`
+
+The core hierarchical recognition module. Determines heading levels using:
+
+1. **TOC-first**: If a DOCX TOC exists (`toc_parser.build_docx_toc_hierarchies`),
+ use it as ground truth for heading levels.
+2. **Regex patterns**: Match numbered headings like `1.2.3`, `第X章`, `（一）`.
+3. **LLM smart parse**: When `smart_title_parse=True`, send candidate headings
+ to the hierarchy model (`HIERARCHY_LLM_MODEL` or `NORMOL_MODEL`) for level
+ assignment.
+4. **Font clustering (PDF)**: K-means on span heights from MinerU `layout.json`
+ to group headings into 5 discrete tiers.
+
+### DOCX Parsing Deep Dive: `doc_parser.py`
+
+```mermaid
+flowchart LR
+ A[load_file_bytes] --> B[iter_block_items]
+ B --> C{Element Type}
+ C -->|CT_P| D[Paragraph + Images]
+ C -->|CT_Tbl| E[Table + Cell Images]
+ C -->|sdt| F[TOC Detection]
+ D --> G[pred_titles → heading levels]
+ G --> H[Build hierarchical tree]
+ E --> I[table2html → HTML]
+ H --> J[get_leaf_dics → flatten]
+ J --> K[postprocess_leaf_dics → LLM summaries]
+ K --> L[convert_doc2dics → DataFrame]
+```
+
+Key logic in `parse_docx()`:
+
+- **`iter_block_items()`**: Iterates OXML body elements, yielding
+ `(ele_num, content, label, meta)` tuples. Labels: `PTXT`, `TABLE`,
+ `IMAGE`, `TOC-AREA`.
+- **Heading stack**: Maintains `headings_stack` with `{heading, content[], level}`
+ dicts. New headings pop the stack to their parent level.
+- **Image dedup**: Uses `perceptual_hash()` for document-level visual dedup.
+ Cached in `_seen_images` dict.
+- **Table handling**: `table2html()` converts python-docx Table to HTML with
+ accurate `rowspan`/`colspan` via direct OXML inspection.
+
+### PDF Parsing: MinerU Pipeline
+
+```mermaid
+flowchart LR
+ PDF[pdf_parser] --> MinerU[MinerU Cloud API]
+ MinerU --> MDFile[Markdown + layout.json]
+ MDFile --> MDParser[md_parser.parse_md]
+ MDParser --> EvalHeadings[eval_md_headings + layout.json]
+ EvalHeadings --> PredTitles[layout_parser.pred_titles]
+ PredTitles --> Chunks[Hierarchical Chunks]
+```
+
+### LLM Models Used
+
+| Task | Config Key | Default Model |
+|:---|:---|:---|
+| Text/table summarization | `NORMOL_MODEL` | `deepseek-chat` |
+| Heading hierarchy recognition | `HIERARCHY_LLM_MODEL` | Falls back to `NORMOL_MODEL` |
+| Image description (VLM) | `IMAGE_MODEL` | `qwen3.5-flash` |
+| Image OCR / Q&A | `IMAGE_MODEL_MAX` | `qwen3.5-flash` |
+| Atlas classification | VLM via `atlas_classifier` | `IMAGE_MODEL` |
+
+---
+
+## Persisted Knowledge Base Schema (On-Disk Output)
+
+After parsing and chunk conversion, results are persisted to `~/.knowhere/{kb_name}/`.
+This on-disk structure is the **authoritative persisted format** — the intermediate
+DataFrame is an internal detail. Below is the complete schema.
+
+### KB-Level Directory Layout
+
+```text
+~/.knowhere/{kb_name}/
+├── knowledge_graph.json # KB-wide graph: file metadata + cross-doc edges
+├── chunk_stats.json # Per-chunk retrieval hit analytics {chunk_id → stats}
+├── {source_file_name}/ # One directory per ingested document
+│ ├── chunks.json # All parsed chunks for this document
+│ ├── doc_nav.json # Hierarchical navigation tree for agentic retrieval
+│ ├── manifest.json # Parse metadata + full heading hierarchy
+│ ├── {source_file_name}.zip # Archived original + parsed assets
+│ ├── images/ # Extracted image assets (PNG/JPG)
+│ ├── tables/ # Extracted table assets (HTML)
+│ ├── preds_3_llm_base.csv # Debug: heading predictions (base LLM pass)
+│ ├── preds_4_llm_final.csv # Debug: heading predictions (final LLM pass)
+│ ├── preds_5_final_output.csv # Debug: final parser DataFrame output
+│ └── toc_hierarchies.json # Debug: extracted TOC structure (DOCX only)
+```
+
+### `knowledge_graph.json` — KB-Wide Graph
+
+```json
+{
+ "version": "2.0",
+ "kb_id": "test_kb",
+ "stats": { "total_files": 3, "total_chunks": 364, "total_cross_file_edges": 0 },
+ "files": {
+ "AI_Security_Report.docx": {
+ "chunks_count": 155,
+ "types": { "image": 13, "table": 1, "text": 141 },
+ "top_keywords": ["model", "security", "ai", "operations", "artificial_intelligence"],
+ "top_summary": "This document includes: Legal Notice, Foreword, 1. Overview, ...",
+ "importance": 0.3,
+ "created_at": "2026-05-09T09:14:12.422208+00:00"
+ }
+ },
+ "edges": []
+}
+```
+
+| Field | Description |
+|:---|:---|
+| `files.{name}.top_keywords` | TF-IDF top keywords across all chunks (used for cross-doc edge scoring) |
+| `files.{name}.top_summary` | Auto-generated outline of top-level headings (injected by `load_nav_top_summary()`) |
+| `files.{name}.importance` | Base importance score (feeds `compute_importance_score()` in ranking) |
+| `edges[]` | Cross-document edges with `{source, target, weight, shared_keywords}` when keyword overlap ≥ 0.8 |
+
+### `chunk_stats.json` — Retrieval Hit Analytics
+
+```json
+{
+ "2e2beffc-90b2-5429-8ee7-3c49260a1204": {
+ "hit_count": 0,
+ "first_hit": null,
+ "last_hit": null,
+ "created_at": "2026-05-09T09:14:12.423011+00:00"
+ }
+}
+```
+
+Keyed by `chunk_id`. `hit_count` and `last_hit` feed into `importance_norm_score`
+for retrieval ranking boost.
+
+### `chunks.json` — Per-Document Chunk Records
+
+The core persisted data. Contains `{"chunks": [...]}` — an ordered array of
+chunk objects. Three chunk types exist:
+
+#### Text Chunk
+
+```json
+{
+ "chunk_id": "d88e4c47-3c48-5bdf-b849-693c00453021",
+ "type": "text",
+ "content": "AI Security Report\n\nThe image displays a tech theme...\n[images/image-1 ai_model.png]\n",
+ "path": "test_kb/AI_Security_Report.docx/1. Overview/1.1 Key Findings",
+ "metadata": {
+ "length": 220,
+ "summary": "",
+ "page_nums": [],
+ "tokens": ["model", "technology", "market", "research", "report"],
+ "keywords": [],
+ "connect_to": [
+ {
+ "target": "2e2beffc-90b2-5429-8ee7-3c49260a1204",
+ "relation": "embeds",
+ "ref": "[images/image-1 ai_model.png]",
+ "position": { "start": 109, "end": 135 }
+ }
+ ]
+ }
+}
+```
+
+#### Image Chunk
+
+```json
+{
+ "chunk_id": "2e2beffc-90b2-5429-8ee7-3c49260a1204",
+ "type": "image",
+ "content": "\nThe image displays a tech theme...\n[images/image-1 ai_model.png]\n",
+ "path": "images/image-1 ai_model.png",
+ "metadata": {
+ "length": 121,
+ "summary": "image-1\nThe image displays a tech theme...",
+ "page_nums": [],
+ "file_path": "images/image-1 ai_model.png",
+ "keywords": [],
+ "tokens": []
+ }
+}
+```
+
+#### Table Chunk
+
+```json
+{
+ "chunk_id": "a5c3d644-479a-51f9-9a54-ff6789c1f6e8",
+ "type": "table",
+ "content": "| Architecture Layer | AI Capabilities |
...
",
+ "path": "tables/table-1 ai_architecture.html",
+ "metadata": {
+ "length": 488,
+ "summary": "table-1\nThe table shows the architecture layers...",
+ "page_nums": [],
+ "file_path": "tables/table-1 ai_architecture.html",
+ "keywords": ["AI_Capabilities", "Security_Engine", "Intelligent_Collaboration"],
+ "tokens": []
+ }
+}
+```
+
+#### Chunk Field Reference
+
+| Field | Type | Description |
+|:---|:---|:---|
+| `chunk_id` | `str` | Deterministic UUID5 hash from content (`gen_str_codes`) — enables cross-doc dedup |
+| `type` | `str` | `"text"` / `"image"` / `"table"` |
+| `content` | `str` | Raw text, VLM description + asset ref, or HTML `` |
+| `path` | `str` | Hierarchical path: `{kb}/{file}/{section1}/{section2}/...` for text; `images/...` or `tables/...` for assets |
+| `metadata.length` | `int` | Character count of content |
+| `metadata.summary` | `str` | LLM summary (images/tables: `"image-N\n{description}"`) |
+| `metadata.tokens` | `list[str]` | Pre-tokenized Chinese terms for BM25 retrieval |
+| `metadata.keywords` | `list[str]` | LLM-extracted keywords (semicolon-split from DataFrame) |
+| `metadata.page_nums` | `list[int]` | Source page numbers (PDF only) |
+| `metadata.file_path` | `str` | Relative asset path for images/tables |
+| `metadata.connect_to` | `list[ConnectionValue]` | Cross-chunk references (text→image/table embeddings) |
+| `metadata.connect_to[].target` | `str` | Target chunk_id |
+| `metadata.connect_to[].relation` | `str` | `"embeds"` (inline asset) or `"related"` |
+| `metadata.connect_to[].ref` | `str` | Original reference string: `"[images/image-1.png]"` |
+| `metadata.connect_to[].position` | `{start, end}` | Character offset of the reference in content |
+
+### `doc_nav.json` — Hierarchical Navigation Tree
+
+Used by agentic retrieval for 2-level section browsing. Structure:
+
+```json
+{
+ "version": "1.0",
+ "file_name": "AI_Security_Report.docx",
+ "stats": { "total_chunks": 155, "text_chunks": 141, "image_chunks": 13, "table_chunks": 1, "max_depth": 4 },
+ "sections": [
+ {
+ "title": "1. Overview",
+ "path": "test_kb/AI_Security_Report.docx/1. Overview",
+ "level": 1,
+ "summary": "This section covers: 1.1 Key Findings, 1.2 Recommendations",
+ "chunk_count": 6,
+ "children": [
+ {
+ "title": "1.1 Key Findings",
+ "path": "test_kb/.../1. Overview/1.1 Key Findings",
+ "level": 2,
+ "summary": "This section covers: Supply Side Perspective, Demand Side Perspective...",
+ "chunk_count": 4,
+ "children": [...]
+ }
+ ]
+ }
+ ],
+ "resources": {
+ "images": [{ "path": "images/image-1 ai_model.png", "summary": "image-1 The image displays a tech theme..." }],
+ "tables": [{ "path": "tables/table-1 ai_architecture.html", "summary": "table-1 The table shows..." }]
+ }
+}
+```
+
+### `manifest.json` — Parse Metadata & Heading Hierarchy
+
+```json
+{
+ "version": "2.0",
+ "job_id": "AI_Security_Report.docx",
+ "source_file_name": "AI_Security_Report.docx",
+ "processing_date": "2026-05-09T07:46:00.395048Z",
+ "statistics": { "total_chunks": 155, "text_chunks": 141, "image_chunks": 13, "table_chunks": 1 },
+ "HIERARCHY": {
+ "Root": {},
+ "1. Overview": {
+ "1.1 Key Findings": { "Supply Side Perspective": {}, "Demand Side Perspective": {} },
+ "1.2 Recommendations": {}
+ },
+ "2. History of AI in Cybersecurity": { "...": {} }
+ }
+}
+```
+
+The `HIERARCHY` field is a nested dict representing the full heading tree
+discovered by `layout_parser.pred_titles()`. Each key is a heading title;
+its value is a dict of child headings (empty `{}` for leaf nodes).
+
+### Intermediate DataFrame (`ALL_DF_COLS`)
+
+Parsers internally produce a `pd.DataFrame` with columns:
+`content, path, type, length, keywords, summary, know_id, tokens, connectto, addtime, page_nums`.
+This is converted to `ChunkPayload` objects via `dataframe_chunk_converter.dataframe_to_chunks()`
+before persisting to `chunks.json`. The DataFrame is a transient internal format;
+debug CSVs (`preds_*.csv`) are saved alongside for troubleshooting.
+
+---
+
+## Stage ④: Publication — Database Schema
+
+### Core Tables
+
+#### `documents`
+
+| Column | Type | Description |
+|:---|:---|:---|
+| `document_id` | `String(36)` PK | `doc_{uuid_hex[:12]}` |
+| `user_id` | `Text` FK → `user.id` | Owner |
+| `namespace` | `String(255)` | Isolation scope (default: `"default"`) |
+| `status` | `String(32)` | `active` / `archived` |
+| `current_job_result_id` | `String(36)` FK | Points to active revision |
+| `source_file_name` | `Text` | Original filename |
+
+#### `document_sections`
+
+| Column | Type | Description |
+|:---|:---|:---|
+| `section_id` | `String(36)` PK | `sec_{uuid_hex[:12]}` |
+| `document_id` | FK → `documents` | Parent document |
+| `job_result_id` | FK → `job_results` | Revision |
+| `parent_section_id` | FK → self | Parent section (tree structure) |
+| `section_path` | `Text` UNIQUE(doc+rev+path) | `"file.docx / Chapter 1 / Section 1.1"` |
+| `section_title` | `Text` | Heading text |
+| `section_level` | `Integer` | Depth in hierarchy (1-based) |
+| `summary` | `Text` | Section summary |
+| `sort_order` | `Integer` | Display order |
+
+#### `document_chunks`
+
+| Column | Type | Description |
+|:---|:---|:---|
+| `id` | `String(36)` PK | `dchk_{uuid_hex[:12]}` |
+| `chunk_id` | `String(64)` | Content hash (deterministic dedup key) |
+| `document_id` | FK → `documents` | Parent document |
+| `section_id` | FK → `document_sections` | Parent section |
+| `chunk_type` | `String(64)` | `text` / `image` / `table` |
+| `content` | `Text` | Chunk content (text/HTML) |
+| `content_search_text` | `Text` | Pre-tokenized for BM25 content channel |
+| `path_search_text` | `Text` | Pre-tokenized for BM25 path channel |
+| `term_search_text` | `Text` | Pre-tokenized for term/grep channel |
+| `content_search_tsv` | `TSVECTOR` (computed) | PostgreSQL GIN index for full-text |
+| `path_search_tsv` | `TSVECTOR` (computed) | PostgreSQL GIN index for path |
+| `source_chunk_path` | `Text` | Original parser path |
+| `file_path` | `Text` | Asset reference (`images/x.jpg`) |
+| `chunk_metadata` | `JSON` | Keywords, tokens, connect_to, etc. |
+| `sort_order` | `Integer` | Display order |
+
+#### `graph_nodes`
+
+| Column | Type | Description |
+|:---|:---|:---|
+| `node_id` | `String(128)` PK | `doc:{document_id}` |
+| `node_kind` | `String(32)` | `document` (only doc-level nodes) |
+| `owner_document_id` | FK → `documents` | Source document |
+| `properties` | `JSON` | `{source_file_name, top_keywords, chunks_count, types, top_summary}` |
+
+#### `graph_edges`
+
+| Column | Type | Description |
+|:---|:---|:---|
+| `edge_id` | `String(160)` PK | `related:{doc_a}<->{doc_b}` (sorted pair) |
+| `edge_kind` | `String(32)` | `related` |
+| `source_node_id` / `target_node_id` | FK → `graph_nodes` | Connected docs |
+| `weight` | `Float` | Keyword overlap score (≥ 0.8 threshold) |
+| `properties` | `JSON` | `{shared_keywords, connection_count}` |
+| `is_directed` | `Boolean` | Always `False` for related edges |
+
+### Publication Logic
+
+`RetrievalPublicationService.publish_document_state()`:
+
+1. **Dedup**: Cross-document content-hash dedup via `_dedup_chunks_by_content`.
+2. **Section tree**: Builds `DocumentSection` tree from chunk paths, creating
+ ancestor sections top-down.
+3. **Search text**: Generates 3 search text channels per chunk:
+ - `content_search_text`: Tokenized content + summary
+ - `path_search_text`: Tokenized file name + section path + summary
+ - `term_search_text`: Raw content + path for substring grep
+4. **Graph**: `DocumentGraphService.publish_document_graph()` creates doc-level
+ `GraphNode` with TF-IDF keywords, then keyword-overlap `GraphEdge`s to peer
+ documents (min 3 shared keywords, score ≥ 0.8).
+
+---
+
+## Stage ⑤: Retrieval Engine
+
+### Entry Point
+
+`shared/services/retrieval/app_service.py` → `run_retrieval_query()`
+
+### Two Retrieval Modes
+
+The system supports two modes controlled by `RETRIEVAL_AGENTIC_ENABLED`:
+
+#### Legacy Mode (3-Channel RRF)
+
+```mermaid
+flowchart LR
+ Q[Query] --> P[Path Channel: BM25 on path_search_text]
+ Q --> C[Content Channel: BM25 on content_search_text]
+ Q --> T[Term Channel: substring on term_search_text]
+ P --> RRF["RRF Fusion (k=60)"]
+ C --> RRF
+ T --> RRF
+ RRF --> Graph[Legacy Graph Routing]
+ Graph --> Rank[Dual-priority ranking]
+ Rank --> Assemble[assemble_retrieval_results]
+```
+
+**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.
+
+**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 have their leaf chunks merged directly into the BFS document tree.
+
+**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).
+
+### Tree Rendering & Hydration
+
+Unlike legacy retrieval which relied on static `hydrate_mode` tags, hydration is now determined dynamically by the `DocTreeNode` structure:
+- **Structural Context (Outlines)**: Sections not drilled into are simply rendered as structural outlines (`title` + `summary`) to guide the LLM.
+- **Leaf Content (Hydration)**: Sections that the LLM explicitly selects for drill-down have their raw chunks (`text`, `image`, `table`) fully hydrated into the `leaf_content` of the tree.
+- **Multi-Modal Inline Embedding**: During hydration, connected inline assets (images/tables) are natively resolved and embedded directly into the text chunk content, supporting multi-modal LLM processing without brittle string-replacement placeholders.
+
+**`_rank_candidates_by_path()`** — Dual-priority ranking:
+
+- When agent results exist: agent_score is primary, discovery_score is tiebreaker
+- Rows with agent_score=0 are demoted to fallback pool
+- Sort key: `(agent_score, discovery_score, dual_hit_flag, importance_norm_score)`
+
+### Result Assembly
+
+`assemble_retrieval_results()`:
+
+1. Filters by `exclude_document_ids` and `exclude_sections`
+2. Filters by `allowed_chunk_types` (data_type parameter)
+3. Hydrates `connect_to` targets (related table chunks inlined into text)
+4. Cleans asset path references from content
+5. Attaches citation: `{document_id, chunk_id, source_file_name, section_path}`
+
+### Small KB Optimization
+
+When `total_chunks <= top_k`, skips the full pipeline and returns all chunks
+directly (router: `small_kb_all`).
+
+### Caching
+
+Results are cached per `(user_id, namespace, query, top_k, filters)` via
+`cache_service`. Cache version is checked before execution; cache is written
+after successful retrieval.
+
+---
+
+## Analytics Tables
+
+#### `retrieval_hit_stats`
+
+Tracks per-chunk and per-document retrieval usage. `hit_count` and `last_hit_at`
+feed into `compute_importance_score()` for ranking boost.
+
+#### `retrieval_runs` / `retrieval_steps`
+
+Append-only agentic retrieval analytics. One `retrieval_runs` row per query,
+with child `retrieval_steps` rows recording each agent action, its input/output,
+latency, and token usage.
+
+---
+
+## Key Implementation Patterns
+
+### Deterministic Chunk IDs
+
+`know_id = gen_str_codes(pure_text)` — SHA-based hash of text content only
+(excludes image/table asset refs). This enables cross-document dedup:
+identical text in different uploads produces the same `chunk_id`.
+
+### Plan-then-Act DOM Mutation
+
+When splitting tables or modifying document structure:
+1. **Pass 1 (Investigate)**: Collect mutation targets into a static plan
+2. **Pass 2 (Execute)**: Apply mutations in **reverse order** to avoid index shifting
+
+### Image Dedup: Perceptual Hash
+
+`perceptual_hash()` computes a visual fingerprint. Images with identical
+hashes are deduplicated within a document, with cached metadata reused.
+
+### Asset Lifecycle (Deterministic UIDs)
+
+`IMAGE_[hash(content+seq)]_IMAGE` — identical images at different positions
+receive unique IDs. Context chaining prevention scans backward past binary
+identifiers to find the nearest valid text.
+
+### LLM Constraints
+
+- **DeepSeek JSON mode**: Requires the word "json" in the prompt when
+ `response_format` is `json_object`
+- **Streaming robustness**: Concatenate `delta.content` only if `not None`
+- **Token pool rotation**: Ali API keys support per-token RPM limits,
+ cooldown, and inline retry with next available token
+
+---
+
+## Development & Debugging
+
+### Local Setup
+
+```bash
+uv sync --all-packages
+cp apps/api/.env.example apps/api/.env
+cp apps/worker/.env.example apps/worker/.env
+./deploy/local-dev/start-dev.sh # PostgreSQL, Redis, LocalStack
+cd apps/api && uv run main.py # API on :5005
+cd apps/worker && uv run worker.py # Celery worker
+```
+
+### Debug Scripts (Worker)
+
+| Script | Purpose |
+|:---|:---|
+| `debug_parse.py` | End-to-end parsing with `MockRedis`, `LOCAL_DEBUG=1` |
+| `debug_hierarchy_llm.py` | Test heading recognition LLM calls |
+| `debug_agentic_e2e.py` | End-to-end agentic retrieval test |
+| `debug_profiler.py` | Document profiler testing |
+| `debug_toc_detection.py` | TOC detection and hierarchy building |
+
+### Quality Checks
+
+```bash
+make lint # Ruff lint
+make lint-fix # Auto-fix safe issues
+make typecheck # Pyright across api, worker, shared
+make check # Both lint + typecheck
+```
+
+### Local Endpoints
+
+| Service | URL |
+|:---|:---|
+| API | `http://localhost:5005` |
+| OpenAPI docs | `http://localhost:5005/docs` |
+| PostgreSQL | `localhost:5432` |
+| Redis | `localhost:6379` |
+| LocalStack (S3) | `http://localhost:4566` |
diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py
index 35c29aa9a..2121cdb4e 100644
--- a/apps/api/app/mcp/retrieval_server.py
+++ b/apps/api/app/mcp/retrieval_server.py
@@ -66,10 +66,15 @@ def to_mcp_query_response(response: dict[str, Any]) -> dict[str, Any]:
result["asset_url"] = row["asset_url"]
results.append(result)
- return {
+ mcp_response: dict[str, Any] = {
"query": response.get("query"),
"results": results,
}
+ # Forward evidence_text for agentic mode
+ if response.get("evidence_text") is not None:
+ mcp_response["evidence_text"] = response["evidence_text"]
+
+ return mcp_response
async def resolve_mcp_user_id(*, ctx: Context | None, db: AsyncSession) -> str:
diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py
index b62fd941d..75986d40e 100644
--- a/apps/api/tests/contract/test_billing_contract.py
+++ b/apps/api/tests/contract/test_billing_contract.py
@@ -1,3 +1,4 @@
+import asyncio
import importlib
import json
from collections.abc import Callable
@@ -123,6 +124,58 @@ async def test_should_initialize_missing_user_balance_during_tier_lookup(
}
+@pytest.mark.asyncio
+async def test_should_initialize_missing_user_balance_once_for_concurrent_requests(
+ api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]],
+) -> None:
+ user_id = f"contract-concurrent-balance-{uuid4().hex[:12]}"
+ api_key = f"sk_contract_{uuid4().hex[:24]}"
+
+ async with api_client_factory() as api_client:
+ await ContractDatabase.insert_user(user_id=user_id)
+ await _insert_api_key_for_user(user_id, api_key)
+ api_client.headers.update({"Authorization": f"Bearer {api_key}"})
+
+ responses = await asyncio.gather(
+ *[api_client.get("/api/v1/billing/credits") for _ in range(8)]
+ )
+
+ balance_row = await ContractDatabase.fetch_one(
+ """
+ SELECT credits_balance, user_tier
+ FROM user_balances
+ WHERE user_id = :user_id
+ """,
+ {"user_id": user_id},
+ )
+ transaction_count_row = await ContractDatabase.fetch_one(
+ """
+ SELECT COUNT(*) AS count
+ FROM credits_transactions
+ WHERE user_id = :user_id
+ AND transaction_type = 'initial_grant'
+ """,
+ {"user_id": user_id},
+ )
+ payment_count_row = await ContractDatabase.fetch_one(
+ """
+ SELECT COUNT(*) AS count
+ FROM payment_records
+ WHERE user_id = :user_id
+ AND payment_type = 'system_grant'
+ """,
+ {"user_id": user_id},
+ )
+
+ assert all(response.status_code == 200 for response in responses)
+ assert all(
+ response.json() == {"credits_balance": 5.0} for response in responses
+ )
+ assert balance_row == {"credits_balance": 5_000_000, "user_tier": "free"}
+ assert transaction_count_row == {"count": 1}
+ assert payment_count_row == {"count": 1}
+
+
@pytest.mark.asyncio
async def test_should_not_register_billing_routes_when_billing_is_disabled(
monkeypatch: MonkeyPatch,
diff --git a/apps/docs/test_github_flow.md b/apps/docs/test_github_flow.md
new file mode 100644
index 000000000..84fd82cd5
--- /dev/null
+++ b/apps/docs/test_github_flow.md
@@ -0,0 +1,4 @@
+# GitHub Flow Test
+
+This is a test file to verify the GitHub workflow simulation (Issue -> Branch -> PR -> Merge).
+It is safe to ignore or delete this file later.
diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py
index 3379b8613..a5c0aee0a 100644
--- a/apps/worker/tests/contract/test_parse_task_contract.py
+++ b/apps/worker/tests/contract/test_parse_task_contract.py
@@ -3,7 +3,9 @@
import json
import shutil
import zipfile
+from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
+from threading import Barrier
from types import SimpleNamespace
from typing import Any
from uuid import uuid4
@@ -566,6 +568,256 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any:
]
+def test_should_initialize_billing_once_for_concurrent_parse_tasks(
+ worker_contract_environment: None,
+ monkeypatch: MonkeyPatch,
+ tmp_path: Path,
+) -> None:
+ (
+ kb_tasks,
+ parse_service,
+ sync_storage_service,
+ engine,
+ sync_job_info_service_cls,
+ sync_job_metadata_service_cls,
+ sync_redis_service_factory,
+ ) = _load_parse_task_modules()
+
+ user_id: str = f"worker-concurrent-user-{uuid4().hex[:12]}"
+ job_ids: list[str] = [f"job_cb_{index}_{uuid4().hex[:12]}" for index in range(2)]
+ source_file_name: str = "contract-concurrent.pdf"
+ s3_keys: dict[str, str] = {
+ job_id: f"uploads/{job_id}.pdf" for job_id in job_ids
+ }
+ job_metadata_by_id: dict[str, dict[str, Any]] = {
+ job_id: _build_pending_file_job_metadata(source_file_name)
+ for job_id in job_ids
+ }
+
+ with engine.begin() as connection:
+ insert_contract_user(connection, user_id=user_id)
+ for job_id in job_ids:
+ insert_contract_job(
+ connection,
+ job_id=job_id,
+ user_id=user_id,
+ status="pending",
+ source_type="file",
+ s3_key=s3_keys[job_id],
+ webhook_enabled=False,
+ job_metadata=job_metadata_by_id[job_id],
+ billing_status="pending",
+ )
+
+ redis_service = sync_redis_service_factory.get_service()
+ for job_id in job_ids:
+ _save_worker_task_cache(
+ job_id=job_id,
+ user_id=user_id,
+ s3_key=s3_keys[job_id],
+ metadata=job_metadata_by_id[job_id],
+ sync_job_info_service_cls=sync_job_info_service_cls,
+ sync_job_metadata_service_cls=sync_job_metadata_service_cls,
+ sync_redis_service_factory=sync_redis_service_factory,
+ )
+
+ _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks)
+ monkeypatch.setattr(kb_tasks.settings, "TMP_PATH", str(tmp_path))
+ monkeypatch.setattr(kb_tasks.settings, "BILLING_ENABLED", True)
+
+ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]:
+ return {
+ "exists": storage_key in s3_keys.values(),
+ "size": _SAMPLE_PDF_PATH.stat().st_size,
+ }
+
+ def fake_generate_download_url(storage_key: str, bucket: str | None) -> dict[str, str]:
+ return {"download_url": f"https://example.test/{storage_key}"}
+
+ def fake_download_s3_file_to_temp(
+ file_url: str, file_ext: str, temp_dir: str
+ ) -> str:
+ assert file_ext == ".pdf"
+ downloaded_path = Path(temp_dir) / f"downloaded{file_ext}"
+ shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path)
+ return str(downloaded_path)
+
+ billing_start_barrier = Barrier(len(job_ids))
+
+ def fake_estimate_page_count(file_path: str) -> int:
+ billing_start_barrier.wait(timeout=10)
+ return 1
+
+ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]:
+ output_dir = (
+ Path(str(kwargs["output_dir"]))
+ / str(kwargs["kb_dir"])
+ / str(kwargs["internal_output_filename"])
+ )
+ output_dir.mkdir(parents=True, exist_ok=True)
+ (output_dir / "full.md").write_text("body", encoding="utf-8")
+
+ file_root = str(kwargs["internal_output_filename"])
+ parsed_rows: list[dict[str, Any]] = [
+ {
+ "content": "chunk body",
+ "path": f"Default_Root/{file_root}/Section/Point",
+ "type": "text",
+ "length": 10,
+ "keywords": "",
+ "summary": "",
+ "know_id": f"{kwargs['job_id']}-chunk-1",
+ "tokens": "",
+ "connectto": "",
+ "addtime": "now",
+ "page_nums": "1",
+ }
+ ]
+ return str(output_dir), pd.DataFrame(parsed_rows)
+
+ class FakeResultStorage:
+ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any:
+ return SimpleNamespace(
+ zip_key=f"results/{job_id}.zip",
+ raw_prefix=f"results/{job_id}/",
+ raw_files={},
+ )
+
+ monkeypatch.setattr(kb_tasks, "verify_s3_file_exists", fake_verify_s3_file_exists)
+ monkeypatch.setattr(
+ sync_storage_service,
+ "verify_s3_file_exists",
+ fake_verify_s3_file_exists,
+ )
+ monkeypatch.setattr(kb_tasks, "generate_download_url", fake_generate_download_url)
+ monkeypatch.setattr(
+ sync_storage_service,
+ "generate_download_url",
+ fake_generate_download_url,
+ )
+ monkeypatch.setattr(kb_tasks, "download_s3_file_to_temp", fake_download_s3_file_to_temp)
+ monkeypatch.setattr(kb_tasks.PageEstimator, "estimate", fake_estimate_page_count)
+ monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse)
+ monkeypatch.setattr(kb_tasks, "get_result_storage", lambda: FakeResultStorage())
+
+ def run_parse_task(job_id: str) -> dict[str, Any]:
+ return dict(kb_tasks.parse_task.run(job_id, user_id, "kb_management"))
+
+ with ThreadPoolExecutor(max_workers=len(job_ids)) as executor:
+ results = list(executor.map(run_parse_task, job_ids))
+
+ expected_credits_charged = int(kb_tasks.settings.MICRO_DOLLARS_PER_PAGE)
+ expected_initial_balance = (
+ int(kb_tasks.settings.FREE_PLAN_INITIAL_CREDITS) * 1_000_000
+ )
+
+ with engine.begin() as connection:
+ job_rows = list(
+ connection.execute(
+ text(
+ """
+ SELECT job_id, status, billing_status, page_count, credits_charged
+ FROM jobs
+ WHERE job_id = ANY(:job_ids)
+ ORDER BY job_id
+ """
+ ),
+ {"job_ids": job_ids},
+ )
+ .mappings()
+ .all()
+ )
+ balance_row = (
+ connection.execute(
+ text(
+ """
+ SELECT credits_balance
+ FROM user_balances
+ WHERE user_id = :user_id
+ """
+ ),
+ {"user_id": user_id},
+ )
+ .mappings()
+ .one()
+ )
+ transaction_rows = list(
+ connection.execute(
+ text(
+ """
+ SELECT transaction_type, COUNT(*) AS count
+ FROM credits_transactions
+ WHERE user_id = :user_id
+ GROUP BY transaction_type
+ ORDER BY transaction_type
+ """
+ ),
+ {"user_id": user_id},
+ )
+ .mappings()
+ .all()
+ )
+ payment_count_row = (
+ connection.execute(
+ text(
+ """
+ SELECT COUNT(*) AS count
+ FROM payment_records
+ WHERE user_id = :user_id
+ AND payment_type = 'system_grant'
+ """
+ ),
+ {"user_id": user_id},
+ )
+ .mappings()
+ .one()
+ )
+
+ metadata_by_job_id = {
+ job_id: sync_job_metadata_service_cls(redis_service).get_metadata(job_id)
+ for job_id in job_ids
+ }
+ result_job_ids = {str(result["job_id"]) for result in results}
+ transaction_counts = {
+ str(row["transaction_type"]): int(row["count"]) for row in transaction_rows
+ }
+
+ assert result_job_ids == set(job_ids)
+ assert all(result["status"] == "success" for result in results)
+ assert [
+ {
+ "job_id": row["job_id"],
+ "status": row["status"],
+ "billing_status": row["billing_status"],
+ "page_count": row["page_count"],
+ "credits_charged": row["credits_charged"],
+ }
+ for row in job_rows
+ ] == [
+ {
+ "job_id": job_id,
+ "status": "done",
+ "billing_status": "charged",
+ "page_count": 1,
+ "credits_charged": expected_credits_charged,
+ }
+ for job_id in sorted(job_ids)
+ ]
+ assert all(
+ metadata_by_job_id[job_id] is not None
+ and metadata_by_job_id[job_id]["billing_status"] == "charged"
+ and metadata_by_job_id[job_id]["billing_amount_micro_dollars"]
+ == expected_credits_charged
+ for job_id in job_ids
+ )
+ assert balance_row == {
+ "credits_balance": expected_initial_balance
+ - (expected_credits_charged * len(job_ids))
+ }
+ assert transaction_counts == {"initial_grant": 1, "usage": len(job_ids)}
+ assert payment_count_row == {"count": 1}
+
+
def test_should_skip_parse_task_when_the_job_is_already_terminal(
worker_contract_environment: None,
monkeypatch: MonkeyPatch,
diff --git a/packages/shared-python/shared/services/billing/credits_service.py b/packages/shared-python/shared/services/billing/credits_service.py
index 566440157..e4c8e03c8 100644
--- a/packages/shared-python/shared/services/billing/credits_service.py
+++ b/packages/shared-python/shared/services/billing/credits_service.py
@@ -11,7 +11,7 @@
from datetime import timedelta
from typing import Any, Dict, Optional
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession
from shared.core.billing import MicroDollar
@@ -93,44 +93,39 @@ async def ensure_user_initialized(
initial_dollars = getattr(settings, "FREE_PLAN_INITIAL_CREDITS", 5)
initial_amount = MicroDollar.from_dollars(initial_dollars).amount
- try:
- async with session.begin_nested():
- # Create balance record
- balance_entry = UserBalance(
- user_id=user_id, credits_balance=initial_amount
- )
- session.add(balance_entry)
-
- # Create initial transaction record
- transaction = CreditsTransaction(
- user_id=user_id,
- credits_amount=initial_amount,
- description="New user registration bonus",
- transaction_type="initial_grant",
- )
- session.add(transaction)
-
- # Create payment record for tracking
- payment = PaymentRecord(
- user_id=user_id,
- payment_type="system_grant",
- amount_cents=0,
- currency="USD",
- status="succeeded",
- credits_amount=initial_amount,
- extra_metadata={"reason": "initial_grant"},
- processed_at=utc_now_naive(),
- )
- session.add(payment)
- await session.flush()
- except IntegrityError:
- existing_balance = await self.repository.get_user_balance(session, user_id)
- if existing_balance:
- logger.info(
- f"User already initialized by concurrent session: user_id={user_id}"
- )
- return
- raise
+ insert_balance_stmt = (
+ pg_insert(UserBalance)
+ .values(user_id=user_id, credits_balance=initial_amount)
+ .on_conflict_do_nothing(index_elements=[UserBalance.user_id])
+ .returning(UserBalance.user_id)
+ )
+ insert_result = await session.execute(insert_balance_stmt)
+ inserted_user_id: str | None = insert_result.scalar_one_or_none()
+ if inserted_user_id is None:
+ return
+
+ # Create initial transaction record only for the request that won creation.
+ transaction: CreditsTransaction = CreditsTransaction(
+ user_id=user_id,
+ credits_amount=initial_amount,
+ description="New user registration bonus",
+ transaction_type="initial_grant",
+ )
+ session.add(transaction)
+
+ # Create payment record for tracking only for the request that won creation.
+ payment: PaymentRecord = PaymentRecord(
+ user_id=user_id,
+ payment_type="system_grant",
+ amount_cents=0,
+ currency="USD",
+ status="succeeded",
+ credits_amount=initial_amount,
+ extra_metadata={"reason": "initial_grant"},
+ processed_at=utc_now_naive(),
+ )
+ session.add(payment)
+ await session.flush()
logger.info(f"User initialized: user_id={user_id}, credits={initial_amount}")
diff --git a/packages/shared-python/shared/services/billing/credits_sync_service.py b/packages/shared-python/shared/services/billing/credits_sync_service.py
index 06b9bc4ed..ba6ddbe38 100644
--- a/packages/shared-python/shared/services/billing/credits_sync_service.py
+++ b/packages/shared-python/shared/services/billing/credits_sync_service.py
@@ -6,7 +6,7 @@
from typing import Any, Dict, Optional
from sqlalchemy import func, select
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.orm import Session
from shared.core.billing import MicroDollar
@@ -41,41 +41,37 @@ def ensure_user_initialized(
initial_dollars = getattr(settings, "FREE_PLAN_INITIAL_CREDITS", 5)
initial_amount = MicroDollar.from_dollars(initial_dollars).amount
- try:
- with session.begin_nested():
- balance_entry = UserBalance(
- user_id=user_id, credits_balance=initial_amount
- )
- session.add(balance_entry)
-
- transaction = CreditsTransaction(
- user_id=user_id,
- credits_amount=initial_amount,
- description="New user registration bonus",
- transaction_type="initial_grant",
- )
- session.add(transaction)
-
- payment = PaymentRecord(
- user_id=user_id,
- payment_type="system_grant",
- amount_cents=0,
- currency="USD",
- status="succeeded",
- credits_amount=initial_amount,
- extra_metadata={"reason": "initial_grant"},
- processed_at=utc_now_naive(),
- )
- session.add(payment)
- session.flush()
- except IntegrityError:
- existing_balance = self.repository.get_user_balance(session, user_id)
- if existing_balance:
- logger.info(
- f"User already initialized by concurrent session: user_id={user_id}"
- )
- return
- raise
+ insert_balance_stmt = (
+ pg_insert(UserBalance)
+ .values(user_id=user_id, credits_balance=initial_amount)
+ .on_conflict_do_nothing(index_elements=[UserBalance.user_id])
+ .returning(UserBalance.user_id)
+ )
+ insert_result = session.execute(insert_balance_stmt)
+ inserted_user_id: str | None = insert_result.scalar_one_or_none()
+ if inserted_user_id is None:
+ return
+
+ transaction: CreditsTransaction = CreditsTransaction(
+ user_id=user_id,
+ credits_amount=initial_amount,
+ description="New user registration bonus",
+ transaction_type="initial_grant",
+ )
+ session.add(transaction)
+
+ payment: PaymentRecord = PaymentRecord(
+ user_id=user_id,
+ payment_type="system_grant",
+ amount_cents=0,
+ currency="USD",
+ status="succeeded",
+ credits_amount=initial_amount,
+ extra_metadata={"reason": "initial_grant"},
+ processed_at=utc_now_naive(),
+ )
+ session.add(payment)
+ session.flush()
logger.info(f"User initialized: user_id={user_id}, credits={initial_amount}")
diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py
index ea127b69d..6bfaef2e8 100644
--- a/packages/shared-python/shared/services/retrieval/agent_navigate.py
+++ b/packages/shared-python/shared/services/retrieval/agent_navigate.py
@@ -3,7 +3,10 @@
import json
import re
-from typing import Any, Sequence
+from typing import Any, Sequence, TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from shared.services.retrieval.agentic.types import DocTreeNode
from loguru import logger
from sqlalchemy import func, select, or_
@@ -33,42 +36,60 @@
Do not include any explanation.
"""
-_VALID_HYDRATE_MODES = frozenset({
- 'outline', 'chunks', 'assets_only', 'image_only', 'table_only',
-})
_SCOPE_NAV_PROMPT = """\
You are a document navigation assistant.
Document: "{doc_name}" (id: {doc_id})
-Current scope: {scope_label}
+{scope_header}
-Below are candidate section paths at this scope level (up to 2 depth levels).
-Indented items are sub-items of the item above.
-Each item shows text/image/table counts.
-Select section paths directly. A selected section path represents the chunks
-under that section subtree; do not ask to drill deeper.
+Below is the document's section tree.
+Sections tagged [SELECT] are within the current scope and may be selected.
+Other sections are shown as structural context only (not selectable).
+Nodes marked [Leaf] have no further sub-sections.
-=== Items ===
+=== Section Tree ===
{items_overview}
-=== End Items ===
+=== End Section Tree ===
User query: {query}
-Select the most relevant section paths (at most {max_select}).
-If NO section path is relevant to the query, you MUST return an empty array []. Do not force-select irrelevant sections.
-Prefer specific sub-items over broad parents when both are listed and the sub-item is sufficient.
+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 [].
+
+ Optional: set "mode" per selection to control what content is retrieved:
+ - "all" (default): retrieve all content types (text, images, tables)
+ - "image": retrieve only image assets from this section
+ - "table": retrieve only table assets from this section
+
+Return ONLY a JSON object:
+{{"selections": [{{"path": "...", "confidence": , "mode": "all"}}, ...]}}
+Do not include any explanation.
+"""
+
+
+_DISCOVERY_SELECT_PROMPT = """\
+You are a document navigation assistant.
+
+Document: "{doc_name}"
+
+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}
-For each selected path, assign a confidence score (0.0 to 1.0) where 1.0 means exactly answers the query and 0.5 means tangentially related.
-Also choose a hydrate_mode:
-- "chunks" (default) return all text/image/table chunks
-- "outline" return only section title + summary, no chunk content
-- "assets_only" return only image and table chunks
-- "image_only" return only image chunks
-- "table_only" return only table chunks
+Select section paths whose content is needed to answer the query.
+If none are relevant, return an EMPTY list [].
-Return ONLY a JSON array:
-[{{"path": "section/path", "confidence": , "hydrate_mode": "chunks"}}, ...]
+Return ONLY a JSON object:
+{{"selections": [{{"path": "...", "confidence": , "mode": "all"}}, ...]}}
Do not include any explanation.
"""
@@ -113,39 +134,6 @@ def _normalize_confidence(value: Any) -> float | None:
return max(0.0, min(parsed, 1.0))
-def _default_confidence_for_rank(rank: int) -> float:
- return round(max(0.25, 0.85 - rank * 0.15), 4)
-
-
-def _parse_chunk_path_selections(text: str) -> list[dict[str, Any]]:
- """Parse chunk path selections from LLM output.
-
- Accepts either a legacy JSON array of strings or a structured array of
- objects with `path`, optional `confidence`, and optional `hydrate_mode`.
- """
- payload = _extract_json_array_payload(text)
- selections: list[dict[str, Any]] = []
- for item in payload:
- if isinstance(item, str):
- path = item.strip()
- if path:
- selections.append({'path': path, 'confidence': None, 'hydrate_mode': 'chunks'})
- continue
- if not isinstance(item, dict):
- continue
- path = str(item.get('path') or item.get('chunk_path') or '').strip()
- if not path:
- continue
- raw_mode = str(item.get('hydrate_mode') or '').strip().lower()
- hydrate_mode = raw_mode if raw_mode in _VALID_HYDRATE_MODES else 'chunks'
- selections.append({
- 'path': path,
- 'confidence': _normalize_confidence(item.get('confidence')),
- 'hydrate_mode': hydrate_mode,
- })
- return selections
-
-
async def _build_knowledge_map_overview(
db: AsyncSession,
*,
@@ -233,13 +221,13 @@ def _format_items_for_llm(
items: list[dict],
max_chars: int = 20000,
) -> tuple[str, bool]:
- """Unified formatting with overflow guard for scope navigation.
+ """Format items with ▸ └ [Leaf] hierarchy for scope navigation.
- Always shows ALL items (L1 + L2). Overflow controls whether
- summaries are included — not which levels are shown.
-
- Normal: path + title + text=N image=I table=T + summary
- Overflow: path + title + text=N image=I table=T (no summary)
+ Supports arbitrary depth levels via absolute ``level`` field.
+ Items with ``show_summary=False`` render title only (structural context).
+ ``[LN]`` tags indicate the absolute document depth of each section.
+ ``[Leaf]`` tags indicate bottom-level sections with no further children.
+ Summaries are included when within budget, dropped on overflow.
Returns (text, overflowed).
"""
@@ -250,34 +238,53 @@ def _format_items_for_llm(
SUMMARY_HEAD_TOKENS = 80
- def _render_line(item: dict, include_summary: bool) -> str:
+ def _render_item(item: dict, include_summary: bool) -> str:
level = item.get('level', 1)
- indent = ' ' if level == 2 else ''
- line = f'{indent}- path="{item["path"]}" title="{item["title"]}"'
- chunk_count = item.get('chunk_count', 0)
- if chunk_count > 0:
- line += f' text={chunk_count}'
- image_count = item.get('image_count', 0)
- if image_count > 0:
- line += f' image={image_count}'
- table_count = item.get('table_count', 0)
- if table_count > 0:
- line += f' table={table_count}'
- if include_summary:
- summary = item.get('summary') or item.get('title', '')
- if summary:
- clipped = truncate_content_preview(summary, head=SUMMARY_HEAD_TOKENS, tail=0)
- line += f'\n{indent} summary: {clipped}'
- return line
-
- # Try full render (with summaries)
- full_lines = [_render_line(item, include_summary=True) for item in items]
+ show = item.get('show_summary', True)
+ is_leaf = item.get('is_leaf', False)
+ leaf_tag = ' [Leaf]' if is_leaf else ''
+ title = item.get('title', '')
+ path = item.get('path', '')
+ summary = item.get('summary') or ''
+
+ # Build chunk count tags (only for current-scope items)
+ counts_str = ''
+ if show:
+ count_parts: list[str] = []
+ chunk_count = item.get('chunk_count', 0)
+ if chunk_count > 0:
+ count_parts.append(f'text={chunk_count}')
+ image_count = item.get('image_count', 0)
+ if image_count > 0:
+ count_parts.append(f'image={image_count}')
+ table_count = item.get('table_count', 0)
+ if table_count > 0:
+ count_parts.append(f'table={table_count}')
+ counts_str = f' [{" ".join(count_parts)}]' if count_parts else ''
+
+ indent = " " * (level - 1)
+ prefix = '▸' if level == 1 else '└'
+ level_tag = f'[L{level}]'
+ select_tag = '[SELECT] ' if show else ''
+
+ lines: list[str] = []
+ lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}" {title}{counts_str}{leaf_tag}')
+
+ if include_summary and show and summary:
+ sub_indent = " " * level
+ clipped = truncate_content_preview(summary, head=SUMMARY_HEAD_TOKENS, tail=0)
+ lines.append(f'{sub_indent}{clipped}')
+
+ return '\n'.join(lines)
+
+ # Try full render (with summaries for show_summary=True items)
+ full_lines = [_render_item(item, include_summary=True) for item in items]
full_text = '\n'.join(full_lines)
if len(full_text) <= max_chars:
return full_text, False
# Overflow: render without summaries
- slim_lines = [_render_line(item, include_summary=False) for item in items]
+ slim_lines = [_render_item(item, include_summary=False) for item in items]
slim_text = '\n'.join(slim_lines)
return slim_text[:max_chars], True
@@ -433,18 +440,23 @@ async def _load_child_sections(
scope_path: str | None = None,
exclude_paths: set[str] | None = None,
) -> list[dict]:
- """Load the next 2 available section depth bands under *scope_path*.
+ """Load the Continuous Context Tree for *scope_path*.
+
+ Returns a flat list sorted by document order, each item:
+ {path, title, summary, chunk_count, image_count, table_count,
+ level, show_summary, is_leaf}
- Returns a flat list sorted by sort_order, each item:
- {path, title, summary, chunk_count, image_count, table_count, level}
+ 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)
- - level=1: nearest available descendant depth under scope
- - level=2: second nearest available descendant depth under scope
- - chunk_count: text chunks under this section (excluding image/table)
- - image_count: image chunks under this section
- - table_count: table chunks under this section
+ When scope_path is None (root), all items are category 2.
+
+ - 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
+ any path matching (exact or subtree) is skipped from category 2
"""
# ── Fetch all sections for this document revision ────────────────────
stmt = (
@@ -468,7 +480,7 @@ async def _load_child_sections(
scope_depth = len(scope_parts)
# Build full section metadata index
- all_sections: dict[str, dict] = {} # path → {title, summary, sort_order, section_id}
+ all_sections: dict[str, dict] = {} # path → {title, summary, sort_order, section_id, parts, depth}
for section_id, title, path, summary, sort_order in section_rows:
if not path:
continue
@@ -483,63 +495,138 @@ async def _load_child_sections(
'depth': len(parts),
}
- # ── Identify the next two real depth bands ───────────────────────────
- #
- # The stored hierarchy is authoritative. Some ingested documents may only
- # expose deeper section rows at root; in that case those rows become this
- # round's relative L1/L2 instead of synthesizing missing ancestors.
- visible_sections: list[tuple[str, dict, int]] = []
- visible_depths: set[int] = set()
+ # ── Build the set of ancestor prefixes for pruning ────────────────────
+ # e.g. scope = "A / B / K" → ancestor_prefixes = {"A", "A / B", "A / B / K"}
+ ancestor_prefixes: set[str] = set()
+ for i in range(1, scope_depth + 1):
+ ancestor_prefixes.add(' / '.join(scope_parts[:i]))
+
+ # ── Classify each section ────────────────────────────────────────────
_excl = exclude_paths or set()
+ items_by_path: dict[str, dict] = {}
+ scope_child_depths: set[int] = set()
+
for path, meta in all_sections.items():
parts = meta['parts']
- if scope_parts and (
- parts[:scope_depth] != scope_parts or len(parts) <= scope_depth
- ):
- continue
- relative_depth = len(parts) - scope_depth
- if relative_depth < 1:
- continue
- # Skip paths already seen in prior revision rounds
- if _excl and any(
- path == ep or path.startswith(ep + ' / ') or ep.startswith(path + ' / ')
- for ep in _excl
- ):
+ depth = meta['depth']
+
+ if scope_depth == 0:
+ # Root scope: everything is a potential child
+ if depth < 1:
+ 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
+ }
continue
- visible_sections.append((path, meta, relative_depth))
- visible_depths.add(relative_depth)
-
- selected_depths = sorted(visible_depths)[:2]
- depth_to_level = {
- depth: idx + 1
- for idx, depth in enumerate(selected_depths)
- }
-
- items_by_path: dict[str, dict] = {}
- for path, meta, relative_depth in visible_sections:
- level = depth_to_level.get(relative_depth)
- if level is None:
+ # --- Non-root scope ---
+
+ # 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?
+ 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.
+ 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,
+ }
continue
- if path not in items_by_path:
+
+ # Category 2: Descendants of scope_path (children to explore)
+ if parts[:scope_depth] == scope_parts and depth > scope_depth:
+ # 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': level,
+ '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
+
+ # Category 3: Everything else → pruned (not added)
+
+ 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)
+ for path in to_remove:
+ del items_by_path[path]
if not items_by_path:
return []
# ── Count chunks per section (text / image / table) ──────────────────
- section_ids = [meta['section_id'] for meta in all_sections.values()]
- if section_ids:
+ # Only count for show_summary=True items (current scope children)
+ scope_item_sids = {item['section_id'] for item in items_by_path.values() if item['show_summary']}
+ # Also need all section_ids for upward aggregation
+ all_section_ids = [meta['section_id'] for meta in all_sections.values()]
+ if all_section_ids and scope_item_sids:
from sqlalchemy import case, literal_column
chunk_stmt = (
select(
@@ -562,7 +649,7 @@ async def _load_child_sections(
)
.where(DocumentChunk.document_id == document_id)
.where(DocumentChunk.job_result_id == job_result_id)
- .where(DocumentChunk.section_id.in_(section_ids))
+ .where(DocumentChunk.section_id.in_(all_section_ids))
.group_by(DocumentChunk.section_id)
)
chunk_rows = (await db.execute(chunk_stmt)).all()
@@ -575,45 +662,242 @@ async def _load_child_sections(
# Build section_id → path mapping for aggregation
sid_to_path = {meta['section_id']: path for path, meta in all_sections.items()}
- # Aggregate chunk counts upward: each item gets counts from itself + descendants
+ # Aggregate chunk counts upward: each show_summary item gets counts from itself + descendants
for sid, (text_c, img_c, tbl_c) in section_id_counts.items():
chunk_path = sid_to_path.get(sid, '')
if not chunk_path:
continue
- # Add to every ancestor item that is in our items_by_path
for item_path, item in items_by_path.items():
+ if not item['show_summary']:
+ continue
if chunk_path == item_path or chunk_path.startswith(item_path + ' / '):
item['chunk_count'] += text_c
item['image_count'] += img_c
item['table_count'] += tbl_c
- # ── Sort: interleave L2 under their L1 parent ────────────────────────
- #
- # Previous sort `(level, sort_order, path)` grouped all L1 first, then
- # all L2. The LLM prompt uses indentation to show L2 as sub-items, so
- # they should appear directly after their L1 parent for readability.
- #
- # Sort key: (parent_sort_order, is_child, own_sort_order)
- # L1 items: (own_sort_order, 0, 0) → primary position
- # L2 items: (parent_sort_order, 1, own) → right after their parent
-
- def _interleave_key(item: dict) -> tuple:
- if item['level'] == 2:
- parts = split_section_path(item['path'])
- if len(parts) >= 2:
- parent_path = ' / '.join(parts[:-1])
- parent_item = items_by_path.get(parent_path)
- if parent_item is not None:
- return (parent_item['sort_order'], 1, item['sort_order'])
- # Orphan L2 (no matching L1 parent in this view): sort by own order
- return (item['sort_order'], 1, item['sort_order'])
- return (item['sort_order'], 0, 0)
-
- sorted_items = sorted(items_by_path.values(), key=_interleave_key)
+ # ── Sort by native document order ─────────────────────────────────────
+ sorted_items = sorted(items_by_path.values(), key=lambda x: x['sort_order'])
# Clean up internal fields
for item in sorted_items:
item.pop('sort_order', None)
item.pop('section_id', None)
+
+ # ── Detect leaf status ────────────────────────────────────────────────
+ # A section is a leaf if no other section in the database for this
+ # document has a path that descends from it.
+ all_section_paths = set(all_sections.keys())
+ for item in sorted_items:
+ item_path = item['path']
+ has_descendants = any(
+ p != item_path and p.startswith(item_path + ' / ')
+ for p in all_section_paths
+ )
+ item['is_leaf'] = not has_descendants
+
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, "mode": str}.
+ """
+ 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]] = []
+ _VALID_MODES = {'all', 'image', 'table'}
+ 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
+ mode = str(item.get('mode') or 'all').strip().lower()
+ if mode not in _VALID_MODES:
+ mode = 'all'
+ selections.append({'path': path, 'confidence': confidence, 'mode': mode})
+
+ return selections
+
+
+# ---------------------------------------------------------------------------
+# Unified document tree rendering (DocTreeNode → single coherent hierarchy)
+# ---------------------------------------------------------------------------
+
+def _render_leaf_chunks(
+ parts: list[str],
+ chunks: list[dict[str, Any]],
+ indent: str,
+ asset_lookup: dict[str, str] | None = None,
+) -> None:
+ """Render hydrated leaf chunks inline with table/image inlining and dedup.
+
+ Uses ``connect_to`` metadata to resolve asset references — the same
+ pattern as ``assemble_retrieval_results``:
+ - **Tables**: inline HTML content at the ``ref`` placeholder
+ - **Images**: inline the ``file_path`` (S3-compatible URL) at the
+ placeholder for multimodal LLMs
+
+ Connected target chunks (images/tables) are expected to already be
+ present in ``chunks`` via ``hydrate_connected_target_rows``.
+ """
+ chunk_by_id: dict[str, dict] = {
+ c.get('chunk_id', ''): c for c in chunks if c.get('chunk_id')
+ }
+ rendered_ids: set[str] = set()
+ for chunk in chunks:
+ cid = chunk.get('chunk_id', '')
+ if cid and cid in rendered_ids:
+ continue
+ if cid:
+ rendered_ids.add(cid)
+
+ chunk_type = (chunk.get('chunk_type') or chunk.get('type') or 'text').strip().lower()
+
+ # Skip standalone image/table chunks — they'll be inlined
+ # via connect_to from their parent text chunk
+ if chunk_type in ('image', 'table'):
+ continue
+
+ content = str(chunk.get('content', '')).strip()
+
+ # Resolve connected assets via connect_to metadata
+ for conn in (chunk.get('chunk_metadata') or {}).get('connect_to') or []:
+ target = chunk_by_id.get(conn.get('target', ''))
+ if not target:
+ continue
+ target_cid = target.get('chunk_id', '')
+ target_type = (target.get('chunk_type') or target.get('type') or '').strip().lower()
+ ref_str = conn.get('ref', '')
+ if not ref_str or ref_str not in content:
+ continue
+
+ if target_cid:
+ rendered_ids.add(target_cid)
+
+ if target_type == 'table':
+ table_html = str(target.get('content', '')).strip()
+ content = content.replace(ref_str, f'\n[表格内容]\n{table_html}\n')
+ elif target_type == 'image':
+ file_path = target.get('file_path') or ''
+ img_desc = str(target.get('content', '')).strip()
+ # Strip self-reference from image description
+ if ref_str in img_desc:
+ img_desc = img_desc.replace(ref_str, '').strip()
+ # Use pre-generated asset URL if available, fall back to file_path
+ asset_url = (asset_lookup or {}).get(target_cid, '') if target_cid else ''
+ display_ref = asset_url or file_path
+ if display_ref:
+ content = content.replace(ref_str, f'\n[图片: {display_ref}]\n{img_desc}\n')
+ elif img_desc:
+ content = content.replace(ref_str, f'\n[图片描述]\n{img_desc}\n')
+
+ for line in content.split('\n'):
+ if line.strip():
+ parts.append(f'{indent}┈ {line}')
+
+
+def render_unified_doc_tree(
+ node: DocTreeNode,
+ doc_name: str,
+ depth: int = 0,
+ asset_lookup: dict[str, str] | None = None,
+) -> str:
+ """Render a DocTreeNode as a single coherent hierarchy.
+
+ Summaries are navigation-only aids and NEVER appear in evidence.
+ The rendered output contains:
+ 1. Structural titles for ALL sections (positioning context)
+ 2. Hydrated chunk content (┈ lines) ONLY for selected leaf paths
+
+ Asset references (tables/images) are resolved via ``connect_to``
+ metadata in hydrated chunks — no separate lookup needed.
+ """
+
+ parts: list[str] = []
+ indent = ' ' * depth
+
+ if depth == 0:
+ parts.append(f'【文档】{doc_name}\n')
+
+ # Track which paths have been rendered via outline_items
+ rendered_paths: set[str] = set()
+
+ # Collect children keys for path-hierarchy dedup:
+ child_prefixes = set(node.children.keys())
+
+ for item in node.outline_items:
+ path = item.get('path', '')
+ title = item.get('title', '')
+ is_leaf = item.get('is_leaf', False)
+ level = item.get('level', 1)
+ leaf_tag = ' [Leaf]' if is_leaf else ''
+
+ # Skip items that belong to a drilled-into child's subtree
+ if any(path.startswith(cp + ' / ') for cp in child_prefixes):
+ continue
+
+ rendered_paths.add(path)
+
+ # Section header (title only — summaries are navigation aids, not evidence)
+ level_tag = f'[L{level}] ' if level else ''
+ if level <= 1:
+ parts.append(f'{indent}▸ {level_tag}{title}{leaf_tag}')
+ else:
+ parts.append(f'{indent}└ {level_tag}{title}{leaf_tag}')
+
+ sub_indent = indent + ' '
+
+ # Case 1: This section was drilled into → show child tree inline
+ if path in node.children:
+ child = node.children[path]
+ child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup)
+ if child_text.strip():
+ parts.append(child_text)
+
+ # Case 2: This is a hydrated leaf → show chunk content inline
+ elif path in node.leaf_content:
+ _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup)
+
+ # Case 3: Unselected → title already rendered above, nothing more needed
+
+ # Render orphan paths: leaf_content and children not covered by outline_items
+ for path in node.leaf_content:
+ if path not in rendered_paths:
+ title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path
+ parts.append(f'{indent}▸ [Leaf] {title}')
+ sub_indent = indent + ' '
+ _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup)
+
+ for path in node.children:
+ if path not in rendered_paths:
+ title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path
+ parts.append(f'{indent}▸ {title} [DrillDown]')
+ child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup)
+ if child_text.strip():
+ parts.append(child_text)
+
+ return '\n'.join(parts)
+
diff --git a/packages/shared-python/shared/services/retrieval/agentic/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/__init__.py
index b34bfdb63..e4b5a0695 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/__init__.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/__init__.py
@@ -1,8 +1,14 @@
"""Agentic retrieval orchestration for Knowhere.
-Upgrades the retrieval pipeline from a fixed sequence to a state/action/observation
-loop with explicit tool selection, budget controls, and trajectory recording.
+Navigate-then-answer loop:
+ Phase 1: Document selection (discovery + KG LLM select)
+ Phase 2: Per-document iterative navigation (scope_navigate_step)
+ Phase 3: attempt_answer → DONE (return answer) or NOT_FOUND → revision
-All tools are thin wrappers around existing retrieval components — no new retrieval
-algorithms, ranking strategies, or prompts are introduced.
+Navigation auto-terminates when the LLM returns empty selections.
+After navigation, attempt_answer is called automatically — its result
+(answer or NOT_FOUND+reason) drives the revision loop.
+
+All tools are thin wrappers around existing retrieval components — no new
+retrieval algorithms, ranking strategies, or prompts are introduced.
"""
diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
index 5893b416a..776244317 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py
@@ -1,12 +1,18 @@
-"""Retrieval Agent orchestrator — the core agent loop.
-
-Runs the state/action/observation cycle using the LLM-driven LLMPolicy
-and agentic tools. The fixed terminal step (hydrate + rank) always
-executes, even if all tools fail — in that case it uses whatever
-discovery rows were collected.
-
-This module never raises exceptions to the caller. Errors are captured
-in the trace and the best available result is returned.
+"""Retrieval Agent orchestrator — navigate-then-answer loop.
+
+Flow:
+ Phase 1: Document selection (bottom_discovery + kg_document_select)
+ Phase 2: Per-document navigation (iterative BFS scope_navigate_step)
+ Phase 3: Render evidence → attempt_answer
+ → DONE (has answer) → return answer + evidence
+ → NOT_FOUND + reason → revision_hint → re-select docs + re-navigate
+ (exclude seen paths) → re-attempt
+ → max_revisions → return best available
+
+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.
"""
from __future__ import annotations
@@ -14,48 +20,128 @@
from typing import Any
from loguru import logger
+from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
+from shared.models.database.document import Document
+
from shared.services.retrieval.agentic.trace import TraceRecorder
from shared.services.retrieval.agentic.types import (
- ActionType,
AgentRunConfig,
AgentState,
+ AgenticResult,
+ CandidateDoc,
+ DocTreeNode,
ToolResult,
)
from shared.services.retrieval.app_service import (
- _hydrate_paths_to_rows,
- _load_chunk_importance_scores,
- _normalize_row_scores,
- _rank_candidates_by_path,
+ generate_retrieval_asset_url,
+ _is_client_result_artifact_ref,
)
from shared.services.retrieval.llm_adapter import LLMFn
+
+
+
+
+def _collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]:
+ """Recursively collect image/table chunks from a doc tree's leaf_content."""
+ media: list[dict[str, Any]] = []
+ for chunks in node.leaf_content.values():
+ for chunk in chunks:
+ ct = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower()
+ if ct in ('image', 'table'):
+ media.append(chunk)
+ for child in node.children.values():
+ media.extend(_collect_media_chunks(child))
+ return media
+
+
+async def _build_asset_url_map(
+ media_chunks: list[dict[str, Any]],
+) -> dict[str, str]:
+ """Generate presigned asset URLs for media chunks.
+
+ Uses the same ``generate_retrieval_asset_url`` as ``_to_public_response``
+ in ``app_service.py`` — no separate logic.
+ """
+ url_map: dict[str, str] = {}
+ for chunk in media_chunks:
+ chunk_id = str(chunk.get('chunk_id') or '').strip()
+ file_path = chunk.get('file_path') or ''
+ job_id = chunk.get('job_id') or ''
+ if not chunk_id or not file_path or not job_id:
+ continue
+ if not _is_client_result_artifact_ref(file_path):
+ continue
+ try:
+ url = await generate_retrieval_asset_url(
+ job_id=str(job_id),
+ artifact_ref=str(file_path),
+ )
+ if url:
+ url_map[chunk_id] = url
+ except Exception as e:
+ logger.warning(f'Failed to generate asset URL for {chunk_id} (ignored): {e}')
+ return url_map
+
+
def _build_config_from_env() -> AgentRunConfig:
"""Read agent config from environment, with sensible defaults."""
return AgentRunConfig(
- max_steps=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_STEPS', '10')),
- max_docs=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_DOCS', '0')),
- max_path_expansions=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_PATH_EXPANSIONS', '2')),
- max_doc_retries=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_DOC_RETRIES', '2')),
+ max_revisions=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_REVISIONS', '2')),
+ max_nav_depth=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_NAV_DEPTH', '3')),
latency_budget_ms=int(os.environ.get('RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS', '12000')),
- min_evidence_paths=int(os.environ.get('RETRIEVAL_AGENTIC_MIN_EVIDENCE_PATHS', '1')),
)
+async def _render_evidence(
+ db: AsyncSession,
+ doc_trees: dict[str, DocTreeNode],
+ doc_id_to_name: dict[str, str],
+) -> str:
+ """Render unified evidence text from doc trees.
+
+ Discovery paths are now handled by ``discovery_select_step`` in Phase 2
+ and merged into doc_trees — no separate fallback needed.
+ """
+ from shared.services.retrieval.agent_navigate import render_unified_doc_tree
+
+ # Build asset URL map for all media chunks (images/tables)
+ # — same pattern as _to_public_response in app_service.py
+ all_media_chunks: list[dict[str, Any]] = []
+ for doc_tree in doc_trees.values():
+ all_media_chunks.extend(_collect_media_chunks(doc_tree))
+ asset_url_map = await _build_asset_url_map(all_media_chunks)
+
+ # Render unified evidence from doc trees
+ evidence_parts: list[str] = []
+ for doc_id, doc_tree in doc_trees.items():
+ if doc_tree.has_content():
+ doc_name = doc_id_to_name.get(doc_id, doc_id)
+ rendered = render_unified_doc_tree(doc_tree, doc_name, asset_lookup=asset_url_map)
+ if rendered.strip():
+ evidence_parts.append(rendered)
+
+ return '\n\n'.join(evidence_parts) if evidence_parts else '(no evidence collected)'
+
+
class RetrievalAgent:
- """Agentic retrieval orchestrator using LLMPolicy.
+ """Agentic retrieval orchestrator — navigate-then-answer loop.
Usage::
agent = RetrievalAgent()
- ranked_rows, router = await agent.run(
+ result = await agent.run(
db, user_id=..., namespace=..., query=..., llm_fn=..., ...
)
+ # result.evidence_text — hierarchical context
+ # result.answer_text — LLM answer (empty if NOT_FOUND after max revisions)
+ # result.referenced_chunks — chunk IDs for hit stats / frontend
- The agent requires a valid ``llm_fn`` to run LLMPolicy. If ``llm_fn``
- is None, the run terminates immediately after bottom_discovery only.
+ The agent requires a valid ``llm_fn`` for LLM-driven navigation.
+ If ``llm_fn`` is None, the run returns discovery-only results.
"""
async def run(
@@ -75,13 +161,16 @@ async def run(
channels: list[str] | None = None,
channel_weights: dict[str, float] | None = None,
config: AgentRunConfig | None = None,
- ) -> tuple[list[dict[str, Any]], str]:
+ ) -> AgenticResult:
"""Run the agentic retrieval pipeline.
- Returns (ranked_rows, router_used). Never raises — errors are
- captured in trace and the best available result is returned.
+ Returns an ``AgenticResult`` containing the rendered evidence
+ text, LLM answer, and referenced chunk IDs. Never raises —
+ errors are captured in trace and the best available result
+ is returned.
"""
- from shared.services.retrieval.agentic.policy import LLMPolicy
+ from shared.services.retrieval.agentic import tools
+ from shared.services.retrieval.agentic.policy import attempt_answer
config = config or _build_config_from_env()
exclude_document_ids = exclude_document_ids or []
@@ -104,23 +193,18 @@ async def run(
logger.info(
f'agentic retrieval START: query="{query[:60]}..." '
- f'top_k={top_k} max_steps={config.max_steps} '
- f'budget={config.latency_budget_ms}ms'
+ f'top_k={top_k} budget={config.latency_budget_ms}ms'
)
if llm_fn is None:
logger.warning('agentic: no llm_fn provided — running discovery-only mode')
- # Build LLMPolicy (needs llm_fn; gracefully handles None)
- policy = LLMPolicy(llm_fn, query=query) if llm_fn else None
-
- # Shared kwargs for tool calls
- tool_kwargs: dict[str, Any] = {
+ # Shared kwargs for bottom_discovery
+ discovery_kwargs: dict[str, Any] = {
'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,
@@ -130,345 +214,417 @@ async def run(
'channel_weights': channel_weights,
}
- # ── Mandatory pre-step: bottom discovery ─────────────────────────────
- # bottom_discovery is always the first action; running it via the LLM
- # policy wastes ~1-2s on a trivial LLM call. We execute it directly
- # and let the LLM loop start from step 2 (kg_document_select etc.).
- logger.info(' agentic: running mandatory bottom_discovery pre-step')
- discovery_result = await self._execute_tool(
- db, ActionType.BOTTOM_DISCOVERY, state, config, **tool_kwargs,
- )
- state.apply(ActionType.BOTTOM_DISCOVERY, discovery_result)
+ # ══════════════════════════════════════════════════════════════════
+ # Phase 1: Discovery + Document Selection
+ # ══════════════════════════════════════════════════════════════════
+ logger.info(' agentic: Phase 1 — discovery + document selection')
+
+ # 1a. Bottom discovery (always runs)
+ discovery_result = await tools.bottom_discovery(db, **discovery_kwargs)
+ state.step_count += 1
+ discovery_rows = discovery_result.payload.get('fused_rows', []) if discovery_result.status != 'error' else []
+ state.discovery_top_doc_ids = discovery_result.payload.get('top_doc_ids', []) if discovery_result.status != 'error' else []
+
if trace_enabled:
trace.record_step(
- ActionType.BOTTOM_DISCOVERY, discovery_result,
- decision_reason='mandatory_pre_step',
+ 'bottom_discovery', discovery_result,
+ decision_reason='phase_1_mandatory',
)
- state.step_count += 1
+
logger.info(
- f' agentic step {state.step_count} (pre-step): action=bottom_discovery '
+ f' agentic step {state.step_count}: bottom_discovery '
f'status={discovery_result.status} latency={discovery_result.latency_ms}ms'
)
- # ── Agent loop (LLM decisions start from here) ────────────────────────
- stop_reason = 'max_steps'
- while state.step_count < config.max_steps:
- if state.elapsed_ms >= config.latency_budget_ms:
- logger.info(
- f' agentic: latency budget hit ({state.elapsed_ms}ms >= {config.latency_budget_ms}ms), '
- f'stopping at step {state.step_count}'
- )
- trace.record_budget_stop('latency')
- stop_reason = 'latency_budget'
- break
+ # 1b. KG document selection (requires LLM)
+ if llm_fn is not None:
+ kg_result = await tools.kg_document_select(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ llm_fn=llm_fn,
+ exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)),
+ )
+ state.step_count += 1
- if policy is None:
- # No LLM: discovery already ran — stop
- stop_reason = 'no_llm_fn'
- break
- else:
- action_type, decision_reason = await policy.decide(state, config)
-
- if action_type is None or action_type == ActionType.DONE:
- stop_reason = 'llm_done' if action_type == ActionType.DONE else 'llm_stop'
- logger.info(
- f' agentic: policy returned {action_type} at step {state.step_count} '
- f'(reason="{decision_reason}"), stopping'
+ if trace_enabled:
+ trace.record_step(
+ 'kg_document_select', kg_result,
+ decision_reason='phase_1_doc_selection',
)
- # Record DONE as a trace step
- if trace_enabled and action_type == ActionType.DONE:
- trace.record_step(
- ActionType.DONE,
- ToolResult(status='done', payload={'reason': decision_reason}),
- decision_reason=decision_reason,
+
+ if kg_result.status == 'selected_docs':
+ for doc_data in kg_result.payload.get('candidate_docs', []):
+ state.selected_docs.append(CandidateDoc(
+ document_id=doc_data.get('document_id', ''),
+ source_file_name=doc_data.get('source_file_name', ''),
+ confidence=doc_data.get('confidence', 0.0),
+ reason=doc_data.get('reason', ''),
+ source=doc_data.get('source', ''),
+ ))
+ state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {}))
+ state.doc_job_map.update(kg_result.payload.get('doc_job_map', {}))
+
+ # If KG returned nothing, use discovery hints
+ if not state.selected_docs and state.discovery_top_doc_ids:
+ hint_ids = [d for d in state.discovery_top_doc_ids if d not in state.ever_explored_doc_ids]
+ if hint_ids:
+ doc_stmt = (
+ select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
+ .where(Document.document_id.in_(hint_ids))
)
- break
+ doc_result = await db.execute(doc_stmt)
+ for did, fname, jrid in doc_result.all():
+ state.selected_docs.append(CandidateDoc(
+ document_id=did,
+ source_file_name=fname or did,
+ confidence=0.5,
+ reason='discovery_hint (KG returned 0)',
+ source='discovery_hint',
+ ))
+ state.doc_id_to_name[did] = fname or did
+ if jrid:
+ state.doc_job_map[did] = jrid
- result = await self._execute_tool(db, action_type, state, config, **tool_kwargs)
- state.apply(action_type, result)
+ logger.info(
+ f' agentic step {state.step_count}: kg_document_select '
+ f'status={kg_result.status} docs={len(state.selected_docs)} '
+ f'latency={kg_result.latency_ms}ms'
+ )
+ # If no LLM or no docs selected, return discovery rows directly
+ if not state.selected_docs:
+ logger.info('agentic: no documents selected — returning discovery results')
+ discovery_refs = [
+ {
+ 'chunk_id': r.get('chunk_id', ''),
+ 'document_id': r.get('document_id', ''),
+ 'chunk_type': r.get('chunk_type', ''),
+ 'section_path': r.get('section_path', ''),
+ 'file_path': r.get('file_path', ''),
+ }
+ for r in discovery_rows[:top_k]
+ if r.get('chunk_id')
+ ]
if trace_enabled:
- trace.record_step(action_type, result, decision_reason=decision_reason)
+ await trace.complete(discovery_rows, 'agentic_discovery_only')
+ return AgenticResult(
+ evidence_text='',
+ answer_text='',
+ referenced_chunks=discovery_refs,
+ router_used='agentic_discovery_only',
+ )
- state.step_count += 1
+ # ══════════════════════════════════════════════════════════════════
+ # Discovery → Navigation integration
+ # Group discovery_rows by document for post-BFS discovery selection
+ # ══════════════════════════════════════════════════════════════════
+ discovery_by_doc: dict[str, list[dict[str, Any]]] = {}
+ for row in discovery_rows:
+ doc_id = row.get('document_id', '')
+ if doc_id:
+ discovery_by_doc.setdefault(doc_id, []).append(row)
+
+ # Auto-register B-class docs (discovery-only, not selected by KG)
+ selected_doc_ids = {d.document_id for d in state.selected_docs}
+ for doc_id in discovery_by_doc:
+ if doc_id not in selected_doc_ids and doc_id not in state.ever_explored_doc_ids:
+ doc_stmt = (
+ select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
+ .where(Document.document_id == doc_id)
+ )
+ doc_result = await db.execute(doc_stmt)
+ row_data = doc_result.first()
+ if row_data:
+ did, fname, jrid = row_data
+ state.selected_docs.append(CandidateDoc(
+ document_id=did,
+ source_file_name=fname or did,
+ confidence=0.4,
+ reason='discovery_auto (not in KG selection)',
+ source='discovery_auto',
+ ))
+ state.doc_id_to_name[did] = fname or did
+ if jrid:
+ state.doc_job_map[did] = jrid
+
+ # ══════════════════════════════════════════════════════════════════
+ # Phase 2 + 3 Loop: Navigate → Render → Attempt Answer → (Revise)
+ # ══════════════════════════════════════════════════════════════════
+ answer_text = ''
+ revision_hint: str | None = None
+ stop_reason = 'max_revisions'
+
+ for round_idx in range(config.max_revisions + 1):
+ if state.elapsed_ms >= config.latency_budget_ms:
+ stop_reason = 'latency_budget'
+ break
+ # ── Phase 2: Per-Document Navigation ────────────────────────
logger.info(
- f' agentic step {state.step_count}: action={action_type.value} '
- f'status={result.status} latency={result.latency_ms}ms '
- f'docs={len(state.selected_docs)} paths={len(state.selected_paths)}'
+ f' agentic: Phase 2 (round {round_idx}) — '
+ f'navigating {len(state.selected_docs)} documents'
)
- # ── Terminal: hydrate + rank + attempt_answer loop ──
- while True:
- ranked_rows = await self._hydrate_and_rank(
- db, state, user_id=user_id, namespace=namespace, top_k=top_k,
- )
+ for doc in state.selected_docs:
+ if state.elapsed_ms >= config.latency_budget_ms:
+ logger.info(' agentic: latency budget hit during Phase 2, stopping')
+ break
- # Include kept rows from prior revision rounds
- if state.kept_path_rows:
- ranked_rows = state.kept_path_rows + ranked_rows
+ job_result_id = state.doc_job_map.get(doc.document_id, '')
+ if not job_result_id:
+ logger.info(f' agentic: skipping doc {doc.document_id} — no job_result_id')
+ state.ever_explored_doc_ids.add(doc.document_id)
+ continue
+
+ doc_name = doc.source_file_name or state.doc_id_to_name.get(doc.document_id, '')
+
+ # B-class docs (discovery_auto) skip BFS, go to discovery_select
+ is_b_class = doc.source == 'discovery_auto'
+
+ if not is_b_class:
+ # Build exclude_paths for this doc from seen_section_keys
+ doc_exclude = {
+ key.split('::', 1)[1]
+ for key in state.seen_section_keys
+ if key.startswith(f'{doc.document_id}::')
+ } if state.seen_section_keys else None
+
+ # BFS queue: (scope_path, parent_node, depth)
+ root = DocTreeNode(scope_path=None)
+ pending: list[tuple[str | None, DocTreeNode, int]] = [(None, root, 0)]
+
+ while pending:
+ if state.elapsed_ms >= config.latency_budget_ms:
+ break
+
+ scope, parent_node, depth = pending.pop(0)
+ if depth >= config.max_nav_depth:
+ continue
+
+ if llm_fn is None:
+ break
+
+ 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=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,
+ )
+ state.step_count += 1
+
+ # Merge step result into parent node
+ parent_node.outline_items = step_node.outline_items
+ parent_node.leaf_content = step_node.leaf_content
+ parent_node.confidence = step_node.confidence
+
+ # 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))
+
+ if trace_enabled:
+ trace.record_step(
+ 'scope_navigate_step', ToolResult(
+ status='navigated' if step_node.has_content() else 'empty',
+ payload={
+ 'document_id': doc.document_id,
+ 'scope': scope or 'root',
+ 'depth': depth,
+ 'outline_count': len(step_node.outline_items),
+ 'leaf_count': len(step_node.leaf_content),
+ 'pending_drills': len(drill_paths),
+ },
+ ),
+ decision_reason=f'nav_r{round_idx}_d{depth}_{doc.source_file_name}',
+ )
+
+ 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'leaves={len(step_node.leaf_content)} '
+ f'drills={len(drill_paths)}'
+ )
+ else:
+ # B-class: no BFS, create empty root
+ root = DocTreeNode(scope_path=None)
+
+ # ── Post-BFS: Discovery selection step ─────────────────────
+ doc_hints = discovery_by_doc.get(doc.document_id, [])
+ if doc_hints and llm_fn is not None and state.elapsed_ms < config.latency_budget_ms:
+ discovery_node = await tools.discovery_select_step(
+ db,
+ document_id=doc.document_id,
+ query=query,
+ llm_fn=llm_fn,
+ user_id=user_id,
+ namespace=namespace,
+ doc_name=doc_name,
+ discovery_hints=doc_hints,
+ )
+ state.step_count += 1
- # Check if we should attempt_answer (need LLM + results + revision budget)
- if (
- policy is None
- or not ranked_rows
- or state.revision_count >= config.max_revisions
- ):
- break
+ if trace_enabled:
+ trace.record_step(
+ 'discovery_select_step', ToolResult(
+ status='selected' if discovery_node.has_content() else 'empty',
+ payload={
+ 'document_id': doc.document_id,
+ 'hints_count': len(doc_hints),
+ 'hydrated_count': len(discovery_node.leaf_content),
+ },
+ ),
+ decision_reason=f'discovery_r{round_idx}_{doc.source_file_name}',
+ )
+
+ # Merge discovery results into BFS tree
+ root.merge(discovery_node)
+
+ # Merge or store doc tree
+ if doc.document_id in state.doc_trees:
+ state.doc_trees[doc.document_id].merge(root)
+ else:
+ state.doc_trees[doc.document_id] = root
+ state.ever_explored_doc_ids.add(doc.document_id)
+
+ # ── Phase 3: Render evidence + attempt_answer ────────────────
+ evidence_text = await _render_evidence(
+ db,
+ state.doc_trees, state.doc_id_to_name,
+ )
- # KG-exhausted guard: if all selected docs have been explored,
- # further revision won't find new content — skip attempt_answer
- all_selected_ids = {d.document_id for d in state.selected_docs}
- unexplored = all_selected_ids - state.ever_explored_doc_ids
- kg_exhausted = len(unexplored) == 0 and len(all_selected_ids) > 0
- if kg_exhausted:
- logger.info(
- f' agentic: KG exhausted — all {len(all_selected_ids)} docs explored, '
- f'skipping attempt_answer'
- )
- stop_reason = 'kg_exhausted'
+ if llm_fn is None:
+ stop_reason = 'no_llm'
break
- # Three-state verdict from LLM
- verdict, verdict_reason = await policy.attempt_answer(
- state, config, ranked_rows,
+ # Auto-trigger attempt_answer
+ status, answer_text, reason = await attempt_answer(
+ llm_fn,
+ query=query,
+ evidence_text=evidence_text,
+ state=state,
+ config=config,
)
- logger.info(
- f' agentic attempt_answer: verdict={verdict} '
- f'revision={state.revision_count}/{config.max_revisions} '
- f'reason="{verdict_reason}"'
- )
-
- if verdict == 'DONE':
- stop_reason = 'attempt_done'
- break
+ state.step_count += 1
- if verdict in ('NOT_SUFFICIENT', 'NOT_FOUND'):
- state.revision_count += 1
- # Save current results as kept rows for next round
- state.kept_path_rows = ranked_rows
- # Record current selected paths as seen
- for p in state.selected_paths:
- doc_id = p.get('document_id', '')
- path = p.get('path', '')
- if doc_id and path:
- state.seen_section_keys.add(f'{doc_id}::{path}')
- # Reset navigation state for re-exploration
- state.selected_paths.clear()
- state.selected_docs.clear() # Bug 1 fix: prevent doc accumulation
- state.pending_doc_index = 0
- state.kg_done = False
- state.discovery_done = False
-
- # Mandatory bottom_discovery pre-step for revision round
- logger.info(
- f' agentic: running mandatory bottom_discovery pre-step '
- f'(revision {state.revision_count})'
- )
- rev_discovery = await self._execute_tool(
- db, ActionType.BOTTOM_DISCOVERY, state, config, **tool_kwargs,
- )
- state.apply(ActionType.BOTTOM_DISCOVERY, rev_discovery)
- if trace_enabled:
- trace.record_step(
- ActionType.BOTTOM_DISCOVERY, rev_discovery,
- decision_reason=f'mandatory_pre_step (revision {state.revision_count})',
- )
- state.step_count += 1
- logger.info(
- f' agentic step {state.step_count} (rev {state.revision_count} pre-step): '
- f'action=bottom_discovery status={rev_discovery.status}'
+ if trace_enabled:
+ trace.record_step(
+ 'attempt_answer', ToolResult(
+ status=status,
+ payload={
+ 'answer_length': len(answer_text),
+ 'reason': reason,
+ },
+ ),
+ decision_reason=f'phase_3_answer_r{round_idx}',
)
- # Re-enter agent loop
- while state.step_count < config.max_steps:
- if state.elapsed_ms >= config.latency_budget_ms:
- stop_reason = 'latency_budget'
- break
+ logger.info(
+ f' agentic: attempt_answer status={status} '
+ f'round={round_idx}/{config.max_revisions} '
+ f'answer_len={len(answer_text)} reason="{reason}"'
+ )
- action_type, decision_reason = await policy.decide(state, config)
- if action_type is None or action_type == ActionType.DONE:
- stop_reason = 'llm_done'
- break
+ if status == 'DONE':
+ stop_reason = 'answer_done'
+ break
- result = await self._execute_tool(
- db, action_type, state, config, **tool_kwargs,
- )
- state.apply(action_type, result)
- if trace_enabled:
- trace.record_step(action_type, result, decision_reason=decision_reason)
- state.step_count += 1
- logger.info(
- f' agentic step {state.step_count} (rev {state.revision_count}): '
- f'action={action_type.value} status={result.status}'
- )
+ # ── NOT_FOUND: prepare revision ──────────────────────────────
+ if round_idx >= config.max_revisions:
+ stop_reason = 'max_revisions'
+ break
- # Loop back to hydrate + attempt_answer
- continue
+ state.revision_count += 1
+ revision_hint = reason
+ logger.info(f' agentic: starting revision {state.revision_count}, hint="{reason}"')
+
+ # Record all explored paths for masking (deepest-first removal)
+ for doc_id, doc_tree in state.doc_trees.items():
+ state.seen_section_keys.update(doc_tree.collect_all_paths(doc_id))
+
+ # Clear doc selection for re-exploration (preserve doc_trees for merge)
+ state.selected_docs.clear()
+
+ # Re-run KG select (allow re-exploring docs for different sections via path masking)
+ kg_result = await tools.kg_document_select(
+ db,
+ user_id=user_id,
+ namespace=namespace,
+ query=query,
+ llm_fn=llm_fn,
+ exclude_document_ids=list(set(exclude_document_ids)),
+ )
+ state.step_count += 1
- # Unknown verdict — treat as DONE
- break
+ if kg_result.status == 'selected_docs':
+ for doc_data in kg_result.payload.get('candidate_docs', []):
+ state.selected_docs.append(CandidateDoc(
+ document_id=doc_data.get('document_id', ''),
+ source_file_name=doc_data.get('source_file_name', ''),
+ confidence=doc_data.get('confidence', 0.0),
+ reason=doc_data.get('reason', ''),
+ source=doc_data.get('source', ''),
+ ))
+ state.doc_id_to_name.update(kg_result.payload.get('doc_id_to_name', {}))
+ state.doc_job_map.update(kg_result.payload.get('doc_job_map', {}))
+
+ if not state.selected_docs:
+ logger.info(' agentic: revision found no new docs — stopping')
+ stop_reason = 'no_new_docs'
+ break
+ # ══════════════════════════════════════════════════════════════════
+ # Final Assembly
+ # ══════════════════════════════════════════════════════════════════
router_used = (
- 'agentic_llm' if state.selected_paths or state.kept_path_rows
+ 'agentic_llm' if any(t.has_content() for t in state.doc_trees.values())
else 'agentic_discovery_only'
)
+ # Collect referenced chunk IDs from all doc trees
+ all_refs: list[dict[str, str]] = []
+ seen_ref_ids: set[str] = set()
+ for doc_tree in state.doc_trees.values():
+ for ref in doc_tree.collect_referenced_ids():
+ cid = ref.get('chunk_id', '')
+ if cid and cid not in seen_ref_ids:
+ seen_ref_ids.add(cid)
+ all_refs.append(ref)
+
+ # Re-render final evidence (may have been updated in last revision)
+ if not evidence_text or evidence_text == '(no evidence collected)':
+ evidence_text = await _render_evidence(
+ db,
+ state.doc_trees, state.doc_id_to_name,
+ )
+
+ result = AgenticResult(
+ evidence_text=evidence_text,
+ answer_text=answer_text,
+ referenced_chunks=all_refs,
+ router_used=router_used,
+ )
+
logger.info(
- f'agentic retrieval DONE: {len(ranked_rows)} results, '
+ f'agentic retrieval DONE: {len(all_refs)} referenced chunks, '
+ f'evidence_text={len(evidence_text)} chars, '
+ f'answer_text={len(answer_text)} chars, '
f'router={router_used}, steps={state.step_count}, '
f'stop_reason={stop_reason}, revisions={state.revision_count}, '
f'{state.elapsed_ms}ms'
)
if trace_enabled:
- await trace.complete(ranked_rows, router_used)
-
- return ranked_rows, router_used
-
- async def _execute_tool(
- self,
- db: AsyncSession,
- action_type: ActionType,
- state: AgentState,
- config: AgentRunConfig,
- **kwargs: Any,
- ) -> ToolResult:
- """Execute a single tool. Never raises."""
- from shared.services.retrieval.agentic import tools
-
- try:
- if action_type == ActionType.BOTTOM_DISCOVERY:
- return await tools.bottom_discovery(db, **kwargs)
-
- elif action_type == ActionType.KG_DOCUMENT_SELECT:
- return await tools.kg_document_select(
- db,
- user_id=kwargs['user_id'],
- namespace=kwargs['namespace'],
- query=kwargs['query'],
- llm_fn=kwargs.get('llm_fn'),
- exclude_document_ids=list(state.excluded_doc_ids | set(kwargs.get('exclude_document_ids', []))),
- )
-
- elif action_type == ActionType.DOCUMENT_PATH_SELECT:
- # Find which document to process next
- doc = None
- if state.pending_doc_index < len(state.selected_docs):
- doc = state.selected_docs[state.pending_doc_index]
-
- if doc is None:
- return ToolResult(status='error', error='no document to process')
-
- job_result_id = state.doc_job_map.get(doc.document_id, '')
- if not job_result_id:
- return ToolResult(
- status='no_confident_match',
- payload={'document_id': doc.document_id, 'reason': 'no job_result_id'},
- )
-
- # Build exclude_paths for this doc from seen_section_keys
- doc_exclude = {
- key.split('::', 1)[1]
- for key in state.seen_section_keys
- if key.startswith(f'{doc.document_id}::')
- } if state.seen_section_keys else None
-
- return await tools.document_path_select(
- db,
- user_id=kwargs['user_id'],
- namespace=kwargs['namespace'],
- query=kwargs['query'],
- llm_fn=kwargs.get('llm_fn'),
- document_id=doc.document_id,
- job_result_id=job_result_id,
- doc_name=doc.source_file_name or state.doc_id_to_name.get(doc.document_id, ''),
- exclude_paths=doc_exclude,
- )
-
- elif action_type == ActionType.GREP_DOCUMENT_DISCOVER:
- return await tools.grep_document_discover(
- db,
- user_id=kwargs['user_id'],
- namespace=kwargs['namespace'],
- query=kwargs['query'],
- exclude_document_ids=list(state.excluded_doc_ids | set(kwargs.get('exclude_document_ids', []))),
- )
-
- elif action_type == ActionType.GRAPH_EXPAND_DOCS:
- doc_ids = [d.document_id for d in state.selected_docs]
- return await tools.graph_expand_docs(
- db,
- user_id=kwargs['user_id'],
- namespace=kwargs['namespace'],
- document_ids=doc_ids,
- )
-
- else:
- return ToolResult(status='error', error=f'unknown action: {action_type}')
-
- except Exception as e:
- logger.warning(f' agentic tool {action_type.value} raised: {e}')
- return ToolResult(status='error', error=str(e))
+ await trace.complete(all_refs, router_used)
- async def _hydrate_and_rank(
- self,
- db: AsyncSession,
- state: AgentState,
- *,
- user_id: str,
- namespace: str,
- top_k: int,
- ) -> list[dict[str, Any]]:
- """Fixed terminal step: hydrate selected paths + rank against discovery.
-
- Reuses _hydrate_paths_to_rows and _rank_candidates_by_path unchanged.
-
- TODO (Token Budget): Replace `ranked_rows[:top_k]` in
- `_rank_candidates_by_path` with token-accumulation truncation
- (tiktoken or character estimate) so the final result set respects
- a configurable LLM context window budget rather than a fixed count.
- """
- try:
-
- # Hydrate agent-selected paths
- navigated_paths: list[dict[str, Any]] = []
- if state.selected_paths:
- navigated_paths = await _hydrate_paths_to_rows(
- db,
- path_selections=state.selected_paths,
- user_id=user_id,
- namespace=namespace,
- )
-
- # Load importance scores for all candidate rows
- all_candidates = state.discovery_paths + navigated_paths
- if all_candidates:
- importance_map = await _load_chunk_importance_scores(
- db, user_id=user_id, namespace=namespace,
- rows=all_candidates,
- )
- for row in all_candidates:
- chunk_id = row.get('chunk_id', '')
- row['importance_raw_score'] = importance_map.get(chunk_id, 0.0)
-
- _normalize_row_scores(
- all_candidates,
- source_field='importance_raw_score',
- target_field='importance_norm_score',
- default=0.0,
- )
-
- # Rank: merge discovery + agent rows
- ranked = _rank_candidates_by_path(
- discovery_rows=state.discovery_paths,
- routed_rows=navigated_paths,
- top_k=top_k,
- )
-
- return ranked
-
- except Exception as e:
- logger.error(f'agentic hydrate_and_rank failed: {e}')
- # Last resort: return raw discovery rows
- return state.discovery_paths[:top_k]
+ return result
diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py
index 12e5afa62..e5a5c0dd4 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/policy.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py
@@ -1,317 +1,131 @@
-"""LLM-driven policy for agentic retrieval.
+"""LLM answer-attempt tool for agentic retrieval.
-Replaces the former RuleBasedPolicy. All control decisions are made by a
-small LLM call so the agent can adapt to query complexity and KB topology
-rather than following a fixed rule tree.
+Provides ``attempt_answer()`` — a single LLM call that tries to answer
+the user's query using the collected evidence.
-Design:
- - ``LLMPolicy.decide()`` is **async** — it calls the LLM to pick the
- next action.
- - On parse failure the agent terminates immediately (no hidden fallback).
- - The prompt is compact: no raw chunk content, only state metadata.
+Returns one of two outcomes:
+ - answer_text (non-empty) → the evidence was sufficient, answer is ready
+ - NOT_FOUND + reason → the evidence was insufficient, triggers a revision
"""
from __future__ import annotations
import json
+import os
import re
from typing import Any
from loguru import logger
-from shared.services.retrieval.agentic.types import ActionType, AgentRunConfig, AgentState
+from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState
from shared.services.retrieval.llm_adapter import LLMFn
-# ── Available actions presented to the LLM ───────────────────────────────────
-# NOTE: BOTTOM_DISCOVERY is intentionally excluded from this list.
-# It is now a mandatory pre-step executed automatically by the orchestrator
-# before the LLM decision loop begins. The LLM should never need to decide
-# whether to run it — doing so wastes one LLM call per run.
-_AVAILABLE_ACTIONS: list[dict[str, Any]] = [
- {
- 'action': ActionType.KG_DOCUMENT_SELECT.value,
- 'description': (
- 'Ask the LLM to select the most relevant documents from the Knowledge Graph '
- 'overview (doc summaries). Use after discovery.'
- ),
- 'when': 'discovery_done is true, kg_done is false',
- },
- {
- 'action': ActionType.DOCUMENT_PATH_SELECT.value,
- 'description': (
- 'Drill into the next pending document\'s section tree and pick relevant '
- 'paths. The document is chosen automatically from selected_docs[pending_doc_index].'
- ),
- 'when': 'kg_done is true, there are pending documents to process',
- },
- {
- 'action': ActionType.GREP_DOCUMENT_DISCOVER.value,
- 'description': (
- 'Run term/grep search across all documents to discover relevant doc IDs '
- 'when KG selection found nothing.'
- ),
- 'when': 'kg_done is true but selected_docs is empty',
- },
- {
- 'action': ActionType.GRAPH_EXPAND_DOCS.value,
- 'description': (
- 'Expand via Knowledge Graph edge relationships to find related documents '
- 'not yet selected.'
- ),
- 'when': 'Optional deepening — use when current results seem insufficient',
- },
- {
- 'action': ActionType.DONE.value,
- 'description': (
- 'Stop the agent loop and proceed to final hydration and ranking. '
- 'Use when you have sufficient evidence paths or the budget is nearly exhausted.'
- ),
- 'when': 'Sufficient paths collected, or further actions would not improve results',
- },
-]
-
-_ACTIONS_BLOCK = '\n'.join(
- f" {i+1}. \"{a['action']}\": {a['description']} [{a['when']}]"
- for i, a in enumerate(_AVAILABLE_ACTIONS)
-)
-
-_POLICY_PROMPT_TEMPLATE = """\
-You are a retrieval agent orchestrating document search for a RAG system.
-Your job: choose the SINGLE best next action given the current state.
-
-QUERY: "{query}"
-
-CURRENT STATE:
-{state_json}
-
-BUDGET: {elapsed_ms}ms elapsed of {budget_ms}ms max | Step {step} of {max_steps}
-
-AVAILABLE ACTIONS:
-{actions_block}
-
-RULES:
-1. Run kg_document_select when discovery_done is true and kg_done is false.
-2. After kg select, run document_path_select for each pending document.
-3. Call done when all pending documents are processed OR you have >= {min_evidence} evidence paths.
-4. Only use grep_document_discover if kg_document_select found 0 documents.
-5. Only use graph_expand_docs if you need more related docs after reviewing results.
-Note: bottom_discovery is already executed automatically before this loop — do NOT attempt to call it.
-
-Return ONLY a JSON object, no markdown, no explanation:
-{{"action": "", "reason": ""}}
-"""
-
-
-def _parse_action_from_response(response: str) -> dict[str, Any] | None:
- """Extract the JSON decision object from the LLM response."""
- # Strip markdown code fences if present
- text = response.strip()
- text = re.sub(r'^```(?:json)?\s*', '', text, flags=re.MULTILINE)
- text = re.sub(r'\s*```$', '', text, flags=re.MULTILINE)
+def _parse_answer_response(text: str) -> dict[str, Any] | None:
+ """Extract a JSON answer object from LLM response text."""
text = text.strip()
-
- # Try to extract {...} block
- match = re.search(r'\{[^{}]*\}', text, re.DOTALL)
+ try:
+ return json.loads(text)
+ except (json.JSONDecodeError, ValueError):
+ pass
+ match = re.search(r'\{.*\}', text, re.DOTALL)
if match:
try:
- return json.loads(match.group(0))
- except json.JSONDecodeError:
+ return json.loads(match.group())
+ except (json.JSONDecodeError, ValueError):
pass
-
- # Try the whole text as JSON
- try:
- return json.loads(text)
- except json.JSONDecodeError:
- return None
-
-
-class LLMPolicy:
- """Async LLM-driven policy. One LLM call per agent step.
-
- Usage::
-
- policy = LLMPolicy(llm_fn)
- action_type, reason = await policy.decide(state, config, query=query)
+ return None
+
+
+async def attempt_answer(
+ llm_fn: LLMFn,
+ *,
+ query: str,
+ evidence_text: str,
+ state: AgentState,
+ config: AgentRunConfig,
+) -> tuple[str, str, str]:
+ """Attempt to answer the query using collected evidence.
+
+ Returns (status, answer_text, reason) where:
+ - status='DONE', answer_text=, reason=''
+ → evidence was sufficient, answer is ready
+ - status='NOT_FOUND', answer_text='', reason=
+ → evidence was insufficient, reason is used as revision_hint
"""
-
- def __init__(self, llm_fn: LLMFn, *, query: str = '') -> None:
- self._llm_fn = llm_fn
- self._query = query
-
- def build_prompt(self, state: AgentState, config: AgentRunConfig) -> str:
- """Build the decision prompt. Public for test inspection."""
- state_data = state.state_summary()
-
- has_pending_docs = state.pending_doc_index < len(state.selected_docs)
- state_data['has_pending_docs'] = has_pending_docs
- state_json = json.dumps(state_data, ensure_ascii=False, indent=2)
-
- allowed_actions = []
- for action in _AVAILABLE_ACTIONS:
- name = action['action']
- if name == ActionType.KG_DOCUMENT_SELECT.value and (not state.discovery_done or state.kg_done):
- continue
- if name == ActionType.DOCUMENT_PATH_SELECT.value and (not state.kg_done or not has_pending_docs):
- continue
- if name == ActionType.GREP_DOCUMENT_DISCOVER.value and (not state.kg_done or len(state.selected_docs) > 0):
- continue
- allowed_actions.append(action)
-
- actions_block = '\n'.join(
- f" {i+1}. \"{a['action']}\": {a['description']} [{a['when']}]"
- for i, a in enumerate(allowed_actions)
- )
-
- return _POLICY_PROMPT_TEMPLATE.format(
- query=self._query,
- state_json=state_json,
- elapsed_ms=state.elapsed_ms,
- budget_ms=config.latency_budget_ms,
- step=state.step_count,
- max_steps=config.max_steps,
- actions_block=actions_block,
- min_evidence=config.min_evidence_paths,
- )
-
- async def decide(
- self,
- state: AgentState,
- config: AgentRunConfig,
- ) -> tuple[ActionType | None, str]:
- """Ask the LLM which action to take next.
-
- Returns ``(action_type, reason)`` or ``(None, reason)`` when the
- LLM says ``done`` or when parsing fails (hard stop).
- """
- prompt = self.build_prompt(state, config)
-
+ prompt = _ATTEMPT_ANSWER_PROMPT.format(
+ query=query,
+ evidence_context=evidence_text,
+ revision_count=state.revision_count,
+ max_revisions=config.max_revisions,
+ )
+
+ verbose = os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', '') == 'true'
+ if verbose:
logger.info(
- f' [LLMPolicy] step={state.step_count} calling LLM '
- f'(state: discovery={state.discovery_done}, kg={state.kg_done}, '
- f'docs={len(state.selected_docs)}, pending={state.pending_doc_index}, '
- f'paths={len(state.selected_paths)})'
+ f'[attempt_answer PROMPT]\n'
+ f'{prompt}'
)
- # ── Verbose prompt logging (controlled by env) ──
- import os
- if os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', 'false') == 'true':
- logger.info(
- f'\n{"="*60}\n'
- f'[LLMPolicy PROMPT step={state.step_count}]\n'
- f'{prompt}\n'
- f'{"="*60}'
- )
-
- raw_response = await self._llm_fn(prompt)
+ raw_response = await llm_fn(prompt)
+ logger.info(f' [attempt_answer] raw={repr(raw_response[:300])}')
+ if verbose:
logger.info(
- f' [LLMPolicy] raw_response={repr(raw_response[:200])}'
+ f'[attempt_answer RESPONSE]\n'
+ f'{raw_response}'
)
- if os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', 'false') == 'true':
- logger.info(
- f'\n{"="*60}\n'
- f'[LLMPolicy RESPONSE step={state.step_count}]\n'
- f'{raw_response}\n'
- f'{"="*60}'
- )
-
- if not raw_response.strip():
- logger.warning(' [LLMPolicy] empty response → stopping agent')
- return None, 'empty LLM response'
-
- parsed = _parse_action_from_response(raw_response)
- if not parsed:
- logger.warning(
- f' [LLMPolicy] could not parse JSON from response: {repr(raw_response[:300])} → stopping agent'
- )
- return None, f'parse_error: {raw_response[:100]}'
+ parsed = _parse_answer_response(raw_response)
+ if not parsed:
+ # Parse error: treat raw text as best-effort answer
+ return 'DONE', raw_response.strip(), 'parse_error — treating raw response as answer'
- action_str = str(parsed.get('action', '')).strip()
- reason = str(parsed.get('reason', '')).strip()
+ status = str(parsed.get('status', 'DONE')).strip().upper()
+ answer = str(parsed.get('answer', '')).strip()
+ reason = str(parsed.get('reason', '')).strip()
- logger.info(f' [LLMPolicy] decided action="{action_str}" reason="{reason}"')
-
- if action_str == ActionType.DONE.value:
- return ActionType.DONE, reason
-
- # Validate against whitelist
- try:
- action_type = ActionType(action_str)
- except ValueError:
- logger.warning(
- f' [LLMPolicy] unknown action "{action_str}" → stopping agent'
- )
- return None, f'unknown_action: {action_str}'
+ if status == 'NOT_FOUND':
+ return 'NOT_FOUND', '', reason or 'LLM returned NOT_FOUND without reason'
- return action_type, reason
-
- async def attempt_answer(
- self,
- state: AgentState,
- config: AgentRunConfig,
- ranked_rows: list[dict[str, Any]],
- ) -> tuple[str, str]:
- """Three-state verdict: is the evidence sufficient?
-
- Returns (verdict, reason) where verdict is one of:
- - 'DONE': evidence is sufficient, stop searching
- - 'NOT_SUFFICIENT': partial match, need more evidence
- - 'NOT_FOUND': no relevant evidence found at all
- """
- # Build evidence summary (paths + previews, not full content)
- evidence_lines: list[str] = []
- for i, row in enumerate(ranked_rows[:20]): # cap to avoid huge prompt
- path = row.get('section_path') or row.get('source_chunk_path') or ''
- content_preview = str(row.get('content', ''))[:150]
- score = round(float(row.get('score', 0.0) or 0.0), 3)
- evidence_lines.append(
- f' {i+1}. path="{path}" score={score}\n'
- f' preview: {content_preview}'
- )
- evidence_text = '\n'.join(evidence_lines) or '(no evidence collected)'
-
- prompt = _ATTEMPT_ANSWER_PROMPT.format(
- query=self._query,
- evidence_count=len(ranked_rows),
- evidence_summary=evidence_text,
- revision_count=state.revision_count,
- max_revisions=config.max_revisions,
- )
-
- raw_response = await self._llm_fn(prompt)
- logger.info(f' [LLMPolicy.attempt_answer] raw={repr(raw_response[:200])}')
-
- parsed = _parse_action_from_response(raw_response)
- if not parsed:
- return 'DONE', 'parse_error — treating as done'
-
- verdict = str(parsed.get('verdict', 'DONE')).strip().upper()
- reason = str(parsed.get('reason', '')).strip()
-
- if verdict not in ('DONE', 'NOT_SUFFICIENT', 'NOT_FOUND'):
- verdict = 'DONE'
-
- return verdict, reason
+ # Any status other than NOT_FOUND → treat as DONE
+ if not answer:
+ answer = reason or '(empty answer)'
+ return 'DONE', answer, ''
_ATTEMPT_ANSWER_PROMPT = """\
-You are evaluating whether the collected evidence can answer the user's query.
+You are a knowledge retrieval assistant. Answer the user's query based
+STRICTLY on the provided evidence. Do NOT use any external knowledge.
QUERY: "{query}"
-EVIDENCE ({evidence_count} items, showing top 20):
-{evidence_summary}
+EVIDENCE CONTEXT:
+The following evidence is organized by document in a unified hierarchy.
+Each document shows its structural outline (section titles + summaries)
+with retrieved content (┈ lines) inline under the relevant sections.
-REVISION: {revision_count} of {max_revisions} revisions used.
-
-Evaluate the evidence and return ONE verdict:
-- "DONE": The evidence is sufficient to answer the query. Use this if the main points are covered.
-- "NOT_SUFFICIENT": Partial match — some relevant info found but key aspects are missing. Only use if more searching could realistically help.
-- "NOT_FOUND": The evidence is completely irrelevant to the query. Only use if nothing matches at all.
+{evidence_context}
-When in doubt, prefer DONE — avoid unnecessary extra search rounds.
+REVISION: {revision_count} of {max_revisions} revisions used.
-Return ONLY a JSON object:
-{{"verdict": "DONE", "reason": "one sentence explanation"}}
+INSTRUCTIONS:
+1. If the evidence contains enough information to answer the query,
+ compose a clear and comprehensive answer. Return:
+ {{"status": "DONE", "answer": ""}}
+
+2. If the evidence does NOT contain sufficient information to answer
+ the query, return NOT_FOUND with a specific reason explaining what
+ information is missing. This reason will be used to guide the next
+ search round. Return:
+ {{"status": "NOT_FOUND", "reason": ""}}
+
+IMPORTANT:
+- Base your judgment ONLY on the actual retrieved content (┈ lines),
+ not just section titles or summaries.
+- Be specific in your NOT_FOUND reason — mention exactly what data,
+ section, or detail you expected but didn't find.
+- When the evidence partially covers the query, still return DONE with
+ the available information and note any gaps in your answer.
+
+Return ONLY a JSON object, no other text.
"""
diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py
index 52c7645d3..c337d2e08 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/tools.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py
@@ -4,10 +4,6 @@
1. Calls existing functions from channels.py, agent_navigate.py, app_service.py
2. Returns a unified ToolResult
3. Never raises — errors are captured in ToolResult.error
-
-No new retrieval algorithms, ranking strategies, or prompts.
-LLM calls inside kg_document_select / document_path_select reuse
-the exact same prompts and parsing logic from agent_navigate.py.
"""
from __future__ import annotations
@@ -19,18 +15,18 @@
from sqlalchemy.ext.asyncio import AsyncSession
from shared.models.database.document import Document
-from shared.services.retrieval.agentic.types import ToolResult
+from shared.services.retrieval.agentic.types import DocTreeNode, ToolResult
from shared.services.retrieval.agent_navigate import (
_build_knowledge_map_overview,
_expand_by_edges,
_format_items_for_llm,
_grep_discover_document_ids,
_load_child_sections,
- _parse_chunk_path_selections,
_parse_json_array,
+ _parse_scope_nav_response,
_SCOPE_NAV_PROMPT,
+ _DISCOVERY_SELECT_PROMPT,
_FILE_SELECT_PROMPT,
- _default_confidence_for_rank,
)
from shared.services.retrieval.app_service import (
_CHANNEL_WEIGHT_CONTENT,
@@ -40,6 +36,7 @@
_merge_same_section_rows,
_normalize_row_scores,
_resolve_allowed_chunk_types,
+ hydrate_connected_target_rows,
merge_channels_rrf,
)
from shared.services.retrieval.channels import content_channel, path_channel, term_channel
@@ -67,11 +64,7 @@ async def bottom_discovery(
internal_recall_k: int | None = None,
**_kwargs: Any,
) -> ToolResult:
- """Run 3-channel BM25 discovery + RRF fusion.
-
- Reuses: channels.path_channel, content_channel, term_channel,
- app_service.merge_channels_rrf, _merge_same_section_rows.
- """
+ """Run 3-channel BM25 discovery + RRF fusion."""
t0 = time.monotonic()
try:
allowed_chunk_types = _resolve_allowed_chunk_types(data_type)
@@ -106,7 +99,7 @@ async def bottom_discovery(
signal_paths=signal_paths, filter_mode=filter_mode,
)
- # RRF fusion — same logic as app_service
+ # RRF fusion
default_weights = {
'path': _CHANNEL_WEIGHT_PATH,
'content': _CHANNEL_WEIGHT_CONTENT,
@@ -178,11 +171,7 @@ async def kg_document_select(
exclude_document_ids: list[str],
**_kwargs: Any,
) -> ToolResult:
- """Select candidate documents from document-level KG.
-
- Reuses: agent_navigate._build_knowledge_map_overview, _parse_json_array.
- Same LLM prompt as agent_navigate._FILE_SELECT_PROMPT.
- """
+ """Select candidate documents from document-level KG."""
t0 = time.monotonic()
try:
overview_text, doc_id_to_name = await _build_knowledge_map_overview(
@@ -204,7 +193,6 @@ async def kg_document_select(
latency_ms=latency,
)
- # LLM file selection — same prompt as agent_navigate
file_prompt = _FILE_SELECT_PROMPT.format(
overview=overview_text, query=query,
)
@@ -239,7 +227,7 @@ async def kg_document_select(
candidate_docs.append({
'document_id': did,
'source_file_name': doc_id_to_name.get(did, ''),
- 'confidence': 0.8,
+ 'confidence': 1.0,
'reason': 'LLM selected from KG overview',
'source': 'kg_llm_select',
})
@@ -261,39 +249,6 @@ async def kg_document_select(
return ToolResult(status='error', error=str(e), latency_ms=latency)
-# ---------------------------------------------------------------------------
-# Tool: document_path_select
-# ---------------------------------------------------------------------------
-
-async def document_path_select(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- query: str,
- llm_fn: LLMFn | None,
- document_id: str,
- job_result_id: str,
- doc_name: str = '',
- max_chunks_per_file: int = 15,
- exclude_paths: set[str] | None = None,
- **_kwargs: Any,
-) -> ToolResult:
- """Document entry point for agentic scope navigation."""
- if llm_fn is None:
- return ToolResult(
- status='no_confident_match',
- payload={'document_id': document_id, 'reason': 'LLM not available'},
- latency_ms=0,
- )
- return await scope_navigate(
- db, document_id=document_id, job_result_id=job_result_id,
- query=query, llm_fn=llm_fn, doc_name=doc_name,
- scope_path=None, max_select=max_chunks_per_file,
- exclude_paths=exclude_paths,
- )
-
-
# ---------------------------------------------------------------------------
# Tool: grep_document_discover
# ---------------------------------------------------------------------------
@@ -307,10 +262,7 @@ async def grep_document_discover(
exclude_document_ids: list[str],
**_kwargs: Any,
) -> ToolResult:
- """Discover documents via term search (GREP).
-
- Reuses: agent_navigate._grep_discover_document_ids.
- """
+ """Discover documents via term search (GREP)."""
t0 = time.monotonic()
try:
grep_doc_ids = await _grep_discover_document_ids(
@@ -326,7 +278,6 @@ async def grep_document_discover(
latency_ms=latency,
)
- # Load doc names and job_result_ids
doc_stmt = (
select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
.where(Document.document_id.in_(grep_doc_ids))
@@ -368,10 +319,7 @@ async def graph_expand_docs(
document_ids: list[str],
**_kwargs: Any,
) -> ToolResult:
- """Expand document set via KG edge traversal.
-
- Reuses: agent_navigate._expand_by_edges.
- """
+ """Expand document set via KG edge traversal."""
t0 = time.monotonic()
try:
expanded_ids = await _expand_by_edges(
@@ -387,7 +335,6 @@ async def graph_expand_docs(
latency_ms=latency,
)
- # Load names and job maps for new docs
doc_stmt = (
select(Document.document_id, Document.source_file_name, Document.current_job_result_id)
.where(Document.document_id.in_(new_ids))
@@ -418,100 +365,242 @@ async def graph_expand_docs(
# ---------------------------------------------------------------------------
-# Tool: scope_navigate (Unified recursive navigation)
+# Tool: scope_navigate_step (single-step navigation)
# ---------------------------------------------------------------------------
-async def scope_navigate(
+_LLM_MODE_TO_HYDRATE: dict[str, str] = {
+ 'all': 'chunks',
+ 'image': 'image_only',
+ 'table': 'table_only',
+}
+
+async def scope_navigate_step(
db: AsyncSession,
*,
document_id: str,
job_result_id: str,
query: str,
llm_fn: LLMFn,
+ user_id: str,
+ namespace: str,
doc_name: str = '',
scope_path: str | None = None,
- max_select: int = 15,
exclude_paths: set[str] | None = None,
-) -> ToolResult:
- """Unified document-internal navigation tool.
-
- 1. Loads 2 levels of child sections under scope_path
- 2. Applies overflow guard (drops summaries if needed)
- 3. LLM selects most relevant items
- 4. Returns selected section paths directly; each path hydrates the
- corresponding section subtree.
+ revision_hint: str | None = None,
+) -> tuple[DocTreeNode, list[dict]]:
+ """Single navigation step — one LLM call, no recursion.
+
+ 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)
"""
- t0 = time.monotonic()
+ from shared.services.retrieval.app_service import _hydrate_paths_to_rows
+
+ empty = DocTreeNode.empty(scope_path)
+
try:
+ # 1. Load continuous context tree
items = await _load_child_sections(
db, document_id, job_result_id, scope_path,
exclude_paths=exclude_paths,
)
if not items:
- latency = int((time.monotonic() - t0) * 1000)
- return ToolResult(
- status='no_items',
- payload={'document_id': document_id, 'scope_path': scope_path},
- latency_ms=latency,
- )
+ return 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)}
+ # 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)'
+ )
prompt = _SCOPE_NAV_PROMPT.format(
doc_name=doc_name or document_id,
doc_id=document_id,
- scope_label=scope_path or 'root',
+ scope_header=scope_header,
items_overview=text,
query=query,
- max_select=max_select,
)
-
- valid_paths = {item['path'] for item in items}
+ if revision_hint:
+ prompt += (
+ f'\n\nIMPORTANT: Previous round feedback: '
+ f'"{revision_hint}". Select specific sections this time.'
+ )
response = await llm_fn(prompt)
- selected = _parse_chunk_path_selections(response)
+ selections = _parse_scope_nav_response(response)
+
+ logger.info(
+ f' scope_navigate_step scope={scope_path or "root"}: '
+ f'selections={len(selections)}, selectable={len(selectable)}, '
+ f'overflowed={overflowed}'
+ )
- accepted: list[dict[str, Any]] = []
- for item in selected:
- path = str(item.get('path') or '').strip()
- if path not in valid_paths:
+ # 4. Build node with LOCAL items only (no ancestors/siblings)
+ node = DocTreeNode(scope_path=scope_path)
+ 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)
+ valid_selections = [
+ s for s in selections
+ if s['path'] in selectable and s['path'] != scope_path
+ ]
+
+ pending: list[dict] = []
+ for sel in valid_selections:
+ path = sel['path']
+ conf = sel.get('confidence', 0.7)
+ item = selectable[path]
+ node.confidence[path] = conf
+
+ if item.get('is_leaf'):
+ # Leaf → hydrate chunks
+ hydrate_mode = _LLM_MODE_TO_HYDRATE.get(
+ str(sel.get('mode', 'all')).strip().lower(), 'chunks'
+ )
+ chunks = await _hydrate_paths_to_rows(
+ db,
+ path_selections=[
+ {'path': path, 'confidence': conf, 'hydrate_mode': hydrate_mode}
+ ],
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ )
+ # Also hydrate connected targets (image/table chunks referenced via connect_to)
+ if chunks:
+ connected = await hydrate_connected_target_rows(
+ db=db,
+ rows=chunks,
+ exclude_document_ids=[],
+ exclude_sections=[],
+ )
+ if connected:
+ chunks = chunks + connected
+ node.leaf_content[path] = chunks
+ else:
+ # Non-leaf → return as pending for orchestrator to queue
+ pending.append(sel)
+
+ return node, pending
+
+ except Exception as e:
+ logger.error(f' scope_navigate_step failed for doc={document_id}: {e}')
+ return empty, []
+
+
+# ---------------------------------------------------------------------------
+# Tool: discovery_select_step (post-navigation discovery selection)
+# ---------------------------------------------------------------------------
+
+_MAX_DISCOVERY_PER_DOC = 3
+
+
+async def discovery_select_step(
+ db: AsyncSession,
+ *,
+ document_id: str,
+ query: str,
+ llm_fn: LLMFn,
+ user_id: str,
+ namespace: str,
+ doc_name: str = '',
+ discovery_hints: list[dict[str, Any]],
+) -> DocTreeNode:
+ """Post-navigation discovery selection step.
+
+ After BFS navigation exhausts for a document, present discovery-found
+ section paths (from bottom_discovery BM25) to the LLM for selection.
+ Selected paths are hydrated as leaf content.
+
+ For B-class documents (discovery-only, not KG-selected), this is the
+ only navigation step — no prior BFS.
+ """
+ from shared.services.retrieval.app_service import _hydrate_paths_to_rows
+
+ node = DocTreeNode(scope_path=None)
+ if not discovery_hints:
+ return node
+
+ # Limit hints per document
+ hints = discovery_hints[:_MAX_DISCOVERY_PER_DOC]
+
+ t0 = time.monotonic()
+ try:
+ # 1. Format hints for LLM
+ hint_lines: list[str] = []
+ hint_by_path: dict[str, dict] = {}
+ for h in hints:
+ sp = h.get('section_path', '')
+ if not sp:
continue
- confidence = item.get('confidence')
- if confidence is None:
- confidence = _default_confidence_for_rank(len(accepted))
- hydrate_mode = item.get('hydrate_mode', 'chunks')
- accepted.append({'path': path, 'confidence': confidence, 'hydrate_mode': hydrate_mode})
- if len(accepted) >= max_select:
- break
+ title = sp.rsplit(' / ', 1)[-1] if ' / ' in sp else sp
+ summary = h.get('summary', '') or ''
+ hint_lines.append(f'▸ path="{sp}" {title} [Leaf]')
+ if summary:
+ clipped = summary[:300]
+ hint_lines.append(f' {clipped}')
+ hint_by_path[sp] = h
+
+ if not hint_lines:
+ return node
+
+ items_text = '\n'.join(hint_lines)
+ prompt = _DISCOVERY_SELECT_PROMPT.format(
+ doc_name=doc_name or document_id,
+ items=items_text,
+ query=query,
+ )
+ response = await llm_fn(prompt)
+ selections = _parse_scope_nav_response(response)
- latency = int((time.monotonic() - t0) * 1000)
-
- if not accepted:
- return ToolResult(
- status='no_confident_match',
- payload={'document_id': document_id, 'reason': 'no path matches query intent'},
- latency_ms=latency,
+ logger.info(
+ f' discovery_select_step doc="{doc_name}": '
+ f'hints={len(hints)} selections={len(selections)}'
+ )
+
+ # 2. Hydrate selected paths
+ valid_selections = [s for s in selections if s['path'] in hint_by_path]
+ for sel in valid_selections:
+ path = sel['path']
+ conf = sel.get('confidence', 0.7)
+ hydrate_mode = _LLM_MODE_TO_HYDRATE.get(
+ str(sel.get('mode', 'all')).strip().lower(), 'chunks'
)
+ node.confidence[path] = conf
+
+ chunks = await _hydrate_paths_to_rows(
+ db,
+ path_selections=[
+ {'path': path, 'confidence': conf, 'hydrate_mode': hydrate_mode}
+ ],
+ user_id=user_id,
+ namespace=namespace,
+ document_id=document_id,
+ )
+ if chunks:
+ connected = await hydrate_connected_target_rows(
+ db=db,
+ rows=chunks,
+ exclude_document_ids=[],
+ exclude_sections=[],
+ )
+ if connected:
+ chunks = chunks + connected
+ node.leaf_content[path] = chunks
+ latency = int((time.monotonic() - t0) * 1000)
logger.info(
- f" agentic.scope_navigate: {len(accepted)} section paths selected, "
- f"status=selected_paths, overflowed={overflowed}"
+ f' discovery_select_step done: hydrated={len(node.leaf_content)} '
+ f'latency={latency}ms'
)
+ return node
- return ToolResult(
- status='selected_paths',
- payload={
- 'document_id': document_id,
- 'selected_paths': accepted,
- 'scope_path': scope_path,
- 'overflowed': overflowed,
- },
- latency_ms=latency,
- )
except Exception as e:
- latency = int((time.monotonic() - t0) * 1000)
- logger.error(f' agentic.scope_navigate failed for doc={document_id}: {e}')
- return ToolResult(
- status='error',
- payload={'document_id': document_id},
- error=str(e),
- latency_ms=latency,
- )
+ logger.error(f' discovery_select_step failed for doc={document_id}: {e}')
+ return node
diff --git a/packages/shared-python/shared/services/retrieval/agentic/trace.py b/packages/shared-python/shared/services/retrieval/agentic/trace.py
index ce320eef6..b504971eb 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/trace.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/trace.py
@@ -15,7 +15,7 @@
from loguru import logger
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.services.retrieval.agentic.types import ActionType, AgentRunConfig, ToolResult
+from shared.services.retrieval.agentic.types import AgentRunConfig, ToolResult
def _now_utc() -> datetime:
@@ -101,7 +101,7 @@ async def create_run(self) -> None:
def record_step(
self,
- action_type: ActionType,
+ action_type: str,
result: ToolResult,
*,
decision_reason: str = '',
@@ -109,7 +109,7 @@ def record_step(
"""Buffer a step record. Flushed on complete()."""
self._steps.append({
'step_index': len(self._steps),
- 'action_type': action_type.value,
+ 'action_type': action_type,
'action_input': {'decision_reason': decision_reason} if decision_reason else {},
'observation_status': result.status,
'observation_payload_keys': list(result.payload.keys()) if result.payload else [],
diff --git a/packages/shared-python/shared/services/retrieval/agentic/types.py b/packages/shared-python/shared/services/retrieval/agentic/types.py
index 4877886de..e9d5f404c 100644
--- a/packages/shared-python/shared/services/retrieval/agentic/types.py
+++ b/packages/shared-python/shared/services/retrieval/agentic/types.py
@@ -8,30 +8,15 @@
import time
from dataclasses import dataclass, field
-from enum import Enum
from typing import Any
-class ActionType(str, Enum):
- """Whitelisted action space for the retrieval agent."""
- BOTTOM_DISCOVERY = 'bottom_discovery'
- KG_DOCUMENT_SELECT = 'kg_document_select'
- DOCUMENT_PATH_SELECT = 'document_path_select'
- GREP_DOCUMENT_DISCOVER = 'grep_document_discover'
- GRAPH_EXPAND_DOCS = 'graph_expand_docs'
- DONE = 'done' # Explicit termination signal from LLMPolicy
-
-
@dataclass
class AgentRunConfig:
"""Budget and limit configuration for a single agent run."""
- max_steps: int = 10
- max_docs: int = 0 # 0 = no limit, LLM decides autonomously
- max_path_expansions: int = 2
- max_doc_retries: int = 2
- max_revisions: int = 2 # max attempt_answer → revise cycles
+ max_revisions: int = 2 # max attempt_answer → revision cycles
+ max_nav_depth: int = 3 # max scope_navigate recursion depth
latency_budget_ms: int = 12000
- min_evidence_paths: int = 1
@dataclass
@@ -47,6 +32,118 @@ class ToolResult:
error: str | None = None
+@dataclass
+class DocTreeNode:
+ """Unified navigation result tree for one document.
+
+ Produced by ``scope_navigate_step``. Captures the full
+ navigation outcome for rendering as a single hierarchy:
+
+ - ``outline_items``: section tree items at this scope level
+ - ``leaf_content``: hydrated chunk rows keyed by section path
+ (leaf selections from LLM)
+ - ``children``: recursive child trees keyed by section path
+ (non-leaf selections, populated by orchestrator BFS queue)
+ - ``confidence``: per-selection confidence for trimming
+ """
+ scope_path: str | None = None
+
+ # Outline items at this level (title + summary)
+ outline_items: list[dict[str, Any]] = field(default_factory=list)
+
+ # Leaf results, keyed by section path:
+ leaf_content: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
+ children: dict[str, 'DocTreeNode'] = field(default_factory=dict)
+
+ # Confidence per selection (for trimming)
+ confidence: dict[str, float] = field(default_factory=dict)
+
+ @staticmethod
+ def empty(scope_path: str | None = None) -> 'DocTreeNode':
+ return DocTreeNode(scope_path=scope_path)
+
+ def has_content(self) -> bool:
+ """Check if this tree has any meaningful content (outline, chunks, or children)."""
+ if self.outline_items:
+ return True
+ if self.leaf_content:
+ return True
+ if self.children:
+ return any(c.has_content() for c in self.children.values())
+ return False
+
+ def collect_all_paths(self, doc_id: str) -> set[str]:
+ """Recursively collect paths of actually-explored sections.
+
+ Only ``leaf_content`` and ``children`` represent sections whose
+ content was retrieved. ``outline_items`` are structural context
+ (titles/summaries shown to the LLM) and must NOT be masked —
+ otherwise revision rounds see 0 candidates and can't re-navigate.
+ """
+ paths: set[str] = set()
+ if self.scope_path:
+ paths.add(f'{doc_id}::{self.scope_path}')
+ for path in self.leaf_content:
+ paths.add(f'{doc_id}::{path}')
+ for path, child in self.children.items():
+ paths.add(f'{doc_id}::{path}')
+ paths.update(child.collect_all_paths(doc_id))
+ return paths
+
+ def flatten_chunk_rows(self) -> list[dict[str, Any]]:
+ """Recursively collect all hydrated chunk rows (document order)."""
+ rows: list[dict[str, Any]] = []
+ for chunks in self.leaf_content.values():
+ rows.extend(chunks)
+ for child in self.children.values():
+ rows.extend(child.flatten_chunk_rows())
+ return rows
+
+ def collect_referenced_ids(self) -> list[dict[str, str]]:
+ """Extract minimal chunk references from all hydrated leaves.
+
+ Returns deduplicated list of {chunk_id, document_id, chunk_type,
+ section_path, file_path, job_id} for hit stats and frontend display.
+ """
+ refs: list[dict[str, str]] = []
+ seen: set[str] = set()
+ for row in self.flatten_chunk_rows():
+ cid = row.get('chunk_id', '')
+ if cid and cid not in seen:
+ seen.add(cid)
+ refs.append({
+ 'chunk_id': cid,
+ 'document_id': row.get('document_id', ''),
+ 'chunk_type': row.get('chunk_type', ''),
+ 'section_path': row.get('section_path', ''),
+ 'file_path': row.get('file_path', ''),
+ 'job_id': row.get('job_id', ''),
+ })
+ return refs
+
+ def merge(self, other: 'DocTreeNode') -> None:
+ """Additive merge for revision cycles.
+
+ Merges outline items, leaf content, children, and confidence from
+ ``other`` into this node. Existing data is preserved; new data is
+ added. For confidence values, the higher value wins.
+ """
+ existing_paths = {item['path'] for item in self.outline_items}
+ for item in other.outline_items:
+ if item.get('path', '') not in existing_paths:
+ self.outline_items.append(item)
+ for path, chunks in other.leaf_content.items():
+ if path not in self.leaf_content:
+ self.leaf_content[path] = chunks
+ for path, child in other.children.items():
+ if path in self.children:
+ self.children[path].merge(child)
+ else:
+ self.children[path] = child
+ for path, conf in other.confidence.items():
+ self.confidence[path] = max(self.confidence.get(path, 0), conf)
+
+
@dataclass
class CandidateDoc:
"""A document selected by kg_document_select."""
@@ -54,199 +151,56 @@ class CandidateDoc:
source_file_name: str = ''
confidence: float = 0.0
reason: str = ''
- source: str = '' # 'kg_llm_select' | 'grep' | 'edge_expand'
+ source: str = '' # 'kg_llm_select' | 'grep' | 'edge_expand' | 'discovery_hint'
+
+
+@dataclass
+class AgenticResult:
+ """Output of agentic retrieval.
+
+ - ``evidence_text``: complete hierarchical context for LLM answering
+ (rendered doc tree with outline + leaf content + inline tables)
+ - ``answer_text``: LLM-generated answer to the query based on the
+ evidence. Empty string when the evidence was insufficient
+ (NOT_FOUND) and max revisions were exhausted.
+ - ``referenced_chunks``: minimal chunk references for hit stats
+ and frontend display (chunk_id, document_id, chunk_type, etc.)
+ - ``router_used``: routing path identifier
+ """
+ evidence_text: str
+ answer_text: str = ''
+ referenced_chunks: list[dict[str, str]] = field(default_factory=list)
+ router_used: str = 'agentic_discovery_only'
@dataclass
class AgentState:
- """Mutable state carried through the agent loop.
+ """Mutable state carried through the 2-phase orchestrator.
- Updated by ``state.apply(action_type, tool_result)`` after each step.
+ Phase 1: Document selection (discovery + KG)
+ Phase 2: Per-document navigation (scope_navigate_step per doc)
+ Phase 3: Assembly + final verdict
"""
# Timing
start_time: float = field(default_factory=time.monotonic)
step_count: int = 0
- # Discovery results (from bottom_discovery)
- discovery_paths: list[dict[str, Any]] = field(default_factory=list)
+ # Phase 1: Discovery
discovery_top_doc_ids: list[str] = field(default_factory=list)
- discovery_done: bool = False
- # KG document selection
+ # Phase 1: KG document selection
selected_docs: list[CandidateDoc] = field(default_factory=list)
- excluded_doc_ids: set[str] = field(default_factory=set)
- doc_retry_count: int = 0
- kg_done: bool = False
- pending_doc_index: int = 0
-
- # Document path selection
- selected_paths: list[dict[str, Any]] = field(default_factory=list)
- path_expansion_count: int = 0
-
- # Last observation (used by policy to decide next action)
- last_observation: ToolResult | None = None
-
- # Context passed through
doc_id_to_name: dict[str, str] = field(default_factory=dict)
doc_job_map: dict[str, str] = field(default_factory=dict)
- # Revision / three-state fields
+ # Phase 2: Per-document navigation results
+ doc_trees: dict[str, DocTreeNode] = field(default_factory=dict) # doc_id → DocTreeNode
+
+ # Revision state
revision_count: int = 0
ever_explored_doc_ids: set[str] = field(default_factory=set)
seen_section_keys: set[str] = field(default_factory=set) # "{doc_id}::{section_path}"
- kept_path_rows: list[dict[str, Any]] = field(default_factory=list)
@property
def elapsed_ms(self) -> int:
return int((time.monotonic() - self.start_time) * 1000)
-
- def state_summary(self) -> dict[str, Any]:
- """Produce a concise state snapshot for the LLMPolicy prompt.
-
- Only includes fields the LLM needs for decision-making — never
- exposes raw chunk content or internal IDs verbatim.
- """
- selected_doc_summaries = [
- {
- 'document_id': d.document_id,
- 'name': d.source_file_name or '(unnamed)',
- 'confidence': round(d.confidence, 2),
- 'source': d.source,
- }
- for d in self.selected_docs
- ]
- selected_path_summaries = [
- {
- 'path': p.get('path', ''),
- 'confidence': round(float(p.get('confidence', 0.0) or 0.0), 2),
- }
- for p in self.selected_paths[:10] # cap to avoid huge prompts
- ]
- last_obs = None
- if self.last_observation:
- last_obs = {
- 'status': self.last_observation.status,
- 'payload_keys': list(self.last_observation.payload.keys()),
- 'error': self.last_observation.error,
- }
- return {
- 'step': self.step_count,
- 'discovery_done': self.discovery_done,
- 'discovery_candidates': len(self.discovery_paths),
- 'discovery_top_doc_ids': self.discovery_top_doc_ids[:5],
- 'kg_done': self.kg_done,
- 'selected_docs': selected_doc_summaries,
- 'pending_doc_index': self.pending_doc_index,
- 'selected_paths_count': len(self.selected_paths),
- 'selected_paths': selected_path_summaries,
- 'doc_retry_count': self.doc_retry_count,
- 'revision_count': self.revision_count,
- 'explored_doc_count': len(self.ever_explored_doc_ids),
- 'kept_rows_count': len(self.kept_path_rows),
- 'last_observation': last_obs,
- }
-
- def apply(self, action_type: ActionType, result: ToolResult) -> None:
- """Update state based on action and its result."""
- self.last_observation = result
-
- if action_type == ActionType.BOTTOM_DISCOVERY:
- self.discovery_done = True
- if result.status != 'error':
- self.discovery_paths = result.payload.get('fused_rows', [])
- self.discovery_top_doc_ids = result.payload.get('top_doc_ids', [])
-
- elif action_type == ActionType.KG_DOCUMENT_SELECT:
- self.kg_done = True
- if result.status == 'selected_docs':
- new_docs = result.payload.get('candidate_docs', [])
- for doc_data in new_docs:
- if isinstance(doc_data, CandidateDoc):
- self.selected_docs.append(doc_data)
- elif isinstance(doc_data, dict):
- self.selected_docs.append(CandidateDoc(
- document_id=doc_data.get('document_id', ''),
- source_file_name=doc_data.get('source_file_name', ''),
- confidence=doc_data.get('confidence', 0.0),
- reason=doc_data.get('reason', ''),
- source=doc_data.get('source', ''),
- ))
- self.doc_id_to_name.update(result.payload.get('doc_id_to_name', {}))
- self.doc_job_map.update(result.payload.get('doc_job_map', {}))
-
- # Merge discovery hints
- existing_ids = {d.document_id for d in self.selected_docs}
- for did in self.discovery_top_doc_ids:
- if did not in existing_ids and did not in self.excluded_doc_ids:
- self.selected_docs.append(CandidateDoc(
- document_id=did,
- source_file_name=self.doc_id_to_name.get(did, ''),
- confidence=0.5, # Lower than LLM's 0.8
- reason='Bottom discovery hit',
- source='discovery_hint',
- ))
- existing_ids.add(did)
-
- elif action_type == ActionType.DOCUMENT_PATH_SELECT:
- if result.status == 'selected_paths':
- doc_id = result.payload.get('document_id', '')
- new_paths = result.payload.get('selected_paths', [])
- for p in new_paths:
- p['document_id'] = doc_id
- self.selected_paths.extend(new_paths)
- self.pending_doc_index += 1
- if doc_id:
- self.ever_explored_doc_ids.add(doc_id)
- elif result.status == 'no_items':
- doc_id = result.payload.get('document_id', '')
- self.pending_doc_index += 1
- if doc_id:
- self.ever_explored_doc_ids.add(doc_id)
- elif result.status == 'need_more_docs':
- failed_doc_id = result.payload.get('document_id', '')
- if failed_doc_id:
- self.excluded_doc_ids.add(failed_doc_id)
- self.doc_retry_count += 1
- self.kg_done = False # allow re-entry to KG select
- elif result.status == 'no_confident_match':
- doc_id = result.payload.get('document_id', '')
- self.pending_doc_index += 1
- if doc_id:
- self.ever_explored_doc_ids.add(doc_id)
- elif result.status == 'error':
- doc_id = result.payload.get('document_id', '')
- self.pending_doc_index += 1
- if doc_id:
- self.ever_explored_doc_ids.add(doc_id)
-
- elif action_type == ActionType.GREP_DOCUMENT_DISCOVER:
- if result.status == 'discovered_docs':
- grep_doc_ids = result.payload.get('document_ids', [])
- doc_id_to_name = result.payload.get('doc_id_to_name', {})
- for did in grep_doc_ids:
- if did not in self.excluded_doc_ids:
- existing_ids = {d.document_id for d in self.selected_docs}
- if did not in existing_ids:
- self.selected_docs.append(CandidateDoc(
- document_id=did,
- source_file_name=doc_id_to_name.get(did, ''),
- source='grep',
- ))
- self.doc_id_to_name.update(doc_id_to_name)
- self.doc_job_map.update(result.payload.get('doc_job_map', {}))
-
- elif action_type == ActionType.GRAPH_EXPAND_DOCS:
- if result.status == 'expanded_docs':
- expanded_ids = result.payload.get('document_ids', [])
- doc_id_to_name = result.payload.get('doc_id_to_name', {})
- for did in expanded_ids:
- if did not in self.excluded_doc_ids:
- existing_ids = {d.document_id for d in self.selected_docs}
- if did not in existing_ids:
- self.selected_docs.append(CandidateDoc(
- document_id=did,
- source_file_name=doc_id_to_name.get(did, ''),
- source='edge_expand',
- ))
- self.doc_id_to_name.update(doc_id_to_name)
- self.doc_job_map.update(result.payload.get('doc_job_map', {}))
diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py
index b91e31f0c..93f6d6d03 100644
--- a/packages/shared-python/shared/services/retrieval/app_service.py
+++ b/packages/shared-python/shared/services/retrieval/app_service.py
@@ -430,6 +430,15 @@ async def _to_public_response(response: dict[str, Any]) -> dict[str, Any]:
'router_used': response.get('router_used'),
'results': [],
}
+
+ # 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:
+ public_response['referenced_chunks'] = response['referenced_chunks']
+
public_results: list[dict[str, Any]] = []
for row in response.get('results', []):
artifact_ref = row.get('file_path')
@@ -450,7 +459,10 @@ async def _to_public_response(response: dict[str, Any]) -> dict[str, Any]:
public_row['asset_url'] = asset_url
elif field in row:
public_row[field] = row[field]
- public_row['source'] = _to_public_source(row)
+ if 'source' in row:
+ public_row['source'] = row['source']
+ else:
+ public_row['source'] = _to_public_source(row)
public_results.append(public_row)
public_response['results'] = public_results
@@ -496,6 +508,52 @@ def _normalize_row_scores(
row[target_field] = round((raw_score - min_score) / denominator, 6)
+def _importance_multiplier(
+ rows: list[dict[str, Any]],
+ *,
+ raw_field: str = 'importance_raw_score',
+ low: float = 0.1,
+ high: float = 2.0,
+) -> None:
+ """Apply adaptive sigmoid-based importance boost to agent/discovery scores.
+
+ Uses median of ``raw_field`` as center and IQR as spread so the curve
+ adapts to any KB size without hard-coded thresholds. When all values
+ are identical (IQR ≈ 0) the multiplier is 1.0 (neutral).
+
+ Output range ``[low, high]`` — default [0.1, 2.0] — is the only
+ configured constant: max 2× boost, min 10%. The function modifies
+ ``agent_score`` and ``discovery_score`` **in place**.
+ """
+ import math
+
+ if not rows:
+ return
+
+ values = sorted(float(r.get(raw_field, 0.0) or 0.0) for r in rows)
+ n = len(values)
+ median = values[n // 2] if n % 2 else (values[n // 2 - 1] + values[n // 2]) / 2
+ q1 = values[n // 4] if n >= 4 else values[0]
+ q3 = values[3 * n // 4] if n >= 4 else values[-1]
+ iqr = q3 - q1
+
+ for row in rows:
+ raw = float(row.get(raw_field, 0.0) or 0.0)
+ if iqr <= 1e-9:
+ mult = 1.0
+ else:
+ z = (raw - median) / iqr
+ s = 1.0 / (1.0 + math.exp(-z))
+ mult = low + (high - low) * s
+ row['importance_multiplier'] = round(mult, 4)
+ row['agent_score'] = round(
+ float(row.get('agent_score', 0.0) or 0.0) * mult, 6,
+ )
+ row['discovery_score'] = round(
+ float(row.get('discovery_score', 0.0) or 0.0) * mult, 6,
+ )
+
+
async def _load_chunk_importance_scores(
db: AsyncSession,
*,
@@ -548,8 +606,7 @@ def _rank_candidates_by_path(
candidate = dict(row)
candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0)
candidate['agent_score'] = 0.0
- candidate['importance_raw_score'] = float(row.get('importance_raw_score', 0.0) or 0.0)
- candidate['importance_norm_score'] = float(row.get('importance_norm_score', 0.0) or 0.0)
+ candidate.setdefault('hydrate_mode', 'chunks')
merged[key] = candidate
insertion_order[key] = counter
counter += 1
@@ -563,22 +620,12 @@ def _rank_candidates_by_path(
candidate = dict(row)
candidate['discovery_score'] = float(row.get('discovery_score', 0.0) or 0.0)
candidate['agent_score'] = routed_agent_score
- candidate['importance_raw_score'] = float(row.get('importance_raw_score', 0.0) or 0.0)
- candidate['importance_norm_score'] = float(row.get('importance_norm_score', 0.0) or 0.0)
merged[key] = candidate
insertion_order[key] = counter
counter += 1
continue
candidate = merged[key]
candidate['agent_score'] = max(float(candidate.get('agent_score', 0.0) or 0.0), routed_agent_score)
- candidate['importance_raw_score'] = max(
- float(candidate.get('importance_raw_score', 0.0) or 0.0),
- float(row.get('importance_raw_score', 0.0) or 0.0),
- )
- candidate['importance_norm_score'] = max(
- float(candidate.get('importance_norm_score', 0.0) or 0.0),
- float(row.get('importance_norm_score', 0.0) or 0.0),
- )
if not candidate.get('source_chunk_path') and row.get('source_chunk_path'):
candidate['source_chunk_path'] = row.get('source_chunk_path')
if not candidate.get('section_path') and row.get('section_path'):
@@ -587,18 +634,16 @@ def _rank_candidates_by_path(
# ── Dual-priority ranking ────────────────────────────────────────────
# When the agent produced results (routed_rows non-empty), rows with
# agent_score=0 are demoted to a fallback pool. Primary sort is by
- # agent_score (LLM confidence, 0-1, cross-round comparable), with
- # discovery_score as tiebreaker only. This avoids the old
- # `max(agent, discovery)` which mixed incompatible score sources.
+ # agent_score (includes importance boost from _importance_multiplier),
+ # with discovery_score as tiebreaker.
has_agent_results = len(routed_rows) > 0
primary_rows: list[dict[str, Any]] = []
fallback_rows: list[dict[str, Any]] = []
for key, row in merged.items():
- discovery_score = float(row.get('discovery_score', 0.0) or 0.0)
agent_score = float(row.get('agent_score', 0.0) or 0.0)
- row['dual_hit_flag'] = 1 if discovery_score > 0.0 and agent_score > 0.0 else 0
+ discovery_score = float(row.get('discovery_score', 0.0) or 0.0)
row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6)
row['score'] = row['evidence_score']
row['_candidate_order'] = insertion_order[key]
@@ -612,8 +657,6 @@ def _sort_key(row):
return (
float(row.get('agent_score', 0.0) or 0.0),
float(row.get('discovery_score', 0.0) or 0.0),
- int(row.get('dual_hit_flag', 0) or 0),
- float(row.get('importance_norm_score', 0.0) or 0.0),
-int(row.get('_candidate_order', 0) or 0),
)
@@ -716,9 +759,14 @@ async def _hydrate_paths_to_rows(
path_selections: list[dict[str, Any]],
user_id: str,
namespace: str,
+ document_id: str | None = None,
) -> list[dict[str, Any]]:
"""Load full chunk rows by section_path or source_chunk_path.
+ When *document_id* is provided the query is scoped to that single
+ document, preventing cross-document collisions on generic paths
+ such as ``Root``.
+
Supports hydrate_mode branching:
- 'chunks' (default): all chunk types under the section subtree
- 'outline': synthetic row from section metadata, no real chunks
@@ -770,6 +818,8 @@ async def _hydrate_paths_to_rows(
.where(Document.status == 'active')
.where(or_(*outline_section_filters))
)
+ if document_id:
+ outline_stmt = outline_stmt.where(Document.document_id == document_id)
outline_result = await db.execute(outline_stmt)
for document, section in outline_result.all():
agent_score = confidence_by_path.get(section.section_path, 0.0)
@@ -820,6 +870,8 @@ async def _hydrate_paths_to_rows(
)
)
)
+ if document_id:
+ stmt = stmt.where(Document.document_id == document_id)
result = await db.execute(stmt)
# Build a map of path → allowed chunk_types based on hydrate_mode
@@ -1048,7 +1100,7 @@ async def run_retrieval_query(
llm_fn = _create_llm()
agent = RetrievalAgent()
- ranked_rows, router_used = await agent.run(
+ agentic_result = await agent.run(
db,
user_id=user_id,
namespace=namespace,
@@ -1063,6 +1115,83 @@ 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
+ enriched_refs: list[dict[str, Any]] = []
+ for ref in agentic_result.referenced_chunks:
+ enriched = dict(ref)
+ chunk_type = _normalize_chunk_type(ref.get('chunk_type'))
+ artifact_ref = ref.get('file_path', '')
+ job_id = ref.get('job_id', '')
+ if chunk_type in _MEDIA_CHUNK_TYPES and job_id and _is_client_result_artifact_ref(artifact_ref):
+ try:
+ asset_url = await generate_retrieval_asset_url(
+ job_id=str(job_id), artifact_ref=str(artifact_ref),
+ )
+ if asset_url:
+ enriched['asset_url'] = asset_url
+ except Exception as e:
+ 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,
+ }
+
+ if cache_version is not None:
+ try:
+ await set_cached_retrieval_query_result(
+ user_id=user_id, namespace=namespace, version=cache_version,
+ query=query, top_k=top_k,
+ exclude_document_ids=exclude_document_ids,
+ exclude_sections=exclude_sections,
+ response=response, **cache_extra,
+ )
+ except Exception as e:
+ logger.warning(f"Failed to write retrieval cache (ignored): {e}")
+
+ try:
+ schedule_retrieval_hit_stats_update(
+ user_id=user_id, namespace=namespace,
+ results=agentic_result.referenced_chunks,
+ )
+ except Exception as e:
+ logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}")
+
+ elapsed_total = round((time.monotonic() - t_start) * 1000)
+ logger.info(
+ 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'{"█" * 70}'
+ )
+
+ return await _to_public_response(response)
else:
# ── LEGACY path (existing code, unchanged) ──
@@ -1208,17 +1337,6 @@ async def run_retrieval_query(
chunk_importance_scores = {}
for row in combined_rows:
row['importance_raw_score'] = float(chunk_importance_scores.get(str(row.get('chunk_id') or ''), 0.0) or 0.0)
- positive_importance = [row['importance_raw_score'] for row in combined_rows if row['importance_raw_score'] > 0.0]
- if positive_importance:
- _normalize_row_scores(
- combined_rows,
- source_field='importance_raw_score',
- target_field='importance_norm_score',
- default=0.5,
- )
- else:
- for row in combined_rows:
- row['importance_norm_score'] = 0.0
ranked_rows = _rank_candidates_by_path(fused_rows, agent_rows, top_k)
if ranked_rows:
@@ -1227,8 +1345,6 @@ async def run_retrieval_query(
logger.info(
' '
f'[{i}] evidence={row.get("evidence_score", 0.0):.4f} '
- f'dual_hit={row.get("dual_hit_flag", 0)} '
- f'importance={row.get("importance_norm_score", 0.0):.4f} '
f'discovery={row.get("discovery_score", 0.0):.4f} '
f'agent={row.get("agent_score", 0.0):.4f} '
f'path={_get_row_path(row)}'
diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py
index 4a8b6e70b..8d9a11c22 100644
--- a/packages/shared-python/shared/services/retrieval/llm_adapter.py
+++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py
@@ -6,13 +6,15 @@
from __future__ import annotations
import asyncio
-from typing import Any, Callable, Coroutine
+from typing import Any, Callable, Coroutine, Union, Sequence, cast
from loguru import logger
from shared.core.config import settings
-LLMFn = Callable[[str], Coroutine[Any, Any, str]]
+# LLMFn accepts either a plain string or a list of ChatCompletionMessageParam
+LLMFnInput = Union[str, Sequence[dict[str, Any]]]
+LLMFn = Callable[[LLMFnInput], Coroutine[Any, Any, str]]
_RETRIEVAL_LLM_TEMPERATURE = 0.1
_RETRIEVAL_LLM_MAX_TOKENS = 2048
@@ -63,14 +65,14 @@ def create_retrieval_llm_fn(
effective_model = model or _resolve_default_model()
- async def llm_fn(prompt: str) -> str:
+ async def llm_fn(prompt: LLMFnInput) -> str:
from shared.utils.OpenAICompatibleClientSync import get_openai_client
client = get_openai_client(model=effective_model)
try:
result = await asyncio.to_thread(
client.chat_completion,
- prompt,
+ cast(Any, prompt),
model=effective_model,
temperature=temperature,
max_tokens=max_tokens,