From e172c66abe3c10223214b4b964144d1593383791 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 12 May 2026 20:06:53 +0800 Subject: [PATCH 1/4] refactor(agentic): clean up dead tool code and explicitly handle self_only mode --- .../services/retrieval/agentic/tools.py | 604 +++++++++++++----- .../shared/services/retrieval/app_service.py | 1 + 2 files changed, 433 insertions(+), 172 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 29d08c015..f862c23cd 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -15,29 +15,29 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document +from shared.services.retrieval.agentic.budget import BudgetExceeded from shared.services.retrieval.agentic.types import DocTreeNode, ToolResult -from shared.services.retrieval.agentic.result_rows import ( - CHANNEL_WEIGHT_CONTENT, - CHANNEL_WEIGHT_PATH, - CHANNEL_WEIGHT_TERM, - INTERNAL_RECALL_K_MULTIPLIER, - hydrate_connected_target_rows, - hydrate_paths_to_rows, - merge_channels_rrf, - merge_same_section_rows, - normalize_row_scores, - resolve_allowed_chunk_types, -) from shared.services.retrieval.agent_navigate import ( _build_knowledge_map_overview, - _expand_by_edges, _format_items_for_llm, - _grep_discover_document_ids, _load_child_sections, _parse_json_array, _parse_scope_nav_response, _SCOPE_NAV_PROMPT, + _DISCOVERY_SELECT_PROMPT, _FILE_SELECT_PROMPT, + _format_budget_block, +) +from shared.services.retrieval.app_service import ( + _CHANNEL_WEIGHT_CONTENT, + _CHANNEL_WEIGHT_PATH, + _CHANNEL_WEIGHT_TERM, + _INTERNAL_RECALL_K_MULTIPLIER, + _merge_same_section_rows, + _normalize_row_scores, + _resolve_allowed_chunk_types, + hydrate_connected_target_rows, + merge_channels_rrf, ) from shared.services.retrieval.channels import content_channel, path_channel, term_channel from shared.services.retrieval.llm_adapter import LLMFn @@ -47,30 +47,6 @@ # Tool: bottom_discovery # --------------------------------------------------------------------------- -_DISCOVERY_SELECT_PROMPT = """\ -You are a document navigation assistant. - -Document: "{doc_name}" - -After navigating the document's section tree, the following section paths -were additionally discovered via keyword and semantic search. -They may contain relevant evidence not found through hierarchical navigation. - -=== Discovery Candidates === -{items} -=== End Discovery Candidates === - -User query: {query} - -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"}}, ...]}} -Do not include any explanation. -""" - - async def bottom_discovery( db: AsyncSession, *, @@ -91,8 +67,8 @@ async def bottom_discovery( """Run 3-channel BM25 discovery + RRF fusion.""" t0 = time.monotonic() try: - allowed_chunk_types = resolve_allowed_chunk_types(data_type) - effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * INTERNAL_RECALL_K_MULTIPLIER + allowed_chunk_types = _resolve_allowed_chunk_types(data_type) + effective_recall_k = internal_recall_k if internal_recall_k is not None else top_k * _INTERNAL_RECALL_K_MULTIPLIER active_channels = set(channels) if channels else {'path', 'content', 'term'} path_rows: list[dict[str, Any]] = [] @@ -125,9 +101,9 @@ async def bottom_discovery( # RRF fusion default_weights = { - 'path': CHANNEL_WEIGHT_PATH, - 'content': CHANNEL_WEIGHT_CONTENT, - 'term': CHANNEL_WEIGHT_TERM, + 'path': _CHANNEL_WEIGHT_PATH, + 'content': _CHANNEL_WEIGHT_CONTENT, + 'term': _CHANNEL_WEIGHT_TERM, } effective_weights = {**default_weights, **(channel_weights or {})} @@ -135,19 +111,19 @@ async def bottom_discovery( weight_list: list[float] = [] if path_rows: channel_lists.append(path_rows) - weight_list.append(effective_weights.get('path', CHANNEL_WEIGHT_PATH)) + weight_list.append(effective_weights.get('path', _CHANNEL_WEIGHT_PATH)) if content_rows: channel_lists.append(content_rows) - weight_list.append(effective_weights.get('content', CHANNEL_WEIGHT_CONTENT)) + weight_list.append(effective_weights.get('content', _CHANNEL_WEIGHT_CONTENT)) if term_rows: channel_lists.append(term_rows) - weight_list.append(effective_weights.get('term', CHANNEL_WEIGHT_TERM)) + weight_list.append(effective_weights.get('term', _CHANNEL_WEIGHT_TERM)) fused_rows = merge_channels_rrf(channel_lists, weight_list, top_k) if channel_lists else [] - fused_rows = merge_same_section_rows(fused_rows) + fused_rows = _merge_same_section_rows(fused_rows) if fused_rows: - normalize_row_scores(fused_rows, source_field='score', target_field='discovery_score', default=0.5) + _normalize_row_scores(fused_rows, source_field='score', target_field='discovery_score', default=0.5) # Extract top document IDs as hints for KG selection doc_id_counts: dict[str, int] = {} @@ -193,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.""" @@ -217,8 +194,20 @@ 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, + budget_block=_format_budget_block(_kwargs.get('budget_snapshot')), ) file_response = await llm_fn(file_prompt) selected_ids = _parse_json_array(file_response) @@ -267,6 +256,8 @@ async def kg_document_select( }, latency_ms=latency, ) + except BudgetExceeded: + raise except Exception as e: latency = int((time.monotonic() - t0) * 1000) logger.error(f' agentic.kg_document_select failed: {e}') @@ -274,129 +265,369 @@ async def kg_document_select( # --------------------------------------------------------------------------- -# Tool: grep_document_discover +# Tool: tool_select_step (lightweight LLM router) # --------------------------------------------------------------------------- -async def grep_document_discover( +_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 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"}} +When budget is TIGHT, prefer NAVIGATE only when more detail is necessary. +When budget is CRITICAL, choose the narrowest terminal tool if assets directly answer the query. +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 + + # Accept a JSON object wrapped in markdown, but do not infer a default tool. + 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 + + upper = text.upper() + if 'FIND_IMAGES' in upper: + return 'FIND_IMAGES' + if 'FIND_TABLES' in upper: + return 'FIND_TABLES' + if 'NAVIGATE' in upper: + return 'NAVIGATE' + return '' + + +async def tool_select_step( db: AsyncSession, *, - user_id: str, - namespace: str, + document_id: str, + job_result_id: str, query: str, - exclude_document_ids: list[str], - **_kwargs: Any, -) -> ToolResult: - """Discover documents via term search (GREP).""" - t0 = time.monotonic() - try: - grep_doc_ids = await _grep_discover_document_ids( - db, user_id=user_id, namespace=namespace, query=query, - exclude_document_ids=exclude_document_ids, - ) + 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, +) -> str: + """Route to the appropriate tool for the current scope. - if not grep_doc_ids: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_docs_found', - payload={'reason': 'GREP found no matching documents'}, - latency_ms=latency, - ) + Returns one of: 'NAVIGATE', 'FIND_IMAGES', 'FIND_TABLES'. - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.document_id.in_(grep_doc_ids)) - ) - doc_result = await db.execute(doc_stmt) - doc_id_to_name: dict[str, str] = {} - doc_job_map: dict[str, str] = {} - for did, fname, jrid in doc_result.all(): - doc_id_to_name[did] = fname or did - if jrid: - doc_job_map[did] = jrid + This is the top-level agent decision — it picks WHICH tool + to invoke. Each tool then handles its own internal logic. - latency = int((time.monotonic() - t0) * 1000) - logger.info(f' agentic.grep_document_discover: {len(grep_doc_ids)} docs found, {latency}ms') - return ToolResult( - status='discovered_docs', - payload={ - 'document_ids': grep_doc_ids, - 'doc_id_to_name': doc_id_to_name, - 'doc_job_map': doc_job_map, - }, - latency_ms=latency, + 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' + + # 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, + 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.' ) - except Exception as e: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.grep_document_discover failed: {e}') - return ToolResult(status='error', error=str(e), latency_ms=latency) + response = await llm_fn(prompt) + + # Parse tool choice + tool = _parse_tool_choice(response) + if not tool: + raise ValueError(f'invalid tool selection response: {response[:200]}') + logger.info( + f' tool_select_step scope={scope_path or "root"}: ' + f'tool={tool} images={total_images} tables={total_tables}' + ) + return tool # --------------------------------------------------------------------------- -# Tool: graph_expand_docs +# Tool: asset_filter_step (programmatic asset extraction) # --------------------------------------------------------------------------- -async def graph_expand_docs( +async def asset_filter_step( db: AsyncSession, *, - user_id: str, - namespace: str, - document_ids: list[str], - **_kwargs: Any, -) -> ToolResult: - """Expand document set via KG edge traversal.""" + 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: - expanded_ids = await _expand_by_edges( - db, document_ids=document_ids, user_id=user_id, namespace=namespace, + # 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) ) - new_ids = [did for did in expanded_ids if did not in document_ids] - - if not new_ids: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_expansion', - payload={'reason': 'no new neighbors found via edges'}, - latency_ms=latency, + 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() - doc_stmt = ( - select(Document.document_id, Document.source_file_name, Document.current_job_result_id) - .where(Document.document_id.in_(new_ids)) + section_path_by_id = {section_id: section_path for section_id, section_path in section_rows} + + # 3. Resolve media → owner text section via unified helper + from shared.services.retrieval.app_service import _resolve_asset_owners_from_rows + + text_stmt = ( + select( + DocumentChunk.section_id, + DocumentChunk.chunk_type, + DocumentChunk.chunk_metadata, + DocumentChunk.source_chunk_path, + ) + .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') ) - doc_result = await db.execute(doc_stmt) - doc_id_to_name: dict[str, str] = {} - doc_job_map: dict[str, str] = {} - for did, fname, jrid in doc_result.all(): - doc_id_to_name[did] = fname or did - if jrid: - doc_job_map[did] = jrid + text_result = await db.execute(text_stmt) + text_row_dicts = [ + { + 'chunk_type': chunk_type, + 'chunk_metadata': metadata or {}, + 'section_id': sid, + 'section_path': section_path_by_id.get(sid, ''), + 'source_chunk_path': scp, + } + for sid, chunk_type, metadata, scp in text_result.all() + ] + owner_by_target_id = _resolve_asset_owners_from_rows(text_row_dicts) + + # Collect connected target IDs for batch-loading + connected_target_ids: set[str] = set(owner_by_target_id.keys()) + + # 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) + + # Owner resolution: prefer connect_to-based owner + owner_info = owner_by_target_id.get(chunk_id) + owner_section_path = owner_info.get('section_path') if owner_info else None + + # Fallback: media's own section_id path, but guard against + # 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 + logger.warning( + f' asset_filter_step: rejecting root-level owner fallback ' + f'chunk_id={chunk_id} section_path={own_section_path}' + ) + own_section_path = None + owner_section_path = own_section_path + + if not owner_section_path: + logger.warning( + f' asset_filter_step unresolved owner: chunk_id={chunk_id} ' + f'file_path={row[3]} scope={scope_path or "root"}' + ) + continue + 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], + 'section_path': owner_section_path, + 'owner_section_path': owner_section_path, + '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' agentic.graph_expand_docs: {len(new_ids)} new docs from edges, {latency}ms') - return ToolResult( - status='expanded_docs', - payload={ - 'document_ids': new_ids, - 'doc_id_to_name': doc_id_to_name, - 'doc_job_map': doc_job_map, - }, - latency_ms=latency, + 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: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.graph_expand_docs failed: {e}') - return ToolResult(status='error', error=str(e), latency_ms=latency) + 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, @@ -411,6 +642,7 @@ async def scope_navigate_step( scope_path: 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. @@ -420,6 +652,8 @@ async def scope_navigate_step( - pending: list of {path, confidence, mode} for non-leaf selections (orchestrator queues these for further drill-down) """ + from shared.services.retrieval.app_service import _hydrate_paths_to_rows + empty = DocTreeNode.empty(scope_path) try: @@ -445,6 +679,7 @@ async def scope_navigate_step( 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, ) @@ -474,6 +709,7 @@ async def scope_navigate_step( ] pending: list[dict] = [] + path_selections = [] for sel in valid_selections: path = sel['path'] conf = sel.get('confidence', 0.7) @@ -481,36 +717,39 @@ 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' - ) - chunks = await hydrate_paths_to_rows( - db, - path_selections=[ - {'path': path, 'confidence': conf, 'hydrate_mode': hydrate_mode} - ], - user_id=user_id, - namespace=namespace, - document_id=document_id, - ) - # Also hydrate connected targets (image/table chunks referenced via connect_to) - if chunks: - connected = await hydrate_connected_target_rows( - db=db, - rows=chunks, - exclude_document_ids=[], - exclude_sections=[], - ) - if connected: - chunks = chunks + connected - node.leaf_content[path] = chunks + path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'chunks'}) else: - # Non-leaf → return as pending for orchestrator to queue pending.append(sel) + # ★ NEW: Also hydrate this node's OWN direct chunks (not descendants) + path_selections.append({'path': path, 'confidence': conf, 'hydrate_mode': 'self_only'}) + + if path_selections: + chunks = await _hydrate_paths_to_rows( + db, + path_selections=path_selections, + user_id=user_id, + namespace=namespace, + document_id=document_id, + ) + if chunks: + connected = await hydrate_connected_target_rows( + db=db, + rows=chunks, + exclude_document_ids=[], + exclude_sections=[], + ) + if connected: + chunks = chunks + connected + 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 + except BudgetExceeded: + raise except Exception as e: logger.error(f' scope_navigate_step failed for doc={document_id}: {e}') return empty, [] @@ -533,6 +772,8 @@ async def discovery_select_step( namespace: str, doc_name: str = '', discovery_hints: list[dict[str, Any]], + revision_hint: str | None = None, + budget_snapshot: dict | None = None, ) -> DocTreeNode: """Post-navigation discovery selection step. @@ -543,6 +784,8 @@ async def discovery_select_step( For B-class documents (discovery-only, not KG-selected), this is the only navigation step — no prior BFS. """ + from shared.services.retrieval.app_service import _hydrate_paths_to_rows + node = DocTreeNode(scope_path=None) if not discovery_hints: return node @@ -571,10 +814,23 @@ 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, + budget_block=_format_budget_block(budget_snapshot), items=items_text, query=query, + revision_context=revision_context, ) response = await llm_fn(prompt) selections = _parse_scope_nav_response(response) @@ -586,19 +842,17 @@ async def discovery_select_step( # 2. Hydrate selected paths valid_selections = [s for s in selections if s['path'] in hint_by_path] + path_selections = [] 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 + path_selections.append({'path': path, 'confidence': conf}) - chunks = await hydrate_paths_to_rows( + if path_selections: + chunks = await _hydrate_paths_to_rows( db, - path_selections=[ - {'path': path, 'confidence': conf, 'hydrate_mode': hydrate_mode} - ], + path_selections=path_selections, user_id=user_id, namespace=namespace, document_id=document_id, @@ -612,7 +866,11 @@ async def discovery_select_step( ) if connected: chunks = chunks + connected - node.leaf_content[path] = chunks + 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]) latency = int((time.monotonic() - t0) * 1000) logger.info( @@ -621,6 +879,8 @@ async def discovery_select_step( ) return node + except BudgetExceeded: + raise except Exception as e: logger.error(f' discovery_select_step failed for doc={document_id}: {e}') return node diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 93f6d6d03..86a12c93b 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -877,6 +877,7 @@ async def _hydrate_paths_to_rows( # Build a map of path → allowed chunk_types based on hydrate_mode _MODE_ALLOWED_TYPES: dict[str, set[str] | None] = { 'chunks': None, # all types + 'self_only': None, # all types, but without descendant filtering 'assets_only': {'image', 'table'}, 'image_only': {'image'}, 'table_only': {'table'}, From feee57cc18cef93f538490dd51ebda60a00d1c15 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 12 May 2026 22:22:02 +0800 Subject: [PATCH 2/4] fix(retrieval): fix NameError in discovery_select_step and typecheck errors --- AGENTS.md | 4 +- .../services/retrieval/agent_navigate.py | 325 +++++++-- .../services/retrieval/agentic/budget.py | 199 +++++ .../retrieval/agentic/orchestrator.py | 681 ++++++++++++++++-- .../services/retrieval/agentic/policy.py | 76 +- .../services/retrieval/agentic/result_rows.py | 448 ------------ .../services/retrieval/agentic/tools.py | 266 +++++-- .../services/retrieval/agentic/trace.py | 6 + .../services/retrieval/agentic/types.py | 50 +- .../shared/services/retrieval/app_service.py | 6 +- .../shared/services/retrieval/llm_adapter.py | 75 +- .../utils/OpenAICompatibleClientSync.py | 49 +- .../shared/utils/token_estimate.py | 65 ++ 13 files changed, 1587 insertions(+), 663 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/budget.py delete mode 100644 packages/shared-python/shared/services/retrieval/agentic/result_rows.py create mode 100644 packages/shared-python/shared/utils/token_estimate.py diff --git a/AGENTS.md b/AGENTS.md index 48a62326b..1bec3165d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -565,7 +565,9 @@ For each selected document, the agent performs a constrained Breadth-First Searc 1. **Scope Navigation**: The document's section tree is dynamically rendered to the LLM. - *Path-Based Hierarchy*: Child nodes are strictly filtered using structural path prefixes (e.g., `child_path.startswith(parent_path + ' / ')`) to maintain structural integrity and eliminate L2 duplicate rendering. - *Visual Constraints*: Actionable drill-down paths are explicitly prefixed with `[SELECT]` tags. The LLM system prompt tightly constrains the model to only pick paths with this tag, preventing redundant re-selection of the current scope. -2. **Discovery Select**: The LLM reviews the specific paths flagged by Phase 1's Bottom Discovery for the current document. Selected discovery paths have their leaf chunks merged directly into the BFS document tree. +2. **Discovery Select**: The LLM reviews the specific paths flagged by Phase 1's Bottom Discovery for the current document. Selected discovery paths are hydrated into leaf chunks (with `job_result_id` dynamically extracted from the chunks) and merged directly into the BFS document tree. + - *Reparenting*: The `DocTreeNode.merge()` process reparents these discovered leaf chunks into the closest matching navigated child node. + - *Orphan Leaves*: Discovered chunks whose paths are not explicitly covered by the BFS `outline_items` are rendered cleanly as `[Leaf]` items (orphans) beneath their appropriate parent, ensuring no relevant data is lost even if the BFS did not explicitly drill into that path. **Phase 3: Verdict & Revision** The combined document tree (BFS Navigation + Discovery) is rendered as unified evidence. The tree naturally displays structural context (outlines) alongside hydrated chunk rows (for selected leaf paths). The LLM attempts to answer the user's query: diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py index 4c3e1080a..446c7d80a 100644 --- a/packages/shared-python/shared/services/retrieval/agent_navigate.py +++ b/packages/shared-python/shared/services/retrieval/agent_navigate.py @@ -21,6 +21,7 @@ _FILE_SELECT_PROMPT = """\ You are a document routing assistant. +{budget_block} Below is a knowledge base overview showing all available documents, their navigation summaries, chunk counts, and media counts. @@ -29,9 +30,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. """ @@ -41,8 +42,9 @@ You are a document navigation assistant. Document: "{doc_name}" (id: {doc_id}) -{scope_header} +{budget_block} +{scope_header} Below is the document's section tree. Sections tagged [SELECT] are within the current scope and may be selected. Other sections are shown as structural context only (not selectable). @@ -58,18 +60,63 @@ - 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. - 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": }}, ...]}} +Do not include any explanation. +""" + + +_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 []. +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. Return ONLY a JSON object: -{{"selections": [{{"path": "...", "confidence": , "mode": "all"}}, ...]}} +{{"selections": [{{"path": "...", "confidence": }}, ...]}} Do not include any explanation. """ +def _format_budget_block(snapshot: dict | None) -> str: + if not snapshot: + return "" + planning = snapshot.get("planning") or {} + context = snapshot.get("context") or {} + return ( + "=== Resource Status ===\n" + f"Planning Budget: {planning.get('status', 'HEALTHY')} " + f"({planning.get('used_pct', 0)}% used)\n" + f"Context Budget: {context.get('status', 'HEALTHY')} " + f"({context.get('used_pct', 0)}% used)\n" + f"KG Coverage: {snapshot.get('explored_chunks', 0)}/" + f"{snapshot.get('total_chunks', 0)} chunks explored\n" + f"Docs Explored: {snapshot.get('explored_docs', 0)}/" + f"{snapshot.get('total_docs', 0)}\n" + "When budget is TIGHT, prefer fewer high-confidence selections over broad exploration. " + "When CRITICAL, be very selective — only pick paths with strong relevance. Return empty if evidence suffices.\n" + "=== End Resource Status ===\n" + ) + + def _extract_json_array_payload(text: str) -> list[Any]: """Best-effort extraction of a JSON array payload from LLM response text.""" text = text.strip() @@ -219,7 +266,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 '' @@ -244,7 +290,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 @@ -455,6 +501,12 @@ async def _load_child_sections( scope_parts = split_section_path(scope) scope_depth = len(scope_parts) + logger.debug( + f' _load_child_sections: scope={scope!r} scope_parts={scope_parts} ' + f'scope_depth={scope_depth} exclude_paths={_excl if (_excl := exclude_paths or set()) else "none"} ' + f'total_sections={len(section_rows)}' + ) + # Build full section metadata index all_sections: dict[str, dict] = {} # path → {title, summary, sort_order, section_id, parts, depth} for section_id, title, path, summary, sort_order in section_rows: @@ -532,7 +584,7 @@ async def _load_child_sections( 'section_id': meta['section_id'], 'show_summary': False, } - else: + 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. @@ -553,12 +605,15 @@ async def _load_child_sections( continue # Category 2: Descendants of scope_path (children to explore) - if parts[:scope_depth] == scope_parts: + is_descendant = parts[:scope_depth] == scope_parts and depth > scope_depth + if is_descendant: # Skip excluded paths - if _excl and any( + 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] = { @@ -574,6 +629,12 @@ async def _load_child_sections( 'show_summary': True, } continue + else: + if depth > scope_depth: + 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) @@ -639,10 +700,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 @@ -651,6 +714,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 @@ -680,7 +816,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 @@ -700,7 +836,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 @@ -710,10 +845,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 @@ -738,25 +870,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 @@ -794,6 +935,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, @@ -818,57 +987,98 @@ def render_unified_doc_tree( if depth == 0: parts.append(f'【文档】{doc_name}\n') - # Track which paths have been rendered via outline_items - rendered_paths: set[str] = set() - # Collect children keys for path-hierarchy dedup: child_prefixes = set(node.children.keys()) + # Helper: min sort_order of a leaf_content entry + def _min_sort(path: str) -> float: + chunks = node.leaf_content.get(path, []) + return min((c.get('sort_order') or float('inf') for c in chunks), default=float('inf')) + + # ── Build a unified render queue ── + # Each entry: (sort_key, render_type, data) + # render_type: 'outline' | 'orphan_leaf' | 'orphan_child' + render_queue: list[tuple[float, str, dict | str]] = [] + + outline_paths: set[str] = set() + # Position counter for outline-only items (no leaf content) to preserve + # their relative ordering among themselves. + outline_position = 0.0 + for item in node.outline_items: path = item.get('path', '') - title = item.get('title', '') - is_leaf = item.get('is_leaf', False) - level = item.get('level', 1) - leaf_tag = ' [Leaf]' if is_leaf else '' - - # Skip items that belong to a drilled-into child's subtree + # Skip items belonging to a drilled-into child's subtree if any(path.startswith(cp + ' / ') for cp in child_prefixes): continue + outline_paths.add(path) - rendered_paths.add(path) - - # Section header (title only — summaries are navigation aids, not evidence) - level_tag = f'[L{level}] ' if level else '' - if level <= 1: - parts.append(f'{indent}▸ {level_tag}{title}{leaf_tag}') + # Determine sort_key: use chunk sort_order if content exists, + # else use a synthetic position to maintain outline ordering. + if path in node.leaf_content or path in node.children: + sort_key = _min_sort(path) if path in node.leaf_content else outline_position else: - parts.append(f'{indent}└ {level_tag}{title}{leaf_tag}') + sort_key = outline_position + outline_position = max(outline_position, sort_key) + 0.001 - sub_indent = indent + ' ' + render_queue.append((sort_key, 'outline', item)) - # Case 1: This section was drilled into → show child tree inline - if path in node.children: - child = node.children[path] - child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup) - if child_text.strip(): - parts.append(child_text) + # Add orphan leaf_content paths (not covered by outline_items) + for path in node.leaf_content: + if path not in outline_paths: + render_queue.append((_min_sort(path), 'orphan_leaf', path)) - # Case 2: This is a hydrated leaf → show chunk content inline - elif path in node.leaf_content: - _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + # Add orphan children (not covered by outline_items) + for path in node.children: + if path not in outline_paths: + render_queue.append((float('inf'), 'orphan_child', path)) + + # Sort by sort_key (stable sort preserves insertion order for ties) + render_queue.sort(key=lambda x: x[0]) + + from typing import cast + + # ── Render the unified queue ── + for _sort_key, rtype, data in render_queue: + if rtype == 'outline': + item = cast(dict, data) + path = item.get('path', '') + title = item.get('title', '') + is_leaf = item.get('is_leaf', False) + level = item.get('level', 1) + leaf_tag = ' [Leaf]' if is_leaf else '' + + level_tag = f'[L{level}] ' if level else '' + if level <= 1: + parts.append(f'{indent}▸ {level_tag}{title}{leaf_tag}') + else: + parts.append(f'{indent}└ {level_tag}{title}{leaf_tag}') + + sub_indent = indent + ' ' - # Case 3: Unselected → title already rendered above, nothing more needed + # Case 1: drilled-into child → render child tree + if path in node.children: + child = node.children[path] + if path in node.leaf_content: + _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + child_text = render_unified_doc_tree(child, doc_name, depth + 1, asset_lookup=asset_lookup) + if child_text.strip(): + parts.append(child_text) - # Render orphan paths: leaf_content and children not covered by outline_items - for path in node.leaf_content: - if path not in rendered_paths: + # Case 2: hydrated leaf → show chunk content + elif path in node.leaf_content: + _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) + + # Case 3: unselected → title only (already rendered above) + + elif rtype == 'orphan_leaf': + path = cast(str, data) title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path parts.append(f'{indent}▸ [Leaf] {title}') sub_indent = indent + ' ' _render_leaf_chunks(parts, node.leaf_content[path], sub_indent, asset_lookup=asset_lookup) - for path in node.children: - if path not in rendered_paths: + elif rtype == 'orphan_child': + path = cast(str, data) title = path.rsplit(' / ', 1)[-1] if ' / ' in path else path parts.append(f'{indent}▸ {title} [DrillDown]') child_text = render_unified_doc_tree(node.children[path], doc_name, depth + 1, asset_lookup=asset_lookup) @@ -876,3 +1086,4 @@ def render_unified_doc_tree( parts.append(child_text) return '\n'.join(parts) + diff --git a/packages/shared-python/shared/services/retrieval/agentic/budget.py b/packages/shared-python/shared/services/retrieval/agentic/budget.py new file mode 100644 index 000000000..1f4940bdd --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/budget.py @@ -0,0 +1,199 @@ +"""Token budget ledger for agentic retrieval runs.""" +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Literal + + +BudgetPoolName = Literal["bootstrap", "planning", "context"] +BudgetStatus = Literal["HEALTHY", "TIGHT", "CRITICAL", "EXHAUSTED"] + + +class BudgetExceeded(Exception): + """Raised when a planned LLM call cannot reserve budget.""" + + +@dataclass +class BudgetPool: + name: BudgetPoolName + capacity: int + used: int = 0 + reserved: int = 0 + + @property + def remaining(self) -> int: + return max(self.capacity - self.used - self.reserved, 0) + + @property + def used_pct(self) -> int: + if self.capacity <= 0: + return 100 + return min(100, int(round((self.used + self.reserved) * 100 / self.capacity))) + + +class BudgetLedger: + """Concurrency-safe ledger with bootstrap/planning/context pools.""" + + def __init__( + self, + *, + total: int, + planning_ratio: float, + bootstrap: int = 2000, + per_doc_min_share: int = 1500, + ) -> None: + total = max(int(total), 1) + bootstrap = max(0, min(int(bootstrap), total)) + remaining = max(total - bootstrap, 0) + planning_ratio = min(max(float(planning_ratio), 0.0), 1.0) + planning_capacity = int(remaining * planning_ratio) + context_capacity = remaining - planning_capacity + + self._lock = asyncio.Lock() + self._pools: dict[BudgetPoolName, BudgetPool] = { + "bootstrap": BudgetPool("bootstrap", bootstrap), + "planning": BudgetPool("planning", planning_capacity), + "context": BudgetPool("context", context_capacity), + } + self._doc_caps: dict[str, int] = {} + self._doc_used: dict[str, int] = {} + self._doc_reserved: dict[str, int] = {} + self._per_doc_min_share = max(int(per_doc_min_share), 0) + self.total_chunks = 0 + self.total_docs = 0 + self.explored_chunks = 0 + self.explored_docs = 0 + self.trimmed_paths: list[dict[str, str]] = [] + + def remaining(self, pool: BudgetPoolName) -> int: + return self._pools[pool].remaining + + def status(self, pool: BudgetPoolName) -> BudgetStatus: + pool_state = self._pools[pool] + if pool_state.remaining <= 0: + return "EXHAUSTED" + used_pct = pool_state.used_pct + if used_pct >= 80: + return "CRITICAL" + if used_pct >= 50: + return "TIGHT" + return "HEALTHY" + + async def allocate_doc_caps(self, doc_chunks: dict[str, int]) -> None: + """Allocate planning soft caps by document chunk counts.""" + async with self._lock: + self._doc_caps.clear() + self._doc_used.clear() + self._doc_reserved.clear() + if not doc_chunks: + return + + planning_capacity = self._pools["planning"].capacity + total_weight = sum(max(int(count), 1) for count in doc_chunks.values()) + for doc_id, count in doc_chunks.items(): + weight = max(int(count), 1) + weighted = int(planning_capacity * weight / total_weight) + self._doc_caps[doc_id] = min( + planning_capacity, + max(self._per_doc_min_share, weighted), + ) + + async def try_reserve( + self, + pool: BudgetPoolName, + est: int, + doc_id: str | None = None, + *, + priority: Literal["normal", "low"] = "normal", + ) -> bool: + est = max(int(est), 0) + if est == 0: + return True + + async with self._lock: + pool_state = self._pools[pool] + if priority == "low" and self.status(pool) == "CRITICAL": + return False + if pool_state.remaining < est: + return False + + pool_state.reserved += est + if pool == "planning" and doc_id: + self._doc_reserved[doc_id] = self._doc_reserved.get(doc_id, 0) + est + return True + + async def commit( + self, + pool: BudgetPoolName, + *, + actual: int, + est: int, + doc_id: str | None = None, + ) -> None: + actual = max(int(actual), 0) + est = max(int(est), 0) + async with self._lock: + pool_state = self._pools[pool] + reserved_delta = min(est, pool_state.reserved) + pool_state.reserved -= reserved_delta + pool_state.used = min(pool_state.capacity, pool_state.used + actual) + + if pool == "planning" and doc_id: + doc_reserved = min(est, self._doc_reserved.get(doc_id, 0)) + if doc_reserved: + self._doc_reserved[doc_id] -= doc_reserved + if self._doc_reserved[doc_id] <= 0: + self._doc_reserved.pop(doc_id, None) + self._doc_used[doc_id] = self._doc_used.get(doc_id, 0) + actual + + async def refund( + self, + pool: BudgetPoolName, + *, + est: int, + doc_id: str | None = None, + ) -> None: + est = max(int(est), 0) + async with self._lock: + pool_state = self._pools[pool] + pool_state.reserved = max(pool_state.reserved - est, 0) + if pool == "planning" and doc_id: + current = self._doc_reserved.get(doc_id, 0) + remaining = max(current - est, 0) + if remaining: + self._doc_reserved[doc_id] = remaining + else: + self._doc_reserved.pop(doc_id, None) + + def mark_explored( + self, + *, + chunks: int = 0, + docs: int = 0, + ) -> None: + self.explored_chunks += max(int(chunks), 0) + self.explored_docs += max(int(docs), 0) + + def snapshot(self) -> dict[str, object]: + snapshot: dict[str, object] = { + name: { + "capacity": pool.capacity, + "used": pool.used, + "reserved": pool.reserved, + "remaining": pool.remaining, + "used_pct": pool.used_pct, + "status": self.status(name), + } + for name, pool in self._pools.items() + } + snapshot.update({ + "total_chunks": self.total_chunks, + "total_docs": self.total_docs, + "explored_chunks": min(self.explored_chunks, self.total_chunks) + if self.total_chunks else self.explored_chunks, + "explored_docs": min(self.explored_docs, self.total_docs) + if self.total_docs else self.explored_docs, + "trimmed_paths": list(self.trimmed_paths), + }) + return snapshot diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index ec42b6b0d..b8b766327 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -17,14 +17,16 @@ from __future__ import annotations import os -from typing import Any +import json +from typing import Any, cast from loguru import logger -from sqlalchemy import select +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document +from shared.models.database.document import Document, DocumentChunk, RetrievalHitStat +from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger, BudgetPoolName from shared.services.retrieval.agentic.trace import TraceRecorder from shared.services.retrieval.agentic.types import ( AgentRunConfig, @@ -34,18 +36,14 @@ DocTreeNode, ToolResult, ) -from shared.services.retrieval.agentic.result_rows import ( +from shared.services.retrieval.app_service import ( generate_retrieval_asset_url, - is_client_result_artifact_ref, + _is_client_result_artifact_ref, ) from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.agentic.policy import attempt_answer -from shared.services.retrieval.agentic.tools import ( - bottom_discovery, - discovery_select_step, - kg_document_select, - scope_navigate_step, -) +from shared.services.retrieval.llm_adapter import current_llm_usage +from shared.services.retrieval.hit_stats_service import compute_importance_score +from shared.utils.token_estimate import estimate_tokens @@ -65,6 +63,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]: @@ -80,7 +86,7 @@ async def _build_asset_url_map( job_id = chunk.get('job_id') or '' if not chunk_id or not file_path or not job_id: continue - if not is_client_result_artifact_ref(file_path): + if not _is_client_result_artifact_ref(file_path): continue try: url = await generate_retrieval_asset_url( @@ -100,9 +106,23 @@ def _build_config_from_env() -> AgentRunConfig: max_revisions=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_REVISIONS', '2')), max_nav_depth=int(os.environ.get('RETRIEVAL_AGENTIC_MAX_NAV_DEPTH', '3')), latency_budget_ms=int(os.environ.get('RETRIEVAL_AGENTIC_LATENCY_BUDGET_MS', '12000')), + token_budget_total=int(os.environ.get('RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL', '40000')), + planning_ratio=float(os.environ.get('RETRIEVAL_AGENTIC_PLANNING_RATIO', '0.5')), + bootstrap_budget=int(os.environ.get('RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET', '2000')), + per_doc_min_share=int(os.environ.get('RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE', '1500')), + inventory_aware=os.environ.get('RETRIEVAL_AGENTIC_INVENTORY_AWARE', 'true') == 'true', ) +def _stringify_llm_input(prompt: Any) -> str: + if isinstance(prompt, str): + return prompt + try: + return json.dumps(prompt, ensure_ascii=False, default=str) + except Exception: + return str(prompt) + + async def _render_evidence( db: AsyncSession, doc_trees: dict[str, DocTreeNode], @@ -134,6 +154,154 @@ async def _render_evidence( return '\n\n'.join(evidence_parts) if evidence_parts else '(no evidence collected)' +async def _load_budget_inventory( + db: AsyncSession, + *, + user_id: str, + namespace: str, + exclude_document_ids: list[str], +) -> tuple[int, int, dict[str, int]]: + stmt = ( + select(Document.document_id, func.count(DocumentChunk.id)) + .join( + DocumentChunk, + (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id), + ) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .group_by(Document.document_id) + ) + if exclude_document_ids: + stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) + + result = await db.execute(stmt) + doc_chunks = {str(doc_id): int(count or 0) for doc_id, count in result.all()} + return sum(doc_chunks.values()), len(doc_chunks), doc_chunks + + +def _iter_leaf_content(node: DocTreeNode): + for path, chunks in node.leaf_content.items(): + yield path, chunks + for child in node.children.values(): + yield from _iter_leaf_content(child) + + +def _collect_confidences(node: DocTreeNode) -> dict[str, float]: + values = dict(node.confidence) + for child in node.children.values(): + for path, score in _collect_confidences(child).items(): + values[path] = max(values.get(path, 0.0), score) + return values + + +def _pop_leaf_path(node: DocTreeNode, path: str) -> bool: + if path in node.leaf_content: + node.leaf_content.pop(path) + return True + for child in node.children.values(): + if _pop_leaf_path(child, path): + return True + return False + + +def _estimate_chunks_tokens(chunks: list[dict[str, Any]]) -> int: + text = '\n'.join(str(chunk.get('content') or '') for chunk in chunks) + return estimate_tokens(text) + + +async def _fetch_importance_norm_scores( + db: AsyncSession, + *, + user_id: str, + namespace: str, + chunk_ids: list[str], +) -> dict[str, float]: + if not chunk_ids: + return {} + stmt = ( + select( + RetrievalHitStat.chunk_id, + RetrievalHitStat.hit_count, + RetrievalHitStat.last_hit_at, + RetrievalHitStat.created_at, + ) + .where(RetrievalHitStat.user_id == user_id) + .where(RetrievalHitStat.namespace == namespace) + .where(RetrievalHitStat.hit_kind == 'chunk') + .where(RetrievalHitStat.chunk_id.in_(chunk_ids)) + ) + result = await db.execute(stmt) + scores: dict[str, float] = {} + for chunk_id, hit_count, last_hit_at, created_at in result.all(): + if chunk_id and last_hit_at and created_at: + scores[str(chunk_id)] = compute_importance_score(hit_count, last_hit_at, created_at) + return scores + + +async def _trim_evidence_to_budget( + db: AsyncSession, + *, + doc_trees: dict[str, DocTreeNode], + doc_id_to_name: dict[str, str], + context_remaining: int, + user_id: str, + namespace: str, + ledger: BudgetLedger | None, + safety_margin: float = 0.9, +) -> str: + full_text = await _render_evidence(db, doc_trees, doc_id_to_name) + target = int(max(context_remaining, 0) * safety_margin) + if target <= 0 or estimate_tokens(full_text) <= target: + return full_text + + candidates: list[tuple[str, str, tuple[float, float, float], int]] = [] + for doc_id, tree in doc_trees.items(): + confidence = _collect_confidences(tree) + for path, chunks in _iter_leaf_content(tree): + chunk_ids = [ + str(chunk.get('chunk_id')) + for chunk in chunks + if chunk.get('chunk_id') + ] + importance = 0.0 + importance_scores = await _fetch_importance_norm_scores( + db, + user_id=user_id, + namespace=namespace, + chunk_ids=chunk_ids, + ) + if importance_scores: + importance = max(importance_scores.values()) + discovery_score = ( + float(chunks[0].get('discovery_score', 0.0) or 0.0) + if chunks else 0.0 + ) + score = (float(confidence.get(path, 0.0) or 0.0), discovery_score, importance) + candidates.append((doc_id, path, score, _estimate_chunks_tokens(chunks))) + + current_estimate = estimate_tokens(full_text) + removed: list[dict[str, str]] = [] + for doc_id, path, _score, token_estimate in sorted( + candidates, + key=lambda item: (item[2], -item[3]), + ): + if current_estimate <= target: + break + if _pop_leaf_path(doc_trees[doc_id], path): + removed.append({'document_id': doc_id, 'path': path}) + current_estimate = max(current_estimate - token_estimate, 0) + + if ledger is not None: + ledger.trimmed_paths.extend(removed) + logger.info( + f' agentic.trim_evidence: removed={len(removed)} ' + f'est_tokens={current_estimate} target={target}' + ) + return await _render_evidence(db, doc_trees, doc_id_to_name) + + class RetrievalAgent: """Agentic retrieval orchestrator — navigate-then-answer loop. @@ -151,6 +319,82 @@ class RetrievalAgent: If ``llm_fn`` is None, the run returns discovery-only results. """ + async def _call_llm_with_budget( + self, + state: AgentState, + llm_fn: LLMFn, + prompt: Any, + *, + pool: BudgetPoolName, + doc_id: str | None = None, + priority: str = 'normal', + ) -> str: + ledger = state.ledger + if ledger is None: + return await llm_fn(prompt) + + prompt_text = _stringify_llm_input(prompt) + est = estimate_tokens(prompt_text) + reserved = await ledger.try_reserve( + pool, + est, + doc_id=doc_id, + priority='low' if priority == 'low' else 'normal', + ) + if not reserved: + raise BudgetExceeded(f'{pool} budget exhausted') + + try: + response = await llm_fn(prompt) + except Exception: + await ledger.refund(pool, est=est, doc_id=doc_id) + raise + + usage = current_llm_usage.get() or {} + actual = int(usage.get('prompt_tokens') or est) + await ledger.commit(pool, actual=actual, est=est, doc_id=doc_id) + return response + + def _budgeted_doc_llm_fn( + self, + state: AgentState, + llm_fn: LLMFn, + *, + doc_id: str, + depth: int, + ) -> LLMFn: + async def _call(prompt): + return await self._call_llm_with_budget( + state, + llm_fn, + prompt, + pool='planning', + doc_id=doc_id, + priority='low' if depth >= 2 else 'normal', + ) + + return _call + + def _budgeted_discovery_llm_fn( + self, + state: AgentState, + llm_fn: LLMFn, + *, + doc_id: str, + low_priority: bool, + ) -> LLMFn: + async def _call(prompt): + return await self._call_llm_with_budget( + state, + llm_fn, + prompt, + pool='planning', + doc_id=doc_id, + priority='low' if low_priority else 'normal', + ) + + return _call + async def run( self, db: AsyncSession, @@ -176,11 +420,33 @@ async def run( errors are captured in trace and the best available result is returned. """ + 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 [] exclude_sections = exclude_sections or [] state = AgentState() + state.ledger = BudgetLedger( + total=config.token_budget_total, + planning_ratio=config.planning_ratio, + bootstrap=config.bootstrap_budget, + per_doc_min_share=config.per_doc_min_share, + ) + total_chunks, total_docs, chunks_count_by_doc = await _load_budget_inventory( + db, + user_id=user_id, + namespace=namespace, + exclude_document_ids=exclude_document_ids, + ) + state.kg_total_chunks = total_chunks + state.kg_total_docs = total_docs + state.ledger.total_chunks = total_chunks + state.ledger.total_docs = total_docs trace = TraceRecorder( db, user_id=user_id, namespace=namespace, query=query, config=config, top_k=top_k, data_type=data_type, @@ -197,12 +463,38 @@ async def run( logger.info( f'agentic retrieval START: query="{query[:60]}..." ' - f'top_k={top_k} budget={config.latency_budget_ms}ms' + f'top_k={top_k} latency_budget={config.latency_budget_ms}ms ' + f'token_budget={config.token_budget_total}' ) if llm_fn is None: logger.warning('agentic: no llm_fn provided — running discovery-only mode') + planning_llm_fn: LLMFn | None = None + bootstrap_llm_fn: LLMFn | None = None + context_llm_fn: LLMFn | None = None + if llm_fn is not None: + base_llm_fn = llm_fn + + async def _planning_llm_call(prompt): + return await self._call_llm_with_budget( + state, base_llm_fn, prompt, pool='planning' + ) + + async def _bootstrap_llm_call(prompt): + return await self._call_llm_with_budget( + state, base_llm_fn, prompt, pool='bootstrap' + ) + + async def _context_llm_call(prompt): + return await self._call_llm_with_budget( + state, base_llm_fn, prompt, pool='context' + ) + + planning_llm_fn = _planning_llm_call + bootstrap_llm_fn = _bootstrap_llm_call + context_llm_fn = _context_llm_call + # Shared kwargs for bottom_discovery discovery_kwargs: dict[str, Any] = { 'user_id': user_id, @@ -224,7 +516,7 @@ async def run( logger.info(' agentic: Phase 1 — discovery + document selection') # 1a. Bottom discovery (always runs) - discovery_result = await bottom_discovery(db, **discovery_kwargs) + discovery_result = await tools.bottom_discovery(db, **discovery_kwargs) state.step_count += 1 discovery_rows = discovery_result.payload.get('fused_rows', []) if discovery_result.status != 'error' else [] state.discovery_top_doc_ids = discovery_result.payload.get('top_doc_ids', []) if discovery_result.status != 'error' else [] @@ -241,15 +533,25 @@ async def run( ) # 1b. KG document selection (requires LLM) - if llm_fn is not None: - kg_result = await kg_document_select( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=llm_fn, - exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)), - ) + if bootstrap_llm_fn is not None: + try: + kg_result = await tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=bootstrap_llm_fn, + exclude_document_ids=list(state.ever_explored_doc_ids | set(exclude_document_ids)), + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(' agentic: bootstrap budget exhausted during document selection') + if trace_enabled: + trace.record_budget_stop('bootstrap_exhausted') + kg_result = ToolResult( + status='no_confident_doc', + payload={'reason': 'bootstrap budget exhausted'}, + ) state.step_count += 1 if trace_enabled: @@ -312,7 +614,11 @@ async def run( if r.get('chunk_id') ] if trace_enabled: - await trace.complete(discovery_rows, 'agentic_discovery_only') + await trace.complete( + discovery_rows, + 'agentic_discovery_only', + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) return AgenticResult( evidence_text='', answer_text='', @@ -353,6 +659,12 @@ async def run( if jrid: state.doc_job_map[did] = jrid + if state.ledger is not None: + await state.ledger.allocate_doc_caps({ + doc.document_id: chunks_count_by_doc.get(doc.document_id, 1) + for doc in state.selected_docs + }) + # ══════════════════════════════════════════════════════════════════ # Phase 2 + 3 Loop: Navigate → Render → Attempt Answer → (Revise) # ══════════════════════════════════════════════════════════════════ @@ -390,11 +702,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) @@ -408,35 +723,178 @@ async def run( if depth >= config.max_nav_depth: continue - if llm_fn is None: + if planning_llm_fn is None: + break + if state.ledger and state.ledger.status('planning') in ('CRITICAL', 'EXHAUSTED'): + logger.info(' agentic: planning budget critical, ending BFS for current doc') break - step_node, drill_paths = await scope_navigate_step( - db, - document_id=doc.document_id, - job_result_id=job_result_id, - query=query, - llm_fn=llm_fn, - user_id=user_id, - namespace=namespace, - doc_name=doc_name, - scope_path=scope, - exclude_paths=doc_exclude, - revision_hint=revision_hint if depth == 0 else None, + doc_llm_fn = self._budgeted_doc_llm_fn( + state, + cast(LLMFn, llm_fn), + doc_id=doc.document_id, + depth=depth, ) + + # ★ Step 1: Tool selection (agent decides which asset tools) + try: + tool_choices = await tools.tool_select_step( + db, + document_id=doc.document_id, + job_result_id=job_result_id, + query=query, + llm_fn=doc_llm_fn, + 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 tool selection') + 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: + 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, + document_id=doc.document_id, + job_result_id=job_result_id, + scope_path=scope, + asset_type=asset_type, + ) + if asset_chunks: + pending_scope_assets.extend(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}' + ) + # ★ 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 - parent_node.leaf_content = step_node.leaf_content + for leaf_path, chunks in step_node.leaf_content.items(): + 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. + parent_node.reparent_leaf_content() + if trace_enabled: trace.record_step( 'scope_navigate_step', ToolResult( @@ -460,23 +918,41 @@ async def run( f'leaves={len(step_node.leaf_content)} ' f'drills={len(drill_paths)}' ) + if state.ledger is not None: + state.ledger.mark_explored( + chunks=sum(len(chunks) for chunks in step_node.leaf_content.values()), + ) else: # B-class: no BFS, create empty root root = DocTreeNode(scope_path=None) # ── Post-BFS: Discovery selection step ───────────────────── doc_hints = discovery_by_doc.get(doc.document_id, []) - if doc_hints and llm_fn is not None and state.elapsed_ms < config.latency_budget_ms: - discovery_node = await discovery_select_step( - db, - document_id=doc.document_id, - query=query, - llm_fn=llm_fn, - user_id=user_id, - namespace=namespace, - doc_name=doc_name, - discovery_hints=doc_hints, + if doc_hints and planning_llm_fn is not None and state.elapsed_ms < config.latency_budget_ms: + doc_discovery_llm_fn = self._budgeted_discovery_llm_fn( + state, + cast(LLMFn, llm_fn), + doc_id=doc.document_id, + low_priority=root.has_content(), ) + try: + discovery_node = await tools.discovery_select_step( + db, + document_id=doc.document_id, + query=query, + llm_fn=doc_discovery_llm_fn, + user_id=user_id, + namespace=namespace, + doc_name=doc_name, + discovery_hints=doc_hints, + revision_hint=revision_hint, + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(' agentic: planning budget exhausted during discovery selection') + if trace_enabled: + trace.record_budget_stop('planning_exhausted') + discovery_node = DocTreeNode(scope_path=None) state.step_count += 1 if trace_enabled: @@ -494,6 +970,10 @@ async def run( # Merge discovery results into BFS tree root.merge(discovery_node) + if state.ledger is not None: + state.ledger.mark_explored( + chunks=sum(len(chunks) for chunks in discovery_node.leaf_content.values()), + ) # Merge or store doc tree if doc.document_id in state.doc_trees: @@ -501,25 +981,58 @@ async def run( else: state.doc_trees[doc.document_id] = root state.ever_explored_doc_ids.add(doc.document_id) + if state.ledger is not None: + state.ledger.mark_explored(docs=1) # ── Phase 3: Render evidence + attempt_answer ──────────────── - evidence_text = await _render_evidence( + evidence_text = await _trim_evidence_to_budget( db, - state.doc_trees, state.doc_id_to_name, + doc_trees=state.doc_trees, + doc_id_to_name=state.doc_id_to_name, + context_remaining=state.ledger.remaining('context') if state.ledger else config.token_budget_total, + user_id=user_id, + namespace=namespace, + ledger=state.ledger, ) - if llm_fn is None: + if context_llm_fn is None: stop_reason = 'no_llm' break - # Auto-trigger attempt_answer - status, answer_text, reason = await attempt_answer( - llm_fn, - query=query, - evidence_text=evidence_text, - state=state, - config=config, - ) + # 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 + ] + + async def vlm_context_call(prompt, _vlm_fn=vlm_fn): + return await self._call_llm_with_budget( + state, cast(LLMFn, _vlm_fn), prompt, pool='context' + ) + + # Auto-trigger attempt_answer (VLM if images present) + try: + status, answer_text, reason = await attempt_answer( + context_llm_fn, + query=query, + evidence_text=evidence_text, + state=state, + config=config, + vlm_fn=vlm_context_call if vlm_fn else None, + image_urls=evidence_image_urls or None, + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(' agentic: context budget exhausted before attempt_answer') + if trace_enabled: + trace.record_budget_stop('context_exhausted') + status, answer_text, reason = 'NOT_FOUND', '', 'context budget exhausted' + stop_reason = 'context_budget' + break state.step_count += 1 if trace_enabled: @@ -560,15 +1073,27 @@ 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) - kg_result = await kg_document_select( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=llm_fn, - exclude_document_ids=list(set(exclude_document_ids)), - ) + # Re-run KG select with revision hint + if bootstrap_llm_fn is None: + stop_reason = 'no_llm' + break + try: + kg_result = await tools.kg_document_select( + db, + user_id=user_id, + namespace=namespace, + query=query, + llm_fn=bootstrap_llm_fn, + exclude_document_ids=list(set(exclude_document_ids)), + revision_hint=revision_hint, + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) + except BudgetExceeded: + logger.info(' agentic: bootstrap budget exhausted during revision doc selection') + if trace_enabled: + trace.record_budget_stop('bootstrap_exhausted') + stop_reason = 'bootstrap_budget' + break state.step_count += 1 if kg_result.status == 'selected_docs': @@ -588,6 +1113,12 @@ async def run( stop_reason = 'no_new_docs' break + if state.ledger is not None: + await state.ledger.allocate_doc_caps({ + doc.document_id: chunks_count_by_doc.get(doc.document_id, 1) + for doc in state.selected_docs + }) + # ══════════════════════════════════════════════════════════════════ # Final Assembly # ══════════════════════════════════════════════════════════════════ @@ -618,6 +1149,8 @@ async def run( answer_text=answer_text, referenced_chunks=all_refs, router_used=router_used, + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + stop_reason=stop_reason, ) logger.info( @@ -630,6 +1163,10 @@ async def run( ) if trace_enabled: - await trace.complete(all_refs, router_used) + await trace.complete( + all_refs, + router_used, + budget_snapshot=state.ledger.snapshot() if state.ledger else None, + ) return result diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index e5a5c0dd4..6d63035e7 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -12,10 +12,13 @@ import json import os import re +from ipaddress import ip_address from typing import Any +from urllib.parse import urlparse from loguru import logger +from shared.services.retrieval.agentic.budget import BudgetExceeded from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState from shared.services.retrieval.llm_adapter import LLMFn @@ -36,6 +39,20 @@ def _parse_answer_response(text: str) -> dict[str, Any] | None: return None +def _is_external_http_url(url: str) -> bool: + parsed = urlparse(str(url)) + if parsed.scheme not in {'http', 'https'} or not parsed.hostname: + return False + host = parsed.hostname.strip().lower() + if host in {'localhost', 'ip6-localhost', 'ip6-loopback'} or host.endswith('.local'): + return False + try: + addr = ip_address(host) + except ValueError: + return True + return not (addr.is_private or addr.is_loopback or addr.is_link_local) + + async def attempt_answer( llm_fn: LLMFn, *, @@ -43,30 +60,79 @@ async def attempt_answer( evidence_text: str, state: AgentState, config: AgentRunConfig, + vlm_fn: LLMFn | None = None, + image_urls: list[str] | None = None, + budget_snapshot: dict | None = None, ) -> tuple[str, str, str]: """Attempt to answer the query using collected evidence. + When ``vlm_fn`` is provided and ``image_urls`` is non-empty, the + answer is generated using the VLM in multimodal message format + (text + image_url parts) so the model can actually *see* chart + images rather than only reading their text descriptions. + + Falls back to ``llm_fn`` (text-only) when VLM is unavailable or + when no images are present. + Returns (status, answer_text, reason) where: - status='DONE', answer_text=, reason='' → evidence was sufficient, answer is ready - status='NOT_FOUND', answer_text='', reason= → evidence was insufficient, reason is used as revision_hint """ - prompt = _ATTEMPT_ANSWER_PROMPT.format( + prompt_text = _ATTEMPT_ANSWER_PROMPT.format( query=query, evidence_context=evidence_text, revision_count=state.revision_count, max_revisions=config.max_revisions, + context_status=((budget_snapshot or {}).get('context') or {}).get('status', 'HEALTHY'), ) verbose = os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', '') == 'true' + + # Decide whether to use VLM (multimodal) or text-only LLM + effective_fn = llm_fn + effective_input: Any = prompt_text + + usable_image_urls = [url for url in image_urls or [] if _is_external_http_url(url)] + if image_urls and len(usable_image_urls) != len(image_urls): + logger.info( + f' [attempt_answer] skipped {len(image_urls) - len(usable_image_urls)} ' + 'non-public image URLs for VLM' + ) + + if vlm_fn and usable_image_urls: + # Build multimodal message: text evidence + image_url parts + content_parts: list[dict[str, Any]] = [ + {'type': 'text', 'text': prompt_text}, + ] + for url in usable_image_urls[:20]: # Cap at 20 images to avoid token overflow + content_parts.append({ + 'type': 'image_url', + 'image_url': {'url': url}, + }) + effective_input = [{'role': 'user', 'content': content_parts}] + effective_fn = vlm_fn + logger.info( + f' [attempt_answer] using VLM with {len(usable_image_urls)} image URLs ' + f'(capped at {min(len(usable_image_urls), 20)})' + ) + if verbose: logger.info( f'[attempt_answer PROMPT]\n' - f'{prompt}' + f'{prompt_text}' ) - raw_response = await llm_fn(prompt) + try: + raw_response = await effective_fn(effective_input) + except BudgetExceeded: + raise + except Exception as exc: + if effective_fn is llm_fn: + raise + logger.warning(f' [attempt_answer] VLM failed, falling back to text LLM: {exc}') + raw_response = await llm_fn(prompt_text) logger.info(f' [attempt_answer] raw={repr(raw_response[:300])}') if verbose: @@ -75,6 +141,9 @@ async def attempt_answer( f'{raw_response}' ) + if not raw_response.strip() and effective_fn is not llm_fn: + return 'NOT_FOUND', '', 'VLM returned empty response for multimodal evidence' + parsed = _parse_answer_response(raw_response) if not parsed: # Parse error: treat raw text as best-effort answer @@ -107,6 +176,7 @@ async def attempt_answer( {evidence_context} REVISION: {revision_count} of {max_revisions} revisions used. +Context budget remaining is {context_status}; the evidence may have been trimmed. INSTRUCTIONS: 1. If the evidence contains enough information to answer the query, diff --git a/packages/shared-python/shared/services/retrieval/agentic/result_rows.py b/packages/shared-python/shared/services/retrieval/agentic/result_rows.py deleted file mode 100644 index b4e39a5f9..000000000 --- a/packages/shared-python/shared/services/retrieval/agentic/result_rows.py +++ /dev/null @@ -1,448 +0,0 @@ -"""Row helpers used by the agentic retrieval pipeline. - -These functions intentionally avoid importing ``app_service`` so the agentic -orchestrator can be imported without creating retrieval package cycles. -""" -from __future__ import annotations - -from typing import Any - -from loguru import logger -from sqlalchemy import and_, or_, select -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.models.database.document import Document, DocumentChunk, DocumentSection -from shared.models.database.job_result import JobResult -from shared.services.retrieval.graph_service import is_excluded_section -from shared.services.retrieval.lexical_text import normalize_section_path -from shared.services.storage.result_storage import get_result_storage - -CHANNEL_WEIGHT_PATH = 1.0 -CHANNEL_WEIGHT_CONTENT = 2.0 -CHANNEL_WEIGHT_TERM = 1.5 -INTERNAL_RECALL_K_MULTIPLIER = 2 -RRF_K = 60 - -DATA_TYPE_ALLOWED_CHUNK_TYPES: dict[int, set[str] | None] = { - 1: None, - 2: {'text'}, - 3: {'image'}, - 4: {'table'}, - 5: {'text', 'image'}, - 6: {'text', 'table'}, -} - -MODE_ALLOWED_TYPES: dict[str, set[str] | None] = { - 'chunks': None, - 'assets_only': {'image', 'table'}, - 'image_only': {'image'}, - 'table_only': {'table'}, -} - - -def resolve_allowed_chunk_types(data_type: int) -> set[str] | None: - return DATA_TYPE_ALLOWED_CHUNK_TYPES.get(data_type) - - -def normalize_chunk_type(raw: str | None) -> str: - return str(raw or '').strip().split('\n', 1)[0].lower() - - -def filter_excluded_rows( - rows: list[dict[str, Any]], - *, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - filtered: list[dict[str, Any]] = [] - excluded_documents = set(exclude_document_ids) - for row in rows: - document_id = row.get('document_id') - if document_id in excluded_documents: - continue - if is_excluded_section( - document_id=document_id, - section_path=row.get('section_path'), - exclude_sections=exclude_sections, - ): - continue - filtered.append(row) - return filtered - - -def iter_connected_target_ids(row: dict[str, Any]) -> list[str]: - metadata = row.get('chunk_metadata') or {} - if not isinstance(metadata, dict): - return [] - - target_ids: list[str] = [] - for item in metadata.get('connect_to') or []: - if not isinstance(item, dict): - continue - target_id = str(item.get('target') or '').strip() - if target_id: - target_ids.append(target_id) - return target_ids - - -async def hydrate_connected_target_rows( - *, - db: AsyncSession | None, - rows: list[dict[str, Any]], - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> list[dict[str, Any]]: - if db is None: - return [] - - existing_chunk_ids = { - str(row.get('chunk_id') or '').strip() - for row in rows - if row.get('chunk_id') - } - target_ids_by_revision: dict[tuple[str, str], set[str]] = {} - for row in rows: - if normalize_chunk_type(row.get('chunk_type')) != 'text': - continue - document_id = str(row.get('document_id') or '').strip() - job_result_id = str(row.get('job_result_id') or '').strip() - if not document_id or not job_result_id: - continue - for target_id in iter_connected_target_ids(row): - if target_id in existing_chunk_ids: - continue - target_ids_by_revision.setdefault((document_id, job_result_id), set()).add(target_id) - - if not target_ids_by_revision: - return [] - - revision_filters = [ - and_( - DocumentChunk.document_id == document_id, - DocumentChunk.job_result_id == job_result_id, - DocumentChunk.chunk_id.in_(sorted(target_ids)), - ) - for (document_id, job_result_id), target_ids in target_ids_by_revision.items() - if target_ids - ] - if not revision_filters: - return [] - - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, DocumentChunk.document_id == Document.document_id) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(or_(*revision_filters)) - .order_by(DocumentChunk.sort_order) - ) - result = await db.execute(stmt) - - hydrated_rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - hydrated_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': 0.0, - '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, - 'sort_order': chunk.sort_order, - } - ) - - return filter_excluded_rows( - hydrated_rows, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - - -def merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: - if not rows: - return rows - groups: dict[str, list[dict[str, Any]]] = {} - order: list[str] = [] - for row in rows: - section_path = row.get('section_path') - if section_path: - key = f"{row.get('document_id', '')}::{section_path}" - else: - key = row.get('chunk_id', '') - if key not in groups: - groups[key] = [] - order.append(key) - groups[key].append(row) - - merged: list[dict[str, Any]] = [] - for key in order: - group = groups[key] - if len(group) == 1: - merged.append(group[0]) - continue - base = dict(group[0]) - base['content'] = '\n'.join(str(row.get('content', '')) for row in group) - base['score'] = max(row.get('score', 0.0) for row in group) - merged.append(base) - return merged - - -def merge_channels_rrf( - channels: list[list[dict[str, Any]]], - weights: list[float], - top_k: int, - k: int = RRF_K, -) -> list[dict[str, Any]]: - score_dict: dict[str, float] = {} - row_by_chunk_id: dict[str, dict[str, Any]] = {} - - for channel_idx, channel_rows in enumerate(channels): - weight = weights[channel_idx] if channel_idx < len(weights) else 1.0 - for rank, row in enumerate(channel_rows): - chunk_id = str(row.get('chunk_id') or '') - if not chunk_id: - continue - rrf_score = weight / (k + rank + 1) - score_dict[chunk_id] = score_dict.get(chunk_id, 0.0) + rrf_score - if chunk_id not in row_by_chunk_id: - row_by_chunk_id[chunk_id] = row - - ranked = sorted(score_dict.items(), key=lambda item: item[1], reverse=True) - results: list[dict[str, Any]] = [] - for chunk_id, fused_score in ranked[:top_k]: - row = row_by_chunk_id[chunk_id] - results.append(dict(row, score=round(fused_score, 6))) - return results - - -def normalize_row_scores( - rows: list[dict[str, Any]], - *, - source_field: str, - target_field: str, - default: float, -) -> None: - if not rows: - return - values = [float(row.get(source_field, 0.0) or 0.0) for row in rows] - min_score = min(values) - max_score = max(values) - if max_score <= 0.0 and min_score <= 0.0: - for row in rows: - row[target_field] = 0.0 - return - if max_score == min_score: - for row in rows: - row[target_field] = default - return - denominator = max_score - min_score - for row in rows: - raw_score = float(row.get(source_field, 0.0) or 0.0) - row[target_field] = round((raw_score - min_score) / denominator, 6) - - -def get_row_path(row: dict[str, Any]) -> str: - return str(row.get('section_path') or row.get('source_chunk_path') or '') - - -async def hydrate_paths_to_rows( - db: AsyncSession, - *, - path_selections: list[dict[str, Any]], - user_id: str, - namespace: str, - document_id: str | None = None, -) -> list[dict[str, Any]]: - if not path_selections: - return [] - - 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() - path = normalize_section_path(raw_path) if raw_path and '/' in raw_path else raw_path - 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 [] - - outline_paths = [path for path in ordered_paths if mode_by_path.get(path) == 'outline'] - chunk_paths = [path for path in ordered_paths if mode_by_path.get(path) != 'outline'] - - rows: list[dict[str, Any]] = [] - - if outline_paths: - outline_section_filters = [ - DocumentSection.section_path == path - for path in outline_paths - ] - - 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)) - ) - if document_id: - outline_stmt = outline_stmt.where(Document.document_id == document_id) - outline_result = await db.execute(outline_stmt) - for document, section in outline_result.all(): - agent_score = confidence_by_path.get(section.section_path, 0.0) - 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', - }) - - 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), - ) - ) - ) - if document_id: - stmt = stmt.where(Document.document_id == document_id) - result = await db.execute(stmt) - - 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 - - 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, - ) - - 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, - }) - - path_order = {path: index for index, path in enumerate(ordered_paths)} - - def get_sort_key(row: dict[str, Any]) -> int: - row_path = get_row_path(row) - if row_path in path_order: - return path_order[row_path] - for path, index in path_order.items(): - if row_path.startswith(f'{path} / '): - return index - return 10**9 - - rows.sort(key=get_sort_key) - hydrated_paths = {get_row_path(row) for row in rows} - resolved_inputs = { - path for path in ordered_paths - if path in hydrated_paths - or any(row_path.startswith(f'{path} / ') for row_path in hydrated_paths) - } - resolved_inputs |= set(outline_paths) - missed = len(ordered_paths) - len(resolved_inputs) - if missed > 0: - missing_paths = [path for path in ordered_paths if path not in resolved_inputs] - logger.warning( - f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' - f'missing[:5]={missing_paths[:5]}' - ) - else: - logger.info(f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved') - return rows - - -async def generate_retrieval_asset_url(*, job_id: str, artifact_ref: str) -> str | None: - return get_result_storage().generate_artifact_url(job_id=job_id, artifact_ref=artifact_ref) - - -def is_client_result_artifact_ref(asset_ref: str | None) -> bool: - return get_result_storage().normalize_artifact_ref(asset_ref) is not None diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index f862c23cd..d11e4da7b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -43,6 +43,104 @@ from shared.services.retrieval.llm_adapter import LLMFn +# --------------------------------------------------------------------------- +# Helper: resolve connected asset → owner text chunk section_path +# --------------------------------------------------------------------------- + +def _build_connected_owner_map(text_chunks: list[dict[str, Any]]) -> dict[str, str]: + """Build target_chunk_id → owner text chunk section_path mapping. + + When text chunks reference images/tables via connect_to metadata, + the referenced assets live in Root section. This map lets us attribute + those assets back to the text chunk's section for correct tree placement. + """ + owner_map: dict[str, str] = {} + for chunk in text_chunks: + if (chunk.get('chunk_type') or 'text') != 'text': + continue + section_path = chunk.get('section_path') or '' + if not section_path: + continue + metadata = chunk.get('chunk_metadata') or {} + if not isinstance(metadata, dict): + continue + for conn in metadata.get('connect_to') or []: + if not isinstance(conn, dict): + continue + target_id = str(conn.get('target') or '').strip() + if target_id and target_id not in owner_map: + owner_map[target_id] = section_path + return owner_map + + +async def _resolve_root_asset_owners( + db: AsyncSession, + *, + document_id: str, + job_result_id: str, + chunks: list[dict[str, Any]], +) -> dict[str, str]: + """Resolve owner section_path for Root-stranded image/table chunks. + + When Root is hydrated directly (e.g. via discovery selection), the + batch contains standalone image/table chunks whose section_path is + 'Root'. ``_build_connected_owner_map`` cannot help because the + referencing text chunks live in other sections outside the batch. + + This function queries the *entire document* for text chunks with + connect_to metadata, using the same logic as + ``_build_connected_owner_map``, to resolve the true owner. + + Returns target_chunk_id → owner_section_path for Root assets only. + Returns empty dict when there are no Root assets (zero DB overhead). + """ + from shared.models.database.document import DocumentChunk, DocumentSection + + root_asset_ids = [ + str(c.get('chunk_id') or '') + for c in chunks + if not c.get('owner_section_path') # skip if already resolved by batch-level owner map + and (c.get('section_path') or '') == 'Root' + and (c.get('chunk_type') or '').lower() in ('image', 'table') + and c.get('chunk_id') + ] + if not root_asset_ids: + return {} + + root_asset_set = set(root_asset_ids) + + # Query all text chunks in this document for connect_to metadata + text_stmt = ( + select( + DocumentChunk.chunk_metadata, + DocumentSection.section_path, + ) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.chunk_type == 'text') + ) + result = await db.execute(text_stmt) + + owner_map: dict[str, str] = {} + for metadata, section_path in result.all(): + if not isinstance(metadata, dict) or not section_path: + continue + for conn in metadata.get('connect_to') or []: + if not isinstance(conn, dict): + continue + target_id = str(conn.get('target') or '').strip() + if target_id in root_asset_set and target_id not in owner_map: + owner_map[target_id] = section_path + + if owner_map: + logger.info( + f' _resolve_root_asset_owners: resolved {len(owner_map)}/{len(root_asset_ids)} ' + f'Root assets to their owner sections' + ) + return owner_map + + # --------------------------------------------------------------------------- # Tool: bottom_discovery # --------------------------------------------------------------------------- @@ -281,68 +379,85 @@ async def kg_document_select( User query: {query} -=== Available Tools === +=== Available Actions === -NAVIGATE +NAVIGATE (always performed) 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"}} -When budget is TIGHT, prefer NAVIGATE only when more detail is necessary. -When budget is CRITICAL, choose the narrowest terminal tool if assets directly answer the query. + 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) -> str: - """Parse tool choice from LLM response. Returns one of NAVIGATE/FIND_IMAGES/FIND_TABLES.""" +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() - _VALID_TOOLS = {'NAVIGATE', 'FIND_IMAGES', 'FIND_TABLES'} + _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): - tool = str(data.get('tool', '')).strip().upper() - if tool in _VALID_TOOLS: - return tool + return _extract_from_data(data) except (ValueError, _json.JSONDecodeError): pass - # Accept a JSON object wrapped in markdown, but do not infer a default tool. + # 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): - tool = str(data.get('tool', '')).strip().upper() - if tool in _VALID_TOOLS: - return tool + 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: - return 'FIND_IMAGES' + result.append('FIND_IMAGES') if 'FIND_TABLES' in upper: - return 'FIND_TABLES' - if 'NAVIGATE' in upper: - return 'NAVIGATE' - return '' + result.append('FIND_TABLES') + return result async def tool_select_step( @@ -357,23 +472,21 @@ async def tool_select_step( exclude_paths: set[str] | None = None, revision_hint: str | None = None, budget_snapshot: dict | None = None, -) -> str: - """Route to the appropriate tool for the current scope. - - Returns one of: 'NAVIGATE', 'FIND_IMAGES', 'FIND_TABLES'. +) -> list[str]: + """Route to the appropriate tools for the current scope. - This is the top-level agent decision — it picks WHICH tool - to invoke. Each tool then handles its own internal logic. + 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 'NAVIGATE' directly. + 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 'NAVIGATE' + return [] # Build lightweight tree summary (titles + counts only, no summaries) summary_lines = [] @@ -397,7 +510,7 @@ async def tool_select_step( 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 + return [] # no assets → skip tool selection, navigate only scope_header = ( f'Current scope: "{scope_path}"' if scope_path @@ -419,15 +532,13 @@ async def tool_select_step( ) response = await llm_fn(prompt) - # Parse tool choice - tool = _parse_tool_choice(response) - if not tool: - raise ValueError(f'invalid tool selection response: {response[:200]}') + # Parse tool choices + asset_tools = _parse_tool_choice(response) logger.info( f' tool_select_step scope={scope_path or "root"}: ' - f'tool={tool} images={total_images} tables={total_tables}' + f'tools={asset_tools or ["NAVIGATE"]} images={total_images} tables={total_tables}' ) - return tool + return asset_tools # --------------------------------------------------------------------------- @@ -498,9 +609,7 @@ async def asset_filter_step( section_path_by_id = {section_id: section_path for section_id, section_path in section_rows} - # 3. Resolve media → owner text section via unified helper - from shared.services.retrieval.app_service import _resolve_asset_owners_from_rows - + # 3. Resolve media → owner text section via connect_to tracing text_stmt = ( select( DocumentChunk.section_id, @@ -524,7 +633,7 @@ async def asset_filter_step( } for sid, chunk_type, metadata, scp in text_result.all() ] - owner_by_target_id = _resolve_asset_owners_from_rows(text_row_dicts) + owner_by_target_id = _build_connected_owner_map(text_row_dicts) # Collect connected target IDs for batch-loading connected_target_ids: set[str] = set(owner_by_target_id.keys()) @@ -574,8 +683,7 @@ async def asset_filter_step( seen_ids.add(chunk_id) # Owner resolution: prefer connect_to-based owner - owner_info = owner_by_target_id.get(chunk_id) - owner_section_path = owner_info.get('section_path') if owner_info else None + owner_section_path = owner_by_target_id.get(chunk_id) # Fallback: media's own section_id path, but guard against # Root / top-level aggregation sections @@ -739,7 +847,31 @@ 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, + job_result_id=job_result_id, + chunks=chunks, + ) + if _root_map: + for c in chunks: + if c.get('owner_section_path'): + continue # already resolved by batch-level owner map + 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') @@ -800,7 +932,7 @@ async def discovery_select_step( hint_by_path: dict[str, dict] = {} for h in hints: sp = h.get('section_path', '') - if not sp: + if not sp or sp == 'Root': continue title = sp.rsplit(' / ', 1)[-1] if ' / ' in sp else sp summary = h.get('summary', '') or '' @@ -865,7 +997,31 @@ async def discovery_select_step( exclude_sections=[], ) if connected: + _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 + _disc_job_result_id = next( + (str(c['job_result_id']) for c in chunks if c.get('job_result_id')), + None, + ) + _root_map = await _resolve_root_asset_owners( + db, + document_id=document_id, + job_result_id=_disc_job_result_id, + chunks=chunks, + ) if _disc_job_result_id else {} + if _root_map: + for c in chunks: + if c.get('owner_section_path'): + continue # already resolved by batch-level owner map + 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') diff --git a/packages/shared-python/shared/services/retrieval/agentic/trace.py b/packages/shared-python/shared/services/retrieval/agentic/trace.py index b504971eb..8a58a243c 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/trace.py +++ b/packages/shared-python/shared/services/retrieval/agentic/trace.py @@ -115,6 +115,7 @@ def record_step( 'observation_payload_keys': list(result.payload.keys()) if result.payload else [], 'latency_ms': result.latency_ms, 'error': result.error, + 'tokens_used': result.tokens_used, 'created_at': _now_utc(), }) @@ -128,6 +129,7 @@ def record_budget_stop(self, reason: str) -> None: 'observation_payload_keys': [], 'latency_ms': 0, 'error': None, + 'tokens_used': 0, 'created_at': _now_utc(), }) @@ -135,6 +137,7 @@ async def complete( self, ranked_rows: list[dict[str, Any]], router_used: str, + budget_snapshot: dict[str, Any] | None = None, ) -> None: """Flush all step records and update the run row. Best-effort.""" if not self._created: @@ -155,6 +158,7 @@ async def complete( observation={ 'status': step_data['observation_status'], 'payload_keys': step_data['observation_payload_keys'], + 'tokens_used': step_data.get('tokens_used', 0), }, latency_ms=step_data['latency_ms'], error=step_data.get('error'), @@ -173,6 +177,8 @@ async def complete( 'step_count': len(self._steps), 'final_doc_ids': doc_ids_in_result, } + if budget_snapshot is not None: + provenance['budget_snapshot'] = budget_snapshot 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 e9d5f404c..498bf08b5 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/types.py @@ -10,6 +10,8 @@ from dataclasses import dataclass, field from typing import Any +from shared.services.retrieval.agentic.budget import BudgetLedger + @dataclass class AgentRunConfig: @@ -17,6 +19,11 @@ class AgentRunConfig: max_revisions: int = 2 # max attempt_answer → revision cycles max_nav_depth: int = 3 # max scope_navigate recursion depth latency_budget_ms: int = 12000 + token_budget_total: int = 40000 + planning_ratio: float = 0.5 + bootstrap_budget: int = 2000 + per_doc_min_share: int = 1500 + inventory_aware: bool = True @dataclass @@ -30,6 +37,7 @@ class ToolResult: payload: dict[str, Any] = field(default_factory=dict) latency_ms: int = 0 error: str | None = None + tokens_used: int = 0 @dataclass @@ -99,6 +107,33 @@ def flatten_chunk_rows(self) -> list[dict[str, Any]]: rows.extend(child.flatten_chunk_rows()) return rows + def add_leaf_chunks(self, path: str, chunks: list[dict[str, Any]]) -> None: + """Merge chunks into a leaf path, deduplicating by (chunk_id, path).""" + if not path or not chunks: + return + existing = self.leaf_content.setdefault(path, []) + seen: set[tuple[str, str]] = { + (str(row.get('chunk_id') or ''), path) + for row in existing + if row.get('chunk_id') + } + for chunk in chunks: + chunk_id = str(chunk.get('chunk_id') or '') + key = (chunk_id, path) + if chunk_id and key in seen: + continue + if chunk_id: + seen.add(key) + existing.append(chunk) + + def reparent_leaf_content(self) -> None: + """Move descendant leaf paths into matching child nodes.""" + for child_path, child in list(self.children.items()): + for leaf_path in list(self.leaf_content.keys()): + if leaf_path == child_path or leaf_path.startswith(child_path + ' / '): + child.add_leaf_chunks(leaf_path, self.leaf_content.pop(leaf_path)) + child.reparent_leaf_content() + def collect_referenced_ids(self) -> list[dict[str, str]]: """Extract minimal chunk references from all hydrated leaves. @@ -133,8 +168,7 @@ def merge(self, other: 'DocTreeNode') -> None: if item.get('path', '') not in existing_paths: self.outline_items.append(item) for path, chunks in other.leaf_content.items(): - if path not in self.leaf_content: - self.leaf_content[path] = chunks + self.add_leaf_chunks(path, chunks) for path, child in other.children.items(): if path in self.children: self.children[path].merge(child) @@ -142,6 +176,7 @@ def merge(self, other: 'DocTreeNode') -> None: self.children[path] = child for path, conf in other.confidence.items(): self.confidence[path] = max(self.confidence.get(path, 0), conf) + self.reparent_leaf_content() @dataclass @@ -166,11 +201,16 @@ class AgenticResult: - ``referenced_chunks``: minimal chunk references for hit stats and frontend display (chunk_id, document_id, chunk_type, etc.) - ``router_used``: routing path identifier + - ``budget_snapshot``: final budget ledger state at run completion + - ``stop_reason``: why the run terminated (answer_done / max_revisions / + latency_budget / context_budget / no_llm / etc.) """ evidence_text: str answer_text: str = '' referenced_chunks: list[dict[str, str]] = field(default_factory=list) router_used: str = 'agentic_discovery_only' + budget_snapshot: dict[str, Any] | None = None + stop_reason: str = '' @dataclass @@ -201,6 +241,12 @@ class AgentState: ever_explored_doc_ids: set[str] = field(default_factory=set) seen_section_keys: set[str] = field(default_factory=set) # "{doc_id}::{section_path}" + # Token budget + KG inventory + ledger: BudgetLedger | None = None + kg_total_chunks: int = 0 + kg_total_docs: int = 0 + explored_chunks: int = 0 + @property def elapsed_ms(self) -> int: return int((time.monotonic() - self.start_time) * 1000) diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 86a12c93b..c066a1c16 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -850,9 +850,13 @@ async def _hydrate_paths_to_rows( # ── Chunk modes: load real chunks with optional type filters ───────── if chunk_paths: section_path_filters = [] + # Separate self_only paths (exact match only, no descendant LIKE) + # from regular chunk paths (exact + descendant subtree match) + self_only_paths = {p for p in chunk_paths if mode_by_path.get(p) == 'self_only'} for path in chunk_paths: section_path_filters.append(DocumentSection.section_path == path) - section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + if path not in self_only_paths: + section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) stmt = ( select(Document, DocumentChunk, DocumentSection, JobResult) diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index 8d9a11c22..c6566d5f3 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio +from contextvars import ContextVar from typing import Any, Callable, Coroutine, Union, Sequence, cast from loguru import logger @@ -15,6 +16,11 @@ # LLMFn accepts either a plain string or a list of ChatCompletionMessageParam LLMFnInput = Union[str, Sequence[dict[str, Any]]] LLMFn = Callable[[LLMFnInput], Coroutine[Any, Any, str]] +LLMUsage = dict[str, int] +current_llm_usage: ContextVar[LLMUsage | None] = ContextVar( + 'current_llm_usage', + default=None, +) _RETRIEVAL_LLM_TEMPERATURE = 0.1 _RETRIEVAL_LLM_MAX_TOKENS = 2048 @@ -69,23 +75,56 @@ async def llm_fn(prompt: LLMFnInput) -> str: from shared.utils.OpenAICompatibleClientSync import get_openai_client client = get_openai_client(model=effective_model) - try: - result = await asyncio.to_thread( - client.chat_completion, - cast(Any, prompt), - model=effective_model, - temperature=temperature, - max_tokens=max_tokens, - ) - return result - except Exception as exc: - logger.warning( - "retrieval: agent LLM call failed (degrading gracefully): " - "model={} error_type={} error={}", - effective_model, - type(exc).__name__, - exc, - ) - return '' + current_llm_usage.set(None) + result, usage = await asyncio.to_thread( + client.chat_completion_with_usage, + cast(Any, prompt), + model=effective_model, + temperature=temperature, + max_tokens=max_tokens, + ) + current_llm_usage.set(usage) + return result 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) + current_llm_usage.set(None) + result, usage = await asyncio.to_thread( + client.chat_completion_with_usage, + cast(Any, prompt), + model=effective_model, + temperature=temperature, + max_tokens=max_tokens, + ) + current_llm_usage.set(usage) + return result + + return vlm_fn diff --git a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py b/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py index a1c5cc3e5..66da76a6d 100644 --- a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py +++ b/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py @@ -22,6 +22,7 @@ from shared.utils.security_utils import mask_api_key LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" +LLMUsage = dict[str, int] _client_cache: Dict[tuple, "OpenAICompatibleClientSync"] = {} _client_cache_lock = threading.Lock() @@ -37,6 +38,21 @@ def _should_mock_llm_calls() -> bool: return bool(getattr(settings, "LLM_MOCK_ENABLED", False)) +def _empty_usage() -> LLMUsage: + return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + +def _extract_usage(response: Any) -> LLMUsage: + usage = getattr(response, "usage", None) + if usage is None: + return _empty_usage() + return { + "prompt_tokens": int(getattr(usage, "prompt_tokens", 0) or 0), + "completion_tokens": int(getattr(usage, "completion_tokens", 0) or 0), + "total_tokens": int(getattr(usage, "total_tokens", 0) or 0), + } + + def _summarize_exception_chain(exc: Exception, *, max_depth: int = 4) -> str: parts: list[str] = [] current: BaseException | None = exc @@ -167,7 +183,7 @@ def _make_ali_pool_call( temperature: float, max_tokens: int, api_kwargs: Dict[str, Any], - ) -> str: + ) -> tuple[str, LLMUsage]: """Acquire a token, make the call, and retry inline on 429.""" from shared.utils.ali_quota_manager import get_ali_quota_manager @@ -198,7 +214,7 @@ def _make_ali_pool_call( internal_message="AI returned empty result", provider=self.default_model, ) - return choices[0].message.content or "" + return choices[0].message.content or "", _extract_usage(response) except openai.RateLimitError as exc: retry_after = _parse_retry_after(exc) quota_manager.mark_rate_limited(lease.token_id, retry_after) @@ -239,7 +255,7 @@ def _make_ali_pool_call( # ------------------------------------------------------------------ - def chat_completion( + def chat_completion_with_usage( self, messages: Union[str, List[ChatCompletionMessageParam]], model: Optional[str] = None, @@ -248,7 +264,7 @@ def chat_completion( top_p: Optional[float] = None, timeout: Optional[int] = None, **kwargs, - ) -> str: + ) -> tuple[str, LLMUsage]: all_messages: List[ChatCompletionMessageParam] if isinstance(messages, list): all_messages = messages # type: ignore[assignment] @@ -285,7 +301,7 @@ def chat_completion( return build_mock_chat_completion_response( messages=all_messages, model_name=effective_model, - ) + ), _empty_usage() # Route through Ali token pool when applicable if self._should_use_ali_pool(): @@ -332,7 +348,7 @@ def chat_completion( ) content = choices[0].message.content or "" - return content + return content, _extract_usage(response) except LLMServiceException: raise except Exception as exc: @@ -348,6 +364,27 @@ def chat_completion( original_exception=exc, ) from exc + def chat_completion( + self, + messages: Union[str, List[ChatCompletionMessageParam]], + model: Optional[str] = None, + temperature: float = 0.1, + max_tokens: int = 4096, + top_p: Optional[float] = None, + timeout: Optional[int] = None, + **kwargs, + ) -> str: + content, _usage = self.chat_completion_with_usage( + messages=messages, + model=model, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + timeout=timeout, + **kwargs, + ) + return content + def _parse_retry_after(exc: openai.RateLimitError) -> int: """Extract Retry-After seconds from a RateLimitError, with sane bounds.""" diff --git a/packages/shared-python/shared/utils/token_estimate.py b/packages/shared-python/shared/utils/token_estimate.py new file mode 100644 index 000000000..3b8cb4b99 --- /dev/null +++ b/packages/shared-python/shared/utils/token_estimate.py @@ -0,0 +1,65 @@ +"""Token estimation helpers for retrieval budgeting. + +The estimator intentionally keeps ``tiktoken`` optional. Production +environments that install it get model-aware counts; other environments use a +conservative mixed Chinese/English heuristic with no extra dependency. +""" +from __future__ import annotations + +import re +from functools import lru_cache + + +_CJK_RE = re.compile(r"[\u4e00-\u9fff]") +_ASCII_WORD_RE = re.compile(r"[A-Za-z0-9_]+") + + +@lru_cache(maxsize=32) +def _get_tiktoken_encoding(model_hint: str | None): + try: + import tiktoken # type: ignore[import-not-found] + except Exception: + return None + + try: + if model_hint: + return tiktoken.encoding_for_model(model_hint) + except Exception: + pass + + try: + return tiktoken.get_encoding("cl100k_base") + except Exception: + return None + + +def _heuristic_estimate(text: str) -> int: + if not text: + return 0 + + zh_chars = len(_CJK_RE.findall(text)) + ascii_chars = sum(len(match.group(0)) for match in _ASCII_WORD_RE.finditer(text)) + other_chars = max(len(text) - zh_chars - ascii_chars, 0) + + mixed_estimate = (zh_chars / 1.5) + (ascii_chars / 4.0) + (other_chars / 3.0) + conservative_floor = len(text) / 2.5 + return max(1, int(max(mixed_estimate, conservative_floor))) + + +def estimate_tokens(text: str, model_hint: str | None = None) -> int: + """Estimate input tokens for ``text``. + + ``model_hint`` is advisory. If no compatible tokenizer is available, the + function falls back to a deterministic heuristic. + """ + if not text: + return 0 + + encoding = _get_tiktoken_encoding(model_hint) + if encoding is not None: + try: + return len(encoding.encode(text)) + except Exception: + pass + + return _heuristic_estimate(text) From 0faae5daec745ca267ecd4c44fead70bee6a14fb Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 12 May 2026 22:26:41 +0800 Subject: [PATCH 3/4] chore: re-trigger CI due to github 429 error From 105e397a19c9c2ff6bd5e24bf31352759522d51a Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Tue, 12 May 2026 22:32:14 +0800 Subject: [PATCH 4/4] build: downgrade codeql-action to v3 to bypass github 429 errors --- .github/workflows/codeql.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 366d9fdf9..8de503391 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -28,7 +28,7 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@v4 + uses: github/codeql-action/init@v3 with: languages: python queries: security-extended,security-and-quality @@ -40,7 +40,7 @@ jobs: python-version: "3.11" - name: Autobuild - uses: github/codeql-action/autobuild@v4 + uses: github/codeql-action/autobuild@v3 - name: Perform CodeQL analysis - uses: github/codeql-action/analyze@v4 + uses: github/codeql-action/analyze@v3