From 84d22668133d6a067f83bb52e113e3efb5ed2f43 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 7 May 2026 22:42:31 +0800 Subject: [PATCH 1/5] feat: enhance graph deletion, sync performance, and zip manifest - Update `remove_document_graph` to use 3-way OR deletion (owner, source, target) and enforce user/namespace scopes. - Optimize graph builder's disk sync performance by removing redundant chunk loads. - Extend `ZipResultService` to include a nested title hierarchy in `manifest.json`. - Simplify hierarchical prompt instructions. --- .../tests/contract/test_documents_contract.py | 181 +++++++++++++++++- .../services/connect_builder/graph_builder.py | 146 ++++++++++++++ .../contract/test_parse_task_contract.py | 11 ++ .../shared/services/ai/prompt_service.py | 19 +- .../services/retrieval/graph_service.py | 21 +- .../services/storage/zip_result_service.py | 32 ++++ 6 files changed, 392 insertions(+), 18 deletions(-) diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index d2898f67e..946cb241d 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -105,6 +105,144 @@ async def _fetch_document(document_id: str) -> dict[str, object]: await engine.dispose() +async def _fetch_graph_counts( + *, + document_id: str, + peer_document_id: str | None = None, +) -> dict[str, int]: + engine = await _create_contract_engine() + document_node_id = f"doc:{document_id}" + try: + async with engine.begin() as connection: + node_count = ( + await connection.execute( + text(""" + SELECT COUNT(*) + FROM graph_nodes + WHERE owner_document_id = :document_id + """), + {"document_id": document_id}, + ) + ).scalar_one() + related_edge_count = ( + await connection.execute( + text(""" + SELECT COUNT(*) + FROM graph_edges + WHERE owner_document_id = :document_id + OR source_node_id = :document_node_id + OR target_node_id = :document_node_id + """), + { + "document_id": document_id, + "document_node_id": document_node_id, + }, + ) + ).scalar_one() + peer_node_count = 0 + if peer_document_id is not None: + peer_node_count = ( + await connection.execute( + text(""" + SELECT COUNT(*) + FROM graph_nodes + WHERE owner_document_id = :peer_document_id + """), + {"peer_document_id": peer_document_id}, + ) + ).scalar_one() + return { + "nodes": int(node_count), + "related_edges": int(related_edge_count), + "peer_nodes": int(peer_node_count), + } + finally: + await engine.dispose() + + +async def _insert_document_graph_fixture( + *, + document_id: str, + job_result_id: str, + peer_document_id: str, + peer_job_result_id: str, + user_id: str = "local-dev-user", + namespace: str = "contract-documents", +) -> None: + engine = await _create_contract_engine() + timestamp = datetime.now(timezone.utc).replace(tzinfo=None) + try: + async with engine.begin() as connection: + await connection.execute( + text(""" + INSERT INTO graph_nodes ( + node_id, + user_id, + namespace, + node_kind, + owner_document_id, + job_result_id, + ref_document_id, + ref_section_id, + properties, + created_at, + updated_at + ) VALUES + (:doc_node_id, :user_id, :namespace, 'document', :document_id, :job_result_id, :document_id, NULL, CAST('{}' AS JSON), :created_at, :updated_at), + (:peer_node_id, :user_id, :namespace, 'document', :peer_document_id, :peer_job_result_id, :peer_document_id, NULL, CAST('{}' AS JSON), :created_at, :updated_at) + """), + { + "doc_node_id": f"doc:{document_id}", + "peer_node_id": f"doc:{peer_document_id}", + "user_id": user_id, + "namespace": namespace, + "document_id": document_id, + "peer_document_id": peer_document_id, + "job_result_id": job_result_id, + "peer_job_result_id": peer_job_result_id, + "created_at": timestamp, + "updated_at": timestamp, + }, + ) + await connection.execute( + text(""" + INSERT INTO graph_edges ( + edge_id, + user_id, + namespace, + edge_kind, + source_node_id, + target_node_id, + owner_document_id, + job_result_id, + is_directed, + weight, + properties, + created_at, + updated_at + ) VALUES + (:owned_edge_id, :user_id, :namespace, 'related', :doc_node_id, :peer_node_id, :document_id, :job_result_id, FALSE, 1.0, CAST('{}' AS JSON), :created_at, :updated_at), + (:incoming_edge_id, :user_id, :namespace, 'related', :peer_node_id, :doc_node_id, :peer_document_id, :peer_job_result_id, FALSE, 1.0, CAST('{}' AS JSON), :created_at, :updated_at) + """), + { + "owned_edge_id": f"edge_{uuid4().hex[:12]}", + "incoming_edge_id": f"edge_{uuid4().hex[:12]}", + "user_id": user_id, + "namespace": namespace, + "doc_node_id": f"doc:{document_id}", + "peer_node_id": f"doc:{peer_document_id}", + "document_id": document_id, + "peer_document_id": peer_document_id, + "job_result_id": job_result_id, + "peer_job_result_id": peer_job_result_id, + "created_at": timestamp, + "updated_at": timestamp, + }, + ) + finally: + await engine.dispose() + + async def _insert_document_revision_with_chunks( *, document_id: str, @@ -662,21 +800,62 @@ async def test_should_archive_a_document_via_the_canonical_archive_route( ], ) -> None: document_id = f"doc_{uuid4().hex[:12]}" + peer_document_id = f"doc_{uuid4().hex[:12]}" async with developer_api_client_factory() as api_client: - await _insert_document(document_id=document_id) + document_revision = await _insert_document_revision_with_chunks( + document_id=document_id, + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "archive-chunk-1", + "chunk_type": "text", + "content": "Archived graph chunk", + "source_chunk_path": "Chapter 1/Archive", + "metadata": {"keywords": ["archive"]}, + } + ], + ) + peer_revision = await _insert_document_revision_with_chunks( + document_id=peer_document_id, + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "peer-chunk-1", + "chunk_type": "text", + "content": "Peer graph chunk", + "source_chunk_path": "Chapter 1/Peer", + "metadata": {"keywords": ["peer"]}, + } + ], + ) + await _insert_document_graph_fixture( + document_id=document_id, + job_result_id=document_revision["job_result_id"], + peer_document_id=peer_document_id, + peer_job_result_id=peer_revision["job_result_id"], + ) response = await api_client.post(f"/api/v1/documents/{document_id}/archive") assert response.status_code == 200 response_json = cast(dict[str, object], response.json()) persisted_document = await _fetch_document(document_id) + graph_counts = await _fetch_graph_counts( + document_id=document_id, + peer_document_id=peer_document_id, + ) assert response_json["document_id"] == document_id assert response_json["status"] == "archived" assert response_json["archived_at"] assert persisted_document["status"] == "archived" assert persisted_document["archived_at"] is not None + assert graph_counts == { + "nodes": 0, + "related_edges": 0, + "peer_nodes": 1, + } @pytest.mark.asyncio diff --git a/apps/worker/app/services/connect_builder/graph_builder.py b/apps/worker/app/services/connect_builder/graph_builder.py index d1a51e3ab..26d5b22c6 100644 --- a/apps/worker/app/services/connect_builder/graph_builder.py +++ b/apps/worker/app/services/connect_builder/graph_builder.py @@ -824,6 +824,22 @@ def _get_stats_path(kb_id: str) -> str: return os.path.join(_get_kb_dir(kb_id), "chunk_stats.json") +def _empty_knowledge_graph(kb_id: str) -> Dict[str, Any]: + """Build an empty v2 knowledge graph for a KB with no local chunks.""" + return { + "version": "2.0", + "updated_at": datetime.now(timezone.utc).isoformat(), + "kb_id": kb_id, + "stats": { + "total_files": 0, + "total_chunks": 0, + "total_cross_file_edges": 0, + }, + "files": {}, + "edges": [], + } + + # ─── Chunk Usage Tracking ───────────────────────────────────────────────────── @@ -1040,6 +1056,113 @@ def _load_all_chunks_from_kb(kb_dir: str) -> List[Dict[str, Any]]: return all_chunks +def _source_files_from_chunks(chunks: List[Dict[str, Any]]) -> set[str]: + """Return the source-file set represented by loaded KB chunks.""" + return { + str(chunk.get("_source_file") or "").strip() + for chunk in chunks + if str(chunk.get("_source_file") or "").strip() + } + + +def _prune_chunk_stats(kb_id: str, chunks: List[Dict[str, Any]]) -> None: + """Remove chunk_stats entries whose chunks no longer exist on disk.""" + stats_path = _get_stats_path(kb_id) + if not os.path.exists(stats_path): + return + + stats = load_chunk_stats(kb_id) + live_chunk_ids = { + str(chunk.get("chunk_id") or chunk.get("know_id", "")) + for chunk in chunks + if chunk.get("chunk_id") or chunk.get("know_id") + } + pruned = {cid: data for cid, data in stats.items() if cid in live_chunk_ids} + if len(pruned) == len(stats): + return + + os.makedirs(os.path.dirname(stats_path), exist_ok=True) + with open(stats_path, "w", encoding="utf-8") as f: + json.dump(pruned, f, ensure_ascii=False, indent=2) + logger.info( + f"📊 Chunk stats pruned: {len(stats) - len(pruned)} stale chunks removed" + ) + + +def sync_knowledge_graph_with_local_files( + kb_id: str, + connect_config: Optional[Dict[str, Any]] = None, + summary_use_llm: bool = False, +) -> Dict[str, Any]: + """Synchronize knowledge_graph.json with current ~/.knowhere/{kb_id} files. + + This is intentionally a no-op when graph files match on-disk document + directories. If a user manually deletes a local parsed document directory, + the graph is rebuilt from remaining chunks and stale chunk_stats entries are + removed. + """ + from app.services.connect_builder.builder import build_connections + from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries + + kb_dir = _get_kb_dir(kb_id) + kg_path = _get_kg_path(kb_id) + os.makedirs(kb_dir, exist_ok=True) + + existing_graph = load_knowledge_graph(kg_path) + chunks_on_disk = _load_all_chunks_from_kb(kb_dir) + disk_files = _source_files_from_chunks(chunks_on_disk) + graph_files = set((existing_graph or {}).get("files", {}).keys()) + + if existing_graph is not None and graph_files == disk_files: + _prune_chunk_stats(kb_id, chunks_on_disk) + return existing_graph + + removed_files = sorted(graph_files - disk_files) + added_files = sorted(disk_files - graph_files) + logger.info( + "📊 Syncing Knowledge Graph with local files: " + f"removed={removed_files}, added={added_files}" + ) + + if not chunks_on_disk: + graph = _empty_knowledge_graph(kb_id) + save_knowledge_graph(graph, kg_path) + _prune_chunk_stats(kb_id, []) + return graph + + try: + file_summaries = enrich_doc_nav_summaries( + kb_dir=kb_dir, + source_file=None, + use_llm=summary_use_llm, + ) + except Exception as e: + logger.warning(f"sync enrich_doc_nav_summaries failed: {e}") + file_summaries = {} + + try: + file_nav_sections = _extract_nav_sections_from_kb(kb_dir, source_file=None) + except Exception as e: + logger.warning(f"sync nav_sections extraction failed: {e}") + file_nav_sections = {} + + stats = load_chunk_stats(kb_id) + connections = build_connections(chunks_on_disk, connect_config) + _merge_related_connections_into_chunks(chunks_on_disk, connections) + _save_chunks_by_source_file(kb_dir, chunks_on_disk) + graph = build_knowledge_graph( + all_chunks=chunks_on_disk, + connections=connections, + kb_id=kb_id, + chunk_stats=stats, + file_summaries=file_summaries, + file_nav_sections=file_nav_sections, + ) + save_knowledge_graph(graph, kg_path) + _prune_chunk_stats(kb_id, chunks_on_disk) + return graph + + # ─── MCP Auto-Registration ─────────────────────────────────────────────────── @@ -1264,8 +1387,31 @@ def build_and_deploy( else None ) existing_graph = load_knowledge_graph(kg_path) + + # Load chunks once and reuse — avoids the double-load where + # sync_knowledge_graph_with_local_files internally calls + # _load_all_chunks_from_kb and then we call it again. + all_on_disk: List[Dict[str, Any]] = [] if existing_graph is not None: all_on_disk = _load_all_chunks_from_kb(kb_dir) + disk_files = _source_files_from_chunks(all_on_disk) + graph_files = set(existing_graph.get("files", {}).keys()) + + if graph_files != disk_files: + # Disk state diverged from graph → full sync rebuild. + # sync will reload chunks internally (it backfills connect_to + # metadata), so we reload afterwards to pick up changes. + existing_graph = sync_knowledge_graph_with_local_files( + kb_id=kb_id, + connect_config=connect_config, + summary_use_llm=summary_use_llm, + ) + all_on_disk = _load_all_chunks_from_kb(kb_dir) + else: + # Files match — fast-path: just prune stale chunk_stats. + _prune_chunk_stats(kb_id, all_on_disk) + + if existing_graph is not None: if not all_on_disk: existing_chunks = extract_chunks_from_graph(existing_graph) else: diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 7fa0ed28e..6fdfb583a 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -297,6 +297,9 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: captured_artifacts["zip_chunks"] = json.loads( zip_file.read("chunks.json") )["chunks"] + captured_artifacts["manifest"] = json.loads( + zip_file.read("manifest.json") + ) return SimpleNamespace( zip_key=f"results/{job_id}.zip", @@ -358,6 +361,14 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: assert captured_artifacts["result_dir"].endswith("Default_Root/contract-parse.pdf") assert captured_artifacts["doc_nav"]["file_name"] == source_file_name assert captured_artifacts["doc_nav"]["sections"][0]["title"] == "公司研究" + assert captured_artifacts["manifest"]["HIERARCHY"] == { + "公司研究": { + "自主可控加强，寒武纪或迎来营收快速放量周期": {}, + }, + "相关研报": { + "要点": {}, + }, + } assert "doc_nav.json" in captured_artifacts["raw_entries"] assert "hierarchy.json" not in captured_artifacts["raw_entries"] assert "hierarchy_slim.json" not in captured_artifacts["raw_entries"] diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index e529ef207..9ea759675 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -399,13 +399,13 @@ def build_prompt(task, texts, query, **kwargs): - Use N in ``[N BODY LINES]`` as a "section bulk" signal when applying the rules below (Rule 6 in particular). - ***Process in THREE steps:*** + ***Process in TWO steps:*** **STEP 1 — Global Pattern Scan (CANDIDATES ONLY)** Enumerate every distinct numbering / structural / semantic granularity pattern that appears on candidate rows and signals hierarchy depth, for example: - - Decimal numbering: "1", "1.1", "1.1.1" → depth increases with dot count - - Enumeration styles: "一、" "（一）" "1、" "①" → shallower to deeper + - Decimal numbering: "1.", "1.1", "1.1.1" → depth increases with dot count + - Enumeration styles: "一、" "（一）" "1、" "①" "1 " → shallower to deeper with increasing numbers - Chapter/section keywords: "Chapter X", "Part X", "第X章", "第X节" - Upper case / lower case differences in candidate headings - Clear semantic granularities or groups of themes @@ -415,11 +415,10 @@ def build_prompt(task, texts, query, **kwargs): **STEP 2 — Assign a level to every candidate (rules in priority order)** A candidate whose preliminary ``level`` is "Not Sure" or any positive integer is **always** open to revision. Pure body text has already been - folded into placeholders, but a candidate that slipped through the - pre-filter **can still be** demoted to level = -1. + folded into placeholders, but a candidate **can still be** demoted to level = -1. - Rule 0 — Global consistency (highest priority among all rules): - Candidates sharing the same structural pattern or semantic granularity MUST receive the + Rule 0 — Global consistency: + Candidates sharing the same structural pattern or semantic granularity SHOULD receive the SAME level across the ENTIRE input. (e.g. every "X.Y" pattern shares one level; every "X.Y.Z" shares a different, deeper level.) @@ -449,12 +448,6 @@ def build_prompt(task, texts, query, **kwargs): Rule 4 — Normalise to start at level 1: The shallowest (the most coarse granularity) heading found MUST be assigned level 1. - **STEP 3 — Consistency check (one pass) before writing output** - Re-scan the level assignments you are about to emit: - - All headings sharing the same pattern (structural or semantic granularity) must share the same level. - - No invalid skips (Rule 1). - If any inconsistency is found, normalise to the most representative level for that pattern. - ***Output requirements*** - Output MUST be a [JSON array] only. - Include ONLY candidate rows you judge to be headings (level >= 1). diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph_service.py index 4efe6ddc8..14a7b24dd 100644 --- a/packages/shared-python/shared/services/retrieval/graph_service.py +++ b/packages/shared-python/shared/services/retrieval/graph_service.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from typing import Any, Iterable, Sequence -from sqlalchemy import delete, select +from sqlalchemy import delete, or_, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session @@ -327,11 +327,24 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d ) def remove_document_graph(self, db: Session, *, scope: GraphScope | None, document_id: str) -> None: - edge_delete = delete(GraphEdge).where(GraphEdge.owner_document_id == document_id) + document_node_id = f"doc:{document_id}" + edge_delete = delete(GraphEdge).where( + or_( + GraphEdge.owner_document_id == document_id, + GraphEdge.source_node_id == document_node_id, + GraphEdge.target_node_id == document_node_id, + ) + ) node_delete = delete(GraphNode).where(GraphNode.owner_document_id == document_id) if scope is not None: - edge_delete = edge_delete.where(GraphEdge.user_id == scope.user_id) - node_delete = node_delete.where(GraphNode.user_id == scope.user_id) + edge_delete = edge_delete.where( + GraphEdge.user_id == scope.user_id, + GraphEdge.namespace == scope.namespace, + ) + node_delete = node_delete.where( + GraphNode.user_id == scope.user_id, + GraphNode.namespace == scope.namespace, + ) db.execute(edge_delete) db.execute(node_delete) db.flush() diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index bfd1513f3..68eb1b0fc 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -90,6 +90,9 @@ def generate_zip_package( ) statistics = self._calculate_statistics(formatted_chunks) + doc_nav: Dict[str, Any] = {} + hierarchy: Dict[str, Any] = {} + # Create ZIP package with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as zip_file: # 1. Generate chunks.json (full version) @@ -134,6 +137,7 @@ def generate_zip_package( # 5. Generate doc_nav.json — unified navigation file try: doc_nav = self._build_doc_nav(formatted_chunks, source_file_name) + hierarchy = self._build_hierarchy_dict(doc_nav.get("sections", [])) doc_nav_json = json.dumps(doc_nav, ensure_ascii=False, indent=2) zip_file.writestr("doc_nav.json", doc_nav_json.encode("utf-8")) logger.info("Added doc_nav.json") @@ -147,6 +151,7 @@ def generate_zip_package( source_file_name=source_file_name, statistics=statistics, job_metadata=job_metadata, + hierarchy=hierarchy, ) manifest_json = json.dumps(manifest, ensure_ascii=False, indent=2) zip_file.writestr("manifest.json", manifest_json.encode("utf-8")) @@ -644,6 +649,7 @@ def _generate_manifest( source_file_name: str, statistics: Dict[str, Any], job_metadata: Dict[str, Any], + hierarchy: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Generate manifest.json""" manifest = { @@ -666,6 +672,7 @@ def _generate_manifest( }, }, "statistics": statistics, + "HIERARCHY": hierarchy or {}, } return manifest @@ -678,6 +685,31 @@ def _calculate_zip_checksum(self, zip_file_path: str) -> str: sha256_hash.update(byte_block) return sha256_hash.hexdigest().lower() + def _build_hierarchy_dict( + self, + sections: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """Build a title-only nested hierarchy from doc_nav sections.""" + hierarchy: Dict[str, Any] = {} + title_counts: Dict[str, int] = {} + + for section in sections: + raw_title = str(section.get("title") or "").strip() + if not raw_title: + continue + + title_counts[raw_title] = title_counts.get(raw_title, 0) + 1 + title = ( + raw_title + if title_counts[raw_title] == 1 + else f"{raw_title} ({title_counts[raw_title]})" + ) + hierarchy[title] = self._build_hierarchy_dict( + section.get("children") or [] + ) + + return hierarchy + def _build_doc_nav( self, formatted_chunks: List[Dict[str, Any]], From 77b612200194d6a8446ff6723da5c7aaa7cb4431 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 8 May 2026 20:44:11 +0800 Subject: [PATCH 2/5] feat(agentic): implement LLMPolicy and unify agentic RAG ranking logic --- apps/worker/app/core/tasks/kb_tasks.py | 45 +- .../services/connect_builder/graph_builder.py | 81 +- .../shared/services/retrieval/__init__.py | 2 - .../services/retrieval/agent_navigate.py | 988 ++++-------------- .../retrieval/agentic/orchestrator.py | 133 ++- .../services/retrieval/agentic/policy.py | 265 ++++- .../services/retrieval/agentic/tools.py | 437 +------- .../services/retrieval/agentic/trace.py | 12 +- .../services/retrieval/agentic/types.py | 92 +- .../shared/services/retrieval/app_service.py | 140 +-- .../services/retrieval/graph_service.py | 21 - .../shared/services/retrieval/lexical_text.py | 25 +- 12 files changed, 694 insertions(+), 1547 deletions(-) diff --git a/apps/worker/app/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/kb_tasks.py index edcadc241..b97353b9c 100644 --- a/apps/worker/app/core/tasks/kb_tasks.py +++ b/apps/worker/app/core/tasks/kb_tasks.py @@ -74,39 +74,6 @@ celery_app = get_celery_app() -def _extract_nav_sections_for_publish(add_dir: str) -> list: - """Extract top-level nav sections from doc_nav.json for chunk metadata injection. - - These sections will be read by graph_service.py when creating GraphNode, - enabling hierarchical navigation at query time. - """ - nav_path = os.path.join(add_dir, "doc_nav.json") - if not os.path.exists(nav_path): - return [] - try: - from shared.utils.text_utils import truncate_content_preview - with open(nav_path, "r", encoding="utf-8") as f: - doc_nav = json.load(f) - sections = [] - for section in doc_nav.get("sections", []): - title = section.get("title", "") - if title.lower() in ("root", "__root__"): - continue - sections.append({ - "title": title, - "path": section.get("path", ""), - "summary": truncate_content_preview( - section.get("summary") or "", head=80, tail=0 - ), - "chunk_count": section.get("chunk_count", 0), - "children_count": len(section.get("children", [])), - }) - return sections - except Exception as _e: - logger.warning( - f'doc_nav extraction failed (non-fatal): add_dir={add_dir!r}, error={_e}' - ) - return [] @celery_app.task( bind=True, @@ -626,7 +593,6 @@ def _parse(job_id: str, user_id: str | None): source_file_name = os.path.basename(source_file_name) document_top_summary = "" - document_nav_sections = [] section_summaries: dict[str, str] = {} if add_dir and source_file_name: if add_contents_df is not None and "path" in add_contents_df.columns: @@ -636,8 +602,6 @@ def _parse(job_id: str, user_id: str | None): source_file_name=str(source_file_name), ) # Enrich non-leaf section summaries (bottom-up aggregation) - # Must run before _extract_nav_sections_for_publish so that - # GraphNode.nav_sections gets populated summaries, not empty strings. try: kb_dir_for_enrich = os.path.dirname(str(add_dir)) summary_use_llm = JobMetadataHelper.get_parsing_param( @@ -652,18 +616,13 @@ def _parse(job_id: str, user_id: str | None): except Exception as _e: logger.warning(f"doc_nav enrichment failed (non-fatal): {_e}") document_top_summary = load_nav_top_summary(str(add_dir), str(source_file_name)) - # Extract nav_sections from doc_nav.json for GraphNode persistence - document_nav_sections = _extract_nav_sections_for_publish(str(add_dir)) - if document_top_summary or document_nav_sections: + if document_top_summary: for chunk in chunks: metadata = chunk.get("metadata") if not isinstance(metadata, dict): metadata = {} chunk["metadata"] = metadata - if document_top_summary: - metadata["document_top_summary"] = document_top_summary - if document_nav_sections: - metadata["document_nav_sections"] = document_nav_sections + metadata["document_top_summary"] = document_top_summary data_id = JobMetadataHelper.get_field(job_metadata, "data_id") diff --git a/apps/worker/app/services/connect_builder/graph_builder.py b/apps/worker/app/services/connect_builder/graph_builder.py index 26d5b22c6..4803034eb 100644 --- a/apps/worker/app/services/connect_builder/graph_builder.py +++ b/apps/worker/app/services/connect_builder/graph_builder.py @@ -618,7 +618,6 @@ def build_knowledge_graph( kb_id: str = "", chunk_stats: Optional[Dict[str, Dict[str, Any]]] = None, file_summaries: Optional[Dict[str, str]] = None, - file_nav_sections: Optional[Dict[str, list]] = None, ) -> Dict[str, Any]: """Build a file-level knowledge graph (v2.0).""" if chunk_stats is None: @@ -654,7 +653,6 @@ def build_knowledge_graph( "types": dict(types_count), "top_keywords": file_keywords.get(fk, []), "top_summary": (file_summaries or {}).get(fk, ""), - "nav_sections": (file_nav_sections or {}).get(fk, []), "importance": _compute_file_importance(cids, chunk_stats), "created_at": datetime.now(timezone.utc).isoformat(), } @@ -690,7 +688,6 @@ def update_knowledge_graph( chunk_stats: Optional[Dict[str, Dict[str, Any]]] = None, file_summaries: Optional[Dict[str, str]] = None, new_connections: Optional[Dict[str, List[Dict[str, Any]]]] = None, - file_nav_sections: Optional[Dict[str, list]] = None, ) -> Dict[str, Any]: """Incrementally update a file-level knowledge graph with new chunks.""" if chunk_stats is None: @@ -776,7 +773,6 @@ def update_knowledge_graph( "types": dict(types_count), "top_keywords": file_keywords.get(fk, []), "top_summary": (file_summaries or {}).get(fk, "") or existing_files.get(fk, {}).get("top_summary", ""), - "nav_sections": (file_nav_sections or {}).get(fk, []) or existing_files.get(fk, {}).get("nav_sections", []), "importance": _compute_file_importance(cids, chunk_stats), "created_at": created_at, } @@ -1140,12 +1136,6 @@ def sync_knowledge_graph_with_local_files( logger.warning(f"sync enrich_doc_nav_summaries failed: {e}") file_summaries = {} - try: - file_nav_sections = _extract_nav_sections_from_kb(kb_dir, source_file=None) - except Exception as e: - logger.warning(f"sync nav_sections extraction failed: {e}") - file_nav_sections = {} - stats = load_chunk_stats(kb_id) connections = build_connections(chunks_on_disk, connect_config) _merge_related_connections_into_chunks(chunks_on_disk, connections) @@ -1156,7 +1146,6 @@ def sync_knowledge_graph_with_local_files( kb_id=kb_id, chunk_stats=stats, file_summaries=file_summaries, - file_nav_sections=file_nav_sections, ) save_knowledge_graph(graph, kg_path) _prune_chunk_stats(kb_id, chunks_on_disk) @@ -1261,63 +1250,7 @@ def _auto_register_mcp() -> None: # ─── doc_nav section extraction for GraphNode persistence ──────────────────── -def _extract_nav_sections_from_kb( - kb_dir: str, - source_file: Optional[str] = None, -) -> Dict[str, list]: - """Extract top-level nav sections from doc_nav.json for GraphNode persistence. - - Returns a dict mapping file_key → list of section summaries. - Each section summary has: title, path, summary (truncated), chunk_count, children_count. - - If source_file is given, only processes that file. Otherwise processes all files. - """ - result: Dict[str, list] = {} - - if source_file: - file_dirs = [os.path.join(kb_dir, source_file)] - else: - file_dirs = [ - os.path.join(kb_dir, d) - for d in os.listdir(kb_dir) - if os.path.isdir(os.path.join(kb_dir, d)) - ] - - for file_dir in file_dirs: - nav_path = os.path.join(file_dir, "doc_nav.json") - if not os.path.exists(nav_path): - continue - - file_key = os.path.basename(file_dir) - try: - with open(nav_path, "r", encoding="utf-8") as f: - doc_nav = json.load(f) - - sections = [] - for section in doc_nav.get("sections", []): - title = section.get("title", "") - # Skip utility sections - if title.lower() in ("root", "__root__"): - continue - from shared.utils.text_utils import truncate_content_preview - sections.append({ - "title": title, - "path": section.get("path", ""), - "summary": truncate_content_preview( - section.get("summary") or "", head=80, tail=0 - ), - "chunk_count": section.get("chunk_count", 0), - "children_count": len(section.get("children", [])), - }) - - if sections: - result[file_key] = sections - logger.info(f" nav_sections extracted: {file_key} → {len(sections)} sections") - - except Exception as e: - logger.warning(f" nav_sections extraction failed for {file_key}: {e}") - return result # ─── One-Stop API ───────────────────────────────────────────────────────────── @@ -1466,16 +1399,8 @@ def build_and_deploy( logger.warning(f"doc_nav summary enrichment failed: {e}") file_summaries = {} - # ── Extract nav_sections from doc_nav.json for GraphNode persistence ── - file_nav_sections: Dict[str, list] = {} - try: - file_nav_sections = _extract_nav_sections_from_kb(kb_dir, source_file) - except Exception as e: - logger.warning(f"doc_nav nav_sections extraction failed: {e}") - # Load chunk_stats for importance calculation stats = load_chunk_stats(kb_id) - stats_chunks: List[Dict[str, Any]] = chunks if existing_graph is None: # ── First build: full ── @@ -1530,7 +1455,6 @@ def build_and_deploy( kb_id=kb_id, chunk_stats=stats, file_summaries=file_summaries, - file_nav_sections=file_nav_sections, ) else: # ── Incremental update ── @@ -1544,10 +1468,8 @@ def build_and_deploy( ) stats_chunks = existing_chunks graph = existing_graph - # Still inject nav_sections and summaries even if chunks unchanged + # Still inject summaries even if chunks unchanged for fk, fdata in graph.get("files", {}).items(): - if file_nav_sections and fk in file_nav_sections: - fdata["nav_sections"] = file_nav_sections[fk] if file_summaries and fk in file_summaries and not fdata.get("top_summary"): fdata["top_summary"] = file_summaries[fk] else: @@ -1574,7 +1496,6 @@ def build_and_deploy( chunk_stats=stats, file_summaries=file_summaries, new_connections=related_connections, - file_nav_sections=file_nav_sections, ) # Save graph diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index 13e7c7302..a5cceddca 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -1,4 +1,3 @@ -from .agent_navigate import agent_navigate from .app_service import list_lexical_chunks, merge_channels_rrf, run_retrieval_query from .cache_service import ( bump_retrieval_namespace_cache_version, @@ -12,7 +11,6 @@ from .llm_adapter import create_retrieval_llm_fn __all__ = [ - "agent_navigate", "create_retrieval_llm_fn", "run_retrieval_query", "list_lexical_chunks", diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py index 6558ecaf8..ad47faf61 100644 --- a/packages/shared-python/shared/services/retrieval/agent_navigate.py +++ b/packages/shared-python/shared/services/retrieval/agent_navigate.py @@ -1,24 +1,8 @@ -"""Agent-driven KG navigation for retrieval — aligned with knowhere-kb. - -Two-stage LLM-driven document routing (mirrors unified_retriever.agent_navigate): - 1. LLM reads a knowledge map overview (file-level metadata) and selects relevant files. - 2. For each file, LLM reads compact chunk previews and selects relevant chunk **paths**. - -Additional KB-aligned mechanisms: - - GREP discovery: term-search hits → include parent document_ids in KG scope. - - Edge expansion: selected documents → follow GraphEdge → include neighbor documents. - -Returns chunk paths (section_path / source_chunk_path), NOT hydrated rows. -The caller (app_service) handles path→row hydration and union with discovery results. - -Falls back gracefully when LLM is unavailable or fails. -""" +"""Shared helpers for agentic KG document routing and scope navigation.""" from __future__ import annotations import json -import math import re -import time from typing import Any, Sequence from loguru import logger @@ -26,12 +10,10 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection, GraphNode, GraphEdge -from shared.services.retrieval.llm_adapter import LLMFn +from shared.services.retrieval.lexical_text import normalize_section_path, split_section_path from shared.utils.text_utils import tokenize_for_retrieval -_CONTENT_PREVIEW_LEN = 120 _MAX_OVERVIEW_FILES = 50 -_MAX_CHUNKS_SLIM_PER_DOC = 80 _FILE_SELECT_PROMPT = """\ You are a document routing assistant. @@ -51,23 +33,27 @@ Do not include any explanation. """ -_CHUNK_SELECT_PROMPT = """\ -You are a document chunk routing assistant. +_SCOPE_NAV_PROMPT = """\ +You are a document navigation assistant. + +Document: "{doc_name}" (id: {doc_id}) +Current scope: {scope_label} -Below are candidate chunks from document "{doc_name}" (id: {doc_id}): +Below are candidate section paths at this scope level (up to 2 depth levels). +Indented items are sub-items of the item above. +Select section paths directly. A selected section path represents the chunks +under that section subtree; do not ask to drill deeper. -=== Chunk Candidates === -{chunks_overview} -=== End Candidates === +=== Items === +{items_overview} +=== End Items === User query: {query} -Select the most relevant chunks (at most {max_chunks}). -Return ONLY a JSON array. Prefer objects with path + confidence, e.g.: -[{{"path": "doc_name/Section A/Subsection B", "confidence": 0.92}}, {{"path": "tables/table-1.html", "confidence": 0.75}}] -You may also return a legacy JSON array of path strings if needed: -["doc_name/Section A/Subsection B", "tables/table-1.html"] -Confidence should be between 0 and 1 and reflect how strongly the path matches the user query. +Select the most relevant section paths (at most {max_select}). +Prefer specific sub-items over broad parents when both are listed and the sub-item is sufficient. +Return ONLY a JSON array: +[{{"path": "section/path", "confidence": 0.9}}, ...] Do not include any explanation. """ @@ -142,37 +128,6 @@ def _parse_chunk_path_selections(text: str) -> list[dict[str, Any]]: return selections -def _keywords_need_repair(keywords: list[str] | None) -> bool: - if not isinstance(keywords, list) or not keywords: - return True - bad = sum(1 for kw in keywords if not kw or len(str(kw)) <= 1 - or re.match(r'^\d+[.,%]*$', str(kw))) - return bad >= len(keywords) * 0.5 - - -def _compute_tfidf_keywords(chunk_metadata_list: list[dict[str, Any]], top_k: int = 10) -> list[str]: - df_count: dict[str, int] = {} - tf_count: dict[str, int] = {} - total = len(chunk_metadata_list) or 1 - for meta in chunk_metadata_list: - if not isinstance(meta, dict): - continue - terms = list(meta.get('tokens', [])) + list(meta.get('keywords', [])) - seen: set[str] = set() - for t in terms: - if not t or len(str(t)) <= 1 or re.match(r'^\d+[.,%]*$', str(t)): - continue - lower = str(t).lower() - tf_count[lower] = tf_count.get(lower, 0) + 1 - if lower not in seen: - df_count[lower] = df_count.get(lower, 0) + 1 - seen.add(lower) - scored = [(term, freq * (math.log(total / (df_count.get(term, 1))) + 1)) - for term, freq in tf_count.items()] - scored.sort(key=lambda x: x[1], reverse=True) - return [s[0] for s in scored[:top_k]] - - async def _build_knowledge_map_overview( db: AsyncSession, *, @@ -256,81 +211,54 @@ def _indent_block(text: str, spaces: int) -> str: return '\n'.join(f'{prefix}{line}' for line in str(text or '').splitlines()) -async def _build_chunks_slim( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, -) -> list[dict[str, str]]: - """Build compact chunk descriptors for LLM chunk path selection. +def _format_items_for_llm( + items: list[dict], + max_chars: int = 20000, +) -> tuple[str, bool]: + """Unified formatting with overflow guard for scope navigation. + + Always shows ALL items (L1 + L2). Overflow controls whether + summaries are included — not which levels are shown. + + Normal: path + title + chunk_count + assets_count + summary + Overflow: path + title + chunk_count + assets_count (no summary) - Aligned with KB's do_get_chunks_slim() + _build_slim_chunk(): - - Each entry has 'path' (section_path or source_chunk_path), 'type', 'preview' - - 'preview' is summary-first (from chunk_metadata.summary), fallback to content[:N] - - LLM will select paths, not chunk_ids + Returns (text, overflowed). """ - stmt = ( - select( - DocumentChunk.chunk_id, - DocumentChunk.chunk_type, - DocumentChunk.content, - DocumentChunk.source_chunk_path, - 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) - .order_by(DocumentChunk.sort_order) - .limit(_MAX_CHUNKS_SLIM_PER_DOC) - ) - result = await db.execute(stmt) - chunks: list[dict[str, str]] = [] - for chunk_id, chunk_type, content, source_chunk_path, chunk_metadata, section_path in result.all(): - # Path: prefer section_path, fallback to source_chunk_path - path = section_path or source_chunk_path or '' - if not path: - continue + from shared.utils.text_utils import truncate_content_preview - # Preview: prefer metadata.summary (aligned with KB _build_slim_chunk) - meta = chunk_metadata or {} - summary = re.sub(r'\s+', ' ', str(meta.get('summary') or '')).strip() - raw_content = re.sub(r'\s+', ' ', str(content or '')).strip() - preview = summary or raw_content[:_CONTENT_PREVIEW_LEN] + if not items: + return '(no items available)', False - entry: dict[str, str] = { - 'path': path, - 'type': chunk_type or 'text', - } - if preview: - entry['preview'] = preview[:_CONTENT_PREVIEW_LEN] - chunks.append(entry) - return chunks + SUMMARY_HEAD_TOKENS = 80 + def _render_line(item: dict, include_summary: bool) -> str: + level = item.get('level', 1) + indent = ' ' if level == 2 else '' + line = f'{indent}- path="{item["path"]}" title="{item["title"]}"' + chunk_count = item.get('chunk_count', 0) + if chunk_count > 0: + line += f' chunks={chunk_count}' + assets = item.get('assets_count', 0) + if assets > 0: + line += f' assets={assets}' + if include_summary: + summary = item.get('summary') or item.get('title', '') + if summary: + clipped = truncate_content_preview(summary, head=SUMMARY_HEAD_TOKENS, tail=0) + line += f'\n{indent} summary: {clipped}' + return line -def _format_chunks_for_llm(chunks: list[dict[str, str]], max_chars: int = 4000) -> str: - """Format compact chunk descriptors for LLM prompt. + # Try full render (with summaries) + full_lines = [_render_line(item, include_summary=True) for item in items] + full_text = '\n'.join(full_lines) + if len(full_text) <= max_chars: + return full_text, False - Shows path (not chunk_id), aligned with KB's _format_chunks_slim(). - """ - if not chunks: - return '(no chunks available)' - - def _render(include_preview: bool) -> str: - lines: list[str] = [] - for c in chunks: - line = f'- [{c["type"]}] path="{c["path"]}"' - if include_preview and c.get('preview'): - line += f' | {c["preview"]}' - if len('\n'.join(lines + [line])) > max_chars: - break - lines.append(line) - return '\n'.join(lines) if lines else '(no chunks available)' - - if len(chunks) > 50: - return _render(include_preview=False) - full = _render(include_preview=True) - return full if len(full) <= max_chars else _render(include_preview=False) + # Overflow: render without summaries + slim_lines = [_render_line(item, include_summary=False) for item in items] + slim_text = '\n'.join(slim_lines) + return slim_text[:max_chars], True # ------------------------------------------------------------------ @@ -473,690 +401,180 @@ async def _expand_by_edges( return ordered -# ------------------------------------------------------------------ -# Main entry point -# ------------------------------------------------------------------ - -async def agent_navigate( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - llm_fn: LLMFn, - max_files: int = 3, - max_chunks_per_file: int = 15, - exclude_document_ids: Sequence[str] = (), -) -> list[dict[str, Any]]: - """Agent-driven KG navigation — returns chunk paths with confidence. - - Aligned with KB unified_retriever.agent_navigate(): - Step 1: LLM selects files from knowledge map overview. - GREP: term-search hit chunks → include their parent documents. - Edge: selected documents → follow edges → include neighbors. - Step 2: For each file, LLM selects chunk paths from compact previews. - - Returns: - List of {"path", "confidence"} objects. - Empty list if no documents or LLM fails. - """ - t0 = time.monotonic() - logger.info('\n' + '=' * 70) - logger.info(' 🧭 AGENT NAVIGATE START') - logger.info(f' query="{query}" max_files={max_files} max_chunks/file={max_chunks_per_file}') - logger.info('=' * 70) - - overview_text, doc_id_to_name = await _build_knowledge_map_overview( - db, user_id=user_id, namespace=namespace, - ) - if overview_text == '(empty)': - logger.info(' ⚠️ No active documents in namespace, skipping agent navigate') - return [] - - logger.info(f'\n 📋 STEP 0: Knowledge Map Overview ({len(doc_id_to_name)} files)') - logger.info(f' {"─" * 60}') - for line in overview_text.split('\n'): - logger.info(f' {line}') - logger.info(f' {"─" * 60}') - - # ── Step 1: LLM selects files ── - logger.info('\n 📄 STEP 1: LLM File Selection') - file_prompt = _FILE_SELECT_PROMPT.format( - overview=overview_text, query=query, - ) - t1 = time.monotonic() - try: - file_response = await llm_fn(file_prompt) - selected_ids = _parse_json_array(file_response) - logger.info(f' LLM raw response: {file_response[:300]}') - logger.info(f' Parsed IDs: {selected_ids}') - except Exception as exc: - logger.error(f' ❌ LLM file selection failed: {exc}') - return [] - - elapsed_file = round((time.monotonic() - t1) * 1000) - - exclude_set = set(exclude_document_ids) - valid_ids = [did for did in selected_ids if did in doc_id_to_name and did not in exclude_set] - - if not valid_ids: - logger.warning( - f' ⚠️ LLM returned no valid files (raw={selected_ids}) in {elapsed_file}ms' - ) - return [] - - logger.info(f' ✅ LLM selected {len(valid_ids)} files in {elapsed_file}ms:') - for did in valid_ids: - logger.info(f' → [{did}] {doc_id_to_name.get(did, "?")}') - - # ── GREP discovery: include parent documents of term-hit chunks ── - logger.info('\n 🔎 STEP 1b: GREP Discovery') - try: - grep_doc_ids = await _grep_discover_document_ids( - db, user_id=user_id, namespace=namespace, query=query, - exclude_document_ids=exclude_document_ids, - ) - logger.info(f' GREP hit document_ids: {grep_doc_ids}') - if grep_doc_ids: - pre_count = len(valid_ids) - for did in grep_doc_ids: - if did not in valid_ids and did in doc_id_to_name: - valid_ids.append(did) - added = len(valid_ids) - pre_count - if added > 0: - logger.info(f' ✅ GREP added {added} new documents') - else: - logger.info(f' ℹ️ GREP found {len(grep_doc_ids)} docs but all already selected') - else: - logger.info(' ℹ️ GREP found no matching documents') - except Exception as exc: - logger.warning(f' ⚠️ GREP discovery failed (ignored): {exc}') - - # ── Edge expansion: include neighbor documents ── - logger.info('\n 🔗 STEP 1c: Edge Expansion') - logger.info(f' Input documents: {valid_ids}') - try: - expanded_ids = await _expand_by_edges( - db, document_ids=valid_ids, user_id=user_id, namespace=namespace, - ) - if len(expanded_ids) > len(valid_ids): - new_from_edges = [d for d in expanded_ids if d not in valid_ids] - logger.info(f' ✅ Edge expansion: {len(valid_ids)} → {len(expanded_ids)} documents') - for did in new_from_edges: - logger.info(f' → added neighbor: [{did}] {doc_id_to_name.get(did, "?")}') - valid_ids = expanded_ids - else: - logger.info(' ℹ️ No new neighbors found via edges') - except Exception as exc: - logger.warning(f' ⚠️ Edge expansion failed (ignored): {exc}') - - logger.info(f'\n 📊 STEP 1 SUMMARY: {len(valid_ids)} documents after all expansions:') - for did in valid_ids: - logger.info(f' [{did}] {doc_id_to_name.get(did, "?")}') - - # ── Step 2: For each file, LLM selects chunk paths ── - logger.info('\n 📑 STEP 2: LLM Chunk Path Selection') - doc_job_map: dict[str, str] = {} - doc_stmt = ( - select(Document.document_id, Document.current_job_result_id) - .where(Document.document_id.in_(valid_ids)) - ) - doc_result = await db.execute(doc_stmt) - for did, jrid in doc_result.all(): - if jrid: - doc_job_map[did] = jrid - - all_selected_paths: list[dict[str, Any]] = [] - seen_paths: set[str] = set() - - for doc_id in valid_ids: - job_result_id = doc_job_map.get(doc_id) - if not job_result_id: - logger.warning(f' ⚠️ doc={doc_id} has no job_result_id, skipping') - continue - - doc_name = doc_id_to_name.get(doc_id, doc_id) - logger.info(f'\n {"─" * 50}') - logger.info(f' 📖 Processing: {doc_name} [{doc_id}]') - - chunks_slim = await _build_chunks_slim( - db, document_id=doc_id, job_result_id=job_result_id, - ) - if not chunks_slim: - logger.info(' ⚠️ No chunks found for this document') - continue - - logger.info(f' chunks_slim: {len(chunks_slim)} entries') - for ci, c in enumerate(chunks_slim[:10]): - logger.info(f' [{ci}] [{c.get("type","?")}] path="{c.get("path","")}" preview="{c.get("preview","")[:80]}"') - if len(chunks_slim) > 10: - logger.info(f' ... and {len(chunks_slim) - 10} more') - - chunks_text = _format_chunks_for_llm(chunks_slim) - chunk_prompt = _CHUNK_SELECT_PROMPT.format( - doc_name=doc_name, - doc_id=doc_id, - chunks_overview=chunks_text, - query=query, - max_chunks=max_chunks_per_file, - ) - - valid_paths = {c['path'] for c in chunks_slim if c.get('path')} - - t2 = time.monotonic() - try: - chunk_response = await llm_fn(chunk_prompt) - logger.info(f' LLM raw response: {chunk_response[:300]}') - parsed_selections = _parse_chunk_path_selections(chunk_response) - except Exception as exc: - logger.error(f' ❌ LLM chunk selection failed: {exc}') - continue +# --------------------------------------------------------------------------- +# Unified scope navigation: load child sections (2-level) +# --------------------------------------------------------------------------- - elapsed_chunk = round((time.monotonic() - t2) * 1000) - accepted: list[dict[str, Any]] = [] - rejected: list[str] = [] - for idx, item in enumerate(parsed_selections): - path = str(item.get('path') or '').strip() - if path not in valid_paths: - if path: - rejected.append(path) - continue - confidence = item.get('confidence') - if confidence is None: - confidence = _default_confidence_for_rank(len(accepted)) - accepted.append({ - 'path': path, - 'confidence': confidence, - }) - if len(accepted) >= max_chunks_per_file: - break - - logger.info(f' ✅ Selected {len(accepted)} paths in {elapsed_chunk}ms:') - for item in accepted: - logger.info(f' → {item["path"]} confidence={item["confidence"]:.4f}') - - if rejected: - logger.warning(f' ⚠️ {len(rejected)} paths rejected (not in valid_paths): {rejected[:5]}') - - for item in accepted: - path = item['path'] - if path in seen_paths: - continue - seen_paths.add(path) - all_selected_paths.append(item) - - elapsed_total = round((time.monotonic() - t0) * 1000) - logger.info(f'\n{"=" * 70}') - logger.info(f' 🧭 AGENT NAVIGATE COMPLETE: {len(all_selected_paths)} paths from {len(valid_ids)} files in {elapsed_total}ms') - for i, item in enumerate(all_selected_paths): - logger.info(f' [{i+1}] {item["path"]} confidence={item["confidence"]:.4f}') - logger.info(f'{"=" * 70}') - return all_selected_paths - - -# ------------------------------------------------------------------ -# doc_nav.json hierarchical navigation helpers -# ------------------------------------------------------------------ - -CHUNK_COUNT_THRESHOLD = 30 - -_NAV_SECTION_PROMPT = """\ -You are a document section navigator. - -Document: "{doc_name}" (id: {doc_id}) - -Below are the sections (up to 2 levels). Indented entries are sub-sections of the item above. -Each entry shows its path, title, chunk count, and a content summary. - -=== Sections === -{sections_overview} -=== End Sections === - -User query: {query} - -Select only the sections most relevant to answering the query. -Prefer more specific sub-sections over their parent when both are listed and the sub-section is sufficient. -Return a JSON array with selected paths and confidence scores (0.0-1.0): -[{{"path": "section/path", "confidence": 0.9}}, ...] -Do not include any explanation. -""" - - -async def _load_nav_sections_from_graph( - db: AsyncSession, - document_id: str, -) -> dict | None: - """Load nav_sections from GraphNode.properties for a document. - - Returns dict with 'sections' list and 'total_chunks' count, - or None if not available. - """ - stmt = ( - select(GraphNode.properties) - .where(GraphNode.owner_document_id == document_id) - .where(GraphNode.node_kind == 'document') - ) - result = await db.execute(stmt) - props = result.scalar_one_or_none() - if not props or not isinstance(props, dict): - return None - nav_sections = props.get('nav_sections') - if not nav_sections: - return None - return { - 'sections': nav_sections, - 'total_chunks': props.get('chunks_count', 0), - } - -async def _load_nav_sections_2level( +async def _load_child_sections( db: AsyncSession, document_id: str, job_result_id: str, -) -> dict | None: - """Load L1 sections from GraphNode + their L2 children from DocumentSection. - - Returns dict with 'sections' (flat list, each entry has a 'level' field) - and 'total_chunks', or None if no nav_sections are available. - Each L2 entry is indented under its L1 parent in the returned list order. - """ - from shared.services.retrieval.lexical_text import section_path_from_chunk_path - from sqlalchemy import select, func - from shared.models.database.document import DocumentSection, DocumentChunk - - nav_data = await _load_nav_sections_from_graph(db, document_id) - if not nav_data: - return None - - l1_sections = nav_data.get('sections', []) - if not l1_sections: - return None - - # Extract the kb prefix from L1 chunk_paths so L2 paths can be normalised - # to the same format. L1 path: "kb_id/filename/Section" - # L2 section_path from DB: "Section / SubSection" - # Normalised L2 chunk_path: "kb_id/filename/Section / SubSection" - kb_prefix = '' - sample_chunk_path = l1_sections[0].get('path', '') - sample_section_path = section_path_from_chunk_path(sample_chunk_path) or '' - if sample_section_path and sample_chunk_path.endswith(sample_section_path): - kb_prefix = sample_chunk_path[: -len(sample_section_path)] - - async def _load_l2_by_path(l1_section_path: str) -> list[dict]: - """Path-based fallback: find direct L2 children when intermediate - DocumentSection nodes are absent (i.e. parent_section_id join returns - nothing but the section_path tree still contains descendants). - - Uses 2 batched queries — one for all descendant sections, one for all - descendant chunks — then aggregates in Python. No N+1 queries. - - Matches section_path entries of the form: - '{l1_section_path} / {child_name}' ← direct L2 - '{l1_section_path} / {child_name} / ...' ← deeper (aggregated up) - """ - prefix = l1_section_path + ' / ' - - # ── Query 1: all descendant sections under this L1 ── - desc_rows = (await db.execute( - select( - DocumentSection.section_id, - DocumentSection.section_path, - DocumentSection.summary, - ) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.section_path.like(prefix + '%')) - .order_by(DocumentSection.sort_order) - )).all() - - if not desc_rows: - return [] - - # Extract unique direct-child names and their canonical paths - sep = ' / ' - child_order: list[str] = [] # preserves first-seen order - child_sections: dict[str, dict] = {} # child_path -> aggregated info - all_desc_section_ids: list = [] - - for sec_id, path, summary in desc_rows: - all_desc_section_ids.append(sec_id) - remainder = path[len(prefix):] - child_name = remainder.split(sep)[0] - child_path = prefix + child_name - if child_path not in child_sections: - child_order.append(child_path) - child_sections[child_path] = { - 'path': child_path, - 'title': child_name, - 'chunk_count': 0, - 'children_count': 0, - 'summary': summary or '', - } - # Count grandchildren (paths with one more ' / ' after child_path) - if sep in path[len(child_path):].lstrip(): - child_sections[child_path]['children_count'] += 1 - - # ── Query 2: chunk counts for all descendant section_ids in one shot ── - chunk_rows = (await db.execute( - select( - DocumentChunk.section_id, - func.count(DocumentChunk.id).label('cnt'), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(all_desc_section_ids)) - .group_by(DocumentChunk.section_id) - )).all() - - # Map section_id → section_path for lookup - sec_id_to_path: dict = {sec_id: path for sec_id, path, _ in desc_rows} - for sec_id, cnt in chunk_rows: - path = sec_id_to_path.get(sec_id, '') - if not path.startswith(prefix): - continue - remainder = path[len(prefix):] - child_name = remainder.split(sep)[0] - child_path = prefix + child_name - if child_path in child_sections: - child_sections[child_path]['chunk_count'] += cnt - - return [child_sections[cp] for cp in child_order] - - - all_sections: list[dict] = [] - for s1 in l1_sections: - entry = dict(s1, level=1) - all_sections.append(entry) - # Load L2 children for sections that have sub-sections - if s1.get('children_count', 0) > 0: - db_path = section_path_from_chunk_path(s1['path']) or s1['path'] - # Try parent_section_id join first (fast, works when intermediate nodes exist) - l2 = await _build_sub_sections_from_db( - db, - document_id=document_id, - job_result_id=job_result_id, - parent_path=db_path, - ) - if not l2: - # Fallback: path-pattern LIKE match (handles absent intermediate nodes) - l2 = await _load_l2_by_path(db_path) - for s2 in l2: - s2['level'] = 2 - # Normalise path: prepend kb_prefix so L2 uses chunk_path format - raw_path = s2.get('path', '') - if kb_prefix and not raw_path.startswith(kb_prefix): - s2['path'] = kb_prefix + raw_path - all_sections.append(s2) - - if not all_sections: - return None - - return { - 'sections': all_sections, - 'total_chunks': nav_data.get('total_chunks', 0), - } - - -async def _build_sub_sections_from_db( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - parent_path: str, + scope_path: str | None = None, ) -> list[dict]: - """Build child section view from DocumentSection table. + """Load the next 2 available section depth bands under *scope_path*. - Finds child sections of the given parent_path and computes chunk - counts via DocumentChunk aggregation. Runs at query time in memory - — no files created. - """ - # Find the parent section by path - parent_stmt = ( - select(DocumentSection.section_id) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.section_path == parent_path) - ) - parent_result = await db.execute(parent_stmt) - parent_section_id = parent_result.scalar_one_or_none() - if parent_section_id is None: - return [] + Returns a flat list sorted by sort_order, each item: + {path, title, summary, chunk_count, assets_count, level} - # Get child sections - children_stmt = ( + - level=1: nearest available descendant depth under scope + - level=2: second nearest available descendant depth under scope + - chunk_count: text chunks under this section (excluding image/table) + - assets_count: image + table chunks under this section + """ + # ── Fetch all sections for this document revision ──────────────────── + stmt = ( select( DocumentSection.section_id, DocumentSection.section_title, DocumentSection.section_path, - DocumentSection.section_level, DocumentSection.summary, + DocumentSection.sort_order, ) .where(DocumentSection.document_id == document_id) .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.parent_section_id == parent_section_id) .order_by(DocumentSection.sort_order) ) - children_result = await db.execute(children_stmt) - children = children_result.all() - - if not children: + section_rows = (await db.execute(stmt)).all() + if not section_rows: return [] - # Get chunk counts per section - section_ids = [row[0] for row in children] - chunk_count_stmt = ( - select( - DocumentChunk.section_id, - func.count(DocumentChunk.id).label('count'), - ) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentChunk.section_id.in_(section_ids)) - .group_by(DocumentChunk.section_id) - ) - chunk_count_result = await db.execute(chunk_count_stmt) - chunk_counts = {row[0]: row[1] for row in chunk_count_result.all()} + scope = normalize_section_path(scope_path) if scope_path else '' + scope_parts = split_section_path(scope) + scope_depth = len(scope_parts) - # Count grandchildren for each child - grandchild_stmt = ( - select( - DocumentSection.parent_section_id, - func.count(DocumentSection.section_id).label('count'), - ) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - .where(DocumentSection.parent_section_id.in_(section_ids)) - .group_by(DocumentSection.parent_section_id) - ) - grandchild_result = await db.execute(grandchild_stmt) - grandchild_counts = {row[0]: row[1] for row in grandchild_result.all()} - - result = [] - for section_id, title, path, level, summary in children: - result.append({ - 'title': title or '', - 'path': path or '', + # Build full section metadata index + all_sections: dict[str, dict] = {} # path → {title, summary, sort_order, section_id} + for section_id, title, path, summary, sort_order in section_rows: + if not path: + continue + path = normalize_section_path(path) + parts = split_section_path(path) + all_sections[path] = { + 'title': title or parts[-1] if parts else path, 'summary': summary or '', - 'chunk_count': chunk_counts.get(section_id, 0), - 'children_count': grandchild_counts.get(section_id, 0), - 'level': level or 1, - }) - return result - - -async def _collect_leaf_paths( - db: AsyncSession, - *, - document_id: str, - job_result_id: str, - section_path: str, -) -> list[dict[str, Any]]: - """Collect all chunk paths under a given section path. - - Returns list of {"path": str, "confidence": float}. - Finds chunks whose section_path starts with the given prefix. - """ - stmt = ( - select(DocumentSection.section_path) - .join(DocumentChunk, DocumentChunk.section_id == DocumentSection.section_id) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - .where(DocumentSection.section_path.like(f'{section_path}%')) - .order_by(DocumentChunk.sort_order) - ) - result = await db.execute(stmt) - paths = [] - seen: set[str] = set() - for (path,) in result.all(): - if path and path not in seen: - seen.add(path) - paths.append({ - 'path': path, - 'confidence': 0.8, - }) - return paths - - -def _format_sections_for_llm( - sections: list[dict], - max_chars: int = 20000, -) -> tuple[str, bool]: - """Format section entries for LLM prompt with 2-level visual indentation. - - Returns: - (text, overflowed) - - text: formatted string to put in the LLM prompt - - overflowed: True when the full L1+L2+summary content did NOT fit in - max_chars. In that case the caller should fall back to the - two-step drill-down: first show L1-only, then show L2 for - selected L1s in a second LLM call. - - Rendering strategy: - - Pass 1: Render ALL L1 sections (path + title only). Guarantees LLM sees - the complete document skeleton even if budget is tight. - - Pass 2: Fill remaining budget with summaries + L2 detail (token-aware). - If the entire section list fits → overflowed=False. - If Pass 2 is cut short → overflowed=True (caller triggers drill-down). - """ - from shared.utils.text_utils import truncate_content_preview - - if not sections: - return '(no sections available)', False - - SUMMARY_HEAD_TOKENS = 80 + 'sort_order': int(sort_order or 0), + 'section_id': section_id, + 'parts': parts, + 'depth': len(parts), + } - # ── Pass 1: L1 skeleton (always included) ───────────────────────────── - l1_lines: list[str] = [] - for s in sections: - if s.get('level', 1) != 1: + # ── Identify the next two real depth bands ─────────────────────────── + # + # The stored hierarchy is authoritative. Some ingested documents may only + # expose deeper section rows at root; in that case those rows become this + # round's relative L1/L2 instead of synthesizing missing ancestors. + visible_sections: list[tuple[str, dict, int]] = [] + visible_depths: set[int] = set() + for path, meta in all_sections.items(): + parts = meta['parts'] + if scope_parts and ( + parts[:scope_depth] != scope_parts or len(parts) <= scope_depth + ): continue - line = f'- path="{s["path"]}" title="{s["title"]}" chunks={s.get("chunk_count", 0)}' - if s.get('children_count', 0) > 0: - line += f' sub_sections={s["children_count"]}' - l1_lines.append(line) - - skeleton = '\n'.join(l1_lines) - if len(skeleton) >= max_chars: - return skeleton[:max_chars], True # extremely large doc - - remaining = max_chars - len(skeleton) - - # ── Pass 2: enrich within remaining budget ───────────────────────────── - enriched_lines: list[str] = [] - total_sections = len(sections) - included = 0 - for s in sections: - level = s.get('level', 1) - indent = ' ' if level == 2 else '' - line = f'{indent}- path="{s["path"]}" title="{s["title"]}" chunks={s.get("chunk_count", 0)}' - if s.get('children_count', 0) > 0: - line += f' sub_sections={s["children_count"]}' - summary = s.get('summary', '') - if summary: - clipped = truncate_content_preview(summary, head=SUMMARY_HEAD_TOKENS, tail=0) - line += f'\n{indent} summary: {clipped}' - candidate = '\n'.join(enriched_lines + [line]) - if len(candidate) > remaining: - break - enriched_lines.append(line) - included += 1 - - overflowed = included < total_sections - result_text = '\n'.join(enriched_lines) if enriched_lines else skeleton - return result_text, overflowed - - -def _format_l1_only_for_llm(sections: list[dict]) -> str: - """Format only L1 sections for the first step of the two-step drill-down. - - Used when the full L1+L2 layout overflows the budget: show LLM just the - top-level chapters (with summaries) so it can pick which ones to drill into. - """ - from shared.utils.text_utils import truncate_content_preview - - lines: list[str] = [] - for s in sections: - if s.get('level', 1) != 1: + relative_depth = len(parts) - scope_depth + if relative_depth < 1: continue - line = f'- path="{s["path"]}" title="{s["title"]}" chunks={s.get("chunk_count", 0)}' - if s.get('children_count', 0) > 0: - line += f' sub_sections={s["children_count"]}' - summary = s.get('summary', '') - if summary: - clipped = truncate_content_preview(summary, head=80, tail=0) - line += f'\n summary: {clipped}' - lines.append(line) - return '\n'.join(lines) if lines else '(no sections available)' + visible_sections.append((path, meta, relative_depth)) + visible_depths.add(relative_depth) + selected_depths = sorted(visible_depths)[:2] + depth_to_level = { + depth: idx + 1 + for idx, depth in enumerate(selected_depths) + } -def _format_l2_for_selected_l1s(sections: list[dict], selected_l1_paths: set[str]) -> str: - """Format L2 children of selected L1 paths for the second drill-down LLM call. + items_by_path: dict[str, dict] = {} - Shows each selected L1 as a header, followed by its L2 children (with summaries). - Used after the LLM has already picked which L1 chapters are relevant. - """ - from shared.utils.text_utils import truncate_content_preview + for path, meta, relative_depth in visible_sections: + level = depth_to_level.get(relative_depth) + if level is None: + continue + if path not in items_by_path: + items_by_path[path] = { + 'path': path, + 'title': meta['title'], + 'summary': meta['summary'], + 'level': level, + 'sort_order': meta['sort_order'], + 'chunk_count': 0, + 'assets_count': 0, + 'section_id': meta['section_id'], + } + + if not items_by_path: + return [] - lines: list[str] = [] - for s in sections: - level = s.get('level', 1) - path = s.get('path', '') - if level == 1: - if path not in selected_l1_paths: - continue - # Show L1 as a context header (not selectable, just for readability) - line = f'## {s["title"]} (path="{path}")' - lines.append(line) - elif level == 2: - # Check if parent L1 is in selected set - parent_match = any( - path.startswith(l1 + '/') or path.startswith(l1 + ' / ') - for l1 in selected_l1_paths + # ── Count chunks per section (text vs assets) ─────────────────────── + section_ids = [meta['section_id'] for meta in all_sections.values()] + if section_ids: + from sqlalchemy import case, literal_column + chunk_stmt = ( + select( + DocumentChunk.section_id, + func.count( + case( + (DocumentChunk.chunk_type.notin_(['image', 'table']), literal_column('1')), + ) + ).label('text_count'), + func.count( + case( + (DocumentChunk.chunk_type.in_(['image', 'table']), literal_column('1')), + ) + ).label('asset_count'), ) - if not parent_match: - continue - line = f' - path="{path}" title="{s["title"]}" chunks={s.get("chunk_count", 0)}' - summary = s.get('summary', '') - if summary: - clipped = truncate_content_preview(summary, head=80, tail=0) - line += f'\n summary: {clipped}' - lines.append(line) - - return '\n'.join(lines) if lines else '(no sub-sections found for selected chapters)' + .where(DocumentChunk.document_id == document_id) + .where(DocumentChunk.job_result_id == job_result_id) + .where(DocumentChunk.section_id.in_(section_ids)) + .group_by(DocumentChunk.section_id) + ) + chunk_rows = (await db.execute(chunk_stmt)).all() + section_id_counts: dict[str, tuple[int, int]] = { + sid: (int(tc), int(ac)) for sid, tc, ac in chunk_rows + } + else: + section_id_counts = {} + # Build section_id → path mapping for aggregation + sid_to_path = {meta['section_id']: path for path, meta in all_sections.items()} + # Aggregate chunk counts upward: each item gets counts from itself + descendants + for sid, (text_c, asset_c) in section_id_counts.items(): + chunk_path = sid_to_path.get(sid, '') + if not chunk_path: + continue + # Add to every ancestor item that is in our items_by_path + for item_path, item in items_by_path.items(): + if chunk_path == item_path or chunk_path.startswith(item_path + ' / '): + item['chunk_count'] += text_c + item['assets_count'] += asset_c + + # ── Sort: interleave L2 under their L1 parent ──────────────────────── + # + # Previous sort `(level, sort_order, path)` grouped all L1 first, then + # all L2. The LLM prompt uses indentation to show L2 as sub-items, so + # they should appear directly after their L1 parent for readability. + # + # Sort key: (parent_sort_order, is_child, own_sort_order) + # L1 items: (own_sort_order, 0, 0) → primary position + # L2 items: (parent_sort_order, 1, own) → right after their parent + + def _interleave_key(item: dict) -> tuple: + if item['level'] == 2: + parts = split_section_path(item['path']) + if len(parts) >= 2: + parent_path = ' / '.join(parts[:-1]) + parent_item = items_by_path.get(parent_path) + if parent_item is not None: + return (parent_item['sort_order'], 1, item['sort_order']) + # Orphan L2 (no matching L1 parent in this view): sort by own order + return (item['sort_order'], 1, item['sort_order']) + return (item['sort_order'], 0, 0) + + sorted_items = sorted(items_by_path.values(), key=_interleave_key) + # Clean up internal fields + for item in sorted_items: + item.pop('sort_order', None) + item.pop('section_id', None) + return sorted_items -def _parse_section_selections(text: str) -> list[dict[str, str]]: - """Parse section selections from LLM output. - Expects JSON array of {path, action} objects. - """ - payload = _extract_json_array_payload(text) - selections: list[dict[str, str]] = [] - for item in payload: - if not isinstance(item, dict): - continue - path = str(item.get('path') or '').strip() - action = str(item.get('action') or 'select').strip().lower() - if path: - selections.append({'path': path, 'action': action}) - return selections diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index d93bc89ea..d7c4d5378 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -1,6 +1,6 @@ """Retrieval Agent orchestrator — the core agent loop. -Runs the state/action/observation cycle using the rule-based policy +Runs the state/action/observation cycle using the LLM-driven LLMPolicy and agentic tools. The fixed terminal step (hydrate + rank) always executes, even if all tools fail — in that case it uses whatever discovery rows were collected. @@ -10,13 +10,13 @@ """ from __future__ import annotations +import json import os from typing import Any from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.policy import RuleBasedPolicy from shared.services.retrieval.agentic.trace import TraceRecorder from shared.services.retrieval.agentic.types import ( ActionType, @@ -46,18 +46,18 @@ def _build_config_from_env() -> AgentRunConfig: class RetrievalAgent: - """Agentic retrieval orchestrator. + """Agentic retrieval orchestrator using LLMPolicy. Usage:: agent = RetrievalAgent() ranked_rows, router = await agent.run( - db, user_id=..., namespace=..., query=..., ... + db, user_id=..., namespace=..., query=..., llm_fn=..., ... ) - """ - def __init__(self, *, policy: RuleBasedPolicy | None = None) -> None: - self.policy = policy or RuleBasedPolicy() + The agent requires a valid ``llm_fn`` to run LLMPolicy. If ``llm_fn`` + is None, the run terminates immediately after bottom_discovery only. + """ async def run( self, @@ -82,6 +82,8 @@ async def run( Returns (ranked_rows, router_used). Never raises — errors are captured in trace and the best available result is returned. """ + from shared.services.retrieval.agentic.policy import LLMPolicy + config = config or _build_config_from_env() exclude_document_ids = exclude_document_ids or [] exclude_sections = exclude_sections or [] @@ -107,6 +109,12 @@ async def run( f'budget={config.latency_budget_ms}ms' ) + if llm_fn is None: + logger.warning('agentic: no llm_fn provided — running discovery-only mode') + + # Build LLMPolicy (needs llm_fn; gracefully handles None) + policy = LLMPolicy(llm_fn, query=query) if llm_fn else None + # Shared kwargs for tool calls tool_kwargs: dict[str, Any] = { 'user_id': user_id, @@ -124,6 +132,7 @@ async def run( } # ── Agent loop ── + stop_reason = 'max_steps' while state.step_count < config.max_steps: if state.elapsed_ms >= config.latency_budget_ms: logger.info( @@ -131,18 +140,40 @@ async def run( f'stopping at step {state.step_count}' ) trace.record_budget_stop('latency') + stop_reason = 'latency_budget' break - action_type = self.policy.decide(state, config) - if action_type is None: - logger.info(f' agentic: policy returned None at step {state.step_count}, stopping') + if policy is None: + # No LLM: run discovery once then stop + if not state.discovery_done: + action_type = ActionType.BOTTOM_DISCOVERY + decision_reason = 'no_llm_fn: discovery only' + else: + stop_reason = 'no_llm_fn' + break + else: + action_type, decision_reason = await policy.decide(state, config) + + if action_type is None or action_type == ActionType.DONE: + stop_reason = 'llm_done' if action_type == ActionType.DONE else 'llm_stop' + logger.info( + f' agentic: policy returned {action_type} at step {state.step_count} ' + f'(reason="{decision_reason}"), stopping' + ) + # Record DONE as a trace step + if trace_enabled and action_type == ActionType.DONE: + trace.record_step( + ActionType.DONE, + ToolResult(status='done', payload={'reason': decision_reason}), + decision_reason=decision_reason, + ) break result = await self._execute_tool(db, action_type, state, config, **tool_kwargs) state.apply(action_type, result) if trace_enabled: - trace.record_step(action_type, result) + trace.record_step(action_type, result, decision_reason=decision_reason) state.step_count += 1 @@ -157,17 +188,20 @@ async def run( db, state, user_id=user_id, namespace=namespace, top_k=top_k, ) - router_used = 'agentic' if state.selected_paths else 'agentic_discovery_only' - - if trace_enabled: - await trace.complete(ranked_rows, router_used) + router_used = ( + 'agentic_llm' if state.selected_paths + else 'agentic_discovery_only' + ) - total_ms = state.elapsed_ms logger.info( f'agentic retrieval DONE: {len(ranked_rows)} results, ' - f'router={router_used}, steps={state.step_count}, {total_ms}ms' + f'router={router_used}, steps={state.step_count}, ' + f'stop_reason={stop_reason}, {state.elapsed_ms}ms' ) + if trace_enabled: + await trace.complete(ranked_rows, router_used) + return ranked_rows, router_used async def _execute_tool( @@ -240,54 +274,6 @@ async def _execute_tool( document_ids=doc_ids, ) - elif action_type == ActionType.NAV_SECTION_SELECT: - if not state.nav_drill_stack: - return ToolResult(status='error', error='nav_drill_stack is empty') - - import asyncio as _asyncio - - # Concurrently navigate ALL pending documents in nav_drill_stack - drill_entries = list(state.nav_drill_stack) # snapshot - nav_tasks = [ - tools.nav_section_select( - db, - user_id=kwargs['user_id'], - namespace=kwargs['namespace'], - query=kwargs['query'], - llm_fn=kwargs.get('llm_fn'), - document_id=entry['document_id'], - job_result_id=state.doc_job_map.get(entry['document_id'], ''), - doc_name=state.doc_id_to_name.get(entry['document_id'], ''), - section_path=entry.get('section_path'), - ) - for entry in drill_entries - ] - nav_results: list[ToolResult] = await _asyncio.gather(*nav_tasks) - - # Merge all results into a single batch ToolResult consumed by state.apply - all_paths: list[dict] = [] - any_selected = False - total_latency = 0 - for res in nav_results: - total_latency = max(total_latency, res.latency_ms) - if res.status == 'selected_paths': - any_selected = True - all_paths.extend(res.payload.get('selected_paths', [])) - - logger.info( - f' agentic.nav_section_select (concurrent): ' - f'{len(drill_entries)} docs → {len(all_paths)} total paths, ' - f'{total_latency}ms' - ) - return ToolResult( - status='selected_paths' if any_selected else 'no_confident_match', - payload={ - 'selected_paths': all_paths, - '_consumed_stack': [e['document_id'] for e in drill_entries], - }, - latency_ms=total_latency, - ) - else: return ToolResult(status='error', error=f'unknown action: {action_type}') @@ -307,12 +293,17 @@ async def _hydrate_and_rank( """Fixed terminal step: hydrate selected paths + rank against discovery. Reuses _hydrate_paths_to_rows and _rank_candidates_by_path unchanged. + + TODO (Token Budget): Replace `ranked_rows[:top_k]` in + `_rank_candidates_by_path` with token-accumulation truncation + (tiktoken or character estimate) so the final result set respects + a configurable LLM context window budget rather than a fixed count. """ try: # Hydrate agent-selected paths - agent_rows: list[dict[str, Any]] = [] + navigated_paths: list[dict[str, Any]] = [] if state.selected_paths: - agent_rows = await _hydrate_paths_to_rows( + navigated_paths = await _hydrate_paths_to_rows( db, path_selections=state.selected_paths, user_id=user_id, @@ -320,7 +311,7 @@ async def _hydrate_and_rank( ) # Load importance scores for all candidate rows - all_candidates = state.discovery_rows + agent_rows + all_candidates = state.discovery_paths + navigated_paths if all_candidates: importance_map = await _load_chunk_importance_scores( db, user_id=user_id, namespace=namespace, @@ -339,8 +330,8 @@ async def _hydrate_and_rank( # Rank: merge discovery + agent rows ranked = _rank_candidates_by_path( - discovery_rows=state.discovery_rows, - routed_rows=agent_rows, + discovery_rows=state.discovery_paths, + routed_rows=navigated_paths, top_k=top_k, ) @@ -349,4 +340,4 @@ async def _hydrate_and_rank( except Exception as e: logger.error(f'agentic hydrate_and_rank failed: {e}') # Last resort: return raw discovery rows - return state.discovery_rows[:top_k] + return state.discovery_paths[:top_k] diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 29a06189b..926ca327d 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -1,50 +1,235 @@ -"""Rule-based policy for agentic retrieval. +"""LLM-driven policy for agentic retrieval. -Determines the next action based on current state and budget constraints. -v1 is purely rule-based — no LLM planning, no learned policy. +Replaces the former RuleBasedPolicy. All control decisions are made by a +small LLM call so the agent can adapt to query complexity and KB topology +rather than following a fixed rule tree. + +Design: + - ``LLMPolicy.decide()`` is **async** — it calls the LLM to pick the + next action. + - On parse failure the agent terminates immediately (no hidden fallback). + - The prompt is compact: no raw chunk content, only state metadata. """ from __future__ import annotations +import json +import re +from typing import Any + +from loguru import logger + from shared.services.retrieval.agentic.types import ActionType, AgentRunConfig, AgentState +from shared.services.retrieval.llm_adapter import LLMFn +# ── Available actions presented to the LLM ─────────────────────────────────── -class RuleBasedPolicy: - """Stateless decision function: (state, config) → next action or None. +_AVAILABLE_ACTIONS: list[dict[str, Any]] = [ + { + 'action': ActionType.BOTTOM_DISCOVERY.value, + 'description': ( + 'Run BM25 3-channel bottom-layer discovery (path / content / term). ' + 'Always call this first to get candidate chunks and top document hints.' + ), + 'when': 'discovery_done is false', + }, + { + 'action': ActionType.KG_DOCUMENT_SELECT.value, + 'description': ( + 'Ask the LLM to select the most relevant documents from the Knowledge Graph ' + 'overview (doc summaries). Use after discovery.' + ), + 'when': 'discovery_done is true, kg_done is false', + }, + { + 'action': ActionType.DOCUMENT_PATH_SELECT.value, + 'description': ( + 'Drill into the next pending document\'s section tree and pick relevant ' + 'paths. The document is chosen automatically from selected_docs[pending_doc_index].' + ), + 'when': 'kg_done is true, there are pending documents to process', + }, + { + 'action': ActionType.GREP_DOCUMENT_DISCOVER.value, + 'description': ( + 'Run term/grep search across all documents to discover relevant doc IDs ' + 'when KG selection found nothing.' + ), + 'when': 'kg_done is true but selected_docs is empty', + }, + { + 'action': ActionType.GRAPH_EXPAND_DOCS.value, + 'description': ( + 'Expand via Knowledge Graph edge relationships to find related documents ' + 'not yet selected.' + ), + 'when': 'Optional deepening — use when current results seem insufficient', + }, + { + 'action': ActionType.DONE.value, + 'description': ( + 'Stop the agent loop and proceed to final hydration and ranking. ' + 'Use when you have sufficient evidence paths or the budget is nearly exhausted.' + ), + 'when': 'Sufficient paths collected, or further actions would not improve results', + }, +] - Returning None means the agent loop should stop and proceed to - the fixed terminal step (hydrate_and_rank). - """ +_ACTIONS_BLOCK = '\n'.join( + f" {i+1}. \"{a['action']}\": {a['description']} [{a['when']}]" + for i, a in enumerate(_AVAILABLE_ACTIONS) +) + +_POLICY_PROMPT_TEMPLATE = """\ +You are a retrieval agent orchestrating document search for a RAG system. +Your job: choose the SINGLE best next action given the current state. + +QUERY: "{query}" + +CURRENT STATE: +{state_json} + +BUDGET: {elapsed_ms}ms elapsed of {budget_ms}ms max | Step {step} of {max_steps} + +AVAILABLE ACTIONS: +{actions_block} + +RULES: +1. Always run bottom_discovery first (if discovery_done is false). +2. After discovery, run kg_document_select (if kg_done is false). +3. After kg select, run document_path_select for each pending document. +4. Call done when all pending documents are processed OR you have >= {min_evidence} evidence paths. +5. Only use grep_document_discover if kg_document_select found 0 documents. +6. Only use graph_expand_docs if you need more related docs after reviewing results. - def decide(self, state: AgentState, config: AgentRunConfig) -> ActionType | None: - # ── Step 1: Always start with bottom_discovery ── - if not state.discovery_done: - return ActionType.BOTTOM_DISCOVERY - - # ── Step 2: KG document selection ── - if not state.kg_done and state.doc_retry_count < config.max_doc_retries: - return ActionType.KG_DOCUMENT_SELECT - - # ── Step 3: Pending doc_nav section navigation (one call per document) ── - if state.nav_drill_stack: - return ActionType.NAV_SECTION_SELECT - - # ── Step 4: Process selected documents one by one ── - if state.selected_docs and state.pending_doc_index < len(state.selected_docs): - # max_docs=0 means no limit — LLM decides autonomously - if config.max_docs == 0 or state.pending_doc_index < config.max_docs: - return ActionType.DOCUMENT_PATH_SELECT - - # ── Step 5: Handle need_more_docs (go back to KG select) ── - if ( - state.last_observation - and state.last_observation.status == 'need_more_docs' - and state.doc_retry_count < config.max_doc_retries - ): - return ActionType.KG_DOCUMENT_SELECT - - # ── Step 6: If no docs selected at all, try GREP as fallback ── - if not state.selected_docs and state.kg_done and state.doc_retry_count == 0: - return ActionType.GREP_DOCUMENT_DISCOVER - - # ── Done: proceed to hydrate_and_rank ── +Return ONLY a JSON object, no markdown, no explanation: +{{"action": "", "reason": ""}} +""" + + +def _parse_action_from_response(response: str) -> dict[str, Any] | None: + """Extract the JSON decision object from the LLM response.""" + # Strip markdown code fences if present + text = response.strip() + text = re.sub(r'^```(?:json)?\s*', '', text, flags=re.MULTILINE) + text = re.sub(r'\s*```$', '', text, flags=re.MULTILINE) + text = text.strip() + + # Try to extract {...} block + match = re.search(r'\{[^{}]*\}', text, re.DOTALL) + if match: + try: + return json.loads(match.group(0)) + except json.JSONDecodeError: + pass + + # Try the whole text as JSON + try: + return json.loads(text) + except json.JSONDecodeError: return None + + +class LLMPolicy: + """Async LLM-driven policy. One LLM call per agent step. + + Usage:: + + policy = LLMPolicy(llm_fn) + action_type, reason = await policy.decide(state, config, query=query) + """ + + def __init__(self, llm_fn: LLMFn, *, query: str = '') -> None: + self._llm_fn = llm_fn + self._query = query + + def build_prompt(self, state: AgentState, config: AgentRunConfig) -> str: + """Build the decision prompt. Public for test inspection.""" + state_data = state.state_summary() + state_json = json.dumps(state_data, ensure_ascii=False, indent=2) + + # Count pending docs + pending_count = len(state.selected_docs) - state.pending_doc_index + + return _POLICY_PROMPT_TEMPLATE.format( + query=self._query, + state_json=state_json, + elapsed_ms=state.elapsed_ms, + budget_ms=config.latency_budget_ms, + step=state.step_count, + max_steps=config.max_steps, + actions_block=_ACTIONS_BLOCK, + min_evidence=config.min_evidence_paths, + ) + + async def decide( + self, + state: AgentState, + config: AgentRunConfig, + ) -> tuple[ActionType | None, str]: + """Ask the LLM which action to take next. + + Returns ``(action_type, reason)`` or ``(None, reason)`` when the + LLM says ``done`` or when parsing fails (hard stop). + """ + prompt = self.build_prompt(state, config) + + logger.info( + f' [LLMPolicy] step={state.step_count} calling LLM ' + f'(state: discovery={state.discovery_done}, kg={state.kg_done}, ' + f'docs={len(state.selected_docs)}, pending={state.pending_doc_index}, ' + f'paths={len(state.selected_paths)})' + ) + + # ── Verbose prompt logging (controlled by env) ── + import os + if os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', 'false') == 'true': + logger.info( + f'\n{"="*60}\n' + f'[LLMPolicy PROMPT step={state.step_count}]\n' + f'{prompt}\n' + f'{"="*60}' + ) + + raw_response = await self._llm_fn(prompt) + + logger.info( + f' [LLMPolicy] raw_response={repr(raw_response[:200])}' + ) + + if os.environ.get('RETRIEVAL_AGENTIC_VERBOSE', 'false') == 'true': + logger.info( + f'\n{"="*60}\n' + f'[LLMPolicy RESPONSE step={state.step_count}]\n' + f'{raw_response}\n' + f'{"="*60}' + ) + + if not raw_response.strip(): + logger.warning(' [LLMPolicy] empty response → stopping agent') + return None, 'empty LLM response' + + parsed = _parse_action_from_response(raw_response) + if not parsed: + logger.warning( + f' [LLMPolicy] could not parse JSON from response: {repr(raw_response[:300])} → stopping agent' + ) + return None, f'parse_error: {raw_response[:100]}' + + action_str = str(parsed.get('action', '')).strip() + reason = str(parsed.get('reason', '')).strip() + + logger.info(f' [LLMPolicy] decided action="{action_str}" reason="{reason}"') + + if action_str == ActionType.DONE.value: + return ActionType.DONE, reason + + # Validate against whitelist + try: + action_type = ActionType(action_str) + except ValueError: + logger.warning( + f' [LLMPolicy] unknown action "{action_str}" → stopping agent' + ) + return None, f'unknown_action: {action_str}' + + return action_type, reason diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 4b3cb6676..a66fa0c5b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -18,24 +18,19 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document +from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.services.retrieval.agentic.types import ToolResult from shared.services.retrieval.agent_navigate import ( - _build_chunks_slim, _build_knowledge_map_overview, - _collect_leaf_paths, _expand_by_edges, - _format_chunks_for_llm, + _format_items_for_llm, _grep_discover_document_ids, - _load_nav_sections_from_graph, - _load_nav_sections_2level, + _load_child_sections, _parse_chunk_path_selections, _parse_json_array, - _CHUNK_SELECT_PROMPT, + _SCOPE_NAV_PROMPT, _FILE_SELECT_PROMPT, - _NAV_SECTION_PROMPT, _default_confidence_for_rank, - CHUNK_COUNT_THRESHOLD, ) from shared.services.retrieval.app_service import ( _CHANNEL_WEIGHT_CONTENT, @@ -283,110 +278,18 @@ async def document_path_select( max_chunks_per_file: int = 15, **_kwargs: Any, ) -> ToolResult: - """Select chunk paths within a single document. - - Reuses: agent_navigate._build_chunks_slim, _format_chunks_for_llm, - _CHUNK_SELECT_PROMPT, _parse_chunk_path_selections. - - Returns one of: - - selected_paths: found relevant paths - - need_nav_drill: large document with doc_nav available, switch to hierarchical mode - - need_more_docs: document not relevant, suggest going back to doc select - - no_confident_match: no paths match above threshold - """ - t0 = time.monotonic() - try: - if llm_fn is None: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': 'LLM not available'}, - latency_ms=latency, - ) - - # Check if doc_nav sections are available for large documents - nav_data = await _load_nav_sections_from_graph(db, document_id) - if nav_data: - total_chunks = nav_data.get('total_chunks', 0) - if total_chunks > CHUNK_COUNT_THRESHOLD: - latency = int((time.monotonic() - t0) * 1000) - logger.info( - f' agentic.document_path_select: doc={document_id} ' - f'total_chunks={total_chunks} > {CHUNK_COUNT_THRESHOLD}, ' - f'switching to nav_section_select' - ) - return ToolResult( - status='need_nav_drill', - payload={ - 'document_id': document_id, - 'total_chunks': total_chunks, - }, - latency_ms=latency, - ) - - chunks_slim = await _build_chunks_slim( - db, document_id=document_id, job_result_id=job_result_id, - ) - if not chunks_slim: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': 'no chunks found'}, - latency_ms=latency, - ) - - chunks_text = _format_chunks_for_llm(chunks_slim) - chunk_prompt = _CHUNK_SELECT_PROMPT.format( - doc_name=doc_name or document_id, - doc_id=document_id, - chunks_overview=chunks_text, - query=query, - max_chunks=max_chunks_per_file, - ) - - valid_paths = {c['path'] for c in chunks_slim if c.get('path')} - chunk_response = await llm_fn(chunk_prompt) - parsed_selections = _parse_chunk_path_selections(chunk_response) - - accepted: list[dict[str, Any]] = [] - for item in parsed_selections: - path = str(item.get('path') or '').strip() - if path not in valid_paths: - continue - confidence = item.get('confidence') - if confidence is None: - confidence = _default_confidence_for_rank(len(accepted)) - accepted.append({'path': path, 'confidence': confidence}) - if len(accepted) >= max_chunks_per_file: - break - - latency = int((time.monotonic() - t0) * 1000) - - if accepted: - logger.info( - f' agentic.document_path_select: {len(accepted)} paths from doc={document_id}, {latency}ms' - ) - return ToolResult( - status='selected_paths', - payload={'document_id': document_id, 'selected_paths': accepted}, - latency_ms=latency, - ) - - # No paths accepted — not relevant + """Document entry point for agentic scope navigation.""" + if llm_fn is None: return ToolResult( status='no_confident_match', - payload={'document_id': document_id, 'reason': 'no path matches query intent'}, - latency_ms=latency, - ) - except Exception as e: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.document_path_select failed for doc={document_id}: {e}') - return ToolResult( - status='error', - payload={'document_id': document_id}, - error=str(e), - latency_ms=latency, + payload={'document_id': document_id, 'reason': 'LLM not available'}, + latency_ms=0, ) + return await scope_navigate( + db, document_id=document_id, job_result_id=job_result_id, + query=query, llm_fn=llm_fn, doc_name=doc_name, + scope_path=None, max_select=max_chunks_per_file, + ) # --------------------------------------------------------------------------- @@ -513,294 +416,55 @@ async def graph_expand_docs( # --------------------------------------------------------------------------- -# Tool: nav_section_select +# Tool: scope_navigate (Unified recursive navigation) # --------------------------------------------------------------------------- -async def nav_section_select( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - llm_fn: LLMFn | None, - document_id: str, - job_result_id: str, - doc_name: str = '', - section_path: str | None = None, # kept for API compat, ignored in 2-level design - **_kwargs: Any, -) -> ToolResult: - """Navigate doc_nav sections with a single 2-level LLM call. - - Shows L1 + L2 sections in one flat overview. LLM picks relevant - sections at either level (no drill/select decision). Each selected - section is immediately expanded to its leaf chunk paths via - _collect_leaf_paths — finest-granularity, deduplicated. - - Falls back to flat chunks_slim selection when no nav sections exist. - """ - t0 = time.monotonic() - try: - if llm_fn is None: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': 'LLM not available'}, - latency_ms=latency, - ) - - # Load L1 + L2 sections in a single pass - nav_data = await _load_nav_sections_2level(db, document_id, job_result_id) - if not nav_data or not nav_data.get('sections'): - logger.info( - f' agentic.nav_section_select: no 2-level sections for ' - f'doc={document_id}, falling back to chunks_slim' - ) - return await _nav_fallback_chunks_slim( - db, - document_id=document_id, - job_result_id=job_result_id, - query=query, - llm_fn=llm_fn, - doc_name=doc_name, - path_scope=None, - ) - - sections = nav_data['sections'] - - # ── Format sections for LLM; detect overflow ────────────────────────── - from shared.services.retrieval.agent_navigate import ( - _format_sections_for_llm, - _format_l1_only_for_llm, - ) - - sections_text, overflowed = _format_sections_for_llm(sections) - valid_paths = {s['path'] for s in sections if s.get('path')} - - def _norm(p: str) -> str: - return p.replace(' / ', '/').replace(' /', '/').replace('/ ', '/') - - norm_to_canonical: dict[str, str] = {_norm(vp): vp for vp in valid_paths} - - def _validate_selections(parsed: list[dict]) -> list[dict]: - """Map LLM path strings to canonical DB paths, drop invalid ones.""" - accepted: list[dict] = [] - seen: set[str] = set() - for item in parsed: - path = str(item.get('path') or '').strip() - canonical = path if path in valid_paths else norm_to_canonical.get(_norm(path)) - if canonical is None: - logger.warning(f' nav_section_select: rejected invalid path: {path}') - continue - if canonical in seen: - continue - seen.add(canonical) - accepted.append({'path': canonical, 'confidence': item.get('confidence') or 0.8}) - return accepted - - accepted: list[dict[str, Any]] = [] - - if not overflowed: - # ── Fast path: everything fits → single LLM call ────────────────── - logger.info(' agentic.nav_section_select: full view (no overflow), single LLM call') - prompt = _NAV_SECTION_PROMPT.format( - doc_name=doc_name or document_id, - doc_id=document_id, - sections_overview=sections_text, - query=query, - ) - response = await llm_fn(prompt) - logger.info(f' agentic.nav_section_select LLM response: {response[:300]}') - accepted = _validate_selections(_parse_chunk_path_selections(response)) - - else: - # ── Overflow path: L1-only LLM selection → mechanical leaf expansion ── - # The document has too many sections to show in one prompt. - # LLM picks the relevant top-level chapters; we then expand every - # selected chapter to ALL its leaf chunk paths without a second LLM - # call — same as the non-overflow path's _collect_leaf_paths behavior. - logger.info( - f' agentic.nav_section_select: overflow — ' - f'L1-only LLM selection for doc={document_id}' - ) - l1_text = _format_l1_only_for_llm(sections) - prompt_l1 = _NAV_SECTION_PROMPT.format( - doc_name=doc_name or document_id, - doc_id=document_id, - sections_overview=( - '[Only top-level chapters are shown below — sub-sections are not listed.]\n' - '[Select from the chapter paths in this list ONLY.]\n' - '[Do NOT return sub-section paths; select at the chapter level.]\n\n' - + l1_text - ), - query=query, - ) - response_l1 = await llm_fn(prompt_l1) - logger.info(f' nav_section_select (overflow) LLM response: {response_l1[:300]}') - - # Validate against L1 paths only. If LLM returns a deeper path - # (e.g. "…/Information extraction/Entity extraction"), truncate it - # to its L1 ancestor ("…/Information extraction") rather than dropping. - l1_valid_paths = {s['path'] for s in sections if s.get('level', 1) == 1} - l1_norm_to_canon = {_norm(p): p for p in l1_valid_paths} - parsed = _parse_chunk_path_selections(response_l1) - - seen_l1: set[str] = set() - for item in parsed: - path = str(item.get('path') or '').strip() - # Try exact L1 match first (with norm fallback) - canonical = ( - path if path in l1_valid_paths - else l1_norm_to_canon.get(_norm(path)) - ) - # If not an exact L1, try truncating to find an L1 ancestor - if canonical is None: - candidate = path - for sep in (' / ', '/'): - while sep in candidate: - candidate = candidate[:candidate.rfind(sep)] - hit = ( - candidate if candidate in l1_valid_paths - else l1_norm_to_canon.get(_norm(candidate)) - ) - if hit: - canonical = hit - logger.info( - f' nav_section_select (overflow): ' - f'truncated "{path}" → L1 "{hit}"' - ) - break - if canonical: - break - if canonical is None: - logger.warning( - f' nav_section_select (overflow): ' - f'no L1 match for "{path}", skipped' - ) - continue - if canonical not in seen_l1: - seen_l1.add(canonical) - accepted.append({ - 'path': canonical, - 'confidence': item.get('confidence') or 0.8, - }) - - # ── Collect latency ────────────────────────────────────────────────── - latency = int((time.monotonic() - t0) * 1000) - - if not accepted: - logger.info( - f' agentic.nav_section_select: no match for doc={document_id}, {latency}ms' - ) - return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': 'no section matches query'}, - latency_ms=latency, - ) - - # ── Expand each selected section to finest-granularity leaf paths ──── - # _collect_leaf_paths uses LIKE-prefix matching, so selecting an L1 - # path like "kb/doc.docx/Information extraction" returns ALL chunks - # under that chapter — L2, L3, leaf chunks — at any depth. - from shared.services.retrieval.lexical_text import section_path_from_chunk_path - - all_leaf_paths: list[dict[str, Any]] = [] - seen_leaf: set[str] = set() - for item in accepted: - db_section_path = section_path_from_chunk_path(item['path']) or item['path'] - leaf_paths = await _collect_leaf_paths( - db, - document_id=document_id, - job_result_id=job_result_id, - section_path=db_section_path, - ) - for lp in leaf_paths: - if lp['path'] not in seen_leaf: - seen_leaf.add(lp['path']) - all_leaf_paths.append(lp) - - mode = 'overflow-L1→all-leaves' if overflowed else 'full' - logger.info( - f' agentic.nav_section_select [{mode}]: {len(accepted)} sections → ' - f'{len(all_leaf_paths)} unique leaf paths, {latency}ms' - ) - - if all_leaf_paths: - return ToolResult( - status='selected_paths', - payload={'document_id': document_id, 'selected_paths': all_leaf_paths}, - latency_ms=latency, - ) - - return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': 'no leaf chunks under selected sections'}, - latency_ms=latency, - ) - - except Exception as e: - latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.nav_section_select failed for doc={document_id}: {e}') - return ToolResult( - status='error', - payload={'document_id': document_id}, - error=str(e), - latency_ms=latency, - ) - - -async def _nav_fallback_chunks_slim( +async def scope_navigate( db: AsyncSession, *, document_id: str, job_result_id: str, query: str, llm_fn: LLMFn, - doc_name: str, - path_scope: str | None = None, + doc_name: str = '', + scope_path: str | None = None, + max_select: int = 15, ) -> ToolResult: - """Fallback: use flat chunks_slim selection when no sections available. - - Builds chunks_slim in-memory from DB (no files), applies path_scope - filter if given, and uses the standard chunk selection prompt. + """Unified document-internal navigation tool. + + 1. Loads 2 levels of child sections under scope_path + 2. Applies overflow guard (drops summaries if needed) + 3. LLM selects most relevant items + 4. Returns selected section paths directly; each path hydrates the + corresponding section subtree. """ t0 = time.monotonic() try: - chunks_slim = await _build_chunks_slim( - db, document_id=document_id, job_result_id=job_result_id, - ) - if not chunks_slim: + items = await _load_child_sections(db, document_id, job_result_id, scope_path) + if not items: latency = int((time.monotonic() - t0) * 1000) return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': 'no chunks found'}, + status='no_items', + payload={'document_id': document_id, 'scope_path': scope_path}, latency_ms=latency, ) - if path_scope: - chunks_slim = [c for c in chunks_slim if c.get('path', '').startswith(path_scope)] - if not chunks_slim: - latency = int((time.monotonic() - t0) * 1000) - return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': f'no chunks under {path_scope}'}, - latency_ms=latency, - ) - - chunks_text = _format_chunks_for_llm(chunks_slim) - prompt = _CHUNK_SELECT_PROMPT.format( + text, overflowed = _format_items_for_llm(items) + prompt = _SCOPE_NAV_PROMPT.format( doc_name=doc_name or document_id, doc_id=document_id, - chunks_overview=chunks_text, + scope_label=scope_path or 'root', + items_overview=text, query=query, - max_chunks=15, + max_select=max_select, ) - valid_paths = {c['path'] for c in chunks_slim if c.get('path')} + valid_paths = {item['path'] for item in items} response = await llm_fn(prompt) - parsed = _parse_chunk_path_selections(response) + selected = _parse_chunk_path_selections(response) accepted: list[dict[str, Any]] = [] - for item in parsed: + for item in selected: path = str(item.get('path') or '').strip() if path not in valid_paths: continue @@ -808,23 +472,36 @@ async def _nav_fallback_chunks_slim( if confidence is None: confidence = _default_confidence_for_rank(len(accepted)) accepted.append({'path': path, 'confidence': confidence}) + if len(accepted) >= max_select: + break latency = int((time.monotonic() - t0) * 1000) - - if accepted: + + if not accepted: return ToolResult( - status='selected_paths', - payload={'document_id': document_id, 'selected_paths': accepted}, + status='no_confident_match', + payload={'document_id': document_id, 'reason': 'no path matches query intent'}, latency_ms=latency, ) + + logger.info( + f" agentic.scope_navigate: {len(accepted)} section paths selected, " + f"status=selected_paths, overflowed={overflowed}" + ) + return ToolResult( - status='no_confident_match', - payload={'document_id': document_id, 'reason': 'no chunk paths matched'}, + status='selected_paths', + payload={ + 'document_id': document_id, + 'selected_paths': accepted, + 'scope_path': scope_path, + 'overflowed': overflowed, + }, latency_ms=latency, ) except Exception as e: latency = int((time.monotonic() - t0) * 1000) - logger.error(f' agentic.nav_fallback failed for doc={document_id}: {e}') + logger.error(f' agentic.scope_navigate failed for doc={document_id}: {e}') return ToolResult( status='error', payload={'document_id': document_id}, diff --git a/packages/shared-python/shared/services/retrieval/agentic/trace.py b/packages/shared-python/shared/services/retrieval/agentic/trace.py index 409a6f301..ce320eef6 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/trace.py +++ b/packages/shared-python/shared/services/retrieval/agentic/trace.py @@ -82,7 +82,7 @@ async def create_run(self) -> None: top_k=self._top_k, data_type=self._data_type, filters=self._filters, - policy_name='rule_based_v1', + policy_name='llm_policy_v1', agentic_enabled=True, cache_hit=False, result_count=0, @@ -99,12 +99,18 @@ async def create_run(self) -> None: except Exception: pass - def record_step(self, action_type: ActionType, result: ToolResult) -> None: + def record_step( + self, + action_type: ActionType, + result: ToolResult, + *, + decision_reason: str = '', + ) -> None: """Buffer a step record. Flushed on complete().""" self._steps.append({ 'step_index': len(self._steps), 'action_type': action_type.value, - 'action_input': {}, + 'action_input': {'decision_reason': decision_reason} if decision_reason else {}, 'observation_status': result.status, 'observation_payload_keys': list(result.payload.keys()) if result.payload else [], 'latency_ms': result.latency_ms, diff --git a/packages/shared-python/shared/services/retrieval/agentic/types.py b/packages/shared-python/shared/services/retrieval/agentic/types.py index bf4eb561d..01054c1e7 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/types.py @@ -17,9 +17,9 @@ class ActionType(str, Enum): BOTTOM_DISCOVERY = 'bottom_discovery' KG_DOCUMENT_SELECT = 'kg_document_select' DOCUMENT_PATH_SELECT = 'document_path_select' - NAV_SECTION_SELECT = 'nav_section_select' GREP_DOCUMENT_DISCOVER = 'grep_document_discover' GRAPH_EXPAND_DOCS = 'graph_expand_docs' + DONE = 'done' # Explicit termination signal from LLMPolicy @dataclass @@ -67,7 +67,7 @@ class AgentState: step_count: int = 0 # Discovery results (from bottom_discovery) - discovery_rows: list[dict[str, Any]] = field(default_factory=list) + discovery_paths: list[dict[str, Any]] = field(default_factory=list) discovery_top_doc_ids: list[str] = field(default_factory=list) discovery_done: bool = False @@ -80,13 +80,8 @@ class AgentState: # Document path selection selected_paths: list[dict[str, Any]] = field(default_factory=list) - agent_rows: list[dict[str, Any]] = field(default_factory=list) path_expansion_count: int = 0 - # doc_nav hierarchical navigation state - nav_drill_stack: list[dict[str, Any]] = field(default_factory=list) - # Each entry: {"document_id": str, "section_path": str | None, "depth": int} - # Last observation (used by policy to decide next action) last_observation: ToolResult | None = None @@ -98,6 +93,49 @@ class AgentState: def elapsed_ms(self) -> int: return int((time.monotonic() - self.start_time) * 1000) + def state_summary(self) -> dict[str, Any]: + """Produce a concise state snapshot for the LLMPolicy prompt. + + Only includes fields the LLM needs for decision-making — never + exposes raw chunk content or internal IDs verbatim. + """ + selected_doc_summaries = [ + { + 'document_id': d.document_id, + 'name': d.source_file_name or '(unnamed)', + 'confidence': round(d.confidence, 2), + 'source': d.source, + } + for d in self.selected_docs + ] + selected_path_summaries = [ + { + 'path': p.get('path', ''), + 'confidence': round(float(p.get('confidence', 0.0) or 0.0), 2), + } + for p in self.selected_paths[:10] # cap to avoid huge prompts + ] + last_obs = None + if self.last_observation: + last_obs = { + 'status': self.last_observation.status, + 'payload_keys': list(self.last_observation.payload.keys()), + 'error': self.last_observation.error, + } + return { + 'step': self.step_count, + 'discovery_done': self.discovery_done, + 'discovery_candidates': len(self.discovery_paths), + 'discovery_top_doc_ids': self.discovery_top_doc_ids[:5], + 'kg_done': self.kg_done, + 'selected_docs': selected_doc_summaries, + 'pending_doc_index': self.pending_doc_index, + 'selected_paths_count': len(self.selected_paths), + 'selected_paths': selected_path_summaries, + 'doc_retry_count': self.doc_retry_count, + 'last_observation': last_obs, + } + def apply(self, action_type: ActionType, result: ToolResult) -> None: """Update state based on action and its result.""" self.last_observation = result @@ -105,7 +143,7 @@ def apply(self, action_type: ActionType, result: ToolResult) -> None: if action_type == ActionType.BOTTOM_DISCOVERY: self.discovery_done = True if result.status != 'error': - self.discovery_rows = result.payload.get('fused_rows', []) + self.discovery_paths = result.payload.get('fused_rows', []) self.discovery_top_doc_ids = result.payload.get('top_doc_ids', []) elif action_type == ActionType.KG_DOCUMENT_SELECT: @@ -126,19 +164,25 @@ def apply(self, action_type: ActionType, result: ToolResult) -> None: self.doc_id_to_name.update(result.payload.get('doc_id_to_name', {})) self.doc_job_map.update(result.payload.get('doc_job_map', {})) + # Merge discovery hints + existing_ids = {d.document_id for d in self.selected_docs} + for did in self.discovery_top_doc_ids: + if did not in existing_ids and did not in self.excluded_doc_ids: + self.selected_docs.append(CandidateDoc( + document_id=did, + source_file_name=self.doc_id_to_name.get(did, ''), + confidence=0.5, # Lower than LLM's 0.8 + reason='Bottom discovery hit', + source='discovery_hint', + )) + existing_ids.add(did) + elif action_type == ActionType.DOCUMENT_PATH_SELECT: if result.status == 'selected_paths': new_paths = result.payload.get('selected_paths', []) self.selected_paths.extend(new_paths) self.pending_doc_index += 1 - elif result.status == 'need_nav_drill': - # Large document with doc_nav available → switch to nav mode - doc_id = result.payload.get('document_id', '') - self.nav_drill_stack.append({ - 'document_id': doc_id, - 'section_path': None, # start from top - 'depth': 0, - }) + elif result.status == 'no_items': self.pending_doc_index += 1 elif result.status == 'need_more_docs': failed_doc_id = result.payload.get('document_id', '') @@ -151,22 +195,6 @@ def apply(self, action_type: ActionType, result: ToolResult) -> None: elif result.status == 'error': self.pending_doc_index += 1 - elif action_type == ActionType.NAV_SECTION_SELECT: - # Concurrent batch nav: all stack entries consumed in one shot. - # The orchestrator embeds '_consumed_stack' listing processed doc IDs. - consumed_ids = set(result.payload.get('_consumed_stack', [])) - if consumed_ids: - self.nav_drill_stack = [ - e for e in self.nav_drill_stack - if e['document_id'] not in consumed_ids - ] - elif self.nav_drill_stack: - # Fallback: legacy single-entry pop (safety net) - self.nav_drill_stack.pop() - if result.status == 'selected_paths': - new_paths = result.payload.get('selected_paths', []) - self.selected_paths.extend(new_paths) - elif action_type == ActionType.GREP_DOCUMENT_DISCOVER: if result.status == 'discovered_docs': grep_doc_ids = result.payload.get('document_ids', []) diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 2e5e5d247..e6ee4a09f 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -12,11 +12,10 @@ from shared.core.database import get_db_context from shared.models.database.document import Document, DocumentChunk, DocumentSection, RetrievalHitStat -from shared.services.retrieval.agent_navigate import agent_navigate from shared.services.retrieval.graph_service import GraphQueryService, is_excluded_section +from shared.services.retrieval.lexical_text import normalize_section_path from shared.services.retrieval.cache_service import get_cached_retrieval_query_result, set_cached_retrieval_query_result from shared.services.retrieval.hit_stats_service import compute_importance_score, record_retrieval_hits -from shared.services.retrieval.llm_adapter import create_retrieval_llm_fn from shared.services.retrieval.channels import path_channel, content_channel, term_channel from shared.services.storage.result_storage import get_result_storage from shared.models.database.job_result import JobResult @@ -790,7 +789,8 @@ async def _hydrate_paths_to_rows( confidence_by_path: dict[str, float] = {} ordered_paths: list[str] = [] for item in path_selections: - path = str(item.get('path') or '').strip() + 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) @@ -802,6 +802,11 @@ async def _hydrate_paths_to_rows( if not ordered_paths: return [] + section_path_filters = [] + for path in ordered_paths: + section_path_filters.append(DocumentSection.section_path == path) + section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + stmt = ( select(Document, DocumentChunk, DocumentSection, JobResult) .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) @@ -813,7 +818,7 @@ async def _hydrate_paths_to_rows( .where(Document.status == 'active') .where( or_( - DocumentSection.section_path.in_(ordered_paths), + *section_path_filters, DocumentChunk.source_chunk_path.in_(ordered_paths), ) ) @@ -829,7 +834,16 @@ async def _hydrate_paths_to_rows( if row_path in seen_paths: continue seen_paths.add(row_path) - agent_score = confidence_by_path.get(row_path, 0.0) + matched_path = row_path + if section and section.section_path not in confidence_by_path: + matched_path = next( + ( + path for path in ordered_paths + if section.section_path == path or section.section_path.startswith(f'{path} / ') + ), + row_path, + ) + agent_score = confidence_by_path.get(matched_path, 0.0) rows.append({ 'document_id': document.document_id, 'chunk_id': chunk.chunk_id, @@ -848,11 +862,24 @@ async def _hydrate_paths_to_rows( 'sort_order': chunk.sort_order, }) - rows.sort(key=lambda r: path_order.get(_get_row_path(r), 10**9)) - missed = len(ordered_paths) - len(rows) + def _row_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, idx in path_order.items(): + if row_path.startswith(f'{path} / '): + return idx + return 10**9 + + rows.sort(key=_row_sort_key) + hydrated_paths = {_get_row_path(r) for r 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) + } + missed = len(ordered_paths) - len(resolved_inputs) if missed > 0: - hydrated_paths = {_get_row_path(r) for r in rows} - missing_paths = [p for p in ordered_paths if p not in hydrated_paths] + missing_paths = [p for p in ordered_paths if p not in resolved_inputs] logger.warning( f' hydrate: {len(rows)}/{len(ordered_paths)} paths resolved (missed={missed}); ' f'missing[:5]={missing_paths[:5]}' @@ -1122,88 +1149,25 @@ async def run_retrieval_query( default=0.5, ) - # ── Agent navigation or lexical graph fallback ── - # Aligned with KB: agent_navigate returns chunk paths with confidence. - logger.info('\n 🧭 PHASE 2: Agent Navigation') + # ── Legacy graph routing ── + logger.info('\n 🧭 PHASE 2: Legacy Graph Routing') router_used = 'discovery_only' - llm_fn = create_retrieval_llm_fn() agent_rows: list[dict[str, Any]] = [] - agent_paths: list[dict[str, Any]] = [] - - if llm_fn is not None: - logger.info(' LLM configured, running agent_navigate...') - t_agent = time.monotonic() - try: - agent_paths = await agent_navigate( - db, - user_id=user_id, - namespace=namespace, - query=query, - llm_fn=llm_fn, - exclude_document_ids=exclude_document_ids, - ) - if agent_paths: - discovery_paths = {_get_row_path(r) for r in fused_rows} - selected_paths = [str(item.get('path') or '') for item in agent_paths if item.get('path')] - new_paths = [path for path in selected_paths if path not in discovery_paths] - overlap_paths = [path for path in selected_paths if path in discovery_paths] - logger.info('\n 🔗 Agent→Discovery union:') - logger.info( - f' agent_paths={len(selected_paths)}, discovery_paths={len(discovery_paths)}, ' - f'overlap_paths={len(overlap_paths)}, new_paths={len(new_paths)}' - ) - if new_paths: - logger.info(' New paths from agent:') - for p in new_paths[:10]: - logger.info(f' → {p}') - if overlap_paths: - logger.info(' Overlap paths reinforced by agent:') - for p in overlap_paths[:10]: - logger.info(f' → {p}') - agent_rows = await _hydrate_paths_to_rows( - db, - path_selections=agent_paths, - user_id=user_id, - namespace=namespace, - ) - logger.info(f' Hydrated {len(agent_rows)} rows from {len(selected_paths)} agent-selected paths') - router_used = 'discovery+agent' - elapsed_agent = round((time.monotonic() - t_agent) * 1000) - logger.info(f' ✅ Agent navigate: {len(agent_paths)} paths ({len(new_paths)} new) in {elapsed_agent}ms') - else: - elapsed_agent = round((time.monotonic() - t_agent) * 1000) - logger.info(f' ⚠️ Agent returned 0 paths in {elapsed_agent}ms, falling back to lexical graph') - agent_rows = await list_graph_routed_chunks( - db, user_id=user_id, namespace=namespace, query=query, - top_k=top_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if agent_rows: - logger.info(f' 📊 Graph fallback: {len(agent_rows)} rows') - except Exception as exc: - logger.error(f' ❌ Agent navigate failed: {exc}, falling back to lexical') - agent_rows = await list_graph_routed_chunks( - db, user_id=user_id, namespace=namespace, query=query, - top_k=top_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if agent_rows: - logger.info(f' 📊 Graph fallback: {len(agent_rows)} rows') - else: - logger.info(' ⚠️ No LLM configured (DS_KEY missing?), using lexical graph routing') - try: - agent_rows = await list_graph_routed_chunks( - db, user_id=user_id, namespace=namespace, query=query, - top_k=top_k, exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - ) - if agent_rows: - logger.info(f' 📊 Graph fallback: {len(agent_rows)} rows') - except Exception as exc: - logger.error(f' ❌ Graph routing failed (ignored): {exc}') - agent_rows = [] - if agent_rows and not agent_paths: + try: + agent_rows = await list_graph_routed_chunks( + db, user_id=user_id, namespace=namespace, query=query, + top_k=top_k, exclude_document_ids=exclude_document_ids, + exclude_sections=exclude_sections, + ) + if agent_rows: + router_used = 'discovery+graph' + logger.info(f' 📊 Graph routing: {len(agent_rows)} rows') + except Exception as exc: + logger.error(f' ❌ Graph routing failed (ignored): {exc}') + agent_rows = [] + + if agent_rows: _normalize_row_scores( agent_rows, source_field='score', diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph_service.py index 14a7b24dd..c8b989d06 100644 --- a/packages/shared-python/shared/services/retrieval/graph_service.py +++ b/packages/shared-python/shared/services/retrieval/graph_service.py @@ -154,23 +154,6 @@ def _extract_document_top_summary( return '' -def _extract_document_nav_sections( - chunk_metadata_list: list[dict[str, Any]], -) -> list[dict[str, Any]]: - """Extract document_nav_sections from chunk metadata. - - The nav_sections list is injected by kb_tasks.py at parse time from - doc_nav.json. Returns the first non-empty list found in chunk metadata. - Each section has: title, path, summary, chunk_count, children_count. - """ - for meta in chunk_metadata_list: - if not isinstance(meta, dict): - continue - nav_sections = meta.get('document_nav_sections') - if isinstance(nav_sections, list) and nav_sections: - return nav_sections - return [] - @dataclass class GraphScope: user_id: str @@ -226,9 +209,6 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d ] top_summary = _extract_document_top_summary(chunk_metadata_list, sections) - # Extract nav_sections from chunk metadata (injected by kb_tasks.py) - nav_sections = _extract_document_nav_sections(chunk_metadata_list) - # ── Clean up old graph data for this document ── self.remove_document_graph(db, scope=GraphScope(user_id=user_id, namespace=namespace), document_id=document_id) @@ -250,7 +230,6 @@ def publish_document_graph(self, db: Session, *, user_id: str, namespace: str, d 'chunks_count': chunks_count, 'types': dict(types_breakdown), 'top_summary': top_summary, - 'nav_sections': nav_sections, }, ) ) diff --git a/packages/shared-python/shared/services/retrieval/lexical_text.py b/packages/shared-python/shared/services/retrieval/lexical_text.py index 79d595a78..42d15366f 100644 --- a/packages/shared-python/shared/services/retrieval/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/lexical_text.py @@ -9,6 +9,27 @@ from shared.utils.text_utils import tokenize_contents_for_retrieval +def normalize_section_path(path: Optional[str]) -> str: + """Return the canonical section path representation used by retrieval.""" + raw = str(path or "").strip() + if not raw: + return "Root" + parts = split_section_path(raw) + if not parts: + return "Root" + return " / ".join(parts) + + +def split_section_path(path: Optional[str]) -> list[str]: + """Split either canonical ``" / "`` paths or raw slash-separated paths.""" + raw = str(path or "").strip() + if not raw or raw == "Root": + return [] + if " / " in raw: + return [p.strip() for p in raw.split(" / ") if p.strip()] + return [p.strip() for p in raw.split("/") if p.strip()] + + def build_lexical_text(value: str) -> str: text = str(value or "").strip() if not text: @@ -44,11 +65,11 @@ def section_path_from_chunk_path(source_path: Optional[str]) -> str: """ if not source_path: return "Root" - parts = [p.strip() for p in source_path.split("/") if p.strip()] + parts = split_section_path(source_path) section_parts = parts[2:] # skip kb_root + filename if not section_parts: return "Root" - return " / ".join(section_parts) + return normalize_section_path(" / ".join(section_parts)) def build_path_lexical_text(source_path: Optional[str]) -> Optional[str]: From 456833462f0b0a36132ddd9cb52f4745fa18f8f7 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 8 May 2026 22:45:33 +0800 Subject: [PATCH 3/5] refactor: implement recursive bottom-up doc_nav summary with virtual document node --- .../connect_builder/summary_builder.py | 279 +++--------------- .../shared-python/shared/utils/text_utils.py | 2 +- 2 files changed, 49 insertions(+), 232 deletions(-) diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 81fde4393..5fecabb71 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -46,11 +46,11 @@ # ─── LLM Interface ─────────────────────────────────────────────────────────── -def _llm_summarize(snippets_text: str, node_name: str) -> str: +def _llm_summarize(snippets_text: str, node_name: str, max_tokens: int = 100) -> str: """ Call LLM to produce a concise summary from aggregated child snippets. - Returns plain text summary (≤100 chars), or "" on failure. + Returns plain text summary, or "" on failure. """ try: from shared.services.ai.prompt_service import build_prompt, _detect_text_language @@ -58,12 +58,12 @@ def _llm_summarize(snippets_text: str, node_name: str) -> str: # Deterministic language lock — see prompt_service._language_directive detected_lang = _detect_text_language(snippets_text) - prompt, temperature, top_p, max_tokens = build_prompt( + prompt, temperature, top_p, _prompt_max_tokens = build_prompt( task="file-summary", texts=snippets_text, query="", paras={ - "max_tokens": 100, + "max_tokens": max_tokens, "node_name": node_name, "lang": detected_lang, }, @@ -90,161 +90,6 @@ def _llm_summarize(snippets_text: str, node_name: str) -> str: return "" -def _normalize_whitespace(text: str) -> str: - return re.sub(r"\s+", " ", str(text or "")).strip() - - -def _normalize_multiline_text(text: str) -> str: - lines = [] - for raw_line in str(text or "").splitlines(): - stripped = raw_line.rstrip() - if not stripped.strip(): - continue - leading_spaces = len(stripped) - len(stripped.lstrip(" ")) - compact = re.sub(r"\s+", " ", stripped.lstrip()) - lines.append((" " * leading_spaces) + compact) - return "\n".join(lines).strip() - - -def _looks_like_title_enum(text: str) -> bool: - normalized = _normalize_whitespace(text).lower() - return any(normalized.startswith(prefix) for prefix in _TITLE_ENUM_PREFIXES) - - -def _truncate_navigation_summary(text: str, max_tokens: int = NAVIGATION_TOP_SUMMARY_MAX_TOKENS) -> str: - """Truncate navigation summary while preserving multiline tree structure. - - Budget is measured in semantic tokens via ``count_cn_en`` (1 Chinese - char = 1 token, 1 English word = 1 token) so the limit is - language-fair. For multiline tree previews, truncation removes - trailing branches. For single-line LLM summaries, the first - sentence is kept; if still too long, ``truncate_text_by_tokens`` - is used (same middle-ellipsis strategy as layout_parser headings). - """ - from shared.utils.text_utils import count_cn_en - from app.services.common.kb_utils import truncate_text_by_tokens - - normalized = _normalize_multiline_text(text) - if not normalized: - return "" - if count_cn_en(normalized) <= max_tokens: - return normalized - - # Multiline content (tree preview) — truncate by keeping lines - lines = normalized.split('\n') - if len(lines) > 1: - kept: List[str] = [] - current_tokens = 0 - for line in lines: - line_tokens = count_cn_en(line) - if kept and current_tokens + line_tokens > max_tokens: - break - kept.append(line) - current_tokens += line_tokens - if kept: - return '\n'.join(kept) - - # Single-line content (LLM summary) — keep first sentence - flat = _normalize_whitespace(normalized) - sentences = re.split(r'(?<=[。！？；.!?;])\s+', flat) - if sentences and sentences[0]: - first = sentences[0] - if count_cn_en(first) <= max_tokens: - return first - - # Absolute fallback: token-aware truncation with middle '...' - return truncate_text_by_tokens(flat, max_tokens, 0, lang_aware=False) - - -def _dedupe_summary_blocks(items: List[str]) -> List[str]: - deduped: List[str] = [] - seen: set[str] = set() - for item in items: - normalized = _normalize_multiline_text(item) - if not normalized: - continue - key = normalized.lower() - if key in seen: - continue - seen.add(key) - deduped.append(normalized) - return deduped - - -def _truncate_tree_title(title: str) -> str: - """Truncate an individual section title for tree preview display. - - Uses the same token-aware truncation as layout_parser headings: - keeps start and end tokens with '...' in the middle. - """ - from app.services.common.kb_utils import truncate_text_by_tokens - - return truncate_text_by_tokens( - title, - _TREE_TITLE_MAX_TOKENS_START, - _TREE_TITLE_MAX_TOKENS_END, - lang_aware=True, - ) - - -# ─── Chunk Lookup ─────────────────────────────────────────────────────────────────── - - -def _build_chunk_lookup( - chunks: List[Dict[str, Any]], -) -> Dict[str, str]: - """ - Build a mapping from the LAST path segment to a snippet string. - - For each chunk, compose a structured snippet: - - If metadata.summary exists → use it directly - - Otherwise → title (node_key) + keywords - """ - lookup: Dict[str, str] = {} - - for chunk in chunks: - path = chunk.get("path", "") - if not path: - continue - - ctype = chunk.get("type", "text") - # Skip image and table chunks — they have their own summaries - # but don't belong to the content tree path hierarchy - if ctype in ("image", "table"): - continue - - # Use the last path segment as tree node key - parts = path.rstrip("/").split("/") - node_key = parts[-1] if parts else "" - if not node_key: - continue - - meta = chunk.get("metadata", {}) - summary = "" - if isinstance(meta, dict): - summary = (meta.get("summary") or "").strip() - - if not summary: - snippet_parts = [node_key] - if isinstance(meta, dict): - kw = meta.get("keywords") - if isinstance(kw, list): - kw_text = ", ".join(str(item).strip() for item in kw if str(item).strip()) - else: - kw_text = str(kw or "").strip() - if kw_text: - snippet_parts.append(f"Keywords: {kw_text}") - summary = "\n".join(snippet_parts) - - if summary: - # Multiple chunks may share the same last segment (e.g. "part 1", "part 2") - if node_key in lookup: - lookup[node_key] = lookup[node_key] + "\n" + summary - else: - lookup[node_key] = summary - - return lookup - # # Uses explicit children arrays: @@ -302,16 +147,15 @@ def ensure_doc_nav_json( def _recursive_summarize_nav( node: Dict[str, Any], - chunk_lookup: Dict[str, str], use_llm: bool = True, + is_top_level: bool = False, ) -> str: """Bottom-up recursive summarization on a doc_nav section node. Operates on the children-array tree structure of doc_nav.json. For each node: - - Leaf (children==[]) → keep existing summary (set during ZIP creation) - or update from chunk_lookup if a better one exists. + - Leaf (children==[]) → keep existing summary (set during ZIP creation). - Non-leaf → recursively summarize children, then aggregate. Writes summary in-place into ``node["summary"]``. @@ -321,43 +165,48 @@ def _recursive_summarize_nav( title = node.get("title", "") if not children: - # Leaf node — check if chunk_lookup has a better summary + # Leaf node — keep existing summary existing = (node.get("summary") or "").strip() - lookup_snippet = chunk_lookup.get(title, "") - if lookup_snippet and (not existing or existing == title): - node["summary"] = lookup_snippet + if existing: + node["summary"] = existing return node.get("summary", "") # Recurse into children child_summaries: List[Tuple[str, str]] = [] for child in children: - child_summary = _recursive_summarize_nav(child, chunk_lookup, use_llm) + child_summary = _recursive_summarize_nav(child, use_llm, is_top_level=False) if child_summary: child_summaries.append((child.get("title", ""), child_summary)) if not child_summaries: return node.get("summary", "") - # Aggregate child summaries + # Aggregate child summaries without hard truncation aggregated_parts = [] for name, summary in child_summaries: - truncated = truncate_content_preview(summary, head=SUMMARY_MAX_LEN, tail=0) - aggregated_parts.append(f"[{name}] {truncated}") + aggregated_parts.append(f"[{name}] {summary}") aggregated_text = "\n".join(aggregated_parts) - if len(child_summaries) <= 1: - result = truncate_content_preview(child_summaries[0][1], head=SUMMARY_MAX_LEN, tail=0) + max_len = NAVIGATION_TOP_SUMMARY_MAX_TOKENS if is_top_level else SUMMARY_MAX_LEN + + if len(child_summaries) <= 1 and not is_top_level: + result = child_summaries[0][1] else: - titles = [name for name, _ in child_summaries] - title_enum = "This section covers: " + ", ".join(titles) + if is_top_level and not use_llm: + titles = [name for name, _ in child_summaries if name.lower() != "root"] + else: + titles = [name for name, _ in child_summaries] + + enum_prefix = "This document includes: " if is_top_level else "This section covers: " + title_enum = enum_prefix + ", ".join(titles) if not use_llm: result = title_enum else: total_len = sum(len(s) for _, s in child_summaries) if total_len > SUMMARY_MAX_LEN: - result = _llm_summarize(aggregated_text, title) + result = _llm_summarize(aggregated_text, title, max_tokens=max_len) if not result: result = title_enum else: @@ -400,55 +249,32 @@ def _check_sections(sections: List[Dict[str, Any]]) -> bool: return _check_sections(sections) -def _build_nav_top_summary(doc_nav: Dict[str, Any]) -> str: +def _build_nav_top_summary( + doc_nav: Dict[str, Any], + use_llm: bool = True +) -> str: """Build navigation-facing top summary from enriched doc_nav.json. Strategy: - - If root section has a high-quality LLM summary → use it. - - Otherwise → build a tree preview from section titles. + Treat all sections as children of a virtual Document node and recursively summarize. """ sections = doc_nav.get("sections", []) if not sections: return "" - # Check for a root-level LLM summary (first section if it's 'Root') - root_section = None - content_sections = [] - for s in sections: - if s.get("title", "").lower() == "root": - root_section = s - else: - content_sections.append(s) - - # Try root summary first - if root_section: - root_summary = _normalize_whitespace(root_section.get("summary", "")) - if root_summary and not _looks_like_title_enum(root_summary): - return _truncate_navigation_summary(root_summary) - - # Build tree preview from section titles - lines: List[str] = [] - excluded = {"root", "images", "tables"} - - def _render_sections(secs: List[Dict], depth: int = 0) -> None: - for sec in secs: - title = sec.get("title", "") - if title.lower() in excluded: - continue - if len(lines) >= NON_LLM_TOP_SUMMARY_MAX_SECTIONS: - break - indent = " " * depth - display = _truncate_tree_title(title) - lines.append(f"{indent}- {display}") - if depth + 1 < NON_LLM_TOP_SUMMARY_MAX_DEPTH: - _render_sections(sec.get("children", []), depth + 1) - - _render_sections(content_sections) - - if lines: - tree_text = "This document includes the following contents:\n" + "\n".join(lines) - return _truncate_navigation_summary(tree_text) - return "" + + virtual_doc_node = { + "title": "Document Overview", + "children": sections + } + + top_summary = _recursive_summarize_nav( + virtual_doc_node, + use_llm=use_llm, + is_top_level=True + ) + + return top_summary def enrich_doc_nav_summaries( @@ -490,31 +316,22 @@ def enrich_doc_nav_summaries( if not force and _doc_nav_has_enriched_summaries(doc_nav): logger.debug(f"Summaries already exist in {DOC_NAV_FILENAME} for {file_name}, skipping") - results[file_name] = _build_nav_top_summary(doc_nav) + results[file_name] = _build_nav_top_summary(doc_nav, use_llm=use_llm) continue - # Load chunks for snippet lookup - chunks_path = os.path.join(file_dir, "chunks.json") - chunks: List[Dict[str, Any]] = [] - if os.path.exists(chunks_path): - with open(chunks_path, "r", encoding="utf-8") as f: - data = json.load(f) - chunks = data.get("chunks", []) - - chunk_lookup = _build_chunk_lookup(chunks) logger.info( f"📝 Enriching {DOC_NAV_FILENAME} summaries for {file_name} " - f"({len(chunk_lookup)} leaf snippets, mode={mode_label})" + f"(mode={mode_label})" ) # Recursively summarize each top-level section for section in doc_nav.get("sections", []): - _recursive_summarize_nav(section, chunk_lookup, use_llm=use_llm) + _recursive_summarize_nav(section, use_llm=use_llm) _save_doc_nav(file_dir, doc_nav) logger.info(f"✅ doc_nav summaries saved for {file_name}") - top_summary = _build_nav_top_summary(doc_nav) + top_summary = _build_nav_top_summary(doc_nav, use_llm=use_llm) results[file_name] = top_summary return results @@ -524,7 +341,7 @@ def load_nav_top_summary(file_dir: str, file_name: str = "") -> str: """Load doc_nav.json and extract the navigation top summary.""" doc_nav = _load_doc_nav(file_dir) if doc_nav is not None: - return _build_nav_top_summary(doc_nav) + return _build_nav_top_summary(doc_nav, use_llm=False) return "" @@ -570,7 +387,7 @@ def _walk(node: Dict[str, Any]) -> None: # This mirrors GraphNode.properties.top_summary and ensures the # DocumentSection Root row has a summary for data completeness. if "Root" not in lookup: - top_summary = _build_nav_top_summary(doc_nav) + top_summary = _build_nav_top_summary(doc_nav, use_llm=False) if top_summary: lookup["Root"] = top_summary diff --git a/packages/shared-python/shared/utils/text_utils.py b/packages/shared-python/shared/utils/text_utils.py index 4423b67f2..dd61bfdd2 100644 --- a/packages/shared-python/shared/utils/text_utils.py +++ b/packages/shared-python/shared/utils/text_utils.py @@ -65,7 +65,7 @@ def count_cn_en(text: str) -> int: def truncate_content_preview( text: str, head: int = 200, - tail: int = 20, + tail: int = 50, ) -> str: """Token-aware content preview truncation. From 76185205a236c86ab4e114a77f4526f75f2703eb Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 8 May 2026 22:50:53 +0800 Subject: [PATCH 4/5] fix: resolve ruff lint errors (unused imports and variables) --- apps/worker/app/core/tasks/kb_tasks.py | 1 - apps/worker/app/services/connect_builder/summary_builder.py | 2 -- .../shared/services/retrieval/agentic/orchestrator.py | 1 - .../shared-python/shared/services/retrieval/agentic/policy.py | 1 - .../shared-python/shared/services/retrieval/agentic/tools.py | 2 +- 5 files changed, 1 insertion(+), 6 deletions(-) diff --git a/apps/worker/app/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/kb_tasks.py index b97353b9c..0529e9bec 100644 --- a/apps/worker/app/core/tasks/kb_tasks.py +++ b/apps/worker/app/core/tasks/kb_tasks.py @@ -5,7 +5,6 @@ All I/O operations use sync services that yield cooperatively under gevent. """ -import json import os from datetime import datetime, timezone diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 5fecabb71..cac4b38bb 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -14,12 +14,10 @@ import json import os -import re from typing import Any, Dict, List, Optional, Tuple from loguru import logger from openai.types.chat import ChatCompletionMessageParam -from shared.utils.text_utils import truncate_content_preview # ─── Constants ──────────────────────────────────────────────────────────────── diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index d7c4d5378..c1cc3000a 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -10,7 +10,6 @@ """ from __future__ import annotations -import json import os from typing import Any diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 926ca327d..7858bcac0 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -148,7 +148,6 @@ def build_prompt(self, state: AgentState, config: AgentRunConfig) -> str: state_json = json.dumps(state_data, ensure_ascii=False, indent=2) # Count pending docs - pending_count = len(state.selected_docs) - state.pending_doc_index return _POLICY_PROMPT_TEMPLATE.format( query=self._query, diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index a66fa0c5b..1463bea61 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -18,7 +18,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import Document from shared.services.retrieval.agentic.types import ToolResult from shared.services.retrieval.agent_navigate import ( _build_knowledge_map_overview, From 2dd0f1c08745d3efcc8f6e9810cf91e901ffd14d Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 8 May 2026 23:03:45 +0800 Subject: [PATCH 5/5] test: update expected top_summary format in unit tests --- apps/worker/tests/contract/test_parse_task_contract.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 6fdfb583a..3379b8613 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -313,13 +313,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") - expected_summary = ( - "This document includes the following contents:\n" - "- 公司研究\n" - " - 自主可控加强，寒武纪或迎来营收快速放量周期\n" - "- 相关研报\n" - " - 要点" - ) + expected_summary = "This document includes: 公司研究, 相关研报" expected_connect_to = [ { "target": "image-1",