diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index 5aa877bf1..a253100c2 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -608,6 +608,21 @@ def iter_block_items(doc_data): # piece-by-piece loses textual overlay and positioning. # Future plan: Use LibreOffice headless conversion to render the entire document # and map the perfectly rendered images back to the layout via text anchors. + # + # Temporary: detect VML-only paragraphs and inject a placeholder so the + # paragraph isn't silently swallowed, leaving its parent section empty. + if not text and not seen_rids: + # No text and no DrawingML images — check for VML content + vml_groups = elem.xpath(".//v:group", namespaces=ns) + vml_images_check = elem.xpath(".//v:imagedata", namespaces=ns) + if vml_groups or vml_images_check: + vml_placeholder = "[VML graphic \u2014 extraction not yet supported]" + yield ele_num, vml_placeholder, "PTXT", None + ele_num += 1 + logger.debug( + f"Injected VML placeholder for paragraph with " + f"{len(vml_groups)} v:group, {len(vml_images_check)} v:imagedata" + ) """ # images (VML: ) — convert to PNG from PIL import Image as PILImage @@ -942,6 +957,13 @@ def convert_doc2dics( for _, row in leaf_dics.iterrows(): key = row["path_identifier"] + # Skip leaf nodes with no actual content (empty heading-only sections) + content_lst = row["content_lst"] + joined = "\n".join(content_lst).strip() + if not joined: + logger.debug(f"Skipping empty leaf node: {key}") + continue + # Build tentative path to check for duplicates tentative_path = doc_name + split_char + key @@ -954,7 +976,7 @@ def convert_doc2dics( path_counter[tentative_path] = 1 path_keys.append((doc_name + split_char + key)) - bottom_content = "\n".join(row["content_lst"]) + bottom_content = joined bottom_tokens = tokenize2stw_remove( [bottom_content], base_llm_paras["stopwords"] ) diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index 33f2d484f..ab85751bc 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -318,7 +318,10 @@ def parse_md( os.makedirs(tb_dir, exist_ok=True) img_dir = os.path.join(output_dir, "images") if os.path.isdir(img_dir): - shutil.rmtree(img_dir) + # Only remove parse_md's own output (image-N-*) from previous runs + for fname in os.listdir(img_dir): + if re.match(r"^image-\d+", fname): + os.remove(os.path.join(img_dir, fname)) os.makedirs(img_dir, exist_ok=True) # initialize vars diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index a5cceddca..a25d45cdc 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -1,4 +1,4 @@ -from .app_service import list_lexical_chunks, merge_channels_rrf, run_retrieval_query +from .app_service import merge_channels_rrf, run_retrieval_query from .cache_service import ( bump_retrieval_namespace_cache_version, get_cached_retrieval_query_result, @@ -13,7 +13,6 @@ __all__ = [ "create_retrieval_llm_fn", "run_retrieval_query", - "list_lexical_chunks", "merge_channels_rrf", "DocumentGraphService", "GraphQueryService", diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py index ad47faf61..ea127b69d 100644 --- a/packages/shared-python/shared/services/retrieval/agent_navigate.py +++ b/packages/shared-python/shared/services/retrieval/agent_navigate.py @@ -33,6 +33,10 @@ Do not include any explanation. """ +_VALID_HYDRATE_MODES = frozenset({ + 'outline', 'chunks', 'assets_only', 'image_only', 'table_only', +}) + _SCOPE_NAV_PROMPT = """\ You are a document navigation assistant. @@ -41,6 +45,7 @@ Below are candidate section paths at this scope level (up to 2 depth levels). Indented items are sub-items of the item above. +Each item shows text/image/table counts. Select section paths directly. A selected section path represents the chunks under that section subtree; do not ask to drill deeper. @@ -51,9 +56,19 @@ User query: {query} Select the most relevant section paths (at most {max_select}). +If NO section path is relevant to the query, you MUST return an empty array []. Do not force-select irrelevant sections. Prefer specific sub-items over broad parents when both are listed and the sub-item is sufficient. + +For each selected path, assign a confidence score (0.0 to 1.0) where 1.0 means exactly answers the query and 0.5 means tangentially related. +Also choose a hydrate_mode: +- "chunks" (default) return all text/image/table chunks +- "outline" return only section title + summary, no chunk content +- "assets_only" return only image and table chunks +- "image_only" return only image chunks +- "table_only" return only table chunks + Return ONLY a JSON array: -[{{"path": "section/path", "confidence": 0.9}}, ...] +[{{"path": "section/path", "confidence": , "hydrate_mode": "chunks"}}, ...] Do not include any explanation. """ @@ -106,7 +121,7 @@ def _parse_chunk_path_selections(text: str) -> list[dict[str, Any]]: """Parse chunk path selections from LLM output. Accepts either a legacy JSON array of strings or a structured array of - objects with `path` and optional `confidence`. + objects with `path`, optional `confidence`, and optional `hydrate_mode`. """ payload = _extract_json_array_payload(text) selections: list[dict[str, Any]] = [] @@ -114,16 +129,19 @@ def _parse_chunk_path_selections(text: str) -> list[dict[str, Any]]: if isinstance(item, str): path = item.strip() if path: - selections.append({'path': path, 'confidence': None}) + selections.append({'path': path, 'confidence': None, 'hydrate_mode': 'chunks'}) continue if not isinstance(item, dict): continue path = str(item.get('path') or item.get('chunk_path') or '').strip() if not path: continue + raw_mode = str(item.get('hydrate_mode') or '').strip().lower() + hydrate_mode = raw_mode if raw_mode in _VALID_HYDRATE_MODES else 'chunks' selections.append({ 'path': path, 'confidence': _normalize_confidence(item.get('confidence')), + 'hydrate_mode': hydrate_mode, }) return selections @@ -220,8 +238,8 @@ def _format_items_for_llm( Always shows ALL items (L1 + L2). Overflow controls whether summaries are included — not which levels are shown. - Normal: path + title + chunk_count + assets_count + summary - Overflow: path + title + chunk_count + assets_count (no summary) + Normal: path + title + text=N image=I table=T + summary + Overflow: path + title + text=N image=I table=T (no summary) Returns (text, overflowed). """ @@ -238,10 +256,13 @@ def _render_line(item: dict, include_summary: bool) -> str: line = f'{indent}- path="{item["path"]}" title="{item["title"]}"' chunk_count = item.get('chunk_count', 0) if chunk_count > 0: - line += f' chunks={chunk_count}' - assets = item.get('assets_count', 0) - if assets > 0: - line += f' assets={assets}' + line += f' text={chunk_count}' + image_count = item.get('image_count', 0) + if image_count > 0: + line += f' image={image_count}' + table_count = item.get('table_count', 0) + if table_count > 0: + line += f' table={table_count}' if include_summary: summary = item.get('summary') or item.get('title', '') if summary: @@ -410,16 +431,20 @@ async def _load_child_sections( document_id: str, job_result_id: str, scope_path: str | None = None, + exclude_paths: set[str] | None = None, ) -> list[dict]: """Load the next 2 available section depth bands under *scope_path*. Returns a flat list sorted by sort_order, each item: - {path, title, summary, chunk_count, assets_count, level} + {path, title, summary, chunk_count, image_count, table_count, level} - level=1: nearest available descendant depth under scope - level=2: second nearest available descendant depth under scope - chunk_count: text chunks under this section (excluding image/table) - - assets_count: image + table chunks under this section + - image_count: image chunks under this section + - table_count: table chunks under this section + - exclude_paths: paths already seen in prior revision rounds; + any path matching (exact or subtree) is skipped """ # ── Fetch all sections for this document revision ──────────────────── stmt = ( @@ -465,6 +490,7 @@ async def _load_child_sections( # round's relative L1/L2 instead of synthesizing missing ancestors. visible_sections: list[tuple[str, dict, int]] = [] visible_depths: set[int] = set() + _excl = exclude_paths or set() for path, meta in all_sections.items(): parts = meta['parts'] if scope_parts and ( @@ -474,6 +500,12 @@ async def _load_child_sections( relative_depth = len(parts) - scope_depth if relative_depth < 1: continue + # Skip paths already seen in prior revision rounds + if _excl and any( + path == ep or path.startswith(ep + ' / ') or ep.startswith(path + ' / ') + for ep in _excl + ): + continue visible_sections.append((path, meta, relative_depth)) visible_depths.add(relative_depth) @@ -497,14 +529,15 @@ async def _load_child_sections( 'level': level, 'sort_order': meta['sort_order'], 'chunk_count': 0, - 'assets_count': 0, + 'image_count': 0, + 'table_count': 0, 'section_id': meta['section_id'], } if not items_by_path: return [] - # ── Count chunks per section (text vs assets) ─────────────────────── + # ── Count chunks per section (text / image / table) ────────────────── section_ids = [meta['section_id'] for meta in all_sections.values()] if section_ids: from sqlalchemy import case, literal_column @@ -518,9 +551,14 @@ async def _load_child_sections( ).label('text_count'), func.count( case( - (DocumentChunk.chunk_type.in_(['image', 'table']), literal_column('1')), + (DocumentChunk.chunk_type == 'image', literal_column('1')), + ) + ).label('image_count'), + func.count( + case( + (DocumentChunk.chunk_type == 'table', literal_column('1')), ) - ).label('asset_count'), + ).label('table_count'), ) .where(DocumentChunk.document_id == document_id) .where(DocumentChunk.job_result_id == job_result_id) @@ -528,8 +566,8 @@ async def _load_child_sections( .group_by(DocumentChunk.section_id) ) chunk_rows = (await db.execute(chunk_stmt)).all() - section_id_counts: dict[str, tuple[int, int]] = { - sid: (int(tc), int(ac)) for sid, tc, ac in chunk_rows + section_id_counts: dict[str, tuple[int, int, int]] = { + sid: (int(tc), int(ic), int(tbc)) for sid, tc, ic, tbc in chunk_rows } else: section_id_counts = {} @@ -538,7 +576,7 @@ async def _load_child_sections( sid_to_path = {meta['section_id']: path for path, meta in all_sections.items()} # Aggregate chunk counts upward: each item gets counts from itself + descendants - for sid, (text_c, asset_c) in section_id_counts.items(): + for sid, (text_c, img_c, tbl_c) in section_id_counts.items(): chunk_path = sid_to_path.get(sid, '') if not chunk_path: continue @@ -546,7 +584,8 @@ async def _load_child_sections( for item_path, item in items_by_path.items(): if chunk_path == item_path or chunk_path.startswith(item_path + ' / '): item['chunk_count'] += text_c - item['assets_count'] += asset_c + item['image_count'] += img_c + item['table_count'] += tbl_c # ── Sort: interleave L2 under their L1 parent ──────────────────────── # diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index c1cc3000a..5893b416a 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -130,7 +130,27 @@ async def run( 'channel_weights': channel_weights, } - # ── Agent loop ── + # ── Mandatory pre-step: bottom discovery ───────────────────────────── + # bottom_discovery is always the first action; running it via the LLM + # policy wastes ~1-2s on a trivial LLM call. We execute it directly + # and let the LLM loop start from step 2 (kg_document_select etc.). + logger.info(' agentic: running mandatory bottom_discovery pre-step') + discovery_result = await self._execute_tool( + db, ActionType.BOTTOM_DISCOVERY, state, config, **tool_kwargs, + ) + state.apply(ActionType.BOTTOM_DISCOVERY, discovery_result) + if trace_enabled: + trace.record_step( + ActionType.BOTTOM_DISCOVERY, discovery_result, + decision_reason='mandatory_pre_step', + ) + state.step_count += 1 + logger.info( + f' agentic step {state.step_count} (pre-step): action=bottom_discovery ' + f'status={discovery_result.status} latency={discovery_result.latency_ms}ms' + ) + + # ── Agent loop (LLM decisions start from here) ──────────────────────── stop_reason = 'max_steps' while state.step_count < config.max_steps: if state.elapsed_ms >= config.latency_budget_ms: @@ -143,13 +163,9 @@ async def run( break if policy is None: - # No LLM: run discovery once then stop - if not state.discovery_done: - action_type = ActionType.BOTTOM_DISCOVERY - decision_reason = 'no_llm_fn: discovery only' - else: - stop_reason = 'no_llm_fn' - break + # No LLM: discovery already ran — stop + stop_reason = 'no_llm_fn' + break else: action_type, decision_reason = await policy.decide(state, config) @@ -182,20 +198,127 @@ async def run( f'docs={len(state.selected_docs)} paths={len(state.selected_paths)}' ) - # ── Fixed terminal step: hydrate + rank ── - ranked_rows = await self._hydrate_and_rank( - db, state, user_id=user_id, namespace=namespace, top_k=top_k, - ) + # ── Terminal: hydrate + rank + attempt_answer loop ── + while True: + ranked_rows = await self._hydrate_and_rank( + db, state, user_id=user_id, namespace=namespace, top_k=top_k, + ) + + # Include kept rows from prior revision rounds + if state.kept_path_rows: + ranked_rows = state.kept_path_rows + ranked_rows + + # Check if we should attempt_answer (need LLM + results + revision budget) + if ( + policy is None + or not ranked_rows + or state.revision_count >= config.max_revisions + ): + break + + # KG-exhausted guard: if all selected docs have been explored, + # further revision won't find new content — skip attempt_answer + all_selected_ids = {d.document_id for d in state.selected_docs} + unexplored = all_selected_ids - state.ever_explored_doc_ids + kg_exhausted = len(unexplored) == 0 and len(all_selected_ids) > 0 + if kg_exhausted: + logger.info( + f' agentic: KG exhausted — all {len(all_selected_ids)} docs explored, ' + f'skipping attempt_answer' + ) + stop_reason = 'kg_exhausted' + break + + # Three-state verdict from LLM + verdict, verdict_reason = await policy.attempt_answer( + state, config, ranked_rows, + ) + logger.info( + f' agentic attempt_answer: verdict={verdict} ' + f'revision={state.revision_count}/{config.max_revisions} ' + f'reason="{verdict_reason}"' + ) + + if verdict == 'DONE': + stop_reason = 'attempt_done' + break + + if verdict in ('NOT_SUFFICIENT', 'NOT_FOUND'): + state.revision_count += 1 + # Save current results as kept rows for next round + state.kept_path_rows = ranked_rows + # Record current selected paths as seen + for p in state.selected_paths: + doc_id = p.get('document_id', '') + path = p.get('path', '') + if doc_id and path: + state.seen_section_keys.add(f'{doc_id}::{path}') + # Reset navigation state for re-exploration + state.selected_paths.clear() + state.selected_docs.clear() # Bug 1 fix: prevent doc accumulation + state.pending_doc_index = 0 + state.kg_done = False + state.discovery_done = False + + # Mandatory bottom_discovery pre-step for revision round + logger.info( + f' agentic: running mandatory bottom_discovery pre-step ' + f'(revision {state.revision_count})' + ) + rev_discovery = await self._execute_tool( + db, ActionType.BOTTOM_DISCOVERY, state, config, **tool_kwargs, + ) + state.apply(ActionType.BOTTOM_DISCOVERY, rev_discovery) + if trace_enabled: + trace.record_step( + ActionType.BOTTOM_DISCOVERY, rev_discovery, + decision_reason=f'mandatory_pre_step (revision {state.revision_count})', + ) + state.step_count += 1 + logger.info( + f' agentic step {state.step_count} (rev {state.revision_count} pre-step): ' + f'action=bottom_discovery status={rev_discovery.status}' + ) + + # Re-enter agent loop + while state.step_count < config.max_steps: + if state.elapsed_ms >= config.latency_budget_ms: + stop_reason = 'latency_budget' + break + + action_type, decision_reason = await policy.decide(state, config) + if action_type is None or action_type == ActionType.DONE: + stop_reason = 'llm_done' + break + + result = await self._execute_tool( + db, action_type, state, config, **tool_kwargs, + ) + state.apply(action_type, result) + if trace_enabled: + trace.record_step(action_type, result, decision_reason=decision_reason) + state.step_count += 1 + logger.info( + f' agentic step {state.step_count} (rev {state.revision_count}): ' + f'action={action_type.value} status={result.status}' + ) + + # Loop back to hydrate + attempt_answer + continue + + # Unknown verdict — treat as DONE + break router_used = ( - 'agentic_llm' if state.selected_paths + 'agentic_llm' if state.selected_paths or state.kept_path_rows else 'agentic_discovery_only' ) logger.info( f'agentic retrieval DONE: {len(ranked_rows)} results, ' f'router={router_used}, steps={state.step_count}, ' - f'stop_reason={stop_reason}, {state.elapsed_ms}ms' + f'stop_reason={stop_reason}, revisions={state.revision_count}, ' + f'{state.elapsed_ms}ms' ) if trace_enabled: @@ -244,6 +367,13 @@ async def _execute_tool( payload={'document_id': doc.document_id, 'reason': 'no job_result_id'}, ) + # Build exclude_paths for this doc from seen_section_keys + doc_exclude = { + key.split('::', 1)[1] + for key in state.seen_section_keys + if key.startswith(f'{doc.document_id}::') + } if state.seen_section_keys else None + return await tools.document_path_select( db, user_id=kwargs['user_id'], @@ -253,6 +383,7 @@ async def _execute_tool( document_id=doc.document_id, job_result_id=job_result_id, doc_name=doc.source_file_name or state.doc_id_to_name.get(doc.document_id, ''), + exclude_paths=doc_exclude, ) elif action_type == ActionType.GREP_DOCUMENT_DISCOVER: @@ -299,6 +430,7 @@ async def _hydrate_and_rank( a configurable LLM context window budget rather than a fixed count. """ try: + # Hydrate agent-selected paths navigated_paths: list[dict[str, Any]] = [] if state.selected_paths: diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 7858bcac0..12e5afa62 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -23,15 +23,11 @@ # ── Available actions presented to the LLM ─────────────────────────────────── +# NOTE: BOTTOM_DISCOVERY is intentionally excluded from this list. +# It is now a mandatory pre-step executed automatically by the orchestrator +# before the LLM decision loop begins. The LLM should never need to decide +# whether to run it — doing so wastes one LLM call per run. _AVAILABLE_ACTIONS: list[dict[str, Any]] = [ - { - 'action': ActionType.BOTTOM_DISCOVERY.value, - 'description': ( - 'Run BM25 3-channel bottom-layer discovery (path / content / term). ' - 'Always call this first to get candidate chunks and top document hints.' - ), - 'when': 'discovery_done is false', - }, { 'action': ActionType.KG_DOCUMENT_SELECT.value, 'description': ( @@ -94,12 +90,12 @@ {actions_block} RULES: -1. Always run bottom_discovery first (if discovery_done is false). -2. After discovery, run kg_document_select (if kg_done is false). -3. After kg select, run document_path_select for each pending document. -4. Call done when all pending documents are processed OR you have >= {min_evidence} evidence paths. -5. Only use grep_document_discover if kg_document_select found 0 documents. -6. Only use graph_expand_docs if you need more related docs after reviewing results. +1. Run kg_document_select when discovery_done is true and kg_done is false. +2. After kg select, run document_path_select for each pending document. +3. Call done when all pending documents are processed OR you have >= {min_evidence} evidence paths. +4. Only use grep_document_discover if kg_document_select found 0 documents. +5. Only use graph_expand_docs if you need more related docs after reviewing results. +Note: bottom_discovery is already executed automatically before this loop — do NOT attempt to call it. Return ONLY a JSON object, no markdown, no explanation: {{"action": "", "reason": ""}} @@ -145,9 +141,26 @@ def __init__(self, llm_fn: LLMFn, *, query: str = '') -> None: def build_prompt(self, state: AgentState, config: AgentRunConfig) -> str: """Build the decision prompt. Public for test inspection.""" state_data = state.state_summary() + + has_pending_docs = state.pending_doc_index < len(state.selected_docs) + state_data['has_pending_docs'] = has_pending_docs state_json = json.dumps(state_data, ensure_ascii=False, indent=2) - # Count pending docs + allowed_actions = [] + for action in _AVAILABLE_ACTIONS: + name = action['action'] + if name == ActionType.KG_DOCUMENT_SELECT.value and (not state.discovery_done or state.kg_done): + continue + if name == ActionType.DOCUMENT_PATH_SELECT.value and (not state.kg_done or not has_pending_docs): + continue + if name == ActionType.GREP_DOCUMENT_DISCOVER.value and (not state.kg_done or len(state.selected_docs) > 0): + continue + allowed_actions.append(action) + + actions_block = '\n'.join( + f" {i+1}. \"{a['action']}\": {a['description']} [{a['when']}]" + for i, a in enumerate(allowed_actions) + ) return _POLICY_PROMPT_TEMPLATE.format( query=self._query, @@ -156,7 +169,7 @@ def build_prompt(self, state: AgentState, config: AgentRunConfig) -> str: budget_ms=config.latency_budget_ms, step=state.step_count, max_steps=config.max_steps, - actions_block=_ACTIONS_BLOCK, + actions_block=actions_block, min_evidence=config.min_evidence_paths, ) @@ -232,3 +245,73 @@ async def decide( return None, f'unknown_action: {action_str}' return action_type, reason + + async def attempt_answer( + self, + state: AgentState, + config: AgentRunConfig, + ranked_rows: list[dict[str, Any]], + ) -> tuple[str, str]: + """Three-state verdict: is the evidence sufficient? + + Returns (verdict, reason) where verdict is one of: + - 'DONE': evidence is sufficient, stop searching + - 'NOT_SUFFICIENT': partial match, need more evidence + - 'NOT_FOUND': no relevant evidence found at all + """ + # Build evidence summary (paths + previews, not full content) + evidence_lines: list[str] = [] + for i, row in enumerate(ranked_rows[:20]): # cap to avoid huge prompt + path = row.get('section_path') or row.get('source_chunk_path') or '' + content_preview = str(row.get('content', ''))[:150] + score = round(float(row.get('score', 0.0) or 0.0), 3) + evidence_lines.append( + f' {i+1}. path="{path}" score={score}\n' + f' preview: {content_preview}' + ) + evidence_text = '\n'.join(evidence_lines) or '(no evidence collected)' + + prompt = _ATTEMPT_ANSWER_PROMPT.format( + query=self._query, + evidence_count=len(ranked_rows), + evidence_summary=evidence_text, + revision_count=state.revision_count, + max_revisions=config.max_revisions, + ) + + raw_response = await self._llm_fn(prompt) + logger.info(f' [LLMPolicy.attempt_answer] raw={repr(raw_response[:200])}') + + parsed = _parse_action_from_response(raw_response) + if not parsed: + return 'DONE', 'parse_error — treating as done' + + verdict = str(parsed.get('verdict', 'DONE')).strip().upper() + reason = str(parsed.get('reason', '')).strip() + + if verdict not in ('DONE', 'NOT_SUFFICIENT', 'NOT_FOUND'): + verdict = 'DONE' + + return verdict, reason + + +_ATTEMPT_ANSWER_PROMPT = """\ +You are evaluating whether the collected evidence can answer the user's query. + +QUERY: "{query}" + +EVIDENCE ({evidence_count} items, showing top 20): +{evidence_summary} + +REVISION: {revision_count} of {max_revisions} revisions used. + +Evaluate the evidence and return ONE verdict: +- "DONE": The evidence is sufficient to answer the query. Use this if the main points are covered. +- "NOT_SUFFICIENT": Partial match — some relevant info found but key aspects are missing. Only use if more searching could realistically help. +- "NOT_FOUND": The evidence is completely irrelevant to the query. Only use if nothing matches at all. + +When in doubt, prefer DONE — avoid unnecessary extra search rounds. + +Return ONLY a JSON object: +{{"verdict": "DONE", "reason": "one sentence explanation"}} +""" diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 1463bea61..52c7645d3 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -276,6 +276,7 @@ async def document_path_select( job_result_id: str, doc_name: str = '', max_chunks_per_file: int = 15, + exclude_paths: set[str] | None = None, **_kwargs: Any, ) -> ToolResult: """Document entry point for agentic scope navigation.""" @@ -289,6 +290,7 @@ async def document_path_select( db, document_id=document_id, job_result_id=job_result_id, query=query, llm_fn=llm_fn, doc_name=doc_name, scope_path=None, max_select=max_chunks_per_file, + exclude_paths=exclude_paths, ) @@ -429,6 +431,7 @@ async def scope_navigate( doc_name: str = '', scope_path: str | None = None, max_select: int = 15, + exclude_paths: set[str] | None = None, ) -> ToolResult: """Unified document-internal navigation tool. @@ -440,7 +443,10 @@ async def scope_navigate( """ t0 = time.monotonic() try: - items = await _load_child_sections(db, document_id, job_result_id, scope_path) + items = await _load_child_sections( + db, document_id, job_result_id, scope_path, + exclude_paths=exclude_paths, + ) if not items: latency = int((time.monotonic() - t0) * 1000) return ToolResult( @@ -471,7 +477,8 @@ async def scope_navigate( confidence = item.get('confidence') if confidence is None: confidence = _default_confidence_for_rank(len(accepted)) - accepted.append({'path': path, 'confidence': confidence}) + hydrate_mode = item.get('hydrate_mode', 'chunks') + accepted.append({'path': path, 'confidence': confidence, 'hydrate_mode': hydrate_mode}) if len(accepted) >= max_select: break diff --git a/packages/shared-python/shared/services/retrieval/agentic/types.py b/packages/shared-python/shared/services/retrieval/agentic/types.py index 01054c1e7..4877886de 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/types.py @@ -29,6 +29,7 @@ class AgentRunConfig: max_docs: int = 0 # 0 = no limit, LLM decides autonomously max_path_expansions: int = 2 max_doc_retries: int = 2 + max_revisions: int = 2 # max attempt_answer → revise cycles latency_budget_ms: int = 12000 min_evidence_paths: int = 1 @@ -89,6 +90,12 @@ class AgentState: doc_id_to_name: dict[str, str] = field(default_factory=dict) doc_job_map: dict[str, str] = field(default_factory=dict) + # Revision / three-state fields + revision_count: int = 0 + ever_explored_doc_ids: set[str] = field(default_factory=set) + seen_section_keys: set[str] = field(default_factory=set) # "{doc_id}::{section_path}" + kept_path_rows: list[dict[str, Any]] = field(default_factory=list) + @property def elapsed_ms(self) -> int: return int((time.monotonic() - self.start_time) * 1000) @@ -133,6 +140,9 @@ def state_summary(self) -> dict[str, Any]: 'selected_paths_count': len(self.selected_paths), 'selected_paths': selected_path_summaries, 'doc_retry_count': self.doc_retry_count, + 'revision_count': self.revision_count, + 'explored_doc_count': len(self.ever_explored_doc_ids), + 'kept_rows_count': len(self.kept_path_rows), 'last_observation': last_obs, } @@ -179,11 +189,19 @@ def apply(self, action_type: ActionType, result: ToolResult) -> None: elif action_type == ActionType.DOCUMENT_PATH_SELECT: if result.status == 'selected_paths': + doc_id = result.payload.get('document_id', '') new_paths = result.payload.get('selected_paths', []) + for p in new_paths: + p['document_id'] = doc_id self.selected_paths.extend(new_paths) self.pending_doc_index += 1 + if doc_id: + self.ever_explored_doc_ids.add(doc_id) elif result.status == 'no_items': + doc_id = result.payload.get('document_id', '') self.pending_doc_index += 1 + if doc_id: + self.ever_explored_doc_ids.add(doc_id) elif result.status == 'need_more_docs': failed_doc_id = result.payload.get('document_id', '') if failed_doc_id: @@ -191,9 +209,15 @@ def apply(self, action_type: ActionType, result: ToolResult) -> None: self.doc_retry_count += 1 self.kg_done = False # allow re-entry to KG select elif result.status == 'no_confident_match': + doc_id = result.payload.get('document_id', '') self.pending_doc_index += 1 + if doc_id: + self.ever_explored_doc_ids.add(doc_id) elif result.status == 'error': + doc_id = result.payload.get('document_id', '') self.pending_doc_index += 1 + if doc_id: + self.ever_explored_doc_ids.add(doc_id) elif action_type == ActionType.GREP_DOCUMENT_DISCOVER: if result.status == 'discovered_docs': diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index e6ee4a09f..b91e31f0c 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -251,89 +251,7 @@ async def assemble_retrieval_results( return assembled -async def list_lexical_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Independent lexical retrieval: path + content channels via ILIKE.""" - recall_k = top_k * _INTERNAL_RECALL_K_MULTIPLIER - excluded_docs = set(exclude_document_ids) - base_stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - ) - if excluded_docs: - base_stmt = base_stmt.where(Document.document_id.notin_(list(excluded_docs))) - - like = f'%{query}%' - content_stmt = base_stmt.where(DocumentChunk.content_lexical_text.ilike(like)).order_by(DocumentChunk.sort_order).limit(recall_k) - path_stmt = base_stmt.where(DocumentChunk.path_lexical_text.ilike(like)).order_by(DocumentChunk.sort_order).limit(recall_k) - - # AsyncSession is stateful and should not be shared across concurrent tasks. - content_result = await db.execute(content_stmt) - path_result = await db.execute(path_stmt) - - def _to_rows(result, channel_score: float) -> list[dict[str, Any]]: - rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - if is_excluded_section(document_id=document.document_id, section_path=section_path, exclude_sections=exclude_sections): - continue - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': channel_score, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - }) - return rows - - content_rows = _to_rows(content_result, _CHANNEL_WEIGHT_CONTENT) - path_rows = _to_rows(path_result, _CHANNEL_WEIGHT_PATH) - return content_rows, path_rows - - -def _grep_search_rows(rows: list[dict[str, Any]], query: str) -> list[dict[str, Any]]: - """Term/grep channel: exact substring matching with scoring from knowhere-kb.""" - import re - query_lower = query.lower().strip() - if not query_lower: - return [] - - units = re.findall(r'[一-鿿]+|[a-zA-Z0-9]+', query_lower) - units = [u for u in units if len(u) > 1] - - scored: list[tuple[float, dict[str, Any]]] = [] - for row in rows: - haystack = (str(row.get('content') or '') + ' ' + str(row.get('section_path') or '')).lower() - if query_lower in haystack: - scored.append((100.0, row)) - elif units: - hit_count = sum(1 for u in units if u in haystack) - if hit_count > 0: - scored.append((float(hit_count), row)) - - scored.sort(key=lambda x: x[0], reverse=True) - return [dict(row, score=score) for score, row in scored] def _merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -666,30 +584,50 @@ def _rank_candidates_by_path( if not candidate.get('section_path') and row.get('section_path'): candidate['section_path'] = row.get('section_path') - ranked_rows: list[dict[str, Any]] = [] + # ── Dual-priority ranking ──────────────────────────────────────────── + # When the agent produced results (routed_rows non-empty), rows with + # agent_score=0 are demoted to a fallback pool. Primary sort is by + # agent_score (LLM confidence, 0-1, cross-round comparable), with + # discovery_score as tiebreaker only. This avoids the old + # `max(agent, discovery)` which mixed incompatible score sources. + has_agent_results = len(routed_rows) > 0 + + primary_rows: list[dict[str, Any]] = [] + fallback_rows: list[dict[str, Any]] = [] + for key, row in merged.items(): discovery_score = float(row.get('discovery_score', 0.0) or 0.0) agent_score = float(row.get('agent_score', 0.0) or 0.0) row['dual_hit_flag'] = 1 if discovery_score > 0.0 and agent_score > 0.0 else 0 - row['evidence_score'] = round(max(discovery_score, agent_score), 6) + row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) row['score'] = row['evidence_score'] row['_candidate_order'] = insertion_order[key] - ranked_rows.append(row) - ranked_rows.sort( - key=lambda row: ( - float(row.get('evidence_score', 0.0) or 0.0), + if has_agent_results and agent_score <= 0.0: + fallback_rows.append(row) + else: + primary_rows.append(row) + + def _sort_key(row): + return ( + float(row.get('agent_score', 0.0) or 0.0), + float(row.get('discovery_score', 0.0) or 0.0), int(row.get('dual_hit_flag', 0) or 0), float(row.get('importance_norm_score', 0.0) or 0.0), - float(row.get('discovery_score', 0.0) or 0.0), - float(row.get('agent_score', 0.0) or 0.0), -int(row.get('_candidate_order', 0) or 0), - ), - reverse=True, - ) + ) + + primary_rows.sort(key=_sort_key, reverse=True) + ranked_rows = primary_rows[:top_k] + + # Back-fill from fallback if primary results are insufficient + if len(ranked_rows) < top_k and fallback_rows: + fallback_rows.sort(key=_sort_key, reverse=True) + ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) + for row in ranked_rows: row.pop('_candidate_order', None) - return ranked_rows[:top_k] + return ranked_rows async def _count_scoped_chunks( @@ -781,12 +719,19 @@ async def _hydrate_paths_to_rows( ) -> list[dict[str, Any]]: """Load full chunk rows by section_path or source_chunk_path. - Used to hydrate agent-selected paths into the standard row format - expected by assemble_retrieval_results(). + Supports hydrate_mode branching: + - 'chunks' (default): all chunk types under the section subtree + - 'outline': synthetic row from section metadata, no real chunks + - 'assets_only': only image + table chunks + - 'image_only': only image chunks + - 'table_only': only table chunks """ if not path_selections: return [] + + # Group selections by hydrate_mode confidence_by_path: dict[str, float] = {} + mode_by_path: dict[str, str] = {} ordered_paths: list[str] = [] for item in path_selections: raw_path = str(item.get('path') or '').strip() @@ -794,73 +739,145 @@ async def _hydrate_paths_to_rows( if not path: continue confidence = float(item.get('confidence', 0.0) or 0.0) + hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() if path not in confidence_by_path: ordered_paths.append(path) confidence_by_path[path] = confidence + mode_by_path[path] = hydrate_mode else: confidence_by_path[path] = max(confidence_by_path[path], confidence) if not ordered_paths: return [] - section_path_filters = [] - for path in ordered_paths: - section_path_filters.append(DocumentSection.section_path == path) - section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + # Separate outline paths from chunk-loading paths + outline_paths = [p for p in ordered_paths if mode_by_path.get(p) == 'outline'] + chunk_paths = [p for p in ordered_paths if mode_by_path.get(p) != 'outline'] - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where( - or_( - *section_path_filters, - DocumentChunk.source_chunk_path.in_(ordered_paths), + rows: list[dict[str, Any]] = [] + + # ── Outline mode: synthesize rows from section metadata ────────────── + if outline_paths: + outline_section_filters = [] + for path in outline_paths: + outline_section_filters.append(DocumentSection.section_path == path) + + outline_stmt = ( + select(Document, DocumentSection) + .join(DocumentSection, (DocumentSection.document_id == Document.document_id) + & (DocumentSection.job_result_id == Document.current_job_result_id)) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(or_(*outline_section_filters)) + ) + outline_result = await db.execute(outline_stmt) + for document, section in outline_result.all(): + agent_score = confidence_by_path.get(section.section_path, 0.0) + summary_text = (section.summary or '').strip() + title_text = (section.section_title or '').strip() + content = f'[Outline] {title_text}' + if summary_text: + content += f'\n{summary_text}' + rows.append({ + 'document_id': document.document_id, + 'chunk_id': f'outline_{section.section_id}', + 'section_id': section.section_id, + 'section_path': section.section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': 'outline', + 'content': content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': None, + 'chunk_metadata': {}, + 'job_result_id': section.job_result_id, + 'job_id': None, + 'source_chunk_path': None, + 'sort_order': section.sort_order, + 'hydrate_mode': 'outline', + }) + + # ── Chunk modes: load real chunks with optional type filters ───────── + if chunk_paths: + section_path_filters = [] + for path in chunk_paths: + section_path_filters.append(DocumentSection.section_path == path) + section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id)) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where( + or_( + *section_path_filters, + DocumentChunk.source_chunk_path.in_(chunk_paths), + ) ) ) - ) - result = await db.execute(stmt) + result = await db.execute(stmt) + + # Build a map of path → allowed chunk_types based on hydrate_mode + _MODE_ALLOWED_TYPES: dict[str, set[str] | None] = { + 'chunks': None, # all types + 'assets_only': {'image', 'table'}, + 'image_only': {'image'}, + 'table_only': {'table'}, + } + + seen_paths: set[str] = set() + for document, chunk, section, job_result in result.all(): + row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' + if row_path in seen_paths: + continue - # Build rows, preserving agent-selected order + # Find which ordered path this row belongs to + matched_path = row_path + if section and section.section_path not in confidence_by_path: + matched_path = next( + ( + path for path in chunk_paths + if section.section_path == path or section.section_path.startswith(f'{path} / ') + ), + row_path, + ) + + # Check chunk_type filter based on hydrate_mode + path_mode = mode_by_path.get(matched_path, 'chunks') + allowed_types = _MODE_ALLOWED_TYPES.get(path_mode) + if allowed_types is not None: + chunk_type_lower = (chunk.chunk_type or '').strip().lower() + if chunk_type_lower not in allowed_types: + continue + + seen_paths.add(row_path) + agent_score = confidence_by_path.get(matched_path, 0.0) + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + 'hydrate_mode': path_mode, + }) + + # ── Sort by agent-selected order ───────────────────────────────────── path_order = {p: idx for idx, p in enumerate(ordered_paths)} - rows: list[dict[str, Any]] = [] - seen_paths: set[str] = set() - for document, chunk, section, job_result in result.all(): - row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' - if row_path in seen_paths: - continue - seen_paths.add(row_path) - matched_path = row_path - if section and section.section_path not in confidence_by_path: - matched_path = next( - ( - path for path in ordered_paths - if section.section_path == path or section.section_path.startswith(f'{path} / ') - ), - row_path, - ) - agent_score = confidence_by_path.get(matched_path, 0.0) - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section.section_path if section else None, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'source_chunk_path': chunk.source_chunk_path, - 'sort_order': chunk.sort_order, - }) def _row_sort_key(row: dict[str, Any]) -> int: row_path = _get_row_path(row) @@ -877,6 +894,8 @@ def _row_sort_key(row: dict[str, Any]) -> int: path for path in ordered_paths if path in hydrated_paths or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) } + # Outline paths are always resolved (synthesized) + resolved_inputs |= set(outline_paths) missed = len(ordered_paths) - len(resolved_inputs) if missed > 0: missing_paths = [p for p in ordered_paths if p not in resolved_inputs] diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph_service.py index c8b989d06..8244a8a72 100644 --- a/packages/shared-python/shared/services/retrieval/graph_service.py +++ b/packages/shared-python/shared/services/retrieval/graph_service.py @@ -47,7 +47,9 @@ def is_excluded_section( for item in exclude_sections: if not isinstance(item, dict): continue - if document_id == str(item.get('document_id') or '').strip() and section_path == str(item.get('section_path') or '').strip(): + exc_doc = str(item.get('document_id') or '').strip() + exc_path = str(item.get('section_path') or '').strip() + if document_id == exc_doc and (section_path == exc_path or section_path.startswith(exc_path + ' / ')): return True return False @@ -368,7 +370,10 @@ async def find_entry_documents( exc_path = str(exc.get('section_path') or '').strip() if exc_doc and exc_path: stmt = stmt.where( - ~((DocumentSection.document_id == exc_doc) & (DocumentSection.section_path == exc_path)) + ~((DocumentSection.document_id == exc_doc) & ( + (DocumentSection.section_path == exc_path) | + DocumentSection.section_path.like(f'{exc_path} / %') + )) ) result = await db.execute(stmt) seen = [row[0] for row in result.all()] diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 3c547c0eb..5322fa809 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -309,28 +309,34 @@ def _publish_document_state_for_job( section_path = section_path_from_chunk_path(source_path) section = sections_by_path.get(section_path) if section is None: - parent_section_id = None path_parts = [p for p in section_path.split(" / ") if p] - if len(path_parts) > 1: - parent_path = " / ".join(path_parts[:-1]) - parent = sections_by_path.get(parent_path) - if parent is not None: - parent_section_id = parent.section_id - section = DocumentSection( - user_id=str(job.user_id), - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - parent_section_id=parent_section_id, - section_path=section_path, - section_title=path_parts[-1] if path_parts else None, - section_level=len(path_parts), - section_metadata={}, - sort_order=len(sections_by_path), - ) - db.add(section) - db.flush() - sections_by_path[section_path] = section + # Ensure all ancestor sections exist (top-down) + for depth in range(1, len(path_parts) + 1): + ancestor_path = " / ".join(path_parts[:depth]) + if ancestor_path in sections_by_path: + continue + ancestor_parent_id = None + if depth > 1: + parent_path = " / ".join(path_parts[:depth - 1]) + parent = sections_by_path.get(parent_path) + if parent is not None: + ancestor_parent_id = parent.section_id + ancestor_section = DocumentSection( + user_id=str(job.user_id), + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + parent_section_id=ancestor_parent_id, + section_path=ancestor_path, + section_title=path_parts[depth - 1], + section_level=depth, + section_metadata={}, + sort_order=len(sections_by_path), + ) + db.add(ancestor_section) + db.flush() + sections_by_path[ancestor_path] = ancestor_section + section = sections_by_path[section_path] chunk_id = chunk.get("chunk_id") or f"chunk_{uuid4().hex[:12]}" section_summary = section.summary if section else None