From eb13044e46aacd9781ec2ddaad6917e9257530e0 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 12 May 2026 15:30:53 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=9B=BE=E8=A1=A8=E5=B7=A5=E5=85=B7?= =?UTF-8?q?=E5=8A=A0=E8=BD=BD=20(Asset=20Tool=20Loading)=20&=20fix=20CodeQ?= =?UTF-8?q?L=20warning?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit implements asset tool loading and also fixes an uninitialized 'evidence_text' variable in orchestrator.py which caused the previous CodeQL PR check failure. --- .../services/retrieval/agent_navigate.py | 148 ++++++-- .../retrieval/agentic/orchestrator.py | 126 ++++++- .../services/retrieval/agentic/policy.py | 44 ++- .../services/retrieval/agentic/tools.py | 355 +++++++++++++++++- .../shared/services/retrieval/llm_adapter.py | 49 +++ 5 files changed, 677 insertions(+), 45 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py index 6bfaef2e8..824cf2e0c 100644 --- a/packages/shared-python/shared/services/retrieval/agent_navigate.py +++ b/packages/shared-python/shared/services/retrieval/agent_navigate.py @@ -29,9 +29,9 @@ === End Overview === User query: {query} - -Based on the query, select all documents that may contain relevant information. -Only skip documents that are clearly irrelevant to the query. +{revision_context} +Based on the query, select documents that may contain relevant information. +If NO document in the knowledge base is relevant to the query, return an EMPTY array []. Return ONLY a JSON array of document IDs, e.g.: ["doc_abc123", "doc_def456"] Do not include any explanation. """ @@ -59,13 +59,8 @@ - Select sections whose content is needed to answer the query. - If the titles and summaries already visible are sufficient (e.g. the query asks for an outline or overview), return an EMPTY list []. - Optional: set "mode" per selection to control what content is retrieved: - - "all" (default): retrieve all content types (text, images, tables) - - "image": retrieve only image assets from this section - - "table": retrieve only table assets from this section - Return ONLY a JSON object: -{{"selections": [{{"path": "...", "confidence": , "mode": "all"}}, ...]}} +{{"selections": [{{"path": "...", "confidence": }}, ...]}} Do not include any explanation. """ @@ -84,12 +79,12 @@ === 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": , "mode": "all"}}, ...]}} +{{"selections": [{{"path": "...", "confidence": }}, ...]}} Do not include any explanation. """ @@ -243,7 +238,6 @@ def _render_item(item: dict, include_summary: bool) -> str: show = item.get('show_summary', True) is_leaf = item.get('is_leaf', False) leaf_tag = ' [Leaf]' if is_leaf else '' - title = item.get('title', '') path = item.get('path', '') summary = item.get('summary') or '' @@ -268,7 +262,7 @@ def _render_item(item: dict, include_summary: bool) -> str: select_tag = '[SELECT] ' if show else '' lines: list[str] = [] - lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}" {title}{counts_str}{leaf_tag}') + lines.append(f'{indent}{prefix} {select_tag}{level_tag} path="{path}"{counts_str}{leaf_tag}') if include_summary and show and summary: sub_indent = " " * level @@ -663,10 +657,12 @@ async def _load_child_sections( sid_to_path = {meta['section_id']: path for path, meta in all_sections.items()} # Aggregate chunk counts upward: each show_summary item gets counts from itself + descendants + # Phase 1: Direct section assignment — counts from chunks directly under each section 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 + for item_path, item in items_by_path.items(): if not item['show_summary']: continue @@ -675,6 +671,79 @@ async def _load_child_sections( item['image_count'] += img_c item['table_count'] += tbl_c + # Phase 2: connect_to reference tracing — Root-level standalone assets + # Images/tables often live in the Root section but are referenced via connect_to + # from text chunks in deeper sections. Trace these references to attribute + # assets to the sections that actually use them. + # + # Algorithm: for each show_summary item, find all text chunks under its subtree, + # collect their connect_to targets, and count how many are image/table chunks. + scope_items_with_zero_assets = [ + item for item in items_by_path.values() + if item['show_summary'] and item['image_count'] == 0 and item['table_count'] == 0 + ] + if scope_items_with_zero_assets: + # Load connect_to metadata for text chunks under all scope sections + scope_section_ids = {item['section_id'] for item in items_by_path.values() if item.get('section_id')} + if scope_section_ids: + from sqlalchemy import literal_column + connect_stmt = ( + select( + DocumentChunk.section_id, + DocumentChunk.chunk_metadata, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(scope_section_ids))) + .where(DocumentChunk.chunk_type == 'text') + ) + connect_result = (await db.execute(connect_stmt)).all() + + # Map section_id → set of connected target chunk_ids + section_target_ids: dict[str, set[str]] = {} + for sec_id, metadata in connect_result: + if not isinstance(metadata, dict): + continue + for conn in metadata.get('connect_to') or []: + target_id = conn.get('target', '') + if target_id: + section_target_ids.setdefault(sec_id, set()).add(target_id) + + if section_target_ids: + # Collect all target chunk_ids and look up their types + all_target_ids = set() + for tids in section_target_ids.values(): + all_target_ids.update(tids) + + target_type_stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_id.in_(list(all_target_ids))) + .where(DocumentChunk.chunk_type.in_(['image', 'table'])) + ) + target_type_result = (await db.execute(target_type_stmt)).all() + target_types: dict[str, str] = {cid: ctype for cid, ctype in target_type_result} + + # Aggregate connected asset counts per section path → upward to items + for sec_id, target_ids in section_target_ids.items(): + ref_path = sid_to_path.get(sec_id, '') + if not ref_path: + continue + ref_img = sum(1 for tid in target_ids if target_types.get(tid) == 'image') + ref_tbl = sum(1 for tid in target_ids if target_types.get(tid) == 'table') + if ref_img == 0 and ref_tbl == 0: + continue + for item_path, item in items_by_path.items(): + if not item['show_summary']: + continue + if ref_path == item_path or ref_path.startswith(item_path + ' / '): + item['image_count'] += ref_img + item['table_count'] += ref_tbl + # ── Sort by native document order ───────────────────────────────────── sorted_items = sorted(items_by_path.values(), key=lambda x: x['sort_order']) # Clean up internal fields @@ -704,7 +773,7 @@ async def _load_child_sections( def _parse_scope_nav_response(text: str) -> list[dict[str, Any]]: """Parse selections JSON from scope navigation LLM response. - Returns list of {"path": str, "confidence": float, "mode": str}. + Returns list of {"path": str, "confidence": float}. """ text = text.strip() # Try direct parse @@ -724,7 +793,6 @@ def _parse_scope_nav_response(text: str) -> list[dict[str, Any]]: return [] selections: list[dict[str, Any]] = [] - _VALID_MODES = {'all', 'image', 'table'} for item in (data.get('selections') or []): if not isinstance(item, dict): continue @@ -734,10 +802,7 @@ def _parse_scope_nav_response(text: str) -> list[dict[str, Any]]: confidence = _normalize_confidence(item.get('confidence')) if confidence is None: confidence = 0.7 - mode = str(item.get('mode') or 'all').strip().lower() - if mode not in _VALID_MODES: - mode = 'all' - selections.append({'path': path, 'confidence': confidence, 'mode': mode}) + selections.append({'path': path, 'confidence': confidence}) return selections @@ -762,25 +827,34 @@ def _render_leaf_chunks( Connected target chunks (images/tables) are expected to already be present in ``chunks`` via ``hydrate_connected_target_rows``. + + Phase 2: After rendering all text chunks, standalone image/table + chunks that were NOT inlined via connect_to are rendered separately. + This handles cases where assets exist at root/section level without + a parent text chunk referencing them. """ chunk_by_id: dict[str, dict] = { c.get('chunk_id', ''): c for c in chunks if c.get('chunk_id') } rendered_ids: set[str] = set() + + # Phase 1: Render text chunks with inline asset resolution for chunk in chunks: cid = chunk.get('chunk_id', '') if cid and cid in rendered_ids: continue - if cid: - rendered_ids.add(cid) chunk_type = (chunk.get('chunk_type') or chunk.get('type') or 'text').strip().lower() - # Skip standalone image/table chunks — they'll be inlined - # via connect_to from their parent text chunk + # Skip standalone image/table chunks — they'll be rendered in Phase 2 + # if not inlined via connect_to from a parent text chunk. + # NOTE: do NOT add to rendered_ids here — Phase 2 needs to see them. if chunk_type in ('image', 'table'): continue + if cid: + rendered_ids.add(cid) + content = str(chunk.get('content', '')).strip() # Resolve connected assets via connect_to metadata @@ -818,6 +892,34 @@ def _render_leaf_chunks( if line.strip(): parts.append(f'{indent}┈ {line}') + # Phase 2: Render standalone image/table chunks not inlined via connect_to + for chunk in chunks: + cid = chunk.get('chunk_id', '') + if cid and cid in rendered_ids: + continue + if cid: + rendered_ids.add(cid) + + chunk_type = (chunk.get('chunk_type') or chunk.get('type') or '').strip().lower() + if chunk_type == 'image': + file_path = chunk.get('file_path') or '' + img_desc = str(chunk.get('content', '')).strip() + asset_url = (asset_lookup or {}).get(cid, '') if cid else '' + display_ref = asset_url or file_path + if display_ref: + parts.append(f'{indent}┈ [图片: {display_ref}]') + if img_desc: + for line in img_desc.split('\n'): + if line.strip(): + parts.append(f'{indent}┈ {line}') + elif chunk_type == 'table': + table_html = str(chunk.get('content', '')).strip() + parts.append(f'{indent}┈ [表格内容]') + if table_html: + for line in table_html.split('\n'): + if line.strip(): + parts.append(f'{indent}┈ {line}') + def render_unified_doc_tree( node: DocTreeNode, diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index 776244317..a757f87a0 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -58,6 +58,14 @@ def _collect_media_chunks(node: DocTreeNode) -> list[dict[str, Any]]: return media +def _collect_media_chunks_all(doc_trees: dict[str, DocTreeNode]) -> list[dict[str, Any]]: + """Collect media chunks from all doc trees.""" + media: list[dict[str, Any]] = [] + for tree in doc_trees.values(): + media.extend(_collect_media_chunks(tree)) + return media + + async def _build_asset_url_map( media_chunks: list[dict[str, Any]], ) -> dict[str, str]: @@ -171,6 +179,9 @@ async def run( """ from shared.services.retrieval.agentic import tools from shared.services.retrieval.agentic.policy import attempt_answer + from shared.services.retrieval.llm_adapter import create_retrieval_vlm_fn + + vlm_fn = create_retrieval_vlm_fn() config = config or _build_config_from_env() exclude_document_ids = exclude_document_ids or [] @@ -353,6 +364,7 @@ async def run( # Phase 2 + 3 Loop: Navigate → Render → Attempt Answer → (Revise) # ══════════════════════════════════════════════════════════════════ answer_text = '' + evidence_text = '' revision_hint: str | None = None stop_reason = 'max_revisions' @@ -385,11 +397,14 @@ async def run( if not is_b_class: # Build exclude_paths for this doc from seen_section_keys - doc_exclude = { + # Starts with revision-carried paths, then accumulates + # leaf paths hydrated during THIS BFS round to prevent + # re-selection in deeper drill-downs. + doc_exclude: set[str] = { 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 + } if state.seen_section_keys else set() # BFS queue: (scope_path, parent_node, depth) root = DocTreeNode(scope_path=None) @@ -406,6 +421,76 @@ async def run( if llm_fn is None: break + # ★ Step 1: Tool selection (agent decides which tool) + tool_choice = await tools.tool_select_step( + db, + document_id=doc.document_id, + job_result_id=job_result_id, + query=query, + llm_fn=llm_fn, + doc_name=doc_name, + scope_path=scope, + exclude_paths=doc_exclude, + revision_hint=revision_hint if depth == 0 else None, + ) + 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_choice, + }, + ), + 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} tool={tool_choice}' + ) + + if tool_choice in ('FIND_IMAGES', 'FIND_TABLES'): + # ★ Terminal: asset filter (no LLM, pure programmatic) + asset_type = 'image' if tool_choice == 'FIND_IMAGES' else 'table' + asset_chunks = await tools.asset_filter_step( + db, + document_id=doc.document_id, + job_result_id=job_result_id, + scope_path=scope, + asset_type=asset_type, + ) + if asset_chunks: + parent_node.leaf_content[scope or 'root'] = asset_chunks + + 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', + 'asset_type': asset_type, + 'chunks_found': len(asset_chunks) if asset_chunks else 0, + }, + ), + decision_reason=f'asset_r{round_idx}_d{depth}_{doc.source_file_name}', + ) + + logger.info( + f' agentic step {state.step_count}: asset_filter_step ' + f'doc="{doc.source_file_name}" scope={scope or "root"} ' + f'type={asset_type} chunks={len(asset_chunks) if asset_chunks else 0}' + ) + # NO further drill-down — terminal action + continue + + # ★ Step 2: NAVIGATE (existing scope_navigate_step) step_node, drill_paths = await tools.scope_navigate_step( db, document_id=doc.document_id, @@ -423,15 +508,32 @@ async def run( # Merge step result into parent node parent_node.outline_items = step_node.outline_items - parent_node.leaf_content = step_node.leaf_content + parent_node.leaf_content.update(step_node.leaf_content) parent_node.confidence = step_node.confidence + # Accumulate hydrated leaf paths into doc_exclude + # so subsequent drill-downs don't re-show them as [SELECT] + for leaf_path in step_node.leaf_content: + 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. + for child_path in list(parent_node.children.keys()): + for leaf_path in list(parent_node.leaf_content.keys()): + if leaf_path.startswith(child_path + ' / '): + parent_node.children[child_path].leaf_content[leaf_path] = \ + parent_node.leaf_content.pop(leaf_path) + if trace_enabled: trace.record_step( 'scope_navigate_step', ToolResult( @@ -471,6 +573,7 @@ async def run( namespace=namespace, doc_name=doc_name, discovery_hints=doc_hints, + revision_hint=revision_hint, ) state.step_count += 1 @@ -507,13 +610,25 @@ async def run( stop_reason = 'no_llm' break - # Auto-trigger attempt_answer + # Collect image URLs from evidence for VLM switch + evidence_image_urls: list[str] = [] + if vlm_fn: + asset_url_map = await _build_asset_url_map( + _collect_media_chunks_all(state.doc_trees), + ) + evidence_image_urls = [ + url for url in asset_url_map.values() if url + ] + + # Auto-trigger attempt_answer (VLM if images present) status, answer_text, reason = await attempt_answer( llm_fn, query=query, evidence_text=evidence_text, state=state, config=config, + vlm_fn=vlm_fn, + image_urls=evidence_image_urls or None, ) state.step_count += 1 @@ -555,7 +670,7 @@ async def run( # Clear doc selection for re-exploration (preserve doc_trees for merge) state.selected_docs.clear() - # Re-run KG select (allow re-exploring docs for different sections via path masking) + # Re-run KG select with revision hint kg_result = await tools.kg_document_select( db, user_id=user_id, @@ -563,6 +678,7 @@ async def run( query=query, llm_fn=llm_fn, exclude_document_ids=list(set(exclude_document_ids)), + revision_hint=revision_hint, ) state.step_count += 1 diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index e5a5c0dd4..ab0087973 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -43,16 +43,26 @@ async def attempt_answer( evidence_text: str, state: AgentState, config: AgentRunConfig, + vlm_fn: LLMFn | None = None, + image_urls: list[str] | None = None, ) -> tuple[str, str, str]: """Attempt to answer the query using collected evidence. + When ``vlm_fn`` is provided and ``image_urls`` is non-empty, the + answer is generated using the VLM in multimodal message format + (text + image_url parts) so the model can actually *see* chart + images rather than only reading their text descriptions. + + Falls back to ``llm_fn`` (text-only) when VLM is unavailable or + when no images are present. + Returns (status, answer_text, reason) where: - status='DONE', answer_text=, reason='' → evidence was sufficient, answer is ready - status='NOT_FOUND', answer_text='', reason= → evidence was insufficient, reason is used as revision_hint """ - prompt = _ATTEMPT_ANSWER_PROMPT.format( + prompt_text = _ATTEMPT_ANSWER_PROMPT.format( query=query, evidence_context=evidence_text, revision_count=state.revision_count, @@ -60,13 +70,35 @@ async def attempt_answer( ) verbose = os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', '') == 'true' + + # Decide whether to use VLM (multimodal) or text-only LLM + effective_fn = llm_fn + effective_input: Any = prompt_text + + if vlm_fn and image_urls: + # Build multimodal message: text evidence + image_url parts + content_parts: list[dict[str, Any]] = [ + {'type': 'text', 'text': prompt_text}, + ] + for url in image_urls[:20]: # Cap at 20 images to avoid token overflow + content_parts.append({ + 'type': 'image_url', + 'image_url': {'url': url}, + }) + effective_input = [{'role': 'user', 'content': content_parts}] + effective_fn = vlm_fn + logger.info( + f' [attempt_answer] using VLM with {len(image_urls)} image URLs ' + f'(capped at {min(len(image_urls), 20)})' + ) + if verbose: logger.info( f'[attempt_answer PROMPT]\n' - f'{prompt}' + f'{prompt_text}' ) - raw_response = await llm_fn(prompt) + raw_response = await effective_fn(effective_input) logger.info(f' [attempt_answer] raw={repr(raw_response[:300])}') if verbose: @@ -75,6 +107,12 @@ async def attempt_answer( f'{raw_response}' ) + # If VLM returned empty (e.g. URL not reachable), fall back to text LLM + if not raw_response.strip() and effective_fn is not llm_fn: + logger.info(' [attempt_answer] VLM returned empty, falling back to text LLM') + raw_response = await llm_fn(prompt_text) + logger.info(f' [attempt_answer] text fallback raw={repr(raw_response[:300])}') + parsed = _parse_answer_response(raw_response) if not parsed: # Parse error: treat raw text as best-effort answer diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index c337d2e08..1098130a6 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -169,6 +169,7 @@ async def kg_document_select( query: str, llm_fn: LLMFn | None, exclude_document_ids: list[str], + revision_hint: str | None = None, **_kwargs: Any, ) -> ToolResult: """Select candidate documents from document-level KG.""" @@ -193,8 +194,19 @@ async def kg_document_select( latency_ms=latency, ) + revision_context = '' + if revision_hint: + revision_context = ( + f'\nIMPORTANT: This is a REVISION round. ' + f'The previous search attempt failed because:\n' + f'"{revision_hint}"\n' + f'Adjust your document selection accordingly. ' + f'If no document can address this, return an EMPTY array [].\n' + ) + file_prompt = _FILE_SELECT_PROMPT.format( overview=overview_text, query=query, + revision_context=revision_context, ) file_response = await llm_fn(file_prompt) selected_ids = _parse_json_array(file_response) @@ -364,15 +376,323 @@ async def graph_expand_docs( 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}" +{scope_header} + +Below is a summary of the current scope's sections: + +{tree_summary} + +User query: {query} + +=== Available Tools === + +NAVIGATE + Drill into specific sections to explore detailed content. + Use when the query requires reading text content in sub-sections. + +FIND_IMAGES + Extract all image/chart/diagram assets under this scope. + Use when the query asks about images, charts, figures, or visual content. + +FIND_TABLES + Extract all table/data assets under this scope. + Use when the query asks about tables, tabular data, or structured data. + +Choose exactly ONE tool. Return ONLY a JSON object: +{{"tool": "NAVIGATE"}} +or {{"tool": "FIND_IMAGES"}} +or {{"tool": "FIND_TABLES"}} +Do not include any explanation. +""" + + +def _parse_tool_choice(text: str) -> str: + """Parse tool choice from LLM response. Returns one of NAVIGATE/FIND_IMAGES/FIND_TABLES.""" + import json as _json + import re as _re + + text = text.strip() + _VALID_TOOLS = {'NAVIGATE', 'FIND_IMAGES', 'FIND_TABLES'} + + # Try JSON parse + try: + data = _json.loads(text) + if isinstance(data, dict): + tool = str(data.get('tool', '')).strip().upper() + if tool in _VALID_TOOLS: + return tool + except (ValueError, _json.JSONDecodeError): + pass + + # Fallback: extract JSON from markdown wrapper + match = _re.search(r'\{.*?\}', text, _re.DOTALL) + if match: + try: + data = _json.loads(match.group()) + if isinstance(data, dict): + tool = str(data.get('tool', '')).strip().upper() + if tool in _VALID_TOOLS: + return tool + except (ValueError, _json.JSONDecodeError): + pass + + # Last resort: keyword match + upper = text.upper() + if 'FIND_IMAGES' in upper: + return 'FIND_IMAGES' + if 'FIND_TABLES' in upper: + return 'FIND_TABLES' + return 'NAVIGATE' + + +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, +) -> str: + """Route to the appropriate tool for the current scope. + + Returns one of: 'NAVIGATE', 'FIND_IMAGES', 'FIND_TABLES'. + + This is the top-level agent decision — it picks WHICH tool + to invoke. Each tool then handles its own internal logic. + + Optimization: if the current scope has no image or table chunks, + skips the LLM call and returns 'NAVIGATE' directly. + """ + items = await _load_child_sections( + db, document_id, job_result_id, scope_path, + exclude_paths=exclude_paths, + ) + if not items: + return 'NAVIGATE' # fallback + + # 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 'NAVIGATE' # no assets → skip tool selection, go straight to navigate + + 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, + 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 choice + tool = _parse_tool_choice(response) + logger.info( + f' tool_select_step scope={scope_path or "root"}: ' + f'tool={tool} images={total_images} tables={total_tables}' + ) + return tool + + +# --------------------------------------------------------------------------- +# Tool: asset_filter_step (programmatic asset extraction) +# --------------------------------------------------------------------------- + +async def asset_filter_step( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + scope_path: str | None, + asset_type: str, # 'image' | 'table' +) -> list[dict[str, Any]]: + """Extract assets from all descendants under scope_path. + + Terminal action — no LLM involved. + Algorithm: load all text chunks under scope → parse connect_to metadata → + batch-load target image/table chunks → return directly. + + Also collects standalone asset chunks (image/table) that exist directly + under the scope but are not referenced via connect_to. + """ + from shared.models.database.document import DocumentChunk, DocumentSection + + t0 = time.monotonic() + try: + # 1. Find all section_ids under scope_path + 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} / %')) + ) + section_result = await db.execute(section_stmt) + section_rows = section_result.all() + section_ids = {row[0] for row in section_rows} + + if not section_ids: + logger.info(f' asset_filter_step: no sections found under scope={scope_path}') + return [] + + # 2. Load target asset chunks directly (standalone assets in the scope) + asset_stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.file_path, + DocumentChunk.section_id, + DocumentChunk.source_chunk_path, + DocumentChunk.chunk_metadata, + DocumentChunk.sort_order, + DocumentChunk.job_result_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(section_ids))) + .where(DocumentChunk.chunk_type == asset_type) + .order_by(DocumentChunk.sort_order) + ) + asset_result = await db.execute(asset_stmt) + asset_rows = asset_result.all() + + # 3. Also find assets referenced via connect_to from text chunks + text_stmt = ( + select( + DocumentChunk.chunk_metadata, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(list(section_ids))) + .where(DocumentChunk.chunk_type == 'text') + ) + text_result = await db.execute(text_stmt) + connected_target_ids: set[str] = set() + for (metadata,) in text_result.all(): + if not isinstance(metadata, dict): + continue + for conn in metadata.get('connect_to') or []: + target_id = conn.get('target', '') + if target_id: + connected_target_ids.add(target_id) + + # Load connected targets that match asset_type + if connected_target_ids: + connected_stmt = ( + select( + DocumentChunk.chunk_id, + DocumentChunk.chunk_type, + DocumentChunk.content, + DocumentChunk.file_path, + DocumentChunk.section_id, + DocumentChunk.source_chunk_path, + DocumentChunk.chunk_metadata, + DocumentChunk.sort_order, + DocumentChunk.job_result_id, + ) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_id.in_(list(connected_target_ids))) + .where(DocumentChunk.chunk_type == asset_type) + .order_by(DocumentChunk.sort_order) + ) + connected_result = await db.execute(connected_stmt) + connected_rows = connected_result.all() + else: + connected_rows = [] + + # 4. Merge and deduplicate + seen_ids: set[str] = set() + chunks: list[dict[str, Any]] = [] + + # Helper to look up job_id from job_result + from shared.models.database.job_result import JobResult + job_stmt = ( + select(JobResult.job_id) + .where(JobResult.id == job_result_id) + ) + job_result_row = await db.execute(job_stmt) + job_id = job_result_row.scalar() or '' + + for row in list(asset_rows) + list(connected_rows): + chunk_id = row[0] + if chunk_id in seen_ids: + continue + seen_ids.add(chunk_id) + chunks.append({ + 'document_id': document_id, + 'chunk_id': chunk_id, + 'chunk_type': row[1], + 'content': row[2], + 'file_path': row[3], + 'section_id': row[4], + 'source_chunk_path': row[5], + 'chunk_metadata': row[6] or {}, + 'sort_order': row[7], + 'job_result_id': job_result_id, + 'job_id': job_id, + }) + + latency = int((time.monotonic() - t0) * 1000) + logger.info( + f' asset_filter_step scope={scope_path or "root"} ' + f'type={asset_type}: {len(chunks)} chunks found, {latency}ms' + ) + return chunks + + except Exception as e: + logger.error(f' asset_filter_step failed: {e}') + return [] + + # --------------------------------------------------------------------------- # Tool: scope_navigate_step (single-step navigation) # --------------------------------------------------------------------------- -_LLM_MODE_TO_HYDRATE: dict[str, str] = { - 'all': 'chunks', - 'image': 'image_only', - 'table': 'table_only', -} async def scope_navigate_step( db: AsyncSession, @@ -459,14 +779,11 @@ async def scope_navigate_step( node.confidence[path] = conf if item.get('is_leaf'): - # Leaf → hydrate chunks - hydrate_mode = _LLM_MODE_TO_HYDRATE.get( - str(sel.get('mode', 'all')).strip().lower(), 'chunks' - ) + # Leaf → hydrate all chunk types chunks = await _hydrate_paths_to_rows( db, path_selections=[ - {'path': path, 'confidence': conf, 'hydrate_mode': hydrate_mode} + {'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'} ], user_id=user_id, namespace=namespace, @@ -511,6 +828,7 @@ async def discovery_select_step( namespace: str, doc_name: str = '', discovery_hints: list[dict[str, Any]], + revision_hint: str | None = None, ) -> DocTreeNode: """Post-navigation discovery selection step. @@ -551,10 +869,22 @@ async def discovery_select_step( return node items_text = '\n'.join(hint_lines) + + revision_context = '' + if revision_hint: + revision_context = ( + f'\nIMPORTANT: This is a REVISION round. ' + f'The previous search attempt failed because:\n' + f'"{revision_hint}"\n' + f'Adjust your selection accordingly. ' + f'If no candidate is relevant, return an EMPTY list [].\n' + ) + prompt = _DISCOVERY_SELECT_PROMPT.format( doc_name=doc_name or document_id, items=items_text, query=query, + revision_context=revision_context, ) response = await llm_fn(prompt) selections = _parse_scope_nav_response(response) @@ -569,15 +899,12 @@ async def discovery_select_step( for sel in valid_selections: path = sel['path'] conf = sel.get('confidence', 0.7) - hydrate_mode = _LLM_MODE_TO_HYDRATE.get( - str(sel.get('mode', 'all')).strip().lower(), 'chunks' - ) node.confidence[path] = conf chunks = await _hydrate_paths_to_rows( db, path_selections=[ - {'path': path, 'confidence': conf, 'hydrate_mode': hydrate_mode} + {'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'} ], user_id=user_id, namespace=namespace, diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index 8d9a11c22..ba5093701 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -89,3 +89,52 @@ async def llm_fn(prompt: LLMFnInput) -> str: return '' return llm_fn + + +def create_retrieval_vlm_fn( + *, + model: str | None = None, + temperature: float = _RETRIEVAL_LLM_TEMPERATURE, + max_tokens: int = 4096, +) -> LLMFn | None: + """Create an async VLM callable for image-aware answer generation. + + Uses the IMAGE_MODEL (e.g. qwen3.5-flash) for multimodal input. + Returns None when the image model is not configured. + + The returned function accepts the same ``LLMFnInput`` type as + ``create_retrieval_llm_fn`` — callers pass either a plain string + or a list of ChatCompletionMessageParam (including image_url parts). + """ + from shared.core.config import settings + + effective_model = model or getattr(settings, 'IMAGE_MODEL', '') or 'qwen3.5-flash' + + if not _has_llm_credentials(): + logger.debug('retrieval: no LLM credentials for VLM, image-aware answering disabled') + return None + + async def vlm_fn(prompt: LLMFnInput) -> str: + from shared.utils.OpenAICompatibleClientSync import get_openai_client + + client = get_openai_client(model=effective_model) + try: + result = await asyncio.to_thread( + client.chat_completion, + cast(Any, prompt), + model=effective_model, + temperature=temperature, + max_tokens=max_tokens, + ) + return result + except Exception as exc: + logger.warning( + "retrieval: VLM call failed (degrading gracefully): " + "model={} error_type={} error={}", + effective_model, + type(exc).__name__, + exc, + ) + return '' + + return vlm_fn