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