diff --git a/apps/api/app/services/rate_limit/tier_service.py b/apps/api/app/services/rate_limit/tier_service.py index 01664e922..91bb73d00 100644 --- a/apps/api/app/services/rate_limit/tier_service.py +++ b/apps/api/app/services/rate_limit/tier_service.py @@ -20,6 +20,7 @@ from shared.models.database.tier_limit import TierLimit from shared.models.database.user_balance import UserBalance +from shared.services.billing.credits_service import CreditsService from shared.services.redis.redis_service import RedisService _DEFAULT_TIER: str = "free" @@ -42,7 +43,13 @@ async def get_tier(user_id: str) -> str: return cached_tier async with get_db_context() as session: - user_tier = await TierService._get_tier_from_db(session, user_id) + try: + user_tier: str = await TierService._get_tier_from_db(session, user_id) + except NotFoundException: + user_tier = await TierService._initialize_missing_user_tier( + session, + user_id, + ) await TierService._set_cached_tier(redis_service, user_id, user_tier) return user_tier @@ -126,6 +133,21 @@ async def _get_tier_from_db(session: AsyncSession, user_id: str) -> str: ) return str(user_tier) + @staticmethod + async def _initialize_missing_user_tier( + session: AsyncSession, + user_id: str, + ) -> str: + """Create missing first-use billing state, then return the user's tier.""" + credits_service: CreditsService = CreditsService() + await credits_service.ensure_user_initialized(session, user_id) + user_tier: str = await TierService._get_tier_from_db(session, user_id) + logger.info( + "Initialized missing user balance during tier lookup: user_id={}", + user_id, + ) + return user_tier + @staticmethod async def _get_cached_tier( redis_service: RedisService, diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index 9a9a48469..b62fd941d 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -1,4 +1,5 @@ import importlib +import json from collections.abc import Callable from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone @@ -10,12 +11,51 @@ from pytest import MonkeyPatch from tests.support.contract_database import ContractDatabase +from shared.utils.api_keys import hash_api_key def _utc_now() -> datetime: return datetime.now(timezone.utc).replace(tzinfo=None) +async def _insert_api_key_for_user(user_id: str, api_key: str) -> None: + timestamp = _utc_now() + api_key_id = f"key_{uuid4().hex[:12]}" + await ContractDatabase.execute( + """ + INSERT INTO api_keys ( + id, + user_id, + key_hash, + key_mask, + name, + enabled_modules, + is_active, + created_at + ) VALUES ( + :id, + :user_id, + :key_hash, + :key_mask, + :name, + CAST(:enabled_modules AS JSON), + :is_active, + :created_at + ) + """, + { + "id": api_key_id, + "user_id": user_id, + "key_hash": hash_api_key(api_key), + "key_mask": f"{api_key[:8]}...{api_key[-4:]}", + "name": f"Contract API Key {user_id}", + "enabled_modules": json.dumps(["all"]), + "is_active": True, + "created_at": timestamp, + }, + ) + + @pytest.mark.asyncio async def test_should_return_the_authenticated_users_initialized_credits_balance( developer_api_client_factory: Callable[ @@ -29,6 +69,60 @@ async def test_should_return_the_authenticated_users_initialized_credits_balance assert response.json() == {"credits_balance": 5.0} +@pytest.mark.asyncio +async def test_should_initialize_missing_user_balance_during_tier_lookup( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], +) -> None: + user_id = f"contract-missing-balance-{uuid4().hex[:12]}" + api_key = f"sk_contract_{uuid4().hex[:24]}" + + async with api_client_factory() as api_client: + await ContractDatabase.insert_user(user_id=user_id) + await _insert_api_key_for_user(user_id, api_key) + api_client.headers.update({"Authorization": f"Bearer {api_key}"}) + + response = await api_client.get("/api/v1/billing/credits") + balance_row = await ContractDatabase.fetch_one( + """ + SELECT credits_balance, user_tier + FROM user_balances + WHERE user_id = :user_id + """, + {"user_id": user_id}, + ) + transaction_row = await ContractDatabase.fetch_one( + """ + SELECT credits_amount, transaction_type + FROM credits_transactions + WHERE user_id = :user_id + AND transaction_type = 'initial_grant' + """, + {"user_id": user_id}, + ) + payment_row = await ContractDatabase.fetch_one( + """ + SELECT credits_amount, payment_type, status + FROM payment_records + WHERE user_id = :user_id + AND payment_type = 'system_grant' + """, + {"user_id": user_id}, + ) + + assert response.status_code == 200 + assert response.json() == {"credits_balance": 5.0} + assert balance_row == {"credits_balance": 5_000_000, "user_tier": "free"} + assert transaction_row == { + "credits_amount": 5_000_000, + "transaction_type": "initial_grant", + } + assert payment_row == { + "credits_amount": 5_000_000, + "payment_type": "system_grant", + "status": "succeeded", + } + + @pytest.mark.asyncio async def test_should_not_register_billing_routes_when_billing_is_disabled( monkeypatch: MonkeyPatch, 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/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/kb_tasks.py index 4cba0c087..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 @@ -74,39 +73,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 +592,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 +601,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 +615,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") @@ -686,6 +644,26 @@ def _parse(job_id: str, user_id: str | None): metadata_service.update_metadata(job_id, processing_timing_updates) job_metadata.update(processing_timing_updates) + # 1.5. Garbage Collection: Remove redundant local media files + try: + from shared.services.retrieval.publication_service import RetrievalPublicationService + + with get_sync_db_context() as db: + job_record = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() + if job_record: + gc_namespace = JobMetadataHelper.get_field(job_metadata, "namespace") or "default" + chunks, dedup_stats = RetrievalPublicationService.garbage_collect_and_dedup_local_media( + db, + job_id=job_id, + user_id=str(job_record.user_id), + namespace=gc_namespace, + add_dir=str(add_dir) if add_dir else "", + chunks=chunks, + ) + except Exception as e: + logger.error(f"[{job_id}] GC failed (non-fatal): {e}") + dedup_stats = None + # Generate ZIP package zip_service = ZipResultService() zip_file_path, checksum, statistics, zip_size = ( @@ -734,6 +712,7 @@ def _parse(job_id: str, user_id: str | None): stored_count=stored_count, delivery_mode="url", section_summaries=section_summaries, + chunk_dedup_stats=dedup_stats, ) logger.info( diff --git a/apps/worker/app/services/connect_builder/graph_builder.py b/apps/worker/app/services/connect_builder/graph_builder.py index d1a51e3ab..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, } @@ -824,6 +820,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 +1052,106 @@ 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 = {} + + 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, + ) + save_knowledge_graph(graph, kg_path) + _prune_chunk_stats(kb_id, chunks_on_disk) + return graph + + # ─── MCP Auto-Registration ─────────────────────────────────────────────────── @@ -1138,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 ───────────────────────────────────────────────────────────── @@ -1264,8 +1320,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: @@ -1320,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 ── @@ -1384,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 ── @@ -1398,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: @@ -1428,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/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 81fde4393..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 ──────────────────────────────────────────────────────────────── @@ -46,11 +44,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 +56,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 +88,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 +145,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 +163,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 +247,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 +314,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 +339,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 +385,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/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index a4838470e..a253100c2 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -1,8 +1,8 @@ # pyright: reportArgumentType=false, reportAttributeAccessIssue=false, reportCallIssue=false, reportGeneralTypeIssues=false -import hashlib import io import json import os +import shutil import zipfile import pandas as pd @@ -15,7 +15,11 @@ remove_spaces, ) from app.services.document_parser.html_parser import table2html -from app.services.document_parser.image_parser import _get_vision_client, ask_image +from app.services.document_parser.image_parser import ( + _get_vision_client, + ask_image, + perceptual_hash, +) from app.services.document_parser.layout_parser import pred_titles from app.services.document_parser.table_parser import sanitize_table_name_from_header from app.services.document_parser.toc_parser import ( @@ -101,10 +105,35 @@ def handle_image( current_heading, img_count, smart_summary=False, + seen_images=None, ): time_stamp = get_str_time() - client = _get_vision_client() + # Document-level dedup: use perceptual hash to catch visually-identical + # images that differ only in compression/metadata + img_hash = perceptual_hash(img_file["data"]) + if seen_images is not None and img_hash in seen_images: + cached = seen_images[img_hash] + headings_stack[-1]["content"].append(cached["image_ref"]) + df_list.append( + [ + cached["image_ref"], + cached["img_path"], + "image", + len(cached["image_ref"]), + "", + cached["img_summary_field"], + cached["temp_uid"], + "", + "", + time_stamp, + "", + ] + ) + logger.debug(f"Skipped duplicate image (hash={img_hash[:12]}...)") + return headings_stack, df_list, False # False = cache hit, don't increment + + client = _get_vision_client() last_context = _find_img_context(headings_stack) # Image index (always present) @@ -150,7 +179,7 @@ def handle_image( img_path = os.path.join(img_dir, f"{img_name}{img_ext}") os.rename(img_raw_path, img_path) # if summary fails, renaming is not applied - temp_uid = gen_str_codes(hashlib.sha256(img_file["data"]).hexdigest()) + temp_uid = gen_str_codes(img_hash) # Build img_summary_field for df_list: image-n + optional summary if img_summary: @@ -183,7 +212,17 @@ def handle_image( "", ] ) - return headings_stack, df_list + + # Cache result for document-level dedup + if seen_images is not None: + seen_images[img_hash] = { + "img_path": img_path, + "image_ref": image_ref, + "img_summary_field": img_summary_field, + "temp_uid": temp_uid, + } + + return headings_stack, df_list, True # True = new image processed def _first_cols_rows(table_block, max_items=10, max_chars=20): @@ -245,6 +284,7 @@ def handle_table( cell_images=None, img_dir=None, img_count=0, + seen_images=None, ): time_stamp = get_str_time() @@ -256,6 +296,31 @@ def handle_table( for (row_idx, col_idx), images in cell_images.items(): descriptions = [] for img_data in images: + # Document-level dedup: perceptual hash for visual duplicates + cell_img_hash = perceptual_hash(img_data["data"]) + if seen_images is not None and cell_img_hash in seen_images: + cached = seen_images[cell_img_hash] + descriptions.append(f"[{cached['img_summary_field']}]") + table_img_entries.append( + [ + cached["image_ref"], + cached["img_path"], + "image", + len(cached["image_ref"]), + "", + cached["img_summary_field"], + cached["temp_uid"], + "", + "", + time_stamp, + "", + ] + ) + logger.debug( + f"Skipped duplicate table cell image (hash={cell_img_hash[:12]}...)" + ) + continue + img_count += 1 img_ext = os.path.splitext(img_data["image_name"])[-1] image_index = f"image-{img_count}" @@ -286,7 +351,7 @@ def handle_table( descriptions.append(f"[{effective_desc}]") # Also add as IMAGE entry in df_list for indexing - temp_uid = gen_str_codes(hashlib.sha256(img_data["data"]).hexdigest()) + temp_uid = gen_str_codes(cell_img_hash) img_summary_field = ( f"{image_index}\n{img_summary}" if img_summary else image_index ) @@ -312,12 +377,22 @@ def handle_table( ] ) + # Cache result for document-level dedup + if seen_images is not None: + seen_images[cell_img_hash] = { + "img_path": relative_img_path, + "image_ref": image_ref, + "img_summary_field": img_summary_field, + "temp_uid": temp_uid, + } + cell_image_map[(row_idx, col_idx)] = " ".join(descriptions) logger.info( f"Extracted {sum(len(v) for v in cell_images.values())} images from table-{table_count + 1} cells" ) + # Generate HTML with image descriptions embedded tb_html_str = table2html( block, cell_image_map=cell_image_map if cell_image_map else None @@ -417,7 +492,10 @@ def iter_block_items(doc_data): "a": "http://schemas.openxmlformats.org/drawingml/2006/main", "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main", + "v": "urn:schemas-microsoft-com:vml", + "o": "urn:schemas-microsoft-com:office:office", } + r_ns = "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}" root = etree.fromstring(xml) body = root.find(".//w:body", namespaces=ns) @@ -500,15 +578,17 @@ def iter_block_items(doc_data): yield ele_num, p_obj or text, label, meta ele_num += 1 - # images + # images (DrawingML: ) + seen_rids = set() blips = elem.xpath(".//a:blip", namespaces=ns) for b in blips: - rid = b.get( - "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed" - ) + rid = b.get(f"{r_ns}embed") + if not rid or rid in seen_rids: + continue target = rel_map.get(rid) if not target or not target.startswith("media/"): continue + seen_rids.add(rid) data = docx.read("word/" + target) yield ( ele_num, @@ -522,6 +602,65 @@ def iter_block_items(doc_data): }, ) ele_num += 1 + + # TODO: Re-evaluate VML group extraction strategy. + # Complex VML composite images () are currently skipped because extracting + # piece-by-piece loses textual overlay and positioning. + # Future plan: Use LibreOffice headless conversion to render the entire document + # and map the perfectly rendered images back to the layout via text anchors. + # + # Temporary: detect VML-only paragraphs and inject a placeholder so the + # paragraph isn't silently swallowed, leaving its parent section empty. + if not text and not seen_rids: + # No text and no DrawingML images — check for VML content + vml_groups = elem.xpath(".//v:group", namespaces=ns) + vml_images_check = elem.xpath(".//v:imagedata", namespaces=ns) + if vml_groups or vml_images_check: + vml_placeholder = "[VML graphic \u2014 extraction not yet supported]" + yield ele_num, vml_placeholder, "PTXT", None + ele_num += 1 + logger.debug( + f"Injected VML placeholder for paragraph with " + f"{len(vml_groups)} v:group, {len(vml_images_check)} v:imagedata" + ) + """ + # images (VML: ) — convert to PNG + from PIL import Image as PILImage + + vml_images = elem.xpath(".//v:imagedata", namespaces=ns) + for v in vml_images: + rid = v.get(f"{r_ns}id") + if not rid or rid in seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + seen_rids.add(rid) + raw_data = docx.read("word/" + target) + # Convert to PNG for uniform downstream handling + try: + pil_img = PILImage.open(io.BytesIO(raw_data)) + png_buf = io.BytesIO() + pil_img.save(png_buf, format="PNG") + png_data = png_buf.getvalue() + except Exception as e: + logger.warning(f"Failed to convert VML image to PNG: {e}") + continue + orig_name = target.split("/")[-1] + png_name = os.path.splitext(orig_name)[0] + ".png" + yield ( + ele_num, + None, + "IMAGE", + { + "image_name": png_name, + "from": "paragraph_vml", + "size": len(png_data), + "data": png_data, + }, + ) + ele_num += 1 + """ map_index += 1 if toc_info["is_field_end"]: @@ -538,15 +677,18 @@ def iter_block_items(doc_data): cell_images = {} # {(row_idx, col_idx): [{'image_name', 'data', 'size'}]} for row_idx, tr in enumerate(elem.findall(".//w:tr", namespaces=ns)): for col_idx, tc in enumerate(tr.findall(".//w:tc", namespaces=ns)): - blips = tc.xpath(".//a:blip", namespaces=ns) + cell_seen_rids = set() imgs_in_cell = [] + # DrawingML images in cell + blips = tc.xpath(".//a:blip", namespaces=ns) for b in blips: - rid = b.get( - "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed" - ) + rid = b.get(f"{r_ns}embed") + if not rid or rid in cell_seen_rids: + continue target = rel_map.get(rid) if not target or not target.startswith("media/"): continue + cell_seen_rids.add(rid) data = docx.read("word/" + target) if ( len(data) < 10 * 1024 @@ -559,6 +701,44 @@ def iter_block_items(doc_data): "size": len(data), } ) + # TODO: VML in tables is temporarily skipped to avoid extracting + # fragmented textless background images. (Same as paragraph VML logic) + """ + # VML images in cell — convert to PNG + from PIL import Image as PILImage + + vml_in_cell = tc.xpath(".//v:imagedata", namespaces=ns) + for v in vml_in_cell: + rid = v.get(f"{r_ns}id") + if not rid or rid in cell_seen_rids: + continue + target = rel_map.get(rid) + if not target or not target.startswith("media/"): + continue + cell_seen_rids.add(rid) + raw_data = docx.read("word/" + target) + try: + pil_img = PILImage.open(io.BytesIO(raw_data)) + png_buf = io.BytesIO() + pil_img.save(png_buf, format="PNG") + png_data = png_buf.getvalue() + except Exception as e: + logger.warning( + f"Failed to convert VML cell image to PNG: {e}" + ) + continue + if len(png_data) < 10 * 1024: + continue + orig_name = target.split("/")[-1] + png_name = os.path.splitext(orig_name)[0] + ".png" + imgs_in_cell.append( + { + "image_name": png_name, + "data": png_data, + "size": len(png_data), + } + ) + """ if imgs_in_cell: cell_images[(row_idx, col_idx)] = imgs_in_cell @@ -607,9 +787,15 @@ def parse_docx( headings_stack = [{"level": -1, "content": doc_structure}] current_heading = "" + # Clean old artifacts to prevent accumulation across debug runs. + # In production each job uses a fresh workspace so rmtree never triggers. tb_dir = os.path.join(output_dir, "tables") + if os.path.isdir(tb_dir): + shutil.rmtree(tb_dir) os.makedirs(tb_dir, exist_ok=True) img_dir = os.path.join(output_dir, "images") + if os.path.isdir(img_dir): + shutil.rmtree(img_dir) os.makedirs(img_dir, exist_ok=True) block_tuples = list(iter_block_items(doc_data)) @@ -671,6 +857,7 @@ def parse_docx( df_list = [] table_count = 0 image_count = 0 + _seen_images: dict[str, dict] = {} # sha256_hex -> cached image info for dedup logger.debug("Parsing docx file... total_blocks={}", len(block_tuples)) for block_tuple in block_tuples: @@ -708,7 +895,7 @@ def parse_docx( if meta and meta.get("size", 0) < 10 * 1024: continue - headings_stack, df_list = handle_image( + headings_stack, df_list, is_new = handle_image( df_list, meta, img_dir, @@ -716,8 +903,10 @@ def parse_docx( current_heading, image_count, llm_paras["summary_image"], + seen_images=_seen_images, ) - image_count += 1 + if is_new: + image_count += 1 current_heading = last_heading_before_block elif label == "TABLE": @@ -734,6 +923,7 @@ def parse_docx( cell_images=meta, img_dir=img_dir, img_count=image_count, + seen_images=_seen_images, ) table_count += 1 current_heading = last_heading_before_block @@ -767,6 +957,13 @@ def convert_doc2dics( for _, row in leaf_dics.iterrows(): key = row["path_identifier"] + # Skip leaf nodes with no actual content (empty heading-only sections) + content_lst = row["content_lst"] + joined = "\n".join(content_lst).strip() + if not joined: + logger.debug(f"Skipping empty leaf node: {key}") + continue + # Build tentative path to check for duplicates tentative_path = doc_name + split_char + key @@ -779,7 +976,7 @@ def convert_doc2dics( path_counter[tentative_path] = 1 path_keys.append((doc_name + split_char + key)) - bottom_content = "\n".join(row["content_lst"]) + bottom_content = joined bottom_tokens = tokenize2stw_remove( [bottom_content], base_llm_paras["stopwords"] ) diff --git a/apps/worker/app/services/document_parser/image_compressor.py b/apps/worker/app/services/document_parser/image_compressor.py index 9893c3de0..892f85ac8 100644 --- a/apps/worker/app/services/document_parser/image_compressor.py +++ b/apps/worker/app/services/document_parser/image_compressor.py @@ -39,16 +39,6 @@ class CompressionStats(NamedTuple): rename_map: dict # {"old_relative_path": "new_relative_path"}, e.g. {"images/foo.png": "images/foo.jpg"} -def _has_transparency(img) -> bool: - """Check whether a PIL Image has meaningful alpha (transparency) data.""" - if img.mode != "RGBA": - return False - # Sample alpha channel — if any pixel has alpha < 250, treat as transparent - alpha = img.getchannel("A") - extrema = alpha.getextrema() - return extrema[0] < 250 - - def compress_output_images( output_dir: str, *, @@ -104,38 +94,9 @@ def compress_output_images( w, h = img.size needs_resize = max(w, h) > max_side is_png = ext == ".png" - has_alpha = is_png and _has_transparency(img) - - if is_png and not has_alpha: - # Opaque PNG → convert to JPEG - if needs_resize: - ratio = max_side / max(w, h) - new_w, new_h = int(w * ratio), int(h * ratio) - img = img.resize((new_w, new_h), Image.LANCZOS) - resized_count += 1 - - # Convert RGBA → RGB for JPEG - if img.mode in ("RGBA", "P", "LA"): - img = img.convert("RGB") - - jpg_path = os.path.splitext(file_path)[0] + ".jpg" - img.save(jpg_path, "JPEG", quality=jpeg_quality, optimize=True) - img.close() - - # Remove original PNG - if jpg_path != file_path: - os.remove(file_path) - # Record rename for downstream reference updates - old_rel = f"images/{filename}" - new_rel = f"images/{os.path.basename(jpg_path)}" - rename_map[old_rel] = new_rel - - converted += 1 - processed += 1 - total_after += os.path.getsize(jpg_path) - elif is_png and has_alpha: - # Transparent PNG → keep as PNG but resize if needed + if is_png: + # Keep as PNG, only resize if needed if needs_resize: ratio = max_side / max(w, h) new_w, new_h = int(w * ratio), int(h * ratio) @@ -188,7 +149,7 @@ def compress_output_images( ratio = total_before / total_after if total_after > 0 else 0 logger.info( f"[image_compressor] Compressed {processed} images " - f"({converted} PNG→JPG, {resized_count} resized), " + f"({resized_count} resized), " f"skipped {skipped}. " f"Size: {total_before / 1024 / 1024:.1f}MB → {total_after / 1024 / 1024:.1f}MB " f"({ratio:.1f}x reduction)" diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py index f5a91c7f8..dc8ffb541 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/image_parser.py @@ -36,6 +36,25 @@ g_img_lock = threading.Lock() +def perceptual_hash(data: bytes) -> str: + """Compute a normalized pixel-data hash for image dedup. + + Word/PDF may embed the same visual image with different compression + or metadata, making raw-byte SHA256 differ. This function decodes + the image, converts to RGBA, and hashes the raw pixel buffer so + that visually-identical images always produce the same digest. + + Falls back to raw-bytes hash when PIL cannot decode the data. + """ + try: + img = Image.open(io.BytesIO(data)) + pixels = img.convert("RGBA").tobytes() + return hashlib.sha256(pixels).hexdigest() + except Exception: + return hashlib.sha256(data).hexdigest() + + + def _get_vision_client() -> OpenAICompatibleClientSync: """Create OpenAI-compatible client for vision models, auto-routing by IMAGE_MODEL name.""" image_model = settings.IMAGE_MODEL or "qwen-vl-plus" diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index 12d467d7a..f308e7f5a 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -1197,7 +1197,7 @@ def hiearchy_llm( model_name=None, max_depth=6, toc_context=None, - max_len=2048, + max_len=8192, task="eval-headings", ): """Apply LLM to analyze the hierarchy of headings diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index c39d42f58..ab85751bc 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -1,8 +1,8 @@ # pyright: reportArgumentType=false, reportAssignmentType=false, reportOptionalIterable=false, reportOptionalMemberAccess=false, reportOptionalOperand=false, reportOptionalSubscript=false -import hashlib import json import os import re +import shutil from pathlib import Path import gevent @@ -22,6 +22,7 @@ _get_vision_client, ask_image, detect_summary_img_md, + perceptual_hash, ) from app.services.document_parser.layout_parser import md_heading_match, pred_titles from app.services.document_parser.stage_profiler import stage_timer @@ -309,10 +310,18 @@ def parse_md( json.dump(toc_hierarchies, f, ensure_ascii=False, indent=2) logger.info(f"Saved TOC hierarchies to {toc_json_path}") - # create local storage + # Clean old artifacts to prevent accumulation across debug runs. + # In production each job uses a fresh workspace so rmtree never triggers. tb_dir = os.path.join(output_dir, "tables") + if os.path.isdir(tb_dir): + shutil.rmtree(tb_dir) os.makedirs(tb_dir, exist_ok=True) img_dir = os.path.join(output_dir, "images") + if os.path.isdir(img_dir): + # Only remove parse_md's own output (image-N-*) from previous runs + for fname in os.listdir(img_dir): + if re.match(r"^image-\d+", fname): + os.remove(os.path.join(img_dir, fname)) os.makedirs(img_dir, exist_ok=True) # initialize vars @@ -332,6 +341,7 @@ def parse_md( img_count = 1 path_counter = {} # Track path occurrences for deduplication deferred_llm_tasks = [] # Collected during loop, executed in parallel after + _seen_images: dict[str, dict] = {} # sha256_hex -> cached image info for dedup # Find layout.json path layout_json_path = os.path.join(output_dir, "layout.json") @@ -474,6 +484,38 @@ def parse_md( img_count += 1 continue + # Document-level dedup: perceptual hash for visual duplicates + with open(source_path, "rb") as f: + img_binary_hash = perceptual_hash(f.read()) + + if img_binary_hash in _seen_images: + cached = _seen_images[img_binary_hash] + content_items.append(cached["img_content"]) + df_list.append( + [ + cached["img_content"], + cached["relative_img_path"], + "image", + len(cached["img_content"]), + "", + cached["img_summary_field"], + cached["temp_uid"], + "", + "", + time_stamp, + str(current_pg_num) if current_pg_num > 0 else "", + ] + ) + logger.debug( + f"Skipped duplicate image (hash={img_binary_hash[:12]}...)" + ) + # Remove unused source file since we reuse the cached image + try: + source_path.unlink() + except OSError: + pass + continue + os.rename(source_path, update_img_path) # Image index (always present) @@ -483,8 +525,6 @@ def parse_md( effective_summary = img_summary or last_context or None # Deterministic know_id: use image binary hash - with open(update_img_path, "rb") as img_f: - img_binary_hash = hashlib.sha256(img_f.read()).hexdigest() temp_uid = gen_str_codes(img_binary_hash) relative_img_path = f"images/{img_name}{img_suffix}" img_ref = build_chunk_ref(relative_img_path) @@ -518,6 +558,15 @@ def parse_md( str(current_pg_num) if current_pg_num > 0 else "", ] ) + + # Cache result for document-level dedup + _seen_images[img_binary_hash] = { + "relative_img_path": relative_img_path, + "img_content": img_content, + "img_summary_field": img_summary_field, + "temp_uid": temp_uid, + } + if base_llm_paras["summary_image"]: # Store img_dir, img_name, img_suffix for post-loop rename (mirrors table deferred task) deferred_llm_tasks.append( diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index bf02fe4d5..3379b8613 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -223,7 +223,15 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "summary": "", "know_id": "kid-1", "tokens": "", - "connectto": "", + "connectto": json.dumps( + [ + { + "target": "table-1", + "relation": "embeds", + "ref": "[tables/table-1.html]", + } + ] + ), "addtime": "now", "page_nums": "1", }, @@ -289,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", @@ -302,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", @@ -350,6 +355,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/billing/credits_service.py b/packages/shared-python/shared/services/billing/credits_service.py index 5bf816f84..566440157 100644 --- a/packages/shared-python/shared/services/billing/credits_service.py +++ b/packages/shared-python/shared/services/billing/credits_service.py @@ -48,11 +48,10 @@ class CreditsService: not ensure the /billing/credits endpoint is called before any credit modification operation. - 2. **Read-only endpoints** (e.g., GET /billing/credits): - MUST be called explicitly by the API route. This is necessary because - `get_balance()` is intentionally kept fast (no initialization check) - for performance. Without this call in the route, first-time users - would see 0 balance instead of their initial credits. + 2. **First-use user flows**: + Called before reading balance data, either from an API route or from + the tier lookup path used by authenticated route guards. `get_balance()` + is intentionally kept fast (no initialization check) for performance. Usage: ------ diff --git a/packages/shared-python/shared/services/chunks/chunk_connections.py b/packages/shared-python/shared/services/chunks/chunk_connections.py new file mode 100644 index 000000000..efb0804f1 --- /dev/null +++ b/packages/shared-python/shared/services/chunks/chunk_connections.py @@ -0,0 +1,222 @@ +"""Build canonical chunk connection metadata.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any, TypeAlias, TypedDict + +from shared.utils.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans + + +class PositionPayload(TypedDict): + start: int + end: int + + +class ConnectionPayload(TypedDict, total=False): + target: str + relation: str + ref: str + position: PositionPayload + score: float + keywords: list[str] + + +RelationshipRef: TypeAlias = str | ChunkRefSpan +ConnectionValue: TypeAlias = str | ConnectionPayload +ConnectionKey: TypeAlias = tuple[str, str, str] +PositionKey: TypeAlias = tuple[str, str] + + +def parse_relationship_refs(type_value: object, content: str) -> list[RelationshipRef]: + parsed_relationships = _parse_type_relationship_refs(type_value) + if parsed_relationships: + return [relationship for relationship in parsed_relationships] + return [span for span in extract_chunk_ref_spans(content)] + + +def build_resource_target_map( + chunks: Sequence[Mapping[str, Any]], + *, + image_files_map: Mapping[str, Mapping[str, Any]] | None = None, + table_files_map: Mapping[str, Mapping[str, Any]] | None = None, +) -> dict[str, str]: + target_map: dict[str, str] = {} + for chunk in chunks: + chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id") or "").strip() + if not chunk_id: + continue + + chunk_type = str(chunk.get("type", "")).strip().split("\n", 1)[0].lower() + if chunk_type not in {"image", "table"}: + continue + + metadata = chunk.get("metadata", {}) + file_path = "" + if isinstance(metadata, dict): + file_path = str(metadata.get("file_path") or "").strip() + if not file_path: + file_map = image_files_map if chunk_type == "image" else table_files_map + file_info = file_map.get(chunk_id) if file_map else None + if file_info: + file_path = str(file_info.get("file_path") or "").strip() + + path_alias = str(chunk.get("path") or "").strip() + aliases = {file_path, path_alias} + for alias in list(aliases): + if alias: + aliases.add(f"[{alias}]") + for alias in aliases: + if alias: + target_map[alias] = chunk_id + return target_map + + +def convert_refs_to_embed_connections( + refs: Sequence[RelationshipRef], target_map: Mapping[str, str] +) -> list[ConnectionPayload]: + connections: list[ConnectionPayload] = [] + for ref in refs: + if isinstance(ref, dict): + ref_text = str(ref.get("ref") or "").strip() + start = ref.get("start") + end = ref.get("end") + else: + ref_text = str(ref or "").strip() + start = None + end = None + if not ref_text: + continue + + target_id = target_map.get(ref_text) + if not target_id and ref_text.startswith("[") and ref_text.endswith("]"): + target_id = target_map.get(ref_text[1:-1].strip()) + if not target_id: + continue + + connection: ConnectionPayload = { + "target": target_id, + "relation": "embeds", + "ref": ref_text, + } + if isinstance(start, int) and isinstance(end, int): + connection["position"] = { + "start": start, + "end": end, + } + connections.append(connection) + return connections + + +def normalize_connect_to_targets( + connects: object, target_map: Mapping[str, str] +) -> list[ConnectionPayload]: + if connects is None or connects == "": + return [] + + raw_items = connects if isinstance(connects, list) else [connects] + normalized: list[ConnectionPayload] = [] + for item in raw_items: + if item is None or item == "": + continue + + if isinstance(item, dict): + target = str(item.get("target") or "").strip() + normalized_target = target_map.get(target, target) + if not normalized_target: + continue + + normalized_item: ConnectionPayload = { + "target": normalized_target, + "relation": str(item.get("relation") or "related"), + } + score = item.get("score") + if isinstance(score, (int, float)): + normalized_item["score"] = float(score) + keywords = item.get("keywords") + if isinstance(keywords, list): + normalized_item["keywords"] = [str(keyword) for keyword in keywords] + ref = item.get("ref") + if ref: + normalized_item["ref"] = str(ref) + position = item.get("position") + if isinstance(position, dict): + start = position.get("start") + end = position.get("end") + if isinstance(start, int) and isinstance(end, int): + normalized_item["position"] = {"start": start, "end": end} + normalized.append(normalized_item) + continue + + target = str(item or "").strip() + normalized_target = target_map.get(target, target) + if normalized_target: + normalized.append( + { + "target": normalized_target, + "relation": "related", + "score": 1.0, + "keywords": [], + } + ) + return normalized + + +def merge_connections( + *connection_lists: Sequence[ConnectionValue], +) -> list[ConnectionValue]: + merged: list[ConnectionValue] = [] + unpositioned_indexes: dict[ConnectionKey, int] = {} + positioned_keys: dict[ConnectionKey, set[PositionKey]] = {} + for connection_list in connection_lists: + for item in connection_list or []: + if not isinstance(item, dict): + continue + key = _get_connection_key(item) + position_key = _get_connection_position_key(item) + if position_key is None: + if key in unpositioned_indexes or key in positioned_keys: + continue + unpositioned_indexes[key] = len(merged) + merged.append(item) + continue + + key_positions = positioned_keys.setdefault(key, set()) + if position_key in key_positions: + continue + key_positions.add(position_key) + unpositioned_index = unpositioned_indexes.pop(key, None) + if unpositioned_index is None: + merged.append(item) + else: + merged[unpositioned_index] = item + return merged + + +def _parse_type_relationship_refs(type_value: object) -> list[str]: + if not isinstance(type_value, str) or "\n" not in type_value: + return [] + lines = [line.strip() for line in type_value.split("\n") if line.strip()] + return [line for line in lines[1:] if line.upper() != "PTXT"] + + +def _get_connection_key(item: ConnectionValue) -> ConnectionKey: + if not isinstance(item, dict): + return ("", "related", "") + return ( + str(item.get("target") or ""), + str(item.get("relation") or "related"), + str(item.get("ref") or ""), + ) + + +def _get_connection_position_key(item: ConnectionValue) -> PositionKey | None: + if not isinstance(item, dict): + return None + position = item.get("position") + if not isinstance(position, dict): + return None + return ( + str(position.get("start", "")), + str(position.get("end", "")), + ) diff --git a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py index 6b2cb7b37..c57301e72 100644 --- a/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py +++ b/packages/shared-python/shared/services/chunks/dataframe_chunk_converter.py @@ -5,13 +5,21 @@ import json import os import uuid -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Sequence from typing import Dict, Literal, Protocol, TypeAlias, TypedDict, Union, cast import pandas as pd from loguru import logger -from shared.utils.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans +from shared.services.chunks.chunk_connections import ( + ConnectionValue, + RelationshipRef, + build_resource_target_map, + convert_refs_to_embed_connections, + merge_connections, + normalize_connect_to_targets, + parse_relationship_refs, +) class _ParserRow(Protocol): @@ -24,28 +32,12 @@ def __len__(self) -> int: ... def iterrows(self) -> Iterable[tuple[object, _ParserRow]]: ... -class PositionPayload(TypedDict): - start: int - end: int - - -class ConnectionPayload(TypedDict, total=False): - target: str - relation: str - ref: str - position: PositionPayload - score: float - keywords: list[str] - - JsonPrimitive: TypeAlias = str | int | float | bool | None JsonValue: TypeAlias = Union[ JsonPrimitive, list["JsonValue"], dict[str, "JsonValue"], ] -RelationshipRef: TypeAlias = str | ChunkRefSpan -ConnectionValue: TypeAlias = str | ConnectionPayload ChunkType: TypeAlias = Literal["text", "image", "table"] @@ -67,7 +59,6 @@ class ChunkPayload(TypedDict): content: str path: str metadata: ChunkMetadata - text: str order: int know_id: str keywords: list[str] @@ -134,15 +125,6 @@ def _safe_parse_tokens(value: object) -> list[str]: return [] -def _safe_parse_relationships(value: object) -> list[str]: - if _is_missing(value): - return [] - if not isinstance(value, str) or "\n" not in value: - return [] - lines = [line.strip() for line in value.split("\n") if line.strip()] - return [line for line in lines[1:] if line.upper() != "PTXT"] - - def _normalize_resource_ref(ref: RelationshipRef) -> str: if isinstance(ref, dict): ref_text = str(ref.get("ref") or "").strip() @@ -197,87 +179,6 @@ def _parse_connect_to(value: object) -> list[ConnectionValue]: ] -def _build_resource_target_map(chunks: Sequence[ChunkPayload]) -> dict[str, str]: - target_map: dict[str, str] = {} - for chunk in chunks: - if chunk["type"] not in {"image", "table"}: - continue - chunk_id = str(chunk["chunk_id"] or chunk["know_id"]).strip() - if not chunk_id: - continue - metadata = chunk["metadata"] - file_path = "" - file_path = str(metadata.get("file_path") or "").strip() - path_alias = chunk["path"].strip() - aliases = {file_path, path_alias} - for alias in list(aliases): - if alias: - aliases.add(f"[{alias}]") - for alias in aliases: - if alias: - target_map[alias] = chunk_id - return target_map - - -def _refs_to_embed_connections( - refs: Sequence[RelationshipRef], target_map: Mapping[str, str] -) -> list[ConnectionPayload]: - connections: list[ConnectionPayload] = [] - for ref in refs: - if isinstance(ref, dict): - ref_text = str(ref.get("ref") or "").strip() - start = ref.get("start") - end = ref.get("end") - else: - ref_text = str(ref or "").strip() - start = None - end = None - if not ref_text: - continue - target_id = target_map.get(ref_text) - if not target_id and ref_text.startswith("[") and ref_text.endswith("]"): - target_id = target_map.get(ref_text[1:-1].strip()) - if not target_id: - continue - connection: ConnectionPayload = { - "target": target_id, - "relation": "embeds", - "ref": ref_text, - } - if isinstance(start, int) and isinstance(end, int): - connection["position"] = { - "start": start, - "end": end, - } - connections.append(connection) - return connections - - -def _merge_connections( - *connection_lists: Sequence[ConnectionValue], -) -> list[ConnectionValue]: - merged: list[ConnectionValue] = [] - seen: set[tuple[str, str, str, str, str]] = set() - for connection_list in connection_lists: - for item in connection_list or []: - if not isinstance(item, dict): - continue - position = item.get("position") - position_data = position if isinstance(position, dict) else {} - key = ( - str(item.get("target") or ""), - str(item.get("relation") or "related"), - str(item.get("ref") or ""), - str(position_data.get("start", "")), - str(position_data.get("end", "")), - ) - if key in seen: - continue - seen.add(key) - merged.append(item) - return merged - - def _parse_page_numbers(value: object) -> list[int]: if _is_missing(value): return [] @@ -305,10 +206,7 @@ def _get_chunk_type(value: object) -> ChunkType: def _get_relationship_refs(type_value: object, content: str) -> list[RelationshipRef]: - parsed_relationships = _safe_parse_relationships(type_value) - if parsed_relationships: - return [relationship for relationship in parsed_relationships] - return [span for span in extract_chunk_ref_spans(content)] + return parse_relationship_refs(type_value, content) def _get_connect_to(metadata: ChunkMetadata) -> list[ConnectionValue]: @@ -390,7 +288,6 @@ def dataframe_to_chunks(df: _ParserDataFrame | None) -> list[Dict[str, JsonValue "content": content, "path": path, "metadata": metadata, - "text": content, "order": index, "know_id": str(know_id), "keywords": metadata["keywords"], @@ -399,18 +296,21 @@ def dataframe_to_chunks(df: _ParserDataFrame | None) -> list[Dict[str, JsonValue } ) - resource_target_map = _build_resource_target_map(chunks) + resource_target_map = build_resource_target_map(chunks) for chunk in chunks: metadata = chunk["metadata"] relationship_refs = metadata.pop("_relationship_refs", []) if chunk["type"] != "text": continue - embed_connections = _refs_to_embed_connections( + embed_connections = convert_refs_to_embed_connections( relationship_refs, resource_target_map ) - metadata["connect_to"] = _merge_connections( + metadata["connect_to"] = merge_connections( embed_connections, - _get_connect_to(metadata), + normalize_connect_to_targets( + _get_connect_to(metadata), + resource_target_map, + ), ) logger.debug(f"DataFrame conversion completed: chunk count={len(chunks)}") diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index f599d71ac..b0b7b197d 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -60,6 +60,7 @@ def finalize_job_success( stored_count: int = 0, delivery_mode: str = "url", section_summaries: Optional[Dict[str, str]] = None, + chunk_dedup_stats: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Finalize a successful job in a single atomic transaction. @@ -82,6 +83,7 @@ def finalize_job_success( inline_payload=inline_payload, result_s3_key=result_s3_key, result_size=zip_size, + chunk_dedup_stats=chunk_dedup_stats, ) normalized_chunks = chunks or [] @@ -100,7 +102,7 @@ def finalize_job_success( chunks=normalized_chunks, ) ) - if published_document_state is not None: + if published_document_state is not None and not published_document_state.get("skipped_all_duplicate"): # Backfill DocumentSection.summary from enriched doc_nav data if section_summaries: self._backfill_section_summaries( @@ -304,6 +306,7 @@ def _upsert_job_result( inline_payload: Optional[Dict[str, Any]] = None, result_s3_key: Optional[str] = None, result_size: Optional[int] = None, + chunk_dedup_stats: Optional[Dict[str, Any]] = None, ) -> JobResult: """Create or update JobResult row.""" result = db.execute(select(JobResult).where(JobResult.job_id == job_id)) @@ -311,17 +314,24 @@ def _upsert_job_result( if existing: existing.delivery_mode = delivery_mode - existing.document_metadata = {} + doc_meta = existing.document_metadata or {} + if chunk_dedup_stats: + doc_meta["chunk_dedup"] = chunk_dedup_stats + existing.document_metadata = doc_meta existing.inline_payload = inline_payload existing.result_s3_key = result_s3_key existing.result_size = result_size db.flush() return existing + doc_meta = {} + if chunk_dedup_stats: + doc_meta["chunk_dedup"] = chunk_dedup_stats + job_result = JobResult( job_id=job_id, delivery_mode=delivery_mode, - document_metadata={}, + document_metadata=doc_meta, inline_payload=inline_payload, result_s3_key=result_s3_key, result_size=result_size, diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index 13e7c7302..a25d45cdc 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -1,5 +1,4 @@ -from .agent_navigate import agent_navigate -from .app_service import list_lexical_chunks, merge_channels_rrf, run_retrieval_query +from .app_service import merge_channels_rrf, run_retrieval_query from .cache_service import ( bump_retrieval_namespace_cache_version, get_cached_retrieval_query_result, @@ -12,10 +11,8 @@ from .llm_adapter import create_retrieval_llm_fn __all__ = [ - "agent_navigate", "create_retrieval_llm_fn", "run_retrieval_query", - "list_lexical_chunks", "merge_channels_rrf", "DocumentGraphService", "GraphQueryService", diff --git a/packages/shared-python/shared/services/retrieval/agent_navigate.py b/packages/shared-python/shared/services/retrieval/agent_navigate.py index 6558ecaf8..ea127b69d 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,42 @@ Do not include any explanation. """ -_CHUNK_SELECT_PROMPT = """\ -You are a document chunk routing assistant. +_VALID_HYDRATE_MODES = frozenset({ + 'outline', 'chunks', 'assets_only', 'image_only', 'table_only', +}) + +_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. +Each item shows text/image/table counts. +Select section paths directly. A selected section path represents the chunks +under that section subtree; do not ask to drill deeper. -=== 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}). +If NO section path is relevant to the query, you MUST return an empty array []. Do not force-select irrelevant sections. +Prefer specific sub-items over broad parents when both are listed and the sub-item is sufficient. + +For each selected path, assign a confidence score (0.0 to 1.0) where 1.0 means exactly answers the query and 0.5 means tangentially related. +Also choose a hydrate_mode: +- "chunks" (default) return all text/image/table chunks +- "outline" return only section title + summary, no chunk content +- "assets_only" return only image and table chunks +- "image_only" return only image chunks +- "table_only" return only table chunks + +Return ONLY a JSON array: +[{{"path": "section/path", "confidence": , "hydrate_mode": "chunks"}}, ...] Do not include any explanation. """ @@ -120,7 +121,7 @@ def _parse_chunk_path_selections(text: str) -> list[dict[str, Any]]: """Parse chunk path selections from LLM output. Accepts either a legacy JSON array of strings or a structured array of - objects with `path` and optional `confidence`. + objects with `path`, optional `confidence`, and optional `hydrate_mode`. """ payload = _extract_json_array_payload(text) selections: list[dict[str, Any]] = [] @@ -128,51 +129,23 @@ def _parse_chunk_path_selections(text: str) -> list[dict[str, Any]]: if isinstance(item, str): path = item.strip() if path: - selections.append({'path': path, 'confidence': None}) + selections.append({'path': path, 'confidence': None, 'hydrate_mode': 'chunks'}) continue if not isinstance(item, dict): continue path = str(item.get('path') or item.get('chunk_path') or '').strip() if not path: continue + raw_mode = str(item.get('hydrate_mode') or '').strip().lower() + hydrate_mode = raw_mode if raw_mode in _VALID_HYDRATE_MODES else 'chunks' selections.append({ 'path': path, 'confidence': _normalize_confidence(item.get('confidence')), + 'hydrate_mode': hydrate_mode, }) return selections -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 +229,57 @@ 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. - 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 + Always shows ALL items (L1 + L2). Overflow controls whether + summaries are included — not which levels are shown. + + Normal: path + title + text=N image=I table=T + summary + Overflow: path + title + text=N image=I table=T (no summary) + + 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' text={chunk_count}' + image_count = item.get('image_count', 0) + if image_count > 0: + line += f' image={image_count}' + table_count = item.get('table_count', 0) + if table_count > 0: + line += f' table={table_count}' + if include_summary: + summary = item.get('summary') or item.get('title', '') + if summary: + 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 +422,198 @@ async def _expand_by_edges( return ordered -# ------------------------------------------------------------------ -# Main entry point -# ------------------------------------------------------------------ +# --------------------------------------------------------------------------- +# Unified scope navigation: load child sections (2-level) +# --------------------------------------------------------------------------- -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 - - 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, + exclude_paths: set[str] | None = None, ) -> list[dict]: - """Build child section view from DocumentSection table. - - Finds child sections of the given parent_path and computes chunk - counts via DocumentChunk aggregation. Runs at query time in memory - — no files created. + """Load the next 2 available section depth bands under *scope_path*. + + Returns a flat list sorted by sort_order, each item: + {path, title, summary, chunk_count, image_count, table_count, level} + + - level=1: nearest available descendant depth under scope + - level=2: second nearest available descendant depth under scope + - chunk_count: text chunks under this section (excluding image/table) + - image_count: image chunks under this section + - table_count: table chunks under this section + - exclude_paths: paths already seen in prior revision rounds; + any path matching (exact or subtree) is skipped """ - # 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 [] - - # Get child sections - children_stmt = ( + # ── 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() + _excl = exclude_paths or 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)' + # Skip paths already seen in prior revision rounds + if _excl and any( + path == ep or path.startswith(ep + ' / ') or ep.startswith(path + ' / ') + for ep in _excl + ): + continue + visible_sections.append((path, meta, relative_depth)) + visible_depths.add(relative_depth) + 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, + 'image_count': 0, + 'table_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 / image / table) ────────────────── + 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 == 'image', literal_column('1')), + ) + ).label('image_count'), + func.count( + case( + (DocumentChunk.chunk_type == 'table', literal_column('1')), + ) + ).label('table_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, int]] = { + sid: (int(tc), int(ic), int(tbc)) for sid, tc, ic, tbc 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, img_c, tbl_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['image_count'] += img_c + item['table_count'] += tbl_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..5893b416a 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. @@ -16,7 +16,6 @@ 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 +45,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 +81,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 +108,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, @@ -123,7 +130,28 @@ async def run( 'channel_weights': channel_weights, } - # ── Agent loop ── + # ── Mandatory pre-step: bottom discovery ───────────────────────────── + # bottom_discovery is always the first action; running it via the LLM + # policy wastes ~1-2s on a trivial LLM call. We execute it directly + # and let the LLM loop start from step 2 (kg_document_select etc.). + logger.info(' agentic: running mandatory bottom_discovery pre-step') + discovery_result = await self._execute_tool( + db, ActionType.BOTTOM_DISCOVERY, state, config, **tool_kwargs, + ) + state.apply(ActionType.BOTTOM_DISCOVERY, discovery_result) + if trace_enabled: + trace.record_step( + ActionType.BOTTOM_DISCOVERY, discovery_result, + decision_reason='mandatory_pre_step', + ) + state.step_count += 1 + logger.info( + f' agentic step {state.step_count} (pre-step): action=bottom_discovery ' + f'status={discovery_result.status} latency={discovery_result.latency_ms}ms' + ) + + # ── Agent loop (LLM decisions start from here) ──────────────────────── + stop_reason = 'max_steps' while state.step_count < config.max_steps: if state.elapsed_ms >= config.latency_budget_ms: logger.info( @@ -131,18 +159,36 @@ 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: discovery already ran — stop + 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 @@ -152,22 +198,132 @@ async def run( f'docs={len(state.selected_docs)} paths={len(state.selected_paths)}' ) - # ── Fixed terminal step: hydrate + rank ── - ranked_rows = await self._hydrate_and_rank( - db, state, user_id=user_id, namespace=namespace, top_k=top_k, - ) + # ── Terminal: hydrate + rank + attempt_answer loop ── + while True: + ranked_rows = await self._hydrate_and_rank( + db, state, user_id=user_id, namespace=namespace, top_k=top_k, + ) - router_used = 'agentic' if state.selected_paths else 'agentic_discovery_only' + # Include kept rows from prior revision rounds + if state.kept_path_rows: + ranked_rows = state.kept_path_rows + ranked_rows - if trace_enabled: - await trace.complete(ranked_rows, router_used) + # Check if we should attempt_answer (need LLM + results + revision budget) + if ( + policy is None + or not ranked_rows + or state.revision_count >= config.max_revisions + ): + break + + # KG-exhausted guard: if all selected docs have been explored, + # further revision won't find new content — skip attempt_answer + all_selected_ids = {d.document_id for d in state.selected_docs} + unexplored = all_selected_ids - state.ever_explored_doc_ids + kg_exhausted = len(unexplored) == 0 and len(all_selected_ids) > 0 + if kg_exhausted: + logger.info( + f' agentic: KG exhausted — all {len(all_selected_ids)} docs explored, ' + f'skipping attempt_answer' + ) + stop_reason = 'kg_exhausted' + break + + # Three-state verdict from LLM + verdict, verdict_reason = await policy.attempt_answer( + state, config, ranked_rows, + ) + logger.info( + f' agentic attempt_answer: verdict={verdict} ' + f'revision={state.revision_count}/{config.max_revisions} ' + f'reason="{verdict_reason}"' + ) + + if verdict == 'DONE': + stop_reason = 'attempt_done' + break + + if verdict in ('NOT_SUFFICIENT', 'NOT_FOUND'): + state.revision_count += 1 + # Save current results as kept rows for next round + state.kept_path_rows = ranked_rows + # Record current selected paths as seen + for p in state.selected_paths: + doc_id = p.get('document_id', '') + path = p.get('path', '') + if doc_id and path: + state.seen_section_keys.add(f'{doc_id}::{path}') + # Reset navigation state for re-exploration + state.selected_paths.clear() + state.selected_docs.clear() # Bug 1 fix: prevent doc accumulation + state.pending_doc_index = 0 + state.kg_done = False + state.discovery_done = False + + # Mandatory bottom_discovery pre-step for revision round + logger.info( + f' agentic: running mandatory bottom_discovery pre-step ' + f'(revision {state.revision_count})' + ) + rev_discovery = await self._execute_tool( + db, ActionType.BOTTOM_DISCOVERY, state, config, **tool_kwargs, + ) + state.apply(ActionType.BOTTOM_DISCOVERY, rev_discovery) + if trace_enabled: + trace.record_step( + ActionType.BOTTOM_DISCOVERY, rev_discovery, + decision_reason=f'mandatory_pre_step (revision {state.revision_count})', + ) + state.step_count += 1 + logger.info( + f' agentic step {state.step_count} (rev {state.revision_count} pre-step): ' + f'action=bottom_discovery status={rev_discovery.status}' + ) + + # Re-enter agent loop + while state.step_count < config.max_steps: + if state.elapsed_ms >= config.latency_budget_ms: + stop_reason = 'latency_budget' + break + + action_type, decision_reason = await policy.decide(state, config) + if action_type is None or action_type == ActionType.DONE: + stop_reason = 'llm_done' + break + + result = await self._execute_tool( + db, action_type, state, config, **tool_kwargs, + ) + state.apply(action_type, result) + if trace_enabled: + trace.record_step(action_type, result, decision_reason=decision_reason) + state.step_count += 1 + logger.info( + f' agentic step {state.step_count} (rev {state.revision_count}): ' + f'action={action_type.value} status={result.status}' + ) + + # Loop back to hydrate + attempt_answer + continue + + # Unknown verdict — treat as DONE + break + + router_used = ( + 'agentic_llm' if state.selected_paths or state.kept_path_rows + 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}, revisions={state.revision_count}, ' + f'{state.elapsed_ms}ms' ) + if trace_enabled: + await trace.complete(ranked_rows, router_used) + return ranked_rows, router_used async def _execute_tool( @@ -211,6 +367,13 @@ async def _execute_tool( payload={'document_id': doc.document_id, 'reason': 'no job_result_id'}, ) + # Build exclude_paths for this doc from seen_section_keys + doc_exclude = { + key.split('::', 1)[1] + for key in state.seen_section_keys + if key.startswith(f'{doc.document_id}::') + } if state.seen_section_keys else None + return await tools.document_path_select( db, user_id=kwargs['user_id'], @@ -220,6 +383,7 @@ async def _execute_tool( document_id=doc.document_id, job_result_id=job_result_id, doc_name=doc.source_file_name or state.doc_id_to_name.get(doc.document_id, ''), + exclude_paths=doc_exclude, ) elif action_type == ActionType.GREP_DOCUMENT_DISCOVER: @@ -240,54 +404,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 +423,18 @@ 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 +442,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 +461,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 +471,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..12e5afa62 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -1,50 +1,317 @@ -"""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. +# NOTE: BOTTOM_DISCOVERY is intentionally excluded from this list. +# It is now a mandatory pre-step executed automatically by the orchestrator +# before the LLM decision loop begins. The LLM should never need to decide +# whether to run it — doing so wastes one LLM call per run. +_AVAILABLE_ACTIONS: list[dict[str, Any]] = [ + { + 'action': ActionType.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. Run kg_document_select when discovery_done is true and kg_done is false. +2. After kg select, run document_path_select for each pending document. +3. Call done when all pending documents are processed OR you have >= {min_evidence} evidence paths. +4. Only use grep_document_discover if kg_document_select found 0 documents. +5. Only use graph_expand_docs if you need more related docs after reviewing results. +Note: bottom_discovery is already executed automatically before this loop — do NOT attempt to call it. + +Return ONLY a JSON object, no markdown, no explanation: +{{"action": "", "reason": ""}} +""" + + +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() - 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 ── + # 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() + + has_pending_docs = state.pending_doc_index < len(state.selected_docs) + state_data['has_pending_docs'] = has_pending_docs + state_json = json.dumps(state_data, ensure_ascii=False, indent=2) + + allowed_actions = [] + for action in _AVAILABLE_ACTIONS: + name = action['action'] + if name == ActionType.KG_DOCUMENT_SELECT.value and (not state.discovery_done or state.kg_done): + continue + if name == ActionType.DOCUMENT_PATH_SELECT.value and (not state.kg_done or not has_pending_docs): + continue + if name == ActionType.GREP_DOCUMENT_DISCOVER.value and (not state.kg_done or len(state.selected_docs) > 0): + continue + allowed_actions.append(action) + + actions_block = '\n'.join( + f" {i+1}. \"{a['action']}\": {a['description']} [{a['when']}]" + for i, a in enumerate(allowed_actions) + ) + + return _POLICY_PROMPT_TEMPLATE.format( + query=self._query, + 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 + + async def attempt_answer( + self, + state: AgentState, + config: AgentRunConfig, + ranked_rows: list[dict[str, Any]], + ) -> tuple[str, str]: + """Three-state verdict: is the evidence sufficient? + + Returns (verdict, reason) where verdict is one of: + - 'DONE': evidence is sufficient, stop searching + - 'NOT_SUFFICIENT': partial match, need more evidence + - 'NOT_FOUND': no relevant evidence found at all + """ + # Build evidence summary (paths + previews, not full content) + evidence_lines: list[str] = [] + for i, row in enumerate(ranked_rows[:20]): # cap to avoid huge prompt + path = row.get('section_path') or row.get('source_chunk_path') or '' + content_preview = str(row.get('content', ''))[:150] + score = round(float(row.get('score', 0.0) or 0.0), 3) + evidence_lines.append( + f' {i+1}. path="{path}" score={score}\n' + f' preview: {content_preview}' + ) + evidence_text = '\n'.join(evidence_lines) or '(no evidence collected)' + + prompt = _ATTEMPT_ANSWER_PROMPT.format( + query=self._query, + evidence_count=len(ranked_rows), + evidence_summary=evidence_text, + revision_count=state.revision_count, + max_revisions=config.max_revisions, + ) + + raw_response = await self._llm_fn(prompt) + logger.info(f' [LLMPolicy.attempt_answer] raw={repr(raw_response[:200])}') + + parsed = _parse_action_from_response(raw_response) + if not parsed: + return 'DONE', 'parse_error — treating as done' + + verdict = str(parsed.get('verdict', 'DONE')).strip().upper() + reason = str(parsed.get('reason', '')).strip() + + if verdict not in ('DONE', 'NOT_SUFFICIENT', 'NOT_FOUND'): + verdict = 'DONE' + + return verdict, reason + + +_ATTEMPT_ANSWER_PROMPT = """\ +You are evaluating whether the collected evidence can answer the user's query. + +QUERY: "{query}" + +EVIDENCE ({evidence_count} items, showing top 20): +{evidence_summary} + +REVISION: {revision_count} of {max_revisions} revisions used. + +Evaluate the evidence and return ONE verdict: +- "DONE": The evidence is sufficient to answer the query. Use this if the main points are covered. +- "NOT_SUFFICIENT": Partial match — some relevant info found but key aspects are missing. Only use if more searching could realistically help. +- "NOT_FOUND": The evidence is completely irrelevant to the query. Only use if nothing matches at all. + +When in doubt, prefer DONE — avoid unnecessary extra search rounds. + +Return ONLY a JSON object: +{{"verdict": "DONE", "reason": "one sentence explanation"}} +""" diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 4b3cb6676..52c7645d3 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -21,21 +21,16 @@ from shared.models.database.document import Document 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, @@ -281,112 +276,22 @@ async def document_path_select( job_result_id: str, doc_name: str = '', max_chunks_per_file: int = 15, + exclude_paths: set[str] | None = None, **_kwargs: Any, ) -> ToolResult: - """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, + exclude_paths=exclude_paths, + ) # --------------------------------------------------------------------------- @@ -513,318 +418,97 @@ 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, + exclude_paths: set[str] | None = None, ) -> 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, + items = await _load_child_sections( + db, document_id, job_result_id, scope_path, + exclude_paths=exclude_paths, ) - if not chunks_slim: + 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 confidence = item.get('confidence') if confidence is None: confidence = _default_confidence_for_rank(len(accepted)) - accepted.append({'path': path, 'confidence': confidence}) + hydrate_mode = item.get('hydrate_mode', 'chunks') + accepted.append({'path': path, 'confidence': confidence, 'hydrate_mode': hydrate_mode}) + if len(accepted) >= max_select: + break 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..4877886de 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 @@ -29,6 +29,7 @@ class AgentRunConfig: max_docs: int = 0 # 0 = no limit, LLM decides autonomously max_path_expansions: int = 2 max_doc_retries: int = 2 + max_revisions: int = 2 # max attempt_answer → revise cycles latency_budget_ms: int = 12000 min_evidence_paths: int = 1 @@ -67,7 +68,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 +81,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 @@ -94,10 +90,62 @@ class AgentState: doc_id_to_name: dict[str, str] = field(default_factory=dict) doc_job_map: dict[str, str] = field(default_factory=dict) + # Revision / three-state fields + revision_count: int = 0 + ever_explored_doc_ids: set[str] = field(default_factory=set) + seen_section_keys: set[str] = field(default_factory=set) # "{doc_id}::{section_path}" + kept_path_rows: list[dict[str, Any]] = field(default_factory=list) + @property def elapsed_ms(self) -> int: return int((time.monotonic() - self.start_time) * 1000) + 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, + 'revision_count': self.revision_count, + 'explored_doc_count': len(self.ever_explored_doc_ids), + 'kept_rows_count': len(self.kept_path_rows), + 'last_observation': last_obs, + } + def apply(self, action_type: ActionType, result: ToolResult) -> None: """Update state based on action and its result.""" self.last_observation = result @@ -105,7 +153,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,20 +174,34 @@ 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': + doc_id = result.payload.get('document_id', '') new_paths = result.payload.get('selected_paths', []) + for p in new_paths: + p['document_id'] = doc_id self.selected_paths.extend(new_paths) self.pending_doc_index += 1 - elif result.status == 'need_nav_drill': - # Large document with doc_nav available → switch to nav mode + if doc_id: + self.ever_explored_doc_ids.add(doc_id) + elif result.status == 'no_items': doc_id = result.payload.get('document_id', '') - self.nav_drill_stack.append({ - 'document_id': doc_id, - 'section_path': None, # start from top - 'depth': 0, - }) self.pending_doc_index += 1 + if doc_id: + self.ever_explored_doc_ids.add(doc_id) elif result.status == 'need_more_docs': failed_doc_id = result.payload.get('document_id', '') if failed_doc_id: @@ -147,25 +209,15 @@ def apply(self, action_type: ActionType, result: ToolResult) -> None: self.doc_retry_count += 1 self.kg_done = False # allow re-entry to KG select elif result.status == 'no_confident_match': + doc_id = result.payload.get('document_id', '') self.pending_doc_index += 1 + if doc_id: + self.ever_explored_doc_ids.add(doc_id) elif result.status == 'error': + doc_id = result.payload.get('document_id', '') self.pending_doc_index += 1 - - 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) + if doc_id: + self.ever_explored_doc_ids.add(doc_id) elif action_type == ActionType.GREP_DOCUMENT_DISCOVER: if result.status == 'discovered_docs': diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 2e5e5d247..b91e31f0c 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 @@ -252,89 +251,7 @@ async def assemble_retrieval_results( return assembled -async def list_lexical_chunks( - db: AsyncSession, - *, - user_id: str, - namespace: str, - query: str, - top_k: int, - exclude_document_ids: list[str], - exclude_sections: list[dict[str, str]], -) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - """Independent lexical retrieval: path + content channels via ILIKE.""" - recall_k = top_k * _INTERNAL_RECALL_K_MULTIPLIER - excluded_docs = set(exclude_document_ids) - - base_stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - ) - if excluded_docs: - base_stmt = base_stmt.where(Document.document_id.notin_(list(excluded_docs))) - like = f'%{query}%' - content_stmt = base_stmt.where(DocumentChunk.content_lexical_text.ilike(like)).order_by(DocumentChunk.sort_order).limit(recall_k) - path_stmt = base_stmt.where(DocumentChunk.path_lexical_text.ilike(like)).order_by(DocumentChunk.sort_order).limit(recall_k) - - # AsyncSession is stateful and should not be shared across concurrent tasks. - content_result = await db.execute(content_stmt) - path_result = await db.execute(path_stmt) - - def _to_rows(result, channel_score: float) -> list[dict[str, Any]]: - rows: list[dict[str, Any]] = [] - for document, chunk, section, job_result in result.all(): - section_path = section.section_path if section else None - if is_excluded_section(document_id=document.document_id, section_path=section_path, exclude_sections=exclude_sections): - continue - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section_path, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': channel_score, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - }) - return rows - - content_rows = _to_rows(content_result, _CHANNEL_WEIGHT_CONTENT) - path_rows = _to_rows(path_result, _CHANNEL_WEIGHT_PATH) - return content_rows, path_rows - - -def _grep_search_rows(rows: list[dict[str, Any]], query: str) -> list[dict[str, Any]]: - """Term/grep channel: exact substring matching with scoring from knowhere-kb.""" - import re - query_lower = query.lower().strip() - if not query_lower: - return [] - - units = re.findall(r'[一-鿿]+|[a-zA-Z0-9]+', query_lower) - units = [u for u in units if len(u) > 1] - - scored: list[tuple[float, dict[str, Any]]] = [] - for row in rows: - haystack = (str(row.get('content') or '') + ' ' + str(row.get('section_path') or '')).lower() - if query_lower in haystack: - scored.append((100.0, row)) - elif units: - hit_count = sum(1 for u in units if u in haystack) - if hit_count > 0: - scored.append((float(hit_count), row)) - - scored.sort(key=lambda x: x[0], reverse=True) - return [dict(row, score=score) for score, row in scored] def _merge_same_section_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -667,30 +584,50 @@ def _rank_candidates_by_path( if not candidate.get('section_path') and row.get('section_path'): candidate['section_path'] = row.get('section_path') - ranked_rows: list[dict[str, Any]] = [] + # ── Dual-priority ranking ──────────────────────────────────────────── + # When the agent produced results (routed_rows non-empty), rows with + # agent_score=0 are demoted to a fallback pool. Primary sort is by + # agent_score (LLM confidence, 0-1, cross-round comparable), with + # discovery_score as tiebreaker only. This avoids the old + # `max(agent, discovery)` which mixed incompatible score sources. + has_agent_results = len(routed_rows) > 0 + + primary_rows: list[dict[str, Any]] = [] + fallback_rows: list[dict[str, Any]] = [] + for key, row in merged.items(): discovery_score = float(row.get('discovery_score', 0.0) or 0.0) agent_score = float(row.get('agent_score', 0.0) or 0.0) row['dual_hit_flag'] = 1 if discovery_score > 0.0 and agent_score > 0.0 else 0 - row['evidence_score'] = round(max(discovery_score, agent_score), 6) + row['evidence_score'] = round(agent_score if has_agent_results else max(discovery_score, agent_score), 6) row['score'] = row['evidence_score'] row['_candidate_order'] = insertion_order[key] - ranked_rows.append(row) - ranked_rows.sort( - key=lambda row: ( - float(row.get('evidence_score', 0.0) or 0.0), + if has_agent_results and agent_score <= 0.0: + fallback_rows.append(row) + else: + primary_rows.append(row) + + def _sort_key(row): + return ( + float(row.get('agent_score', 0.0) or 0.0), + float(row.get('discovery_score', 0.0) or 0.0), int(row.get('dual_hit_flag', 0) or 0), float(row.get('importance_norm_score', 0.0) or 0.0), - float(row.get('discovery_score', 0.0) or 0.0), - float(row.get('agent_score', 0.0) or 0.0), -int(row.get('_candidate_order', 0) or 0), - ), - reverse=True, - ) + ) + + primary_rows.sort(key=_sort_key, reverse=True) + ranked_rows = primary_rows[:top_k] + + # Back-fill from fallback if primary results are insufficient + if len(ranked_rows) < top_k and fallback_rows: + fallback_rows.sort(key=_sort_key, reverse=True) + ranked_rows.extend(fallback_rows[:top_k - len(ranked_rows)]) + for row in ranked_rows: row.pop('_candidate_order', None) - return ranked_rows[:top_k] + return ranked_rows async def _count_scoped_chunks( @@ -782,77 +719,186 @@ async def _hydrate_paths_to_rows( ) -> list[dict[str, Any]]: """Load full chunk rows by section_path or source_chunk_path. - Used to hydrate agent-selected paths into the standard row format - expected by assemble_retrieval_results(). + Supports hydrate_mode branching: + - 'chunks' (default): all chunk types under the section subtree + - 'outline': synthetic row from section metadata, no real chunks + - 'assets_only': only image + table chunks + - 'image_only': only image chunks + - 'table_only': only table chunks """ if not path_selections: return [] + + # Group selections by hydrate_mode confidence_by_path: dict[str, float] = {} + mode_by_path: dict[str, str] = {} ordered_paths: list[str] = [] for item in path_selections: - 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) + hydrate_mode = str(item.get('hydrate_mode') or 'chunks').strip().lower() if path not in confidence_by_path: ordered_paths.append(path) confidence_by_path[path] = confidence + mode_by_path[path] = hydrate_mode else: confidence_by_path[path] = max(confidence_by_path[path], confidence) if not ordered_paths: return [] - stmt = ( - select(Document, DocumentChunk, DocumentSection, JobResult) - .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) - & (DocumentChunk.job_result_id == Document.current_job_result_id)) - .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) - .join(JobResult, JobResult.id == DocumentChunk.job_result_id) - .where(Document.user_id == user_id) - .where(Document.namespace == namespace) - .where(Document.status == 'active') - .where( - or_( - DocumentSection.section_path.in_(ordered_paths), - DocumentChunk.source_chunk_path.in_(ordered_paths), + # Separate outline paths from chunk-loading paths + outline_paths = [p for p in ordered_paths if mode_by_path.get(p) == 'outline'] + chunk_paths = [p for p in ordered_paths if mode_by_path.get(p) != 'outline'] + + rows: list[dict[str, Any]] = [] + + # ── Outline mode: synthesize rows from section metadata ────────────── + if outline_paths: + outline_section_filters = [] + for path in outline_paths: + outline_section_filters.append(DocumentSection.section_path == path) + + outline_stmt = ( + select(Document, DocumentSection) + .join(DocumentSection, (DocumentSection.document_id == Document.document_id) + & (DocumentSection.job_result_id == Document.current_job_result_id)) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where(or_(*outline_section_filters)) + ) + outline_result = await db.execute(outline_stmt) + for document, section in outline_result.all(): + agent_score = confidence_by_path.get(section.section_path, 0.0) + summary_text = (section.summary or '').strip() + title_text = (section.section_title or '').strip() + content = f'[Outline] {title_text}' + if summary_text: + content += f'\n{summary_text}' + rows.append({ + 'document_id': document.document_id, + 'chunk_id': f'outline_{section.section_id}', + 'section_id': section.section_id, + 'section_path': section.section_path, + 'source_file_name': document.source_file_name, + 'chunk_type': 'outline', + 'content': content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': None, + 'chunk_metadata': {}, + 'job_result_id': section.job_result_id, + 'job_id': None, + 'source_chunk_path': None, + 'sort_order': section.sort_order, + 'hydrate_mode': 'outline', + }) + + # ── Chunk modes: load real chunks with optional type filters ───────── + if chunk_paths: + section_path_filters = [] + for path in chunk_paths: + section_path_filters.append(DocumentSection.section_path == path) + section_path_filters.append(DocumentSection.section_path.like(f'{path} / %')) + + stmt = ( + select(Document, DocumentChunk, DocumentSection, JobResult) + .join(DocumentChunk, (DocumentChunk.document_id == Document.document_id) + & (DocumentChunk.job_result_id == Document.current_job_result_id)) + .outerjoin(DocumentSection, DocumentSection.section_id == DocumentChunk.section_id) + .join(JobResult, JobResult.id == DocumentChunk.job_result_id) + .where(Document.user_id == user_id) + .where(Document.namespace == namespace) + .where(Document.status == 'active') + .where( + or_( + *section_path_filters, + DocumentChunk.source_chunk_path.in_(chunk_paths), + ) ) ) - ) - result = await db.execute(stmt) + result = await db.execute(stmt) + + # Build a map of path → allowed chunk_types based on hydrate_mode + _MODE_ALLOWED_TYPES: dict[str, set[str] | None] = { + 'chunks': None, # all types + 'assets_only': {'image', 'table'}, + 'image_only': {'image'}, + 'table_only': {'table'}, + } + + seen_paths: set[str] = set() + for document, chunk, section, job_result in result.all(): + row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' + if row_path in seen_paths: + continue + + # Find which ordered path this row belongs to + matched_path = row_path + if section and section.section_path not in confidence_by_path: + matched_path = next( + ( + path for path in chunk_paths + if section.section_path == path or section.section_path.startswith(f'{path} / ') + ), + row_path, + ) + + # Check chunk_type filter based on hydrate_mode + path_mode = mode_by_path.get(matched_path, 'chunks') + allowed_types = _MODE_ALLOWED_TYPES.get(path_mode) + if allowed_types is not None: + chunk_type_lower = (chunk.chunk_type or '').strip().lower() + if chunk_type_lower not in allowed_types: + continue - # Build rows, preserving agent-selected order + seen_paths.add(row_path) + agent_score = confidence_by_path.get(matched_path, 0.0) + rows.append({ + 'document_id': document.document_id, + 'chunk_id': chunk.chunk_id, + 'section_id': chunk.section_id, + 'section_path': section.section_path if section else None, + 'source_file_name': document.source_file_name, + 'chunk_type': chunk.chunk_type, + 'content': chunk.content, + 'score': agent_score, + 'agent_score': agent_score, + 'file_path': chunk.file_path, + 'chunk_metadata': chunk.chunk_metadata or {}, + 'job_result_id': chunk.job_result_id, + 'job_id': job_result.job_id if job_result else None, + 'source_chunk_path': chunk.source_chunk_path, + 'sort_order': chunk.sort_order, + 'hydrate_mode': path_mode, + }) + + # ── Sort by agent-selected order ───────────────────────────────────── path_order = {p: idx for idx, p in enumerate(ordered_paths)} - rows: list[dict[str, Any]] = [] - seen_paths: set[str] = set() - for document, chunk, section, job_result in result.all(): - row_path = (section.section_path if section else None) or chunk.source_chunk_path or '' - if row_path in seen_paths: - continue - seen_paths.add(row_path) - agent_score = confidence_by_path.get(row_path, 0.0) - rows.append({ - 'document_id': document.document_id, - 'chunk_id': chunk.chunk_id, - 'section_id': chunk.section_id, - 'section_path': section.section_path if section else None, - 'source_file_name': document.source_file_name, - 'chunk_type': chunk.chunk_type, - 'content': chunk.content, - 'score': agent_score, - 'agent_score': agent_score, - 'file_path': chunk.file_path, - 'chunk_metadata': chunk.chunk_metadata or {}, - 'job_result_id': chunk.job_result_id, - 'job_id': job_result.job_id if job_result else None, - 'source_chunk_path': chunk.source_chunk_path, - 'sort_order': chunk.sort_order, - }) - 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) + } + # Outline paths are always resolved (synthesized) + resolved_inputs |= set(outline_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 +1168,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 4efe6ddc8..8244a8a72 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 @@ -47,7 +47,9 @@ def is_excluded_section( for item in exclude_sections: if not isinstance(item, dict): continue - if document_id == str(item.get('document_id') or '').strip() and section_path == str(item.get('section_path') or '').strip(): + exc_doc = str(item.get('document_id') or '').strip() + exc_path = str(item.get('section_path') or '').strip() + if document_id == exc_doc and (section_path == exc_path or section_path.startswith(exc_path + ' / ')): return True return False @@ -154,23 +156,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 +211,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 +232,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, }, ) ) @@ -327,11 +308,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() @@ -376,7 +370,10 @@ async def find_entry_documents( exc_path = str(exc.get('section_path') or '').strip() if exc_doc and exc_path: stmt = stmt.where( - ~((DocumentSection.document_id == exc_doc) & (DocumentSection.section_path == exc_path)) + ~((DocumentSection.document_id == exc_doc) & ( + (DocumentSection.section_path == exc_path) | + DocumentSection.section_path.like(f'{exc_path} / %') + )) ) result = await db.execute(stmt) seen = [row[0] for row in result.all()] diff --git a/packages/shared-python/shared/services/retrieval/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]: diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 0bc95b163..5322fa809 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -8,6 +8,7 @@ from __future__ import annotations +from collections import defaultdict from datetime import datetime, timezone from typing import Any, Dict, List, Optional from uuid import uuid4 @@ -35,6 +36,141 @@ def utc_now_naive() -> datetime: class RetrievalPublicationService: + + # ── Chunk-level content-hash dedup ────────────────────────────────── + # Mirrors graph_builder._dedup_chunks_by_content but operates on + # the DB document_chunks table instead of local knowledge_graph.json. + + @staticmethod + def _collect_existing_chunk_id_map( + db: Session, + *, + user_id: str, + namespace: str, + ) -> Dict[str, str]: + """Return {chunk_id -> document_id} for all active document chunks + in the given (user_id, namespace) scope. + + Only considers chunks belonging to the *current* revision of each + active document (Document.current_job_result_id == DocumentChunk.job_result_id). + """ + rows = db.execute( + select(DocumentChunk.chunk_id, DocumentChunk.document_id) + .join( + Document, + (Document.document_id == DocumentChunk.document_id) + & (Document.current_job_result_id == DocumentChunk.job_result_id), + ) + .where( + Document.user_id == user_id, + Document.namespace == namespace, + Document.status == "active", + ) + ).all() + return {row[0]: row[1] for row in rows} + + @staticmethod + def _dedup_chunks_by_content( + new_chunks: List[Dict[str, Any]], + existing_chunk_map: Dict[str, str], + ) -> tuple[List[Dict[str, Any]], Dict[str, int]]: + """Filter new_chunks: discard any whose chunk_id already exists. + + Uses the same deterministic know_id (content-hash) comparison as + graph_builder._dedup_chunks_by_content. + + Returns: + (deduped_chunks, overlap_by_document) + - deduped_chunks: chunks whose chunk_id is NOT in existing_chunk_map + - overlap_by_document: {document_id: count} of skipped chunks per + existing document (for observability logging) + """ + overlap_by_document: Dict[str, int] = defaultdict(int) + deduped: List[Dict[str, Any]] = [] + skipped = 0 + + for chunk in new_chunks: + cid = str(chunk.get("chunk_id") or chunk.get("know_id", "")) + if cid and cid in existing_chunk_map: + skipped += 1 + overlap_by_document[existing_chunk_map[cid]] += 1 + else: + deduped.append(chunk) + + if skipped > 0: + logger.warning( + f"📊 DB chunk dedup: {skipped}/{len(new_chunks)} duplicate chunks " + f"skipped (by chunk_id), {len(deduped)} new chunks to insert. " + f"Overlap by document: {dict(overlap_by_document)}" + ) + return deduped, dict(overlap_by_document) + + @classmethod + def garbage_collect_and_dedup_local_media( + cls, + db: Session, + *, + job_id: str, + user_id: str, + namespace: str, + add_dir: str, + chunks: List[Dict[str, Any]], + ) -> tuple[List[Dict[str, Any]], Dict[str, Any]]: + """ + Deduplicates chunks against the DB and physically deletes associated redundant + media files (images/tables) from the local add_dir before ZIP packaging. + Returns the deduplicated chunks. + """ + import os + + logger.info(f"[{job_id}] Starting local GC for redundant media files in namespace: {namespace}...") + try: + existing_map = cls._collect_existing_chunk_id_map( + db, user_id=user_id, namespace=namespace + ) + deduped_chunks, overlap = cls._dedup_chunks_by_content(chunks, existing_map) + + stats = { + "total_incoming": len(chunks), + "duplicates_skipped": len(chunks) - len(deduped_chunks), + "new_chunks_inserted": len(deduped_chunks), + "overlap_by_document": overlap, + } + + if len(deduped_chunks) < len(chunks): + active_paths = set() + for c in deduped_chunks: + fp = c.get("metadata", {}).get("file_path") or c.get("file_path") + if fp: + active_paths.add(fp) + + deleted_count = 0 + if add_dir and os.path.exists(add_dir): + for c in chunks: + fp = c.get("metadata", {}).get("file_path") or c.get("file_path") + if fp and fp not in active_paths: + abs_path = os.path.join(add_dir, fp) + if os.path.exists(abs_path): + os.remove(abs_path) + deleted_count += 1 + + logger.info(f"[{job_id}] GC complete: permanently removed {deleted_count} redundant local media files.") + return deduped_chunks, stats + else: + logger.info(f"[{job_id}] GC complete: no redundant chunks found.") + return chunks, stats + except Exception as e: + logger.error(f"[{job_id}] GC failed (non-fatal): {e}") + stats = { + "total_incoming": len(chunks), + "duplicates_skipped": 0, + "new_chunks_inserted": len(chunks), + "overlap_by_document": {}, + } + return chunks, stats + + # ── Public API ────────────────────────────────────────────────────── + def get_existing_document_scope( self, db: Session, @@ -65,7 +201,7 @@ def publish_document_state( job_id: str, job_result_id: str, chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, str]]: + ) -> Optional[Dict[str, Any]]: job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() if not job: logger.warning(f"Job not found for document publication: {job_id}") @@ -85,13 +221,29 @@ def _publish_document_state_for_job( job: Job, job_result_id: str, chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, str]]: + ) -> Optional[Dict[str, Any]]: - metadata = job.job_metadata or {} - namespace = metadata.get("namespace") - document_id = metadata.get("document_id") - source_file_name = metadata.get("source_file_name") or metadata.get("file_name") + job_metadata = job.job_metadata or {} + namespace = job_metadata.get("namespace") or "default" + document_id = job_metadata.get("document_id") + source_file_name = job_metadata.get("source_file_name") or job_metadata.get("file_name") + deduped_chunks = chunks + + # If ALL chunks are duplicates → skip document creation entirely + if not deduped_chunks: + logger.warning( + f"⏭️ All chunks are duplicates of existing documents. " + f"Skipping document creation for job_id={job.job_id}." + ) + return { + "user_id": str(job.user_id), + "namespace": namespace, + "document_id": None, + "skipped_all_duplicate": True, + } + + # ── Document upsert (original logic, but only for deduped chunks) ── document = None if document_id: document = db.execute( @@ -107,7 +259,7 @@ def _publish_document_state_for_job( document = Document( document_id=document_id or f"doc_{uuid4().hex[:12]}", user_id=str(job.user_id), - namespace=namespace or "default", + namespace=namespace, status="active", current_job_result_id=job_result_id, source_file_name=source_file_name, @@ -149,35 +301,42 @@ def _publish_document_state_for_job( .where(DocumentSection.job_result_id == job_result_id) ) + # ── Insert only deduped (non-duplicate) chunks ────────────────── sections_by_path: Dict[str, DocumentSection] = {} - for index, chunk in enumerate(chunks): - metadata = chunk.get("metadata") or {} - source_path = metadata.get("path") or chunk.get("path") + for index, chunk in enumerate(deduped_chunks): + chunk_metadata = chunk.get("metadata") or {} + source_path = chunk_metadata.get("path") or chunk.get("path") section_path = section_path_from_chunk_path(source_path) section = sections_by_path.get(section_path) if section is None: - parent_section_id = None path_parts = [p for p in section_path.split(" / ") if p] - if len(path_parts) > 1: - parent_path = " / ".join(path_parts[:-1]) - parent = sections_by_path.get(parent_path) - if parent is not None: - parent_section_id = parent.section_id - section = DocumentSection( - user_id=str(job.user_id), - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - parent_section_id=parent_section_id, - section_path=section_path, - section_title=path_parts[-1] if path_parts else None, - section_level=len(path_parts), - section_metadata={}, - sort_order=len(sections_by_path), - ) - db.add(section) - db.flush() - sections_by_path[section_path] = section + # Ensure all ancestor sections exist (top-down) + for depth in range(1, len(path_parts) + 1): + ancestor_path = " / ".join(path_parts[:depth]) + if ancestor_path in sections_by_path: + continue + ancestor_parent_id = None + if depth > 1: + parent_path = " / ".join(path_parts[:depth - 1]) + parent = sections_by_path.get(parent_path) + if parent is not None: + ancestor_parent_id = parent.section_id + ancestor_section = DocumentSection( + user_id=str(job.user_id), + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + parent_section_id=ancestor_parent_id, + section_path=ancestor_path, + section_title=path_parts[depth - 1], + section_level=depth, + section_metadata={}, + sort_order=len(sections_by_path), + ) + db.add(ancestor_section) + db.flush() + sections_by_path[ancestor_path] = ancestor_section + section = sections_by_path[section_path] chunk_id = chunk.get("chunk_id") or f"chunk_{uuid4().hex[:12]}" section_summary = section.summary if section else None @@ -209,8 +368,8 @@ def _publish_document_state_for_job( ), term_search_text=build_term_search_text(chunk, path_text=path_text), source_chunk_path=source_path, - file_path=metadata.get("file_path") or chunk.get("file_path"), - chunk_metadata=metadata, + file_path=chunk_metadata.get("file_path") or chunk.get("file_path"), + chunk_metadata=chunk_metadata, sort_order=chunk.get("order", index), ) ) 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 2ae93b1f3..68eb1b0fc 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -13,7 +13,13 @@ from loguru import logger from PIL import Image -from shared.utils.chunk_refs import extract_chunk_ref_spans +from shared.services.chunks.chunk_connections import ( + build_resource_target_map, + convert_refs_to_embed_connections, + merge_connections, + normalize_connect_to_targets, + parse_relationship_refs, +) from shared.utils.text_utils import truncate_content_preview import pandas as pd @@ -84,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) @@ -128,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") @@ -141,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")) @@ -201,163 +212,11 @@ def _format_chunks( table_files_map: Dict[str, Dict[str, Any]], ) -> List[Dict[str, Any]]: """Convert chunks data to ZIP specification format""" - - def safe_parse_rels(type_val): - """Safely parse in-document relationship fields from type metadata.""" - rels = [] - if type_val and isinstance(type_val, str): - if "\n" in type_val: - lines = [ - line.strip() for line in type_val.split("\n") if line.strip() - ] - rels.extend([line for line in lines[1:] if line.upper() != "PTXT"]) - return rels if rels else [] - - def build_resource_target_map() -> Dict[str, str]: - """Build ref/path -> chunk_id aliases for image and table chunks.""" - target_map: Dict[str, str] = {} - for chunk in chunks: - chunk_id = str( - chunk.get("chunk_id") or chunk.get("know_id") or "" - ).strip() - if not chunk_id: - continue - chunk_type_str = ( - str(chunk.get("type", "")).strip().split("\n", 1)[0].lower() - ) - if chunk_type_str not in {"image", "table"}: - continue - metadata = chunk.get("metadata", {}) - file_path = "" - if isinstance(metadata, dict): - file_path = str(metadata.get("file_path") or "").strip() - if not file_path: - file_info = ( - image_files_map.get(chunk_id) - if chunk_type_str == "image" - else table_files_map.get(chunk_id) - ) - if file_info: - file_path = str(file_info.get("file_path") or "").strip() - path_alias = str(chunk.get("path") or "").strip() - aliases = {file_path, path_alias} - for alias in list(aliases): - if alias: - aliases.add(f"[{alias}]") - for alias in aliases: - if alias: - target_map[alias] = chunk_id - return target_map - - def normalize_connect_to( - connects, target_map: Dict[str, str] - ) -> List[Dict[str, Any]]: - """Normalize connect_to to chunk_id-based entries.""" - if not connects: - return [] - - raw_items = connects if isinstance(connects, list) else [connects] - normalized = [] - for item in raw_items: - if not item: - continue - - if isinstance(item, dict): - target = str(item.get("target") or "").strip() - normalized_target = target_map.get(target, target) - if not normalized_target: - continue - - normalized_item = { - "target": normalized_target, - "relation": item.get("relation", "related"), - } - if "score" in item: - normalized_item["score"] = item.get("score", 1.0) - if "keywords" in item: - normalized_item["keywords"] = item.get("keywords", []) - if "ref" in item and item.get("ref"): - normalized_item["ref"] = item.get("ref") - if "position" in item and isinstance(item.get("position"), dict): - normalized_item["position"] = item.get("position") - normalized.append(normalized_item) - continue - - item_str = str(item).strip() - if not item_str: - continue - normalized_target = target_map.get(item_str, item_str) - normalized.append( - { - "target": normalized_target, - "relation": "related", - "score": 1.0, - "keywords": [], - } - ) - - return normalized - - def refs_to_embed_connections( - refs: List[Any], target_map: Dict[str, str] - ) -> List[Dict[str, Any]]: - """Convert resource refs to connect_to embeds entries.""" - normalized = [] - for ref in refs: - if isinstance(ref, dict): - ref_str = str(ref.get("ref") or "").strip() - start = ref.get("start") - end = ref.get("end") - else: - ref_str = str(ref or "").strip() - start = None - end = None - if not ref_str: - continue - target_id = target_map.get(ref_str) - if not target_id and ref_str.startswith("[") and ref_str.endswith("]"): - target_id = target_map.get(ref_str[1:-1].strip()) - if not target_id: - continue - connection: Dict[str, Any] = { - "target": target_id, - "relation": "embeds", - "ref": ref_str, - } - if isinstance(start, int) and isinstance(end, int): - connection["position"] = { - "start": start, - "end": end, - } - normalized.append(connection) - return normalized - - def merge_connections( - *connection_lists: List[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Merge connect_to entries while keeping stable order.""" - merged: List[Dict[str, Any]] = [] - seen = set() - for connection_list in connection_lists: - for item in connection_list or []: - if not isinstance(item, dict): - continue - position = item.get("position") - position_data = position if isinstance(position, dict) else {} - key = ( - str(item.get("target") or ""), - str(item.get("relation") or "related"), - str(item.get("ref") or ""), - str(position_data.get("start", "")), - str(position_data.get("end", "")), - ) - if key in seen: - continue - seen.add(key) - merged.append(item) - return merged - - resource_target_map = build_resource_target_map() + resource_target_map = build_resource_target_map( + chunks, + image_files_map=image_files_map, + table_files_map=table_files_map, + ) formatted = [] for chunk in chunks: @@ -404,16 +263,14 @@ def merge_connections( ) # Convert in-text resource refs into embeds edges. - relationship_refs = safe_parse_rels( - chunk.get("type_raw") or chunk_type_str + relationship_refs = parse_relationship_refs( + chunk.get("type_raw") or chunk_type_str, + str(content), ) - if not relationship_refs: - relationship_refs = extract_chunk_ref_spans(content) - - embed_connections = refs_to_embed_connections( + embed_connections = convert_refs_to_embed_connections( relationship_refs, resource_target_map ) - related_connections = normalize_connect_to( + related_connections = normalize_connect_to_targets( existing_metadata.get("connect_to") or chunk.get("connect_to") or chunk.get("connectto"), @@ -792,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 = { @@ -814,6 +672,7 @@ def _generate_manifest( }, }, "statistics": statistics, + "HIERARCHY": hierarchy or {}, } return manifest @@ -826,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]], diff --git a/packages/shared-python/shared/utils/text_utils.py b/packages/shared-python/shared/utils/text_utils.py index df77471d1..dd61bfdd2 100644 --- a/packages/shared-python/shared/utils/text_utils.py +++ b/packages/shared-python/shared/utils/text_utils.py @@ -20,6 +20,7 @@ class _JiebaModule(Protocol): def lcut(self, sentence: str) -> list[str]: ... + def cut(self, sentence: str) -> list[str]: ... warnings.filterwarnings( @@ -64,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. @@ -199,7 +200,10 @@ def _tokenize_english_segment(text: str) -> list[str]: def _tokenize_cjk_segment(text: str) -> list[str]: if not text.strip(): return [] - return list(_jieba.lcut(text)) + try: + return list(_jieba.lcut(text)) + except AttributeError: + return list(_jieba.cut(text)) def _resolve_retrieval_stopwords( @@ -298,7 +302,10 @@ def tokenize2stw_remove(contents: List[str], stopwords: Optional[List[str]] = No for content in contents: # Pre-clean: remove IMAGE_/TABLE_ markers and reference labels content = _CHUNK_MARKER_RE.sub('', content) - raw_tokens = _jieba.lcut(content) + try: + raw_tokens = _jieba.lcut(content) + except AttributeError: + raw_tokens = list(_jieba.cut(content)) # Filter: keep only tokens with meaningful characters (Chinese/English/numbers) tokens = [t for t in raw_tokens if _is_meaningful_token(t)] # Remove stopwords