diff --git a/apps/api/.env.example b/apps/api/.env.example index 54c9a9709..4a01c6d8d 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -85,6 +85,14 @@ NORMOL_MODEL=deepseek-chat HIERARCHY_LLM_MODEL=qwen3.6-flash IMAGE_MODEL=qwen3.5-flash IMAGE_MODEL_MAX=qwen3.5-flash +RETRIEVAL_DECOMPOSITION_ENABLED=false +RETRIEVAL_PLANNER_MODEL= +RETRIEVAL_PLANNER_THINKING_BUDGET=4000 +RETRIEVAL_DECOMPOSITION_MAX_STEPS=5 +RETRIEVAL_WALLET_TOTAL_BUDGET=200000 +RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET=40000 +RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET=6000 +RETRIEVAL_WORKFLOW_PARALLEL_MAX=3 # 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 e0106f0bc..aaee0b7bc 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -52,6 +52,11 @@ class RetrievalQueryRequest(BaseModel): internal_recall_k: int | None = Field( None, ge=1, description="Override per-channel recall count" ) + enable_decomposition: bool | None = Field( + None, + description="Deprecated: agentic mode now always uses workflow decomposition. This field is ignored.", + deprecated=True, + ) @field_validator("channels") @classmethod @@ -63,7 +68,35 @@ def validate_channels(cls, v: list[str]) -> list[str]: return v -@router.post("/query") +class WorkflowStepResponse(BaseModel): + step_id: str + sub_query: str + step_kind: Literal["retrieve", "synthesize"] + depends_on: list[str] + output_role: str + status: Literal["done", "skipped", "error", "budget_stop"] + answer_text: str + evidence_text: str | None = None + referenced_chunks: list[dict] = Field(default_factory=list) + budget_snapshot: dict | None = None + child_run_id: str | None = None + + +class RetrievalQueryResponse(BaseModel): + namespace: str + query: str + router_used: str + answer_text: str | None = None + referenced_chunks: list[dict] = Field(default_factory=list) + results: list[dict] = Field(default_factory=list) + plan: dict | None = None + steps: list[WorkflowStepResponse] | None = None + final_strategy_used: str | None = None + wallet_snapshot: dict | None = None + planner_snapshot: dict | None = None + + +@router.post("/query", response_model=RetrievalQueryResponse) async def query_retrieval( payload: RetrievalQueryRequest, current_user: CurrentUser = Depends(with_current_user), @@ -85,4 +118,5 @@ async def query_retrieval( rerank=payload.rerank, threshold=payload.threshold, internal_recall_k=payload.internal_recall_k, + enable_decomposition=payload.enable_decomposition, ) diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 2deca35fe..fe528c236 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -182,6 +182,13 @@ async def test_should_return_empty_results_for_an_empty_query( "query": "", "router_used": "empty_query_filtered", "results": [], + "answer_text": None, + "final_strategy_used": None, + "plan": None, + "planner_snapshot": None, + "referenced_chunks": [], + "steps": None, + "wallet_snapshot": None, } diff --git a/apps/worker/.env.example b/apps/worker/.env.example index 49bb7aae4..6c46d9df7 100644 --- a/apps/worker/.env.example +++ b/apps/worker/.env.example @@ -85,6 +85,14 @@ NORMOL_MODEL=deepseek-chat HIERARCHY_LLM_MODEL=deepseek-chat IMAGE_MODEL=qwen3.5-flash IMAGE_MODEL_MAX=qwen3.5-flash +RETRIEVAL_DECOMPOSITION_ENABLED=false +RETRIEVAL_PLANNER_MODEL= +RETRIEVAL_PLANNER_THINKING_BUDGET=4000 +RETRIEVAL_DECOMPOSITION_MAX_STEPS=5 +RETRIEVAL_WALLET_TOTAL_BUDGET=200000 +RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET=40000 +RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET=6000 +RETRIEVAL_WORKFLOW_PARALLEL_MAX=3 # Required for specific features: billing and analytics BILLING_ENABLED=false diff --git a/packages/shared-python/shared/core/config/ai.py b/packages/shared-python/shared/core/config/ai.py index 2b25b5571..2164df3d2 100644 --- a/packages/shared-python/shared/core/config/ai.py +++ b/packages/shared-python/shared/core/config/ai.py @@ -35,6 +35,38 @@ class AIConfig(BaseModel): default="qwen3.5-flash", description="Higher-capability image model for OCR and ask-image Q&A", ) + RETRIEVAL_DECOMPOSITION_ENABLED: bool = Field( + default=False, + description="Enable query-decomposition workflow before agentic retrieval.", + ) + RETRIEVAL_PLANNER_MODEL: str = Field( + default="", + description="Reasoning-capable model used by the workflow query planner.", + ) + RETRIEVAL_PLANNER_THINKING_BUDGET: int = Field( + default=4000, + description="Token budget for the query planner thinking call.", + ) + RETRIEVAL_DECOMPOSITION_MAX_STEPS: int = Field( + default=5, + description="Maximum number of planned workflow steps.", + ) + RETRIEVAL_WALLET_TOTAL_BUDGET: int = Field( + default=200000, + description="Total workflow token wallet for decomposed retrieval.", + ) + RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET: int = Field( + default=40000, + description="Default token budget issued to each retrieve step.", + ) + RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET: int = Field( + default=6000, + description="Default token budget issued to each synthesize step.", + ) + RETRIEVAL_WORKFLOW_PARALLEL_MAX: int = Field( + default=3, + description="Maximum concurrent workflow steps in the same DAG batch.", + ) # Runtime LLM controls. LLM_MOCK_ENABLED: bool = Field( diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 8b5ba0a23..41bed45bb 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -380,6 +380,9 @@ class RetrievalRun(Base): result_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) final_doc_ids: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True) result_provenance: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + parent_run_id: Mapped[Optional[str]] = mapped_column(String(36), nullable=True, index=True) + workflow_step_id: Mapped[Optional[str]] = mapped_column(String(64), nullable=True) + workflow_plan: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True) latency_ms: Mapped[int] = mapped_column(Integer, nullable=False, default=0) token_count: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) error: Mapped[Optional[str]] = mapped_column(Text, nullable=True) diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index a25d45cdc..832e81eb6 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -8,10 +8,11 @@ ) from .graph_service import DocumentGraphService, GraphQueryService, GraphScope from .hit_stats_service import record_retrieval_hits -from .llm_adapter import create_retrieval_llm_fn +from .llm_adapter import create_retrieval_llm_fn, create_retrieval_planner_fn __all__ = [ "create_retrieval_llm_fn", + "create_retrieval_planner_fn", "run_retrieval_query", "merge_channels_rrf", "DocumentGraphService", diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py index 3f7a2d316..d5b1bed78 100644 --- a/packages/shared-python/shared/services/retrieval/agent_navigate.py +++ b/packages/shared-python/shared/services/retrieval/agent_navigate.py @@ -38,9 +38,35 @@ """ -_SCOPE_NAV_PROMPT = """\ + +_DISCOVERY_SELECT_PROMPT = """\ You are a document navigation assistant. +Document: "{doc_name}" + +{budget_block} +After navigating the document's section tree, the following section paths +were additionally discovered via keyword and semantic search. +They may contain relevant evidence not found through hierarchical navigation. + +=== Discovery Candidates === +{items} +=== End Discovery Candidates === + +User query: {query} +{revision_context} +Select section paths whose content is needed to answer the query. +If none are relevant, return an EMPTY list []. + +Return ONLY a JSON object: +{{"selections": [{{"path": "...", "confidence": }}, ...]}} +Do not include any explanation. +""" + + +_ACTION_PROMPT = """\ +You are a document navigation agent. + Document: "{doc_name}" (id: {doc_id}) {budget_block} @@ -56,44 +82,106 @@ User query: {query} -Select sections to drill into for more detailed content. -- You may ONLY select sections marked with [SELECT]. Do NOT select any other sections. -- Select sections whose content is needed to answer the query. -- If the titles and summaries already visible are sufficient (e.g. the query asks for an outline or overview), return an EMPTY list []. -- When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. -- When budget is CRITICAL, be very selective — only pick paths with strong relevance. Return [] if current evidence already suffices. +=== Available Actions === + +Choose ONE action: + +NAVIGATE — Drill into selected sections for detailed content. + Consider this when the query targets specific topics and you need deeper text evidence. + Select one or more [SELECT] sections. + +STOP — Current scope evidence is sufficient. No further drill-down. + Consider this when: + - The query asks for an outline, overview, or summary + - The query is broad/global, the tree section can fulfill it without drilling into individual sections. + - You have already collected enough evidence at this level. + +{tools_block} + +When action is NAVIGATE, provide selections: +- You may ONLY select sections marked with [SELECT]. + +When action is STOP, selections must be empty. Return ONLY a JSON object: -{{"selections": [{{"path": "...", "confidence": }}, ...]}} +{{"action": "NAVIGATE", "tools": [...], "selections": [{{"path": "...", "confidence": }}, ...]}} +or +{{"action": "STOP", "tools": [...], "selections": []}} Do not include any explanation. """ -_DISCOVERY_SELECT_PROMPT = """\ -You are a document navigation assistant. +def _parse_action_response(text: str) -> dict: + """Parse the unified action response from LLM. -Document: "{doc_name}" + Returns dict with keys: + action: 'NAVIGATE' | 'STOP' + tools: list[str] (subset of FIND_IMAGES, FIND_TABLES) + selections: list[dict] (each has 'path' and optional 'confidence') -{budget_block} -After navigating the document's section tree, the following section paths -were additionally discovered via keyword and semantic search. -They may contain relevant evidence not found through hierarchical navigation. + When action is STOP, selections are forced to empty. + """ + import json as _json + import re as _re -=== Discovery Candidates === -{items} -=== End Discovery Candidates === + text = text.strip() + _ASSET_TOOLS = {'FIND_IMAGES', 'FIND_TABLES'} + default = {'action': 'NAVIGATE', 'tools': [], 'selections': []} -User query: {query} -{revision_context} -Select section paths whose content is needed to answer the query. -If none are relevant, return an EMPTY list []. -When budget is TIGHT, prefer fewer high-confidence candidates. -When budget is CRITICAL, be very selective — only pick paths with strong relevance. Return [] if evidence suffices. + def _extract(data: dict) -> dict: + action = str(data.get('action', 'NAVIGATE')).strip().upper() + if action not in ('NAVIGATE', 'STOP'): + action = 'NAVIGATE' -Return ONLY a JSON object: -{{"selections": [{{"path": "...", "confidence": }}, ...]}} -Do not include any explanation. -""" + tools_val = data.get('tools') or [] + if isinstance(tools_val, list): + tools = [str(t).strip().upper() for t in tools_val if str(t).strip().upper() in _ASSET_TOOLS] + else: + tools = [] + + # STOP → no selections allowed + if action == 'STOP': + return {'action': action, 'tools': tools, 'selections': []} + + selections_val = data.get('selections') or [] + selections = [] + if isinstance(selections_val, list): + for s in selections_val: + if isinstance(s, dict) and s.get('path'): + conf = _normalize_confidence(s.get('confidence', 0.7)) + selections.append({'path': str(s['path']), 'confidence': conf or 0.7}) + + return {'action': action, 'tools': tools, 'selections': selections} + + # Try JSON parse + try: + data = _json.loads(text) + if isinstance(data, dict): + return _extract(data) + except (ValueError, _json.JSONDecodeError): + pass + + # Try extracting JSON from markdown fences + fence_match = _re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, _re.DOTALL) + if fence_match: + try: + data = _json.loads(fence_match.group(1).strip()) + if isinstance(data, dict): + return _extract(data) + except (ValueError, _json.JSONDecodeError): + pass + + # Try finding a JSON object anywhere + brace_match = _re.search(r'\{.*\}', text, _re.DOTALL) + if brace_match: + try: + data = _json.loads(brace_match.group()) + if isinstance(data, dict): + return _extract(data) + except (ValueError, _json.JSONDecodeError): + pass + + return default def _format_budget_block(snapshot: dict | None) -> str: @@ -287,7 +375,7 @@ def _render_item(item: dict, include_summary: bool) -> str: indent = " " * (level - 1) prefix = '▸' if level == 1 else '└' level_tag = f'[L{level}]' - select_tag = '[SELECT] ' if show else '' + select_tag = '[SELECT] ' if item.get('selectable', False) else '' lines: list[str] = [] lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}') @@ -459,7 +547,7 @@ async def _load_child_sections( db: AsyncSession, document_id: str, job_result_id: str, - scope_path: str | None = None, + scope_path: str | list[str] | None = None, exclude_paths: set[str] | None = None, ) -> list[dict]: """Load the Continuous Context Tree for *scope_path*. @@ -468,17 +556,16 @@ async def _load_child_sections( {path, title, summary, chunk_count, image_count, table_count, level, show_summary, is_leaf} - The tree contains three categories of nodes: - 1. Ancestors of scope_path + their siblings → show_summary=False (title only) - 2. Children of scope_path (2 depth bands) → show_summary=True (with summary) - 3. Everything else → pruned (not returned) - - When scope_path is None (root), all items are category 2. + scope_path can be: + - None: root scope, all items are selectable (2 depth bands). + - str: single scope, descendants are selectable. + - list[str]: multi-scope, descendants of ALL paths are selectable + simultaneously — used when the LLM selected multiple drill-down + paths in the previous step. - level: absolute depth in the document (1-based) - show_summary: controls whether _format_items_for_llm renders summary - - exclude_paths: paths already seen in prior revision rounds; - any path matching (exact or subtree) is skipped from category 2 + - exclude_paths: paths already hydrated; skipped from selectable items """ # ── Fetch all sections for this document revision ──────────────────── stmt = ( @@ -497,12 +584,20 @@ async def _load_child_sections( if not section_rows: return [] - scope = normalize_section_path(scope_path) if scope_path else '' - scope_parts = split_section_path(scope) - scope_depth = len(scope_parts) + # ── Normalize scope(s) ─────────────────────────────────────────────── + # Multi-scope: list of paths to expand simultaneously + if isinstance(scope_path, list): + scope_list = [normalize_section_path(p) for p in scope_path] + elif scope_path: + scope_list = [normalize_section_path(scope_path)] + else: + scope_list = [] # root + + # For logging, derive representative scope info + scope_depth = len(split_section_path(scope_list[0])) if scope_list else 0 logger.debug( - f' _load_child_sections: scope={scope!r} scope_parts={scope_parts} ' + f' _load_child_sections: scopes={scope_list or ["root"]} ' f'scope_depth={scope_depth} exclude_paths={_excl if (_excl := exclude_paths or set()) else "none"} ' f'total_sections={len(section_rows)}' ) @@ -524,134 +619,105 @@ async def _load_child_sections( } # ── Build the set of ancestor prefixes for pruning ──────────────────── - # e.g. scope = "A / B / K" → ancestor_prefixes = {"A", "A / B", "A / B / K"} + # For multi-scope, union all ancestor prefixes from all scope paths ancestor_prefixes: set[str] = set() - for i in range(1, scope_depth + 1): - ancestor_prefixes.add(' / '.join(scope_parts[:i])) + for sp in scope_list: + sp_parts = split_section_path(sp) + for i in range(1, len(sp_parts) + 1): + ancestor_prefixes.add(' / '.join(sp_parts[:i])) # ── Classify each section ──────────────────────────────────────────── _excl = exclude_paths or set() items_by_path: dict[str, dict] = {} - scope_child_depths: set[int] = set() + # Per-scope depth bands: track child depths separately per scope + per_scope_child_depths: dict[str, set[int]] = {sp: set() for sp in scope_list} if scope_list else {} + root_child_depths: set[int] = set() # used when scope_list is empty (root) + + def _make_item(path: str, meta: dict, show_summary: bool) -> dict: + return { + 'path': path, + 'title': meta['title'], + 'summary': meta['summary'], + 'level': meta['depth'], + 'sort_order': meta['sort_order'], + 'chunk_count': 0, + 'image_count': 0, + 'table_count': 0, + 'section_id': meta['section_id'], + 'show_summary': show_summary, + } + + def _is_excluded(path: str) -> bool: + return bool(_excl and any( + path == ep or path.startswith(ep + ' / ') for ep in _excl + )) for path, meta in all_sections.items(): parts = meta['parts'] depth = meta['depth'] - if scope_depth == 0: + if not scope_list: # Root scope: everything is a potential child - if depth < 1: + if depth < 1 or _is_excluded(path): continue - # Skip excluded paths - if _excl and any( - path == ep or path.startswith(ep + ' / ') - for ep in _excl - ): - continue - scope_child_depths.add(depth) - items_by_path[path] = { - 'path': path, - 'title': meta['title'], - 'summary': meta['summary'], - 'level': depth, - 'sort_order': meta['sort_order'], - 'chunk_count': 0, - 'image_count': 0, - 'table_count': 0, - 'section_id': meta['section_id'], - 'show_summary': True, # will be refined after depth band selection - } + root_child_depths.add(depth) + items_by_path[path] = _make_item(path, meta, show_summary=True) continue - # --- Non-root scope --- + # --- Non-root scope(s) --- + # Check if this path is a descendant of ANY scope in scope_list + matched_scope: str | None = None + for sp in scope_list: + sp_parts = split_section_path(sp) + sp_depth = len(sp_parts) + if depth > sp_depth and parts[:sp_depth] == sp_parts: + matched_scope = sp + break + + if matched_scope: + # Category 2: descendant of a scope path → selectable + if _is_excluded(path): + continue + per_scope_child_depths[matched_scope].add(depth) + items_by_path[path] = _make_item(path, meta, show_summary=True) + continue - # Category 1: Ancestors and their siblings (structural context) - # A node is an ancestor/sibling if its depth <= scope_depth AND - # its parent prefix matches the scope's ancestry chain. - if depth <= scope_depth: - # Check: is this node in the ancestry chain or a sibling of one? + # Category 1: structural context (ancestors of scope paths only) + # Only show nodes that are on the ancestor chain of a scope path. + # Non-scope siblings (e.g. 法律声明, 前言 when navigating into + # chapters 一~六) are pruned to reduce token waste and prevent + # summary overflow in _format_items_for_llm. + max_scope_depth = max(len(split_section_path(sp)) for sp in scope_list) + if depth <= max_scope_depth: if depth == 1: - # All L1 nodes are either the ancestor or its siblings - items_by_path[path] = { - 'path': path, - 'title': meta['title'], - 'summary': meta['summary'], - 'level': depth, - 'sort_order': meta['sort_order'], - 'chunk_count': 0, - 'image_count': 0, - 'table_count': 0, - 'section_id': meta['section_id'], - 'show_summary': False, - } - elif depth <= scope_depth: - # For deeper ancestors/siblings: their parent must be in the - # ancestor chain. e.g. "A / C" is a sibling of "A / B" only - # if "A" is an ancestor of scope. + if path in ancestor_prefixes: + items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) + else: parent_prefix = ' / '.join(parts[:-1]) if parent_prefix in ancestor_prefixes: - items_by_path[path] = { - 'path': path, - 'title': meta['title'], - 'summary': meta['summary'], - 'level': depth, - 'sort_order': meta['sort_order'], - 'chunk_count': 0, - 'image_count': 0, - 'table_count': 0, - 'section_id': meta['section_id'], - 'show_summary': False, - } + items_by_path.setdefault(path, _make_item(path, meta, show_summary=False)) continue - # Category 2: Descendants of scope_path (children to explore) - # depth > scope_depth is guaranteed by the continue at line above - is_descendant = parts[:scope_depth] == scope_parts - if is_descendant: - # Skip excluded paths - is_excluded = _excl and any( - path == ep or path.startswith(ep + ' / ') - for ep in _excl - ) - if is_excluded: - logger.debug(f' _load_child_sections: EXCLUDED descendant path={path!r}') - continue - scope_child_depths.add(depth) - items_by_path[path] = { - 'path': path, - 'title': meta['title'], - 'summary': meta['summary'], - 'level': depth, - 'sort_order': meta['sort_order'], - 'chunk_count': 0, - 'image_count': 0, - 'table_count': 0, - 'section_id': meta['section_id'], - 'show_summary': True, - } - continue - else: - logger.debug( - f' _load_child_sections: NOT descendant path={path!r} ' - f'parts[:scope_depth]={parts[:scope_depth]} != scope_parts={scope_parts}' - ) - - # Category 3: Everything else → pruned (not added) + # Category 3: pruned if not items_by_path: return [] - # ── Limit children to 2 depth bands (relative to scope) ───────────── - if scope_child_depths: - if scope_depth == 0: - allowed_depths = sorted(scope_child_depths)[:2] - else: - allowed_depths = sorted(scope_child_depths)[:2] - allowed_set = set(allowed_depths) - to_remove = [] - for path, item in items_by_path.items(): - if item['show_summary'] and item['level'] not in allowed_set: - to_remove.append(path) + # ── Limit children to 2 depth bands (relative to each scope) ──────── + allowed_set: set[int] = set() + if scope_list: + for sp, depths in per_scope_child_depths.items(): + if depths: + allowed_set.update(sorted(depths)[:2]) + else: + if root_child_depths: + allowed_set.update(sorted(root_child_depths)[:2]) + + if allowed_set: + to_remove = [ + path for path, item in items_by_path.items() + if item['show_summary'] and item['level'] not in allowed_set + ] for path in to_remove: del items_by_path[path] @@ -806,48 +872,30 @@ async def _load_child_sections( ) item['is_leaf'] = not has_descendants + # ── Assign selectability ────────────────────────────────────────────── + # Rule: in the 2-band window, only the DEEPER band is selectable. + # Leaf nodes at the shallower band are still selectable (no children + # to drill into). Structural context (show_summary=False) is never + # selectable. + if allowed_set: + shallowest_band = min(allowed_set) + for item in sorted_items: + if not item.get('show_summary', True): + # Structural context → never selectable + item['selectable'] = False + elif item['level'] == shallowest_band and not item.get('is_leaf', False): + # Shallowest band, non-leaf → grouping header, not selectable + item['selectable'] = False + else: + item['selectable'] = True + else: + for item in sorted_items: + item['selectable'] = item.get('show_summary', True) + return sorted_items # --------------------------------------------------------------------------- -# LLM response parser (for scope_navigate) -# --------------------------------------------------------------------------- - -def _parse_scope_nav_response(text: str) -> list[dict[str, Any]]: - """Parse selections JSON from scope navigation LLM response. - - Returns list of {"path": str, "confidence": float}. - """ - text = text.strip() - # Try direct parse - try: - data = json.loads(text) - except (json.JSONDecodeError, ValueError): - # Extract JSON object from markdown wrapper - match = re.search(r'\{.*\}', text, re.DOTALL) - if not match: - return [] - try: - data = json.loads(match.group()) - except (json.JSONDecodeError, ValueError): - return [] - - if not isinstance(data, dict): - return [] - - selections: list[dict[str, Any]] = [] - for item in (data.get('selections') or []): - if not isinstance(item, dict): - continue - path = str(item.get('path') or '').strip() - if not path: - continue - confidence = _normalize_confidence(item.get('confidence')) - if confidence is None: - confidence = 0.7 - selections.append({'path': path, 'confidence': confidence}) - - return selections # --------------------------------------------------------------------------- diff --git a/packages/shared-python/shared/services/retrieval/agentic/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/__init__.py index e4b5a0695..e59b26254 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/__init__.py +++ b/packages/shared-python/shared/services/retrieval/agentic/__init__.py @@ -2,10 +2,11 @@ Navigate-then-answer loop: Phase 1: Document selection (discovery + KG LLM select) - Phase 2: Per-document iterative navigation (scope_navigate_step) + Phase 2: Per-document iterative navigation (navigate_step — unified action) Phase 3: attempt_answer → DONE (return answer) or NOT_FOUND → revision -Navigation auto-terminates when the LLM returns empty selections. +Each navigate_step decides action (NAVIGATE/STOP), optional asset tools, +and section selections in a single LLM call. STOP terminates drill-down. After navigation, attempt_answer is called automatically — its result (answer or NOT_FOUND+reason) drives the revision loop. diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index b8b766327..77bf80bca 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -2,7 +2,7 @@ Flow: Phase 1: Document selection (bottom_discovery + kg_document_select) - Phase 2: Per-document navigation (iterative BFS scope_navigate_step) + Phase 2: Per-document navigation (iterative BFS via navigate_step) Phase 3: Render evidence → attempt_answer → DONE (has answer) → return answer + evidence → NOT_FOUND + reason → revision_hint → re-select docs + re-navigate @@ -10,9 +10,10 @@ → max_revisions → return best available The orchestrator drives navigation via an iterative BFS queue per document, -calling scope_navigate_step at each level. Navigation auto-terminates when -the LLM returns empty selections. After navigation completes, attempt_answer -is called automatically — no separate verdict step needed. +calling navigate_step at each level. Each navigate_step is a single LLM call +that decides action (NAVIGATE/STOP), asset tools (FIND_IMAGES/FIND_TABLES), +and section selections. STOP terminates the drill-down for that scope. +After navigation completes, attempt_answer is called automatically. """ from __future__ import annotations @@ -100,6 +101,117 @@ async def _build_asset_url_map( return url_map +def _collect_all_leaf_paths(node: DocTreeNode) -> set[str]: + """Recursively collect all leaf_content keys across the entire tree.""" + paths = set(node.leaf_content.keys()) + for child in node.children.values(): + paths.update(_collect_all_leaf_paths(child)) + return paths + + +def _collect_visible_paths(node: DocTreeNode) -> set[str]: + """Collect all outline_items paths across the entire tree. + + These are sections that are "visible" in the rendered tree (shown to the + LLM during navigation) even if no chunks have been hydrated into them yet. + Used as fallback targets for asset reconciliation. + """ + paths = {item['path'] for item in node.outline_items if item.get('path')} + for child in node.children.values(): + paths.update(_collect_visible_paths(child)) + return paths + + +def _find_closest_ancestor(path: str, target_paths: set[str]) -> str | None: + """Walk up a section path to find the closest ancestor in target_paths. + + Example: path="kb/file/Ch1/S1.1/S1.1.1", target_paths={"kb/file/Ch1/S1.1"} + → returns "kb/file/Ch1/S1.1" + + Uses the ' / ' separator convention from the section path format. + """ + parts = path.split(' / ') + # Walk from most specific to least specific (skip the full path itself) + for i in range(len(parts) - 1, 0, -1): + ancestor = ' / '.join(parts[:i]) + if ancestor in target_paths: + return ancestor + return None + + +def _reconcile_deferred_assets( + tree: DocTreeNode, + pending_assets: list[dict], +) -> None: + """Place collected assets into the tree based on final navigated paths. + + Called ONCE after the entire BFS + discovery merge completes for a + document. Asset placement uses a two-tier strategy: + + 1. **Exact match**: If the asset's ``owner_section_path`` matches a + leaf_content key, place directly (existing behavior). + 2. **Closest visible ancestor**: If exact match fails, walk up the + owner_section_path hierarchy to find the nearest ancestor that + appears in either leaf_content or outline_items. This handles + the case where the LLM stopped navigation early (e.g. at root) + but still requested images/tables — assets at L3 get attributed + to the visible L2 section on their path. + """ + final_paths = _collect_all_leaf_paths(tree) + visible_paths = _collect_visible_paths(tree) + all_target_paths = final_paths | visible_paths + + if not all_target_paths: + return + + # Collect existing chunk_ids to avoid duplicates + existing_ids = { + str(row.get('chunk_id') or '') + for row in tree.flatten_chunk_rows() + if row.get('chunk_id') + } + + placed = 0 + ancestor_placed = 0 + for asset in pending_assets: + chunk_id = str(asset.get('chunk_id') or '') + if chunk_id and chunk_id in existing_ids: + continue # already in tree via hydrate_connected_target_rows + + owner_path = ( + asset.get('owner_section_path') + or asset.get('section_path') + ) + if not owner_path: + continue + + # Tier 1: exact match in leaf_content or visible outline + target_path = owner_path if owner_path in all_target_paths else None + + # Tier 2: closest visible ancestor fallback + if target_path is None: + target_path = _find_closest_ancestor(owner_path, all_target_paths) + if target_path: + ancestor_placed += 1 + + if target_path is None: + continue # no visible ancestor → discard + + # Place into root; reparent_leaf_content will move to correct child + tree.add_leaf_chunks(target_path, [asset]) + if chunk_id: + existing_ids.add(chunk_id) + placed += 1 + + if placed: + tree.reparent_leaf_content() + logger.info( + f' deferred asset reconcile: {placed}/{len(pending_assets)} ' + f'assets placed into {len(final_paths)} leaf + {len(visible_paths)} visible paths ' + f'(ancestor_fallback={ancestor_placed})' + ) + + def _build_config_from_env() -> AgentRunConfig: """Read agent config from environment, with sensible defaults.""" return AgentRunConfig( @@ -412,6 +524,9 @@ async def run( channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, config: AgentRunConfig | None = None, + ledger: BudgetLedger | None = None, + parent_run_id: str | None = None, + workflow_step_id: str | None = None, ) -> AgenticResult: """Run the agentic retrieval pipeline. @@ -431,7 +546,7 @@ async def run( exclude_sections = exclude_sections or [] state = AgentState() - state.ledger = BudgetLedger( + state.ledger = ledger or BudgetLedger( total=config.token_budget_total, planning_ratio=config.planning_ratio, bootstrap=config.bootstrap_budget, @@ -455,6 +570,8 @@ async def run( 'exclude_sections': exclude_sections, 'signal_paths': signal_paths, }, + parent_run_id=parent_run_id, + workflow_step_id=workflow_step_id, ) trace_enabled = os.environ.get('RETRIEVAL_AGENTIC_TRACE_ENABLED', 'true') == 'true' @@ -711,9 +828,11 @@ async def _context_llm_call(prompt): if key.startswith(f'{doc.document_id}::') } if state.seen_section_keys else set() - # BFS queue: (scope_path, parent_node, depth) + # BFS queue: (scope_path(s), parent_node, depth) + # scope can be: None (root), str, or list[str] (multi-scope) root = DocTreeNode(scope_path=None) - pending: list[tuple[str | None, DocTreeNode, int]] = [(None, root, 0)] + pending: list[tuple[str | list[str] | None, DocTreeNode, int]] = [(None, root, 0)] + doc_pending_assets: list[dict] = [] # deferred asset reconcile while pending: if state.elapsed_ms >= config.latency_budget_ms: @@ -736,14 +855,16 @@ async def _context_llm_call(prompt): depth=depth, ) - # ★ Step 1: Tool selection (agent decides which asset tools) + # ★ Unified navigate step (supports multi-scope batching) try: - tool_choices = await tools.tool_select_step( + action, asset_tools, step_node, drill_paths = await tools.navigate_step( db, document_id=doc.document_id, job_result_id=job_result_id, query=query, llm_fn=doc_llm_fn, + user_id=user_id, + namespace=namespace, doc_name=doc_name, scope_path=scope, exclude_paths=doc_exclude, @@ -751,37 +872,18 @@ async def _context_llm_call(prompt): budget_snapshot=state.ledger.snapshot() if state.ledger else None, ) except BudgetExceeded: - logger.info(' agentic: planning budget exhausted during tool selection') + logger.info(' agentic: planning budget exhausted during navigation') if trace_enabled: trace.record_budget_stop('planning_exhausted') break state.step_count += 1 - if trace_enabled: - trace.record_step( - 'tool_select_step', ToolResult( - status='selected', - payload={ - 'document_id': doc.document_id, - 'scope': scope or 'root', - 'depth': depth, - 'tool_choice': tool_choices or ['NAVIGATE'], - }, - ), - decision_reason=f'tool_r{round_idx}_d{depth}_{doc.source_file_name}', - ) - - logger.info( - f' agentic step {state.step_count}: tool_select_step ' - f'doc="{doc.source_file_name}" scope={scope or "root"} ' - f'depth={depth} tools={tool_choices or ["NAVIGATE"]}' - ) - - pending_scope_assets: list[dict] = [] - for asset_tool in tool_choices: + # ★ Asset collection (deferred reconcile) — runs if LLM selected tools + # scope is passed directly: None (root), str, or list[str] (multi-scope). + # asset_filter_step handles all forms natively. + for asset_tool in asset_tools: if asset_tool not in ('FIND_IMAGES', 'FIND_TABLES'): continue - # ★ Asset collection (programmatic extraction) asset_type = 'image' if asset_tool == 'FIND_IMAGES' else 'table' asset_chunks = await tools.asset_filter_step( db, @@ -791,15 +893,16 @@ async def _context_llm_call(prompt): asset_type=asset_type, ) if asset_chunks: - pending_scope_assets.extend(asset_chunks) + doc_pending_assets.extend(asset_chunks) + scope_display = scope if isinstance(scope, list) else (scope or 'root') if trace_enabled: trace.record_step( 'asset_filter_step', ToolResult( status='filtered' if asset_chunks else 'empty', payload={ 'document_id': doc.document_id, - 'scope': scope or 'root', + 'scope': scope_display, 'asset_type': asset_type, 'chunks_found': len(asset_chunks) if asset_chunks else 0, }, @@ -809,33 +912,9 @@ async def _context_llm_call(prompt): logger.info( f' agentic step {state.step_count}: asset_filter_step ' - f'doc="{doc.source_file_name}" scope={scope or "root"} ' + f'doc="{doc.source_file_name}" scope={scope_display} ' f'type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}' ) - # ★ Fallthrough: always proceed to NAVIGATE. - - # ★ Step 2: NAVIGATE (existing scope_navigate_step) - try: - step_node, drill_paths = await tools.scope_navigate_step( - db, - document_id=doc.document_id, - job_result_id=job_result_id, - query=query, - llm_fn=doc_llm_fn, - user_id=user_id, - namespace=namespace, - doc_name=doc_name, - scope_path=scope, - exclude_paths=doc_exclude, - revision_hint=revision_hint if depth == 0 else None, - budget_snapshot=state.ledger.snapshot() if state.ledger else None, - ) - except BudgetExceeded: - logger.info(' agentic: planning budget exhausted during navigation') - if trace_enabled: - trace.record_budget_stop('planning_exhausted') - break - state.step_count += 1 # Merge step result into parent node parent_node.outline_items = step_node.outline_items @@ -843,66 +922,34 @@ async def _context_llm_call(prompt): parent_node.add_leaf_chunks(leaf_path, chunks) parent_node.confidence = step_node.confidence - # ★ Step 2.5: Reconcile pending assets with navigated leaf content - if pending_scope_assets: - # Collect all chunk_ids already present in any leaf_content - existing_ids = { - str(row.get('chunk_id') or '') - for row in parent_node.flatten_chunk_rows() - if row.get('chunk_id') - } - # Filter out assets already inlined - supplementary = [ - a for a in pending_scope_assets - if str(a.get('chunk_id') or '') not in existing_ids - ] - if supplementary: - # Only place assets into sections already in the - # navigated tree. If the owner path isn't part of - # the tree, fall back to the current scope. - navigated_paths = set(parent_node.leaf_content.keys()) | set(parent_node.children.keys()) - for asset in supplementary: - owner_path = ( - asset.get('owner_section_path') - or asset.get('section_path') - or scope - ) - if owner_path and owner_path in navigated_paths: - parent_node.add_leaf_chunks(str(owner_path), [asset]) - elif scope: - parent_node.add_leaf_chunks(str(scope), [asset]) - # Accumulate hydrated leaf paths into doc_exclude - # so subsequent drill-downs don't re-show them as [SELECT]. - # Skip paths that are pending drill-down (non-leaf hybrid nodes - # hydrated via self_only) — their children must remain selectable. drill_path_set = {sel['path'] for sel in drill_paths} for leaf_path in step_node.leaf_content: if leaf_path not in drill_path_set: doc_exclude.add(leaf_path) - # Queue non-leaf selections for further drill-down - for sel in drill_paths: - child = DocTreeNode(scope_path=sel['path']) - parent_node.children[sel['path']] = child - pending.append((sel['path'], child, depth + 1)) - - # Re-parent leaf paths that belong to a child's subtree. - # When the LLM selects both a parent section (non-leaf) - # and one of its children (leaf) at the same depth, the - # leaf chunks are stored on the parent node. Move them - # into the child node so render_unified_doc_tree nests - # them correctly instead of rendering orphans at root. + # Queue non-leaf selections as a SINGLE batched item + # (all drill paths expand simultaneously in the next call) + if drill_paths: + for sel in drill_paths: + child = DocTreeNode(scope_path=sel['path']) + parent_node.children[sel['path']] = child + batch_scope = [sel['path'] for sel in drill_paths] + pending.append((batch_scope, parent_node, depth + 1)) + + # Re-parent leaf paths that belong to a child's subtree parent_node.reparent_leaf_content() if trace_enabled: trace.record_step( - 'scope_navigate_step', ToolResult( - status='navigated' if step_node.has_content() else 'empty', + 'navigate_step', ToolResult( + status=f'{action.lower()}' + (' (content)' if step_node.has_content() else ''), payload={ 'document_id': doc.document_id, - 'scope': scope or 'root', + 'scope': scope if isinstance(scope, str) else (scope or 'root'), 'depth': depth, + 'action': action, + 'asset_tools': asset_tools, 'outline_count': len(step_node.outline_items), 'leaf_count': len(step_node.leaf_content), 'pending_drills': len(drill_paths), @@ -911,10 +958,12 @@ async def _context_llm_call(prompt): decision_reason=f'nav_r{round_idx}_d{depth}_{doc.source_file_name}', ) + scope_log = scope if isinstance(scope, str) else (', '.join(scope) if scope else 'root') logger.info( - f' agentic step {state.step_count}: scope_navigate_step ' - f'doc="{doc.source_file_name}" scope={scope or "root"} ' - f'depth={depth} outline={len(step_node.outline_items)} ' + f' agentic step {state.step_count}: navigate_step ' + f'doc="{doc.source_file_name}" scope={scope_log} ' + f'depth={depth} action={action} tools={asset_tools} ' + f'outline={len(step_node.outline_items)} ' f'leaves={len(step_node.leaf_content)} ' f'drills={len(drill_paths)}' ) @@ -975,6 +1024,38 @@ async def _context_llm_call(prompt): chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()), ) + # ── Deferred asset reconcile ────────────────────────────── + # Assets were collected across all BFS depths but NOT placed + # into the tree yet. Now that the final navigated paths are + # known (BFS + discovery), filter and place only those assets + # whose owner path matches a navigated leaf. + if not is_b_class and doc_pending_assets: + # Inject doc file name as a visible root-level path, but + # ONLY when BFS stopped at root (no children = STOP action). + if doc_name and not root.children and not any( + item.get('path') == doc_name for item in root.outline_items + ): + root.outline_items.insert(0, {'path': doc_name, 'level': 0}) + _reconcile_deferred_assets(root, doc_pending_assets) + if trace_enabled: + trace.record_step( + 'deferred_asset_reconcile', ToolResult( + status='reconciled', + payload={ + 'document_id': doc.document_id, + 'pending_count': len(doc_pending_assets), + 'placed_count': sum( + 1 for a in doc_pending_assets + if str(a.get('chunk_id') or '') in { + str(r.get('chunk_id') or '') + for r in root.flatten_chunk_rows() + } + ), + }, + ), + decision_reason=f'deferred_reconcile_r{round_idx}_{doc.source_file_name}', + ) + # Merge or store doc tree if doc.document_id in state.doc_trees: state.doc_trees[doc.document_id].merge(root) diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index d11e4da7b..4774d2491 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -22,8 +22,8 @@ _format_items_for_llm, _load_child_sections, _parse_json_array, - _parse_scope_nav_response, - _SCOPE_NAV_PROMPT, + _parse_action_response, + _ACTION_PROMPT, _DISCOVERY_SELECT_PROMPT, _FILE_SELECT_PROMPT, _format_budget_block, @@ -362,184 +362,6 @@ async def kg_document_select( return ToolResult(status='error', error=str(e), latency_ms=latency) -# --------------------------------------------------------------------------- -# Tool: tool_select_step (lightweight LLM router) -# --------------------------------------------------------------------------- - -_TOOL_SELECT_PROMPT = """\ -You are a document navigation agent. - -Document: "{doc_name}" - -{budget_block} -{scope_header} -Below is a summary of the current scope's sections: - -{tree_summary} - -User query: {query} - -=== Available Actions === - -NAVIGATE (always performed) - Drill into specific sections to explore detailed content. - This action always runs — you do not need to select it. - -FIND_IMAGES (optional, additive) - Also extract image/chart/diagram assets under this scope. - Select this when the query asks about images, charts, figures, or visual content. - -FIND_TABLES (optional, additive) - Also extract table/data assets under this scope. - Select this when the query asks about tables, tabular data, or structured data. - -You may select ZERO, ONE, or BOTH optional actions. -Navigation always happens regardless of your selection. - -Return ONLY a JSON object: -{{"tools": []}} — navigate only, no extra assets -{{"tools": ["FIND_IMAGES"]}} — navigate + extract images -{{"tools": ["FIND_TABLES"]}} — navigate + extract tables -{{"tools": ["FIND_IMAGES", "FIND_TABLES"]}} — navigate + extract both -When budget is TIGHT, prefer fewer extra actions. -When budget is CRITICAL, return empty tools unless assets directly answer the query. -Do not include any explanation. -""" - - -def _parse_tool_choice(text: str) -> list[str]: - """Parse tool choices from LLM response. - - Returns a list of selected tools (subset of FIND_IMAGES, FIND_TABLES). - NAVIGATE is always implicit — an empty list means "navigate only". - """ - import json as _json - import re as _re - - text = text.strip() - _ASSET_TOOLS = {'FIND_IMAGES', 'FIND_TABLES'} - - def _extract_from_data(data: dict) -> list[str]: - # New format: {"tools": [...]} - tools_val = data.get('tools') - if isinstance(tools_val, list): - return [str(t).strip().upper() for t in tools_val if str(t).strip().upper() in _ASSET_TOOLS] - # Legacy format: {"tool": "..."} - tool_val = str(data.get('tool', '')).strip().upper() - if tool_val in _ASSET_TOOLS: - return [tool_val] - if tool_val == 'NAVIGATE': - return [] - return [] - - # Try JSON parse - try: - data = _json.loads(text) - if isinstance(data, dict): - return _extract_from_data(data) - except (ValueError, _json.JSONDecodeError): - pass - - # Accept a JSON object wrapped in markdown - match = _re.search(r'\{.*?\}', text, _re.DOTALL) - if match: - try: - data = _json.loads(match.group()) - if isinstance(data, dict): - return _extract_from_data(data) - except (ValueError, _json.JSONDecodeError): - pass - - # Fallback: scan for tool names in raw text - upper = text.upper() - result: list[str] = [] - if 'FIND_IMAGES' in upper: - result.append('FIND_IMAGES') - if 'FIND_TABLES' in upper: - result.append('FIND_TABLES') - return result - - -async def tool_select_step( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - query: str, - llm_fn: LLMFn, - doc_name: str = '', - scope_path: str | None = None, - exclude_paths: set[str] | None = None, - revision_hint: str | None = None, - budget_snapshot: dict | None = None, -) -> list[str]: - """Route to the appropriate tools for the current scope. - - Returns a list of asset tools to run (FIND_IMAGES, FIND_TABLES). - NAVIGATE always runs implicitly after any asset extraction. - - Optimization: if the current scope has no image or table chunks, - skips the LLM call and returns an empty list (navigate only). - """ - items = await _load_child_sections( - db, document_id, job_result_id, scope_path, - exclude_paths=exclude_paths, - ) - if not items: - return [] - - # Build lightweight tree summary (titles + counts only, no summaries) - summary_lines = [] - for item in items: - if not item.get('show_summary'): - continue - title = item.get('title', '') - img = item.get('image_count', 0) - tbl = item.get('table_count', 0) - txt = item.get('chunk_count', 0) - counts = f'text={txt}' - if img > 0: - counts += f' image={img}' - if tbl > 0: - counts += f' table={tbl}' - summary_lines.append(f'- {title} [{counts}]') - - tree_summary = '\n'.join(summary_lines) or '(empty)' - - # Check if scope has ANY images or tables — skip prompt if none - total_images = sum(i.get('image_count', 0) for i in items) - total_tables = sum(i.get('table_count', 0) for i in items) - if total_images == 0 and total_tables == 0: - return [] # no assets → skip tool selection, navigate only - - scope_header = ( - f'Current scope: "{scope_path}"' if scope_path - else 'Current scope: root (document top level)' - ) - prompt = _TOOL_SELECT_PROMPT.format( - doc_name=doc_name or document_id, - scope_header=scope_header, - budget_block=_format_budget_block(budget_snapshot), - tree_summary=tree_summary, - query=query, - ) - if revision_hint: - prompt += ( - f'\n\nIMPORTANT: This is a REVISION round. ' - f'The previous search attempt failed because:\n' - f'"{revision_hint}"\n' - f'Adjust your tool selection accordingly.' - ) - response = await llm_fn(prompt) - - # Parse tool choices - asset_tools = _parse_tool_choice(response) - logger.info( - f' tool_select_step scope={scope_path or "root"}: ' - f'tools={asset_tools or ["NAVIGATE"]} images={total_images} tables={total_tables}' - ) - return asset_tools - # --------------------------------------------------------------------------- # Tool: asset_filter_step (programmatic asset extraction) @@ -550,7 +372,7 @@ async def asset_filter_step( *, document_id: str, job_result_id: str, - scope_path: str | None, + scope_path: str | list[str] | None, asset_type: str, # 'image' | 'table' ) -> list[dict[str, Any]]: """Extract assets from all descendants under scope_path. @@ -561,22 +383,36 @@ async def asset_filter_step( Also collects standalone asset chunks (image/table) that exist directly under the scope but are not referenced via connect_to. + + scope_path can be: + - None: root scope (entire document) + - str: single scope path + - list[str]: multiple scope paths (queried simultaneously) """ from shared.models.database.document import DocumentChunk, DocumentSection t0 = time.monotonic() try: - # 1. Find all section_ids under scope_path + # 1. Find all section_ids under scope_path(s) + # Normalize scope to list for uniform handling + scope_list = ( + scope_path if isinstance(scope_path, list) + else [scope_path] if scope_path + else [] + ) + section_stmt = ( select(DocumentSection.section_id, DocumentSection.section_path) .where(DocumentSection.document_id == document_id) .where(DocumentSection.job_result_id == job_result_id) ) - if scope_path: - section_stmt = section_stmt.where( - (DocumentSection.section_path == scope_path) | - (DocumentSection.section_path.like(f'{scope_path} / %')) - ) + if scope_list: + from sqlalchemy import or_ + scope_filters = [] + for sp in scope_list: + scope_filters.append(DocumentSection.section_path == sp) + scope_filters.append(DocumentSection.section_path.like(f'{sp} / %')) + section_stmt = section_stmt.where(or_(*scope_filters)) section_result = await db.execute(section_stmt) section_rows = section_result.all() section_ids = {row[0] for row in section_rows} @@ -635,6 +471,20 @@ async def asset_filter_step( ] owner_by_target_id = _build_connected_owner_map(text_row_dicts) + # Replace synthetic "Root" owner with the document's source_file_name. + # Root is a hybrid node whose real path is the file name (e.g. + # "32_安全大模型技术与市场研究报告_1.docx"); the DB stores the + # synthetic label "Root" which cannot match any outline node. + if any(v == 'Root' for v in owner_by_target_id.values()): + doc_stmt = select(Document.source_file_name).where( + Document.document_id == document_id + ) + doc_file_name = (await db.execute(doc_stmt)).scalar() or '' + if doc_file_name: + for tid in list(owner_by_target_id): + if owner_by_target_id[tid] == 'Root': + owner_by_target_id[tid] = doc_file_name + # Collect connected target IDs for batch-loading connected_target_ids: set[str] = set(owner_by_target_id.keys()) @@ -689,8 +539,9 @@ async def asset_filter_step( # Root / top-level aggregation sections if not owner_section_path: own_section_path = section_path_by_id.get(row[4]) - if own_section_path and ' / ' not in own_section_path: - # Reject document-root level sections as fallback owners + if own_section_path and own_section_path == 'Root': + # Reject only the synthetic Root aggregation label; + # legitimate L1 sections (e.g. "前言") are valid owners. logger.warning( f' asset_filter_step: rejecting root-level owner fallback ' f'chunk_id={chunk_id} section_path={own_section_path}' @@ -731,13 +582,11 @@ async def asset_filter_step( logger.error(f' asset_filter_step failed: {e}') return [] - # --------------------------------------------------------------------------- -# Tool: scope_navigate_step (single-step navigation) +# Tool: navigate_step (unified action — merges tool_select + scope_navigate) # --------------------------------------------------------------------------- - -async def scope_navigate_step( +async def navigate_step( db: AsyncSession, *, document_id: str, @@ -747,73 +596,149 @@ async def scope_navigate_step( user_id: str, namespace: str, doc_name: str = '', - scope_path: str | None = None, + scope_path: str | list[str] | None = None, exclude_paths: set[str] | None = None, revision_hint: str | None = None, budget_snapshot: dict | None = None, -) -> tuple[DocTreeNode, list[dict]]: - """Single navigation step — one LLM call, no recursion. +) -> tuple[str, list[str], DocTreeNode, list[dict]]: + """Unified navigation step — one LLM call for action + tools + selections. + + scope_path can be: + - None: root scope + - str: single scope to drill into + - list[str]: multiple scopes to expand simultaneously Returns: - - node: DocTreeNode with outline_items (current scope local items only) - and leaf_content (hydrated leaf selections) - - pending: list of {path, confidence, mode} for non-leaf selections - (orchestrator queues these for further drill-down) + - action: 'STOP' | 'NAVIGATE' + - asset_tools: list of asset tools to run (FIND_IMAGES, FIND_TABLES) + - node: DocTreeNode with outline_items and leaf_content + - pending: list of {path, confidence} for non-leaf drill-downs (empty when STOP) """ from shared.services.retrieval.app_service import _hydrate_paths_to_rows - empty = DocTreeNode.empty(scope_path) + # Normalize scope for internal use + scope_paths: list[str] = ( + scope_path if isinstance(scope_path, list) + else [scope_path] if scope_path + else [] + ) + # Set of scope path strings (for filtering selections) + scope_path_set = set(scope_paths) + + empty = DocTreeNode.empty(scope_paths[0] if scope_paths else None) try: - # 1. Load continuous context tree + # 1. Load continuous context tree (supports multi-scope) items = await _load_child_sections( db, document_id, job_result_id, scope_path, exclude_paths=exclude_paths, ) if not items: - return empty, [] + return 'STOP', [], empty, [] - # 2. Build selectable index (only current-scope items with summary) - selectable = {item['path']: item for item in items if item.get('show_summary', True)} + # 2. Build selectable index + selectable = {item['path']: item for item in items if item.get('selectable', False)} - # 3. Format full tree and call LLM - text, overflowed = _format_items_for_llm(items) - scope_header = ( - f'Current scope: navigating into "{scope_path}"' - if scope_path else - 'Current scope: root (document top level)' + # 3. Count ALL image/table chunks under the scope subtree(s) + from shared.models.database.document import DocumentChunk, DocumentSection + from sqlalchemy import func as sa_func + + scope_section_stmt = ( + select(DocumentSection.section_id) + .where(DocumentSection.document_id == document_id) + .where(DocumentSection.job_result_id == job_result_id) ) - prompt = _SCOPE_NAV_PROMPT.format( + if scope_paths: + from sqlalchemy import or_ + scope_filters = [] + for sp in scope_paths: + scope_filters.append(DocumentSection.section_path == sp) + scope_filters.append(DocumentSection.section_path.like(f'{sp} / %')) + scope_section_stmt = scope_section_stmt.where(or_(*scope_filters)) + scope_section_ids = await db.execute(scope_section_stmt) + all_section_ids = [r[0] for r in scope_section_ids.all()] + + total_images = 0 + total_tables = 0 + if all_section_ids: + count_stmt = ( + select( + DocumentChunk.chunk_type, + sa_func.count(DocumentChunk.id), + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(all_section_ids)) + .where(DocumentChunk.chunk_type.in_(['image', 'table'])) + .group_by(DocumentChunk.chunk_type) + ) + count_result = await db.execute(count_stmt) + for chunk_type, cnt in count_result.all(): + if chunk_type == 'image': + total_images = cnt + elif chunk_type == 'table': + total_tables = cnt + + tools_block = '' + if total_images > 0 or total_tables > 0: + tools_lines = ['\nOptional asset tools (usable with NAVIGATE or STOP):\n'] + if total_images > 0: + tools_lines.append( + f' FIND_IMAGES — Extract all image/chart assets under this scope ({total_images} available).\n' + ) + if total_tables > 0: + tools_lines.append( + f' FIND_TABLES — Extract all table/data assets under this scope ({total_tables} available).\n' + ) + tools_block = ''.join(tools_lines) + + # 4. Format tree and build prompt + text, overflowed = _format_items_for_llm(items) + if not scope_paths: + scope_header = 'Current scope: root (document top level)' + elif len(scope_paths) == 1: + scope_header = f'Current scope: navigating into "{scope_paths[0]}"' + else: + scope_header = f'Current scope: navigating into {len(scope_paths)} sections' + prompt = _ACTION_PROMPT.format( doc_name=doc_name or document_id, doc_id=document_id, scope_header=scope_header, budget_block=_format_budget_block(budget_snapshot), items_overview=text, query=query, + tools_block=tools_block, ) if revision_hint: prompt += ( f'\n\nIMPORTANT: Previous round feedback: ' - f'"{revision_hint}". Select specific sections this time.' + f'"{revision_hint}". Adjust your selections accordingly.' ) + + # 5. Single LLM call response = await llm_fn(prompt) - selections = _parse_scope_nav_response(response) + parsed = _parse_action_response(response) + action = parsed['action'] + asset_tools = parsed['tools'] + selections = parsed['selections'] + scope_label = ', '.join(scope_paths) if scope_paths else 'root' logger.info( - f' scope_navigate_step scope={scope_path or "root"}: ' - f'selections={len(selections)}, selectable={len(selectable)}, ' + f' navigate_step scope={scope_label}: ' + f'action={action} tools={asset_tools} ' + f'selections={len(selections)} selectable={len(selectable)} ' f'overflowed={overflowed}' ) - # 4. Build node with LOCAL items only (no ancestors/siblings) - node = DocTreeNode(scope_path=scope_path) + # 6. Build node with LOCAL items only + node = DocTreeNode(scope_path=scope_paths[0] if scope_paths else None) local_items = [item for item in items if item.get('show_summary', True)] node.outline_items = local_items - # 5. Dispatch selections (guard: never re-select scope_path itself) + # 7. Dispatch selections (only present when action == NAVIGATE) valid_selections = [ s for s in selections - if s['path'] in selectable and s['path'] != scope_path + if s['path'] in selectable and s['path'] not in scope_path_set ] pending: list[dict] = [] @@ -827,8 +752,8 @@ async def scope_navigate_step( if item.get('is_leaf'): path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'}) else: - pending.append(sel) - # ★ NEW: Also hydrate this node's OWN direct chunks (not descendants) + # Non-leaf → will be batched into a single next call + pending.append({'path': path, 'confidence': conf}) path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'self_only'}) if path_selections: @@ -847,17 +772,12 @@ async def scope_navigate_step( exclude_sections=[], ) if connected: - # Resolve owner_section_path for connected assets: - # map target_chunk_id → the section_path of the text chunk - # that references it via connect_to. _owner_map = _build_connected_owner_map(chunks) for c in connected: if not c.get('owner_section_path'): c['owner_section_path'] = _owner_map.get(str(c.get('chunk_id') or '')) chunks = chunks + connected - # Resolve Root-stranded assets to their true owner sections - # via document-wide connect_to lookup _root_map = await _resolve_root_asset_owners( db, document_id=document_id, @@ -867,24 +787,23 @@ async def scope_navigate_step( if _root_map: for c in chunks: if c.get('owner_section_path'): - continue # already resolved by batch-level owner map + continue cid = str(c.get('chunk_id') or '') if cid in _root_map: c['owner_section_path'] = _root_map[cid] for chunk in chunks: - # Distribute chunk to its real path or fallback to the selection path real_path = chunk.get('owner_section_path') or chunk.get('section_path') or chunk.get('source_chunk_path') if real_path: node.add_leaf_chunks(str(real_path), [chunk]) - return node, pending + return action, asset_tools, node, pending except BudgetExceeded: raise except Exception as e: - logger.error(f' scope_navigate_step failed for doc={document_id}: {e}') - return empty, [] + logger.error(f' navigate_step failed for doc={document_id}: {e}') + return 'STOP', [], empty, [] # --------------------------------------------------------------------------- @@ -927,13 +846,15 @@ async def discovery_select_step( t0 = time.monotonic() try: - # 1. Format hints for LLM + # 1. Format hints for LLM (deduplicate by section_path) hint_lines: list[str] = [] hint_by_path: dict[str, dict] = {} for h in hints: sp = h.get('section_path', '') if not sp or sp == 'Root': continue + if sp in hint_by_path: + continue # skip duplicate section_path title = sp.rsplit(' / ', 1)[-1] if ' / ' in sp else sp summary = h.get('summary', '') or '' hint_lines.append(f'▸ path="{sp}" {title} [Leaf]') @@ -965,7 +886,9 @@ async def discovery_select_step( revision_context=revision_context, ) response = await llm_fn(prompt) - selections = _parse_scope_nav_response(response) + # Parse {"selections": [...]} response — reuse action parser's extraction + parsed = _parse_action_response(response) + selections = parsed.get('selections', []) logger.info( f' discovery_select_step doc="{doc_name}": ' diff --git a/packages/shared-python/shared/services/retrieval/agentic/trace.py b/packages/shared-python/shared/services/retrieval/agentic/trace.py index 8a58a243c..9c48e7a30 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/trace.py +++ b/packages/shared-python/shared/services/retrieval/agentic/trace.py @@ -50,6 +50,9 @@ def __init__( top_k: int = 10, data_type: int = 1, filters: dict[str, Any] | None = None, + parent_run_id: str | None = None, + workflow_step_id: str | None = None, + workflow_plan: dict[str, Any] | None = None, ) -> None: self._db = db self._run_id = f'aret_{uuid4().hex[:12]}' @@ -60,6 +63,9 @@ def __init__( self._top_k = top_k self._data_type = data_type self._filters = filters or {} + self._parent_run_id = parent_run_id + self._workflow_step_id = workflow_step_id + self._workflow_plan = workflow_plan self._steps: list[dict[str, Any]] = [] self._start_time = time.monotonic() self._created = False @@ -86,6 +92,9 @@ async def create_run(self) -> None: agentic_enabled=True, cache_hit=False, result_count=0, + parent_run_id=self._parent_run_id, + workflow_step_id=self._workflow_step_id, + workflow_plan=self._workflow_plan, latency_ms=0, created_at=_now_utc(), ) @@ -179,6 +188,10 @@ async def complete( } if budget_snapshot is not None: provenance['budget_snapshot'] = budget_snapshot + if self._parent_run_id: + provenance['parent_run_id'] = self._parent_run_id + if self._workflow_step_id: + provenance['workflow_step_id'] = self._workflow_step_id stmt = ( update(RetrievalRun) diff --git a/packages/shared-python/shared/services/retrieval/agentic/types.py b/packages/shared-python/shared/services/retrieval/agentic/types.py index 498bf08b5..a448a5f6b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/types.py @@ -44,7 +44,7 @@ class ToolResult: class DocTreeNode: """Unified navigation result tree for one document. - Produced by ``scope_navigate_step``. Captures the full + Produced by ``navigate_step``. Captures the full navigation outcome for rendering as a single hierarchy: - ``outline_items``: section tree items at this scope level @@ -218,7 +218,7 @@ class AgentState: """Mutable state carried through the 2-phase orchestrator. Phase 1: Document selection (discovery + KG) - Phase 2: Per-document navigation (scope_navigate_step per doc) + Phase 2: Per-document navigation (navigate_step per doc) Phase 3: Assembly + final verdict """ # Timing diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index c066a1c16..dc14c055a 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -982,6 +982,7 @@ async def run_retrieval_query( rerank: bool = False, threshold: float = 0.0, internal_recall_k: int | None = None, + enable_decomposition: bool | None = None, # deprecated: now always uses workflow ) -> dict[str, Any]: """Checkerboard retrieval: 3 independent channels -> RRF -> agent/graph union -> assembly.""" t_start = time.monotonic() @@ -1015,6 +1016,8 @@ async def run_retrieval_query( rerank=rerank, threshold=threshold, internal_recall_k=internal_recall_k, + # Always True: agentic mode now always routes through workflow + decomposition_enabled=True, ) cache_version: int | None = None @@ -1096,22 +1099,22 @@ async def run_retrieval_query( logger.info(f' ✅ Small KB: {len(results)} results in {elapsed_total}ms') return await _to_public_response(response) - # ══ Route: agentic vs legacy ══ + # ══ Route: agentic (unified workflow) vs legacy ══ _agentic_enabled = os.environ.get('RETRIEVAL_AGENTIC_ENABLED', 'false') == 'true' if _agentic_enabled: - # ── AGENTIC path (all errors self-contained, no fallback to legacy) ── - from shared.services.retrieval.agentic.orchestrator import RetrievalAgent - from shared.services.retrieval.llm_adapter import create_retrieval_llm_fn as _create_llm - - llm_fn = _create_llm() - agent = RetrievalAgent() - agentic_result = await agent.run( + # ── Unified agentic path via WorkflowOrchestrator ── + # Simple queries: planner returns a single-step plan (no decomposition). + # Complex queries: planner returns a multi-step plan with synthesize. + # Both go through the same code path. + from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator + + workflow = WorkflowOrchestrator() + workflow_result = await workflow.run( db, user_id=user_id, namespace=namespace, query=query, top_k=top_k, - llm_fn=llm_fn, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, data_type=data_type, @@ -1120,11 +1123,10 @@ async def run_retrieval_query( channels=channels, channel_weights=channel_weights, ) - router_used = agentic_result.router_used - # Generate asset URLs for media chunks in referenced_chunks + # Enrich referenced_chunks with asset URLs (images/tables) enriched_refs: list[dict[str, Any]] = [] - for ref in agentic_result.referenced_chunks: + for ref in workflow_result.referenced_chunks: enriched = dict(ref) chunk_type = _normalize_chunk_type(ref.get('chunk_type')) artifact_ref = ref.get('file_path', '') @@ -1140,30 +1142,9 @@ async def run_retrieval_query( logger.warning(f'Failed to generate agentic asset URL (ignored): {e}') enriched_refs.append(enriched) - # Build backward-compatible results[] from referenced_chunks - # (minimal: chunk_id + document_id + chunk_type + section_path) - results = [ - { - 'chunk_id': ref.get('chunk_id'), - 'document_id': ref.get('document_id'), - 'chunk_type': ref.get('chunk_type'), - 'source': { - 'document_id': ref.get('document_id'), - 'section_path': ref.get('section_path'), - }, - } - for ref in enriched_refs - ] - - response = { - "namespace": namespace, - "query": query, - "router_used": router_used, - "results": results, - "evidence_text": agentic_result.evidence_text, - "answer_text": agentic_result.answer_text, - "referenced_chunks": enriched_refs, - } + response = workflow_result.to_api_response() + # Override referenced_chunks with enriched versions + response['referenced_chunks'] = enriched_refs if cache_version is not None: try: @@ -1180,7 +1161,7 @@ async def run_retrieval_query( try: schedule_retrieval_hit_stats_update( user_id=user_id, namespace=namespace, - results=agentic_result.referenced_chunks, + results=enriched_refs, ) except Exception as e: logger.warning(f"Failed to trigger retrieval hit stats update (ignored): {e}") @@ -1190,14 +1171,15 @@ async def run_retrieval_query( f'\n{"█" * 70}\n' f' ✅ AGENTIC RETRIEVAL COMPLETE: ' f'{len(enriched_refs)} chunks | ' - f'evidence={len(agentic_result.evidence_text)} chars | ' - f'answer={len(agentic_result.answer_text)} chars | ' - f'router={router_used} | {elapsed_total}ms\n' + f'answer={len(workflow_result.answer_text)} chars | ' + f'router={workflow_result.router_used} | {elapsed_total}ms\n' f'{"█" * 70}' ) return await _to_public_response(response) + else: + # ── LEGACY path (existing code, unchanged) ── # ── Channel execution ── diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index b7c2bd88b..bb165eb41 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -8,6 +8,7 @@ from shared.services.redis import RedisServiceFactory _RETRIEVAL_CACHE_TTL_SECONDS = 300 +_WORKFLOW_PLAN_CACHE_TTL_SECONDS = 600 _VERSION_FALLBACK = 0 @@ -42,6 +43,7 @@ def _cache_shape_digest( rerank: bool = False, threshold: float = 0.0, internal_recall_k: int | None = None, + decomposition_enabled: bool | None = None, ) -> str: normalized_excludes = sorted(exclude_document_ids) normalized_sections = _normalize_exclude_sections(exclude_sections) @@ -55,6 +57,7 @@ def _cache_shape_digest( str(rerank), str(threshold), str(internal_recall_k), + str(decomposition_enabled), ] ) payload = f"{query}|{top_k}|{'|'.join(normalized_excludes)}|{'|'.join(normalized_sections)}|{extra}" @@ -178,3 +181,37 @@ async def set_cached_retrieval_query_result( response, ex=_RETRIEVAL_CACHE_TTL_SECONDS, ) + + +def _workflow_plan_cache_key(*, user_id: str, namespace: str, query: str) -> str: + digest = hashlib.sha256(query.encode("utf-8")).hexdigest() + return f"retrieval:workflow:plan:{user_id}:{namespace}:{digest}" + + +async def get_cached_workflow_plan( + *, + user_id: str, + namespace: str, + query: str, +) -> dict[str, Any] | None: + redis_service = RedisServiceFactory.get_service() + cached = await redis_service.get( + _workflow_plan_cache_key(user_id=user_id, namespace=namespace, query=query), + default=None, + ) + return cached if isinstance(cached, dict) else None + + +async def set_cached_workflow_plan( + *, + user_id: str, + namespace: str, + query: str, + plan: dict[str, Any], +) -> None: + redis_service = RedisServiceFactory.get_service() + await redis_service.set( + _workflow_plan_cache_key(user_id=user_id, namespace=namespace, query=query), + plan, + ex=_WORKFLOW_PLAN_CACHE_TTL_SECONDS, + ) diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index c6566d5f3..c5e592b5e 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -54,6 +54,21 @@ def _resolve_default_model() -> str: return getattr(settings, 'NORMOL_MODEL', None) or 'deepseek-chat' +def _resolve_planner_model(*, thinking: bool) -> str: + configured = getattr(settings, 'RETRIEVAL_PLANNER_MODEL', '') or '' + if configured: + return configured + if getattr(settings, 'DS_KEY', ''): + return 'deepseek-reasoner' if thinking else 'deepseek-chat' + if getattr(settings, 'ALI_API_KEYS', ''): + return 'qwq-32b-preview' if thinking else 'qwen-plus' + if getattr(settings, 'GLM_API_KEY', ''): + return 'glm-4-plus' if thinking else 'glm-4-flash' + if getattr(settings, 'GPT_API_KEY', ''): + return 'o3-mini' if thinking else 'gpt-4o-mini' + return getattr(settings, 'NORMOL_MODEL', None) or 'deepseek-chat' + + def create_retrieval_llm_fn( *, model: str | None = None, @@ -89,6 +104,37 @@ async def llm_fn(prompt: LLMFnInput) -> str: return llm_fn +def create_retrieval_planner_fn( + *, + thinking: bool = True, + model: str | None = None, + max_tokens: int = 8192, +) -> LLMFn | None: + """Create a reasoning-capable LLM callable for query planning.""" + if not _has_llm_credentials(): + logger.debug('retrieval: no LLM credentials configured, workflow planner disabled') + return None + + effective_model = model or _resolve_planner_model(thinking=thinking) + + async def llm_fn(prompt: LLMFnInput) -> str: + from shared.utils.OpenAICompatibleClientSync import get_openai_client + + client = get_openai_client(model=effective_model) + current_llm_usage.set(None) + result, usage = await asyncio.to_thread( + client.chat_completion_with_usage, + cast(Any, prompt), + model=effective_model, + temperature=0.0, + max_tokens=max_tokens, + ) + current_llm_usage.set(usage) + return result + + return llm_fn + + def create_retrieval_vlm_fn( *, model: str | None = None, diff --git a/packages/shared-python/shared/services/retrieval/workflow/__init__.py b/packages/shared-python/shared/services/retrieval/workflow/__init__.py new file mode 100644 index 000000000..67f896d3d --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/__init__.py @@ -0,0 +1,13 @@ +"""Query-decomposition retrieval workflow.""" +from .types import PlannedStep, QueryPlan, StepResult, WorkflowResult +from .wallet import BudgetWallet +from .orchestrator import WorkflowOrchestrator + +__all__ = [ + "BudgetWallet", + "PlannedStep", + "QueryPlan", + "StepResult", + "WorkflowResult", + "WorkflowOrchestrator", +] diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py new file mode 100644 index 000000000..9c5a20aff --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -0,0 +1,382 @@ +"""Workflow orchestrator for decomposed retrieval queries.""" +from __future__ import annotations + +import asyncio +import os +import time +from typing import Any +from uuid import uuid4 + +from loguru import logger +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.database import get_db_context +from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.orchestrator import RetrievalAgent +from shared.services.retrieval.agentic.types import AgenticResult +from shared.services.retrieval.cache_service import ( + get_cached_workflow_plan, + set_cached_workflow_plan, +) +from shared.services.retrieval.llm_adapter import ( + create_retrieval_llm_fn, + create_retrieval_planner_fn, +) +from shared.services.retrieval.workflow.planner import QueryPlanner +from shared.services.retrieval.workflow.synthesizer import compose_final_answer, synthesize_step +from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult, WorkflowResult +from shared.services.retrieval.workflow.wallet import BudgetWallet + + +class WorkflowOrchestrator: + """Plan and execute a query workflow DAG.""" + + def __init__(self) -> None: + self.parent_run_id = f'wret_{uuid4().hex[:12]}' + + async def run( + self, + db: AsyncSession, + *, + user_id: str, + namespace: str, + query: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int = 1, + signal_paths: list[str] | None = None, + filter_mode: str = 'delete', + channels: list[str] | None = None, + channel_weights: dict[str, float] | None = None, + llm_fn=None, + ) -> WorkflowResult: + t0 = time.monotonic() + llm_fn = llm_fn or create_retrieval_llm_fn() + planner_llm = create_retrieval_planner_fn(thinking=True) + planner_budget = _env_int('RETRIEVAL_PLANNER_THINKING_BUDGET', 4000) + wallet_total = _env_int('RETRIEVAL_WALLET_TOTAL_BUDGET', 200000) + per_retrieve = _env_int('RETRIEVAL_WALLET_PER_RETRIEVE_STEP_BUDGET', 40000) + per_synthesize = _env_int('RETRIEVAL_WALLET_PER_SYNTHESIZE_STEP_BUDGET', 6000) + max_steps = _env_int('RETRIEVAL_DECOMPOSITION_MAX_STEPS', 5) + + planner_ledger = BudgetLedger( + total=planner_budget, + planning_ratio=0.0, + bootstrap=planner_budget, + per_doc_min_share=0, + ) + plan = await self._load_or_plan( + user_id=user_id, + namespace=namespace, + query=query, + planner_llm=planner_llm, + planner_ledger=planner_ledger, + max_steps=max_steps, + wallet_total=wallet_total, + per_retrieve=per_retrieve, + ) + + wallet = BudgetWallet( + total=wallet_total, + per_retrieve_step_default=per_retrieve, + per_synthesize_step_default=per_synthesize, + ) + ledgers = await wallet.allocate(plan) + results_by_id: dict[str, StepResult] = {} + sem = asyncio.Semaphore(_env_int('RETRIEVAL_WORKFLOW_PARALLEL_MAX', 3)) + + for batch in plan.topological_batches(): + await asyncio.gather( + *[ + self._run_step( + db, + step=step, + ledger=ledgers[step.id], + results_by_id=results_by_id, + semaphore=sem, + user_id=user_id, + namespace=namespace, + top_k=step.top_k or top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=step.data_type or data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + llm_fn=llm_fn, + ) + for step in batch + ] + ) + for step in batch: + await wallet.reclaim(step.id, ledgers[step.id]) + + answer_text = compose_final_answer(plan, results_by_id) + ordered_results = [results_by_id[step.id] for step in plan.steps if step.id in results_by_id] + referenced_chunks = _dedupe_references( + ref for step_result in ordered_results for ref in step_result.referenced_chunks + ) + api_results = _references_to_results(referenced_chunks) + elapsed_ms = int((time.monotonic() - t0) * 1000) + logger.info( + 'workflow retrieval DONE: steps={} refs={} answer_chars={} elapsed={}ms', + len(ordered_results), + len(referenced_chunks), + len(answer_text), + elapsed_ms, + ) + return WorkflowResult( + namespace=namespace, + query=query, + router_used='workflow_decomposed' if len(plan.steps) > 1 else 'workflow_single_step', + answer_text=answer_text, + plan=plan, + steps=ordered_results, + referenced_chunks=referenced_chunks, + results=api_results, + final_strategy_used=plan.final_strategy, + wallet_snapshot=wallet.snapshot(), + planner_snapshot=planner_ledger.snapshot(), + parent_run_id=self.parent_run_id, + ) + + async def _load_or_plan( + self, + *, + user_id: str, + namespace: str, + query: str, + planner_llm, + planner_ledger: BudgetLedger, + max_steps: int, + wallet_total: int, + per_retrieve: int, + ) -> QueryPlan: + try: + cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query) + if cached: + return QueryPlan.from_dict(cached, original_query=query) + except Exception as exc: + logger.warning(f'workflow plan cache read failed (ignored): {exc}') + + planner = QueryPlanner( + llm_fn=planner_llm, + planner_ledger=planner_ledger, + max_steps=max_steps, + total_budget=wallet_total, + per_step_budget=per_retrieve, + ) + plan = await planner.plan(query=query) + try: + await set_cached_workflow_plan( + user_id=user_id, + namespace=namespace, + query=query, + plan=plan.to_dict(), + ) + except Exception as exc: + logger.warning(f'workflow plan cache write failed (ignored): {exc}') + return plan + + async def _run_step( + self, + db: AsyncSession, + *, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + semaphore: asyncio.Semaphore, + user_id: str, + namespace: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int, + signal_paths: list[str] | None, + filter_mode: str, + channels: list[str] | None, + channel_weights: dict[str, float] | None, + llm_fn, + ) -> None: + async with semaphore: + if step.step_kind == 'synthesize': + await self._run_synthesize_step(step, ledger, results_by_id, llm_fn) + return + await self._run_retrieve_step( + db, + step=step, + ledger=ledger, + results_by_id=results_by_id, + user_id=user_id, + namespace=namespace, + top_k=top_k, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + llm_fn=llm_fn, + ) + + async def _run_retrieve_step( + self, + db: AsyncSession, + *, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + user_id: str, + namespace: str, + top_k: int, + exclude_document_ids: list[str], + exclude_sections: list[dict[str, str]], + data_type: int, + signal_paths: list[str] | None, + filter_mode: str, + channels: list[str] | None, + channel_weights: dict[str, float] | None, + llm_fn, + ) -> None: + try: + # AsyncSession is not safe for concurrent use. Workflow steps may + # run in the same topological batch, so each retrieve step opens an + # isolated session and leaves the parent session untouched. + async with get_db_context() as step_db: + agentic_result = await RetrievalAgent().run( + step_db, + user_id=user_id, + namespace=namespace, + query=step.sub_query, + top_k=top_k, + llm_fn=llm_fn, + exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + data_type=data_type, + signal_paths=signal_paths, + filter_mode=filter_mode, + channels=channels, + channel_weights=channel_weights, + ledger=ledger, + parent_run_id=self.parent_run_id, + workflow_step_id=step.id, + ) + results_by_id[step.id] = _step_result_from_agentic(step, agentic_result) + except Exception as exc: + logger.exception(f'workflow retrieve step failed: step_id={step.id}') + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status='error', + error=str(exc), + budget_snapshot=ledger.snapshot(), + ) + + async def _run_synthesize_step( + self, + step: PlannedStep, + ledger: BudgetLedger, + results_by_id: dict[str, StepResult], + llm_fn, + ) -> None: + if llm_fn is None: + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status='skipped', + answer_text='', + error='llm unavailable for synthesis', + budget_snapshot=ledger.snapshot(), + ) + return + prior = {dep: results_by_id[dep] for dep in step.depends_on if dep in results_by_id} + try: + answer = await synthesize_step(step, prior_results=prior, llm_fn=llm_fn, ledger=ledger) + refs = _dedupe_references( + ref for result in prior.values() for ref in result.referenced_chunks + ) + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status='done', + answer_text=answer, + referenced_chunks=refs, + budget_snapshot=ledger.snapshot(), + ) + except Exception as exc: + results_by_id[step.id] = StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status='budget_stop' if 'budget' in str(exc).lower() else 'error', + answer_text='(budget exhausted)' if 'budget' in str(exc).lower() else '', + error=str(exc), + budget_snapshot=ledger.snapshot(), + ) + + +def _step_result_from_agentic(step: PlannedStep, result: AgenticResult) -> StepResult: + status = 'budget_stop' if 'budget' in (result.stop_reason or '') else 'done' + return StepResult( + step_id=step.id, + sub_query=step.sub_query, + step_kind=step.step_kind, + depends_on=step.depends_on, + output_role=step.output_role, + status=status, # type: ignore[arg-type] + answer_text=result.answer_text, + evidence_text=result.evidence_text, + referenced_chunks=result.referenced_chunks, + budget_snapshot=result.budget_snapshot, + router_used=result.router_used, + stop_reason=result.stop_reason, + ) + + +def _dedupe_references(refs) -> list[dict[str, Any]]: + seen: set[str] = set() + out: list[dict[str, Any]] = [] + for ref in refs: + chunk_id = str(ref.get('chunk_id') or '') + key = chunk_id or str(ref) + if key in seen: + continue + seen.add(key) + out.append(dict(ref)) + return out + + +def _references_to_results(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + 'chunk_id': ref.get('chunk_id'), + 'document_id': ref.get('document_id'), + 'chunk_type': ref.get('chunk_type'), + 'source': { + 'document_id': ref.get('document_id'), + 'section_path': ref.get('section_path'), + }, + } + for ref in refs + ] + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, str(default))) + except (TypeError, ValueError): + return default diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py new file mode 100644 index 000000000..e34713db6 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -0,0 +1,239 @@ +"""Query planner for decomposed retrieval workflows.""" +from __future__ import annotations + +import json +import re +import time +from typing import Any + +from loguru import logger + +from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger +from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage +from shared.services.retrieval.workflow.types import FinalStrategy, OutputRole, PlannedStep, QueryPlan, StepKind +from shared.utils.token_estimate import estimate_tokens + + +_PLAN_SCHEMA = { + "reasoning_summary": "