diff --git a/.github/workflows/build-images.yml b/.github/workflows/build-images.yml index 66a5b255..4819a65a 100644 --- a/.github/workflows/build-images.yml +++ b/.github/workflows/build-images.yml @@ -546,6 +546,8 @@ jobs: API_WEBHOOK_ENDPOINT: https://api-staging.knowhereto.ai/v1/internal/s3-events SNS_TOPIC_ARN: arn:aws:sns:us-east-1:107424103509:knowhere-staging-s3-events QSTASH_CALLBACK_BASE_URL: https://api-staging.knowhereto.ai/api/v1 + API_CPU: "256" + API_MEMORY: "1024" WORKER_CPU: "2048" WORKER_MEMORY: "4096" shell: bash @@ -772,6 +774,8 @@ jobs: API_WEBHOOK_ENDPOINT: https://api.knowhereto.ai/v1/internal/s3-events SNS_TOPIC_ARN: arn:aws:sns:us-east-1:107424103509:knowhere-prod-s3-events QSTASH_CALLBACK_BASE_URL: https://api.knowhereto.ai/api/v1 + API_CPU: "512" + API_MEMORY: "2048" WORKER_CPU: "2048" WORKER_MEMORY: "4096" shell: bash diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index e5a4bad6..63941a1f 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -6,7 +6,7 @@ from uuid import uuid4 from httpx import AsyncClient -from sqlalchemy import select +from sqlalchemy import Engine, event, select from shared.models.database.document import DocumentMapUnit from shared.services.retrieval.publication_content import ( @@ -88,6 +88,91 @@ async def test_classic_route_maps_winning_unit_to_one_chunk( } +async def test_classic_route_uses_token_hash_lookup_for_frequency_query( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-token-hash-{identifier}" + statements: list[str] = [] + + def capture_frequency_query( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if ( + "FROM document_map_unit_tokens" in statement + and "frequency" in statement + and "token_hash" in statement + ): + statements.append(statement) + + event.listen(Engine, "before_cursor_execute", capture_frequency_query) + try: + async with developer_api_client_factory() as api_client: + await _publish_document( + namespace=namespace, + source_file_name="token-hash.pdf", + chunks=[ + { + "chunk_id": f"token-hash-{identifier}", + "type": "text", + "content": "token hash lookup marker", + "path": "token-hash.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"token-hash-filler-a-{identifier}", + "type": "text", + "content": "unrelated filler a", + "path": "token-hash.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"token-hash-filler-b-{identifier}", + "type": "text", + "content": "unrelated filler b", + "path": "token-hash.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + { + "chunk_id": f"token-hash-filler-c-{identifier}", + "type": "text", + "content": "unrelated filler c", + "path": "token-hash.pdf/Root/Section/c", + "order": 4, + "metadata": {}, + }, + ], + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "token hash lookup", + "top_k": 1, + "use_agentic": False, + }, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_frequency_query) + + assert response.status_code == 200 + assert statements + assert "token_hash = ANY" in statements[-1] + assert "token = ANY" not in statements[-1] + assert "matching_tokens AS MATERIALIZED" in statements[-1] + assert "FROM matching_tokens" in statements[-1] + + async def test_classic_route_image_filter_scores_only_units_with_images( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] diff --git a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py index 30994dad..a920bb39 100644 --- a/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py +++ b/apps/api/tests/contract/test_retrieval_lazy_tree_contract.py @@ -5,6 +5,9 @@ from shared.services.retrieval.nav.nav_hierarchy import NodeMeta, ProviderToolSpace from shared.services.retrieval.nav.nav_knowhere import KnowhereProvider, SectionRow from shared.services.retrieval.nav.nav_map_scores import _walk_tree +from shared.services.retrieval.nav._compat import Chunk +from shared.services.retrieval.nav.nav_compose import pack_nav_evidence +from shared.services.retrieval.nav.nav_types import NavConfig, NavState class _MetadataForbiddenProvider(KnowhereProvider): @@ -31,3 +34,72 @@ def test_tree_walk_reads_children_and_titles_without_materializing_metadata() -> assert children == {"root": ["child"], "child": []} assert leaves == {"child"} assert titles == {"root": "Root", "child": "Child"} + + +def test_evidence_pack_reads_titles_without_materializing_subtree_metadata() -> None: + class _CountingMetadataProvider(_MetadataForbiddenProvider): + metadata_calls = 0 + + def node_meta(self, section_id: str) -> NodeMeta: + self.metadata_calls += 1 + return super().node_meta(section_id) + + provider = _CountingMetadataProvider( + doc_id="doc", + sections=[ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("child", "root", "Root / Child", "Child", 1, "", 1), + ], + units=(), + ) + toolspace = ProviderToolSpace(provider) + state = NavState(doc_id="doc", query="child") + chunk = Chunk( + node_id="child", + doc_id="doc", + text="evidence", + line_ids=(1,), + section_id="child", + ) + + result = pack_nav_evidence( + [(chunk, 1.0)], + toolspace, + state, + NavConfig(), + budget_chars=100, + ) + + assert result.evidence_text == "[E1]\n[§ Child]\nevidence" + assert provider.metadata_calls == 0 + + +def test_evidence_pack_identifies_header_owners_from_parent_chain() -> None: + provider = KnowhereProvider( + doc_id="doc", + sections=[ + SectionRow("root", None, "Root", "Root", 0, "", 0), + SectionRow("parent", "root", "Root / Parent", "Parent", 1, "", 1), + SectionRow( + "child", "parent", "Root / Parent / Child", "Child", 2, "", 2 + ), + ], + units=(), + ) + toolspace = ProviderToolSpace(provider) + state = NavState(doc_id="doc", query="child") + chunks = [ + Chunk("parent", "doc", "parent evidence", (1,), "parent"), + Chunk("child", "doc", "child evidence", (2,), "child"), + ] + + result = pack_nav_evidence( + [(chunks[0], 1.0), (chunks[1], 0.9)], + toolspace, + state, + NavConfig(), + budget_chars=200, + ) + + assert result.kept_chunks == [chunks[1]] + assert result.evidence_text == "[E1]\n[§ Child]\nchild evidence" diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index e95efd23..b4b047d8 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -15,6 +15,7 @@ ) from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace from shared.services.retrieval.nav._compat import Chunk, EpisodeResult +from shared.services.retrieval.nav import nav_knowhere from shared.services.retrieval.nav.nav_map_scores import ( build_score_units, compute_corpus_map_and_unit_scores, @@ -41,6 +42,78 @@ _USER_ID = "local-dev-user" +def test_read_only_score_loader_drives_frequency_lookup_from_token_hash( + monkeypatch, +) -> None: + document_id = "doc-frequency" + job_result_id = "revision-frequency" + executions: list[tuple[str, object]] = [] + + class FakeCursor: + def __init__(self) -> None: + self.rows: list[tuple[object, ...]] = [] + + def execute(self, statement: str, parameters: object = None) -> None: + executions.append((statement, parameters)) + if "document_map_unit_indexes" in statement: + self.rows = [(document_id, job_result_id, 1, 1, 0.0, 0.0)] + elif "FROM document_map_units AS units" in statement: + self.rows = [("unit-frequency", document_id, "chunk-frequency", "section-frequency", 1, 1)] + elif "FROM document_map_unit_tokens" in statement: + self.rows = [("unit-frequency", "path", "retrieval", 1)] + else: + self.rows = [] + + def fetchall(self) -> list[tuple[object, ...]]: + return self.rows + + def close(self) -> None: + return None + + class FakeConnection: + def __init__(self) -> None: + self.cursor_instance = FakeCursor() + + def set_session(self, *, readonly: bool, autocommit: bool) -> None: + assert readonly is True + assert autocommit is True + + def cursor(self) -> FakeCursor: + return self.cursor_instance + + def close(self) -> None: + return None + + connection = FakeConnection() + monkeypatch.setattr(nav_knowhere, "_connect", lambda _dsn: connection) + store = nav_knowhere.ReadOnlyChunkStore( + dsn="postgresql://test", + revisions={document_id: job_result_id}, + ) + + corpus = store.load_persisted_score_corpus( + [document_id], + {document_id: ["section-frequency"]}, + ["retrieval"], + ) + + assert corpus is not None + frequency_executions = [ + (statement, parameters) + for statement, parameters in executions + if "FROM document_map_unit_tokens" in statement + ] + assert len(frequency_executions) == 1 + statement, parameters = frequency_executions[0] + assert "map_unit_id = ANY" in statement + assert "token_hash = ANY" in statement + assert isinstance(parameters, list) + assert parameters[0] == ["unit-frequency"] + assert parameters[1] == [ + "6e51d6a3d90b6a3243d38e6da6b3f31f49867c1360beba83da8ca9630f9672c7" + ] + + class _IncompleteIndexStore: """Minimal lazy store whose missing index returns no persisted scores.""" diff --git a/apps/worker/app/services/document_agent/structure/null_page_react.py b/apps/worker/app/services/document_agent/structure/null_page_react.py index e4a4fd16..ad022a5f 100644 --- a/apps/worker/app/services/document_agent/structure/null_page_react.py +++ b/apps/worker/app/services/document_agent/structure/null_page_react.py @@ -36,10 +36,11 @@ def react_budget() -> int: Ordered query strategy (follow this order; skip a step only if already tried or not applicable to the given title / parent path). Pattern-level only — do not invent document-specific titles: -1. Derive the search line from the given title by removing leading number / - letter / punctuation prefixes and trailing metadata qualifiers (document - identifiers/codes, revision labels, and similar). Keep the semantic title - body. Prefer that body over a metadata-only query when both are present. +1. Derive the semantic title body from the given title: strip leading + number / letter / punctuation prefixes and trailing metadata qualifiers + (document identifiers/codes, revision labels, and similar). The given title + may be a merged heading (semantic body plus a code); still reduce it to the + semantic body first. Grep that body before any metadata-only query. 2. When the parent path indicates appendices/annexes (or the title is a lettered appendix-style entry): grep the structural form "Appendix " using the letter taken from the title. Prefer this diff --git a/apps/worker/app/services/document_agent/structure/toc_anchoring.py b/apps/worker/app/services/document_agent/structure/toc_anchoring.py index 3e8b57f8..c217bc33 100644 --- a/apps/worker/app/services/document_agent/structure/toc_anchoring.py +++ b/apps/worker/app/services/document_agent/structure/toc_anchoring.py @@ -195,15 +195,16 @@ def _try_outline_anchoring_route( if (judge_result.payload or {}).get("choice") != OUTLINE_CHOICE: return False - toc_with_level: list[dict[str, Any]] = [] - for entry in kept: - row: dict[str, Any] = { + # Judge digest used full ``kept`` (incl. null pages). Anchoring / null-page + # ReAct only consume resolvable outline destinations (``paged_kept``). + toc_with_level: list[dict[str, Any]] = [ + { "heading": entry["heading"], "level": entry["level"], + "physical_page": int(entry["page"]), } - if entry.get("page") is not None: - row["physical_page"] = int(entry["page"]) - toc_with_level.append(row) + for entry in paged_kept + ] hierarchy = { "source": "pdf_outline", diff --git a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py index dc272cbc..1d877988 100644 --- a/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py +++ b/apps/worker/app/services/document_agent/tools/find_toc_anchor_pages.py @@ -83,12 +83,30 @@ def _merge_keyword_split_lines( return merged +def _match_toc_keyword_in_line(normalized_line: str) -> str | None: + """Return the longest TOC keyword contained in a normalized line.""" + if not normalized_line: + return None + return next( + ( + keyword + for keyword in sorted(TOC_KEYWORDS, key=len, reverse=True) + if keyword in normalized_line + ), + None, + ) + + def _find_toc_text_matches(lines: list[str]) -> list[dict[str, Any]]: - """Match TOC keywords only as whole lines after keyword-split repair.""" + """Match TOC keywords as line-level containment after keyword-split repair. + + Cross-line handling stays keyword-internal only (e.g. 目+录). Hit rule is + ``keyword in normalized_line`` (longest match wins), not whole-line equality. + """ matches: list[dict[str, Any]] = [] for raw_line, start_idx, end_idx in _merge_keyword_split_lines(lines): - keyword = normalize_match_text(raw_line) - if keyword not in TOC_KEYWORDS: + keyword = _match_toc_keyword_in_line(normalize_match_text(raw_line)) + if keyword is None: continue matches.append( { @@ -205,8 +223,9 @@ def _filter_recurring_elements( @register_tool( name="find.toc_anchor_pages", description=( - "Scan full PDF page text for whole-line TOC keywords, filter recurring " - "navigation elements, then render candidate PNGs for VLM confirmation." + "Scan full PDF page text for line-level TOC keywords (containment after " + "keyword-split repair), filter recurring navigation elements, then " + "render candidate PNGs for VLM confirmation." ), preconditions=(has_page_labels, has_page_full_text), ) diff --git a/apps/worker/app/services/document_agent/tools/probe_outline.py b/apps/worker/app/services/document_agent/tools/probe_outline.py index 99093ca6..0a701412 100644 --- a/apps/worker/app/services/document_agent/tools/probe_outline.py +++ b/apps/worker/app/services/document_agent/tools/probe_outline.py @@ -1,4 +1,4 @@ -"""probe.outline: read PDF bookmarks via get_toc and build a pruned tree.""" +"""probe.outline: read PDF bookmarks via get_toc and build a nested tree.""" from __future__ import annotations @@ -18,8 +18,12 @@ def _normalize_page(raw: Any) -> int | None: return page if page > 0 else None -def _flat_toc_to_forest(entries: list[list[Any]]) -> list[dict[str, Any]]: - """Convert flat ``[level, title, page]`` rows into a nested forest.""" +def build_outline_forest(entries: list[list[Any]]) -> list[dict[str, Any]]: + """Convert flat ``[level, title, page]`` rows into a nested forest. + + No-page destinations (``page <= 0``) become ``page=None`` and are retained; + printed TOC likewise keeps entries whose printed pages are out of range. + """ roots: list[dict[str, Any]] = [] stack: list[dict[str, Any]] = [] for row in entries: @@ -45,30 +49,6 @@ def _flat_toc_to_forest(entries: list[list[Any]]) -> list[dict[str, Any]]: return roots -def prune_outline_forest(nodes: list[dict[str, Any]]) -> list[dict[str, Any]]: - """Keep no-page parents when descendants have pages; drop no-page leaves/subtrees.""" - kept: list[dict[str, Any]] = [] - for node in nodes: - children = prune_outline_forest(list(node.get("children") or [])) - page = node.get("page") - if page is None and not children: - # No-page leaf, or entire no-page subtree after child prune. - continue - kept.append( - { - "title": node["title"], - "level": node["level"], - "page": page, - "children": children, - } - ) - return kept - - -def build_outline_forest(entries: list[list[Any]]) -> list[dict[str, Any]]: - return prune_outline_forest(_flat_toc_to_forest(entries)) - - def _count_nodes(nodes: list[dict[str, Any]]) -> int: total = 0 for node in nodes: @@ -79,8 +59,8 @@ def _count_nodes(nodes: list[dict[str, Any]]) -> int: @register_tool( name="probe.outline", description=( - "Read PDF bookmark outline via get_toc(simple=True) and return a pruned tree. " - "No-page parents are kept when children have pages; no-page leaves/subtrees are dropped." + "Read PDF bookmark outline via get_toc(simple=True) and return a nested " + "tree. No-page destinations are kept as page=null." ), parameters={ "type": "object", diff --git a/apps/worker/scripts/page_memory/_debug_pm_shared.py b/apps/worker/scripts/page_memory/_debug_pm_shared.py index 43c913d4..e9a7d72e 100644 --- a/apps/worker/scripts/page_memory/_debug_pm_shared.py +++ b/apps/worker/scripts/page_memory/_debug_pm_shared.py @@ -604,6 +604,9 @@ def run_stage1_toc( persist_anatomy_map(coordinator.ctx, {}) profile_path = out_dir / DOC_PROFILE_FILENAME write_debug_json(profile_path, anatomy.to_dict()) + # Stage-2 resumes outline separately: anatomy map does not carry it. + outline_roots = list(coordinator.blackboard.pdf_outline_roots or []) + persist_pdf_outline_roots(out_dir, outline_roots) # Canonical profile is at package root; drop nested duplicate. try: (out_dir / "_doc_agent" / "anatomy_map.json").unlink() @@ -620,6 +623,7 @@ def run_stage1_toc( payload={ "toc_pages": list(getattr(anatomy.toc_result, "toc_pages", []) or []), "region_count": len(list(anatomy.toc_hierarchies or [])), + "outline_root_count": len(outline_roots), "skip_toc_anchoring": True, }, ) @@ -752,6 +756,7 @@ def pipeline_state_path(out_dir: Path) -> Path: STAGE0_STATE_NAME = "stage0_state.json" PAGE_TEXT_CACHE_NAME = "page_full_text_cache.json" +PDF_OUTLINE_ROOTS_NAME = "pdf_outline_roots.json" def stage0_state_path(out_dir: Path) -> Path: @@ -762,6 +767,80 @@ def page_text_cache_path(out_dir: Path) -> Path: return out_dir / "_doc_agent" / PAGE_TEXT_CACHE_NAME +def pdf_outline_roots_path(out_dir: Path) -> Path: + """Stage-1 outline forest consumed by Stage-2 ``run_toc_anchoring``.""" + return out_dir / "_doc_agent" / PDF_OUTLINE_ROOTS_NAME + + +def persist_pdf_outline_roots(out_dir: Path, roots: list[Any] | None) -> Path: + """Write Stage-1 ``probe.outline`` forest for Stage-2 resume.""" + path = pdf_outline_roots_path(out_dir) + write_debug_json(path, list(roots or [])) + return path + + +def load_pdf_outline_roots(out_dir: Path) -> list[dict[str, Any]] | None: + """Load Stage-1 outline forest, or ``None`` when the artifact is missing.""" + path = pdf_outline_roots_path(out_dir) + if not path.exists(): + return None + data = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(data, list): + raise ValueError(f"pdf_outline_roots must be a list: {path}") + return [row for row in data if isinstance(row, dict)] + + +def restore_pdf_outline_roots_for_anchoring(coordinator, out_dir: Path) -> None: + """Put Stage-1 outline on the blackboard (re-probe + persist if missing). + + Production keeps ``probe.outline`` on the same blackboard through + ``run_toc_anchoring``. Staged debug must restore that forest explicitly. + """ + outline_roots = load_pdf_outline_roots(out_dir) + if outline_roots is not None: + coordinator.blackboard.pdf_outline_roots = outline_roots + logger.info( + " restored pdf_outline_roots from Stage-1 ({} roots)", + len(outline_roots), + ) + return + + from app.services.document_agent.registry import REGISTRY + + logger.warning(" Stage-1 pdf_outline_roots missing; re-running probe.outline") + REGISTRY.dispatch("probe.outline", coordinator.ctx, {}) + persist_pdf_outline_roots( + out_dir, list(coordinator.blackboard.pdf_outline_roots or []) + ) + + +def load_stage1_into_coordinator_for_anchoring( + coordinator, + out_dir: Path, + anatomy, +) -> None: + """Resume Stage-0 + Stage-1 TOC state for production ``run_toc_anchoring``. + + Loads page text/features from Stage-0, TOC extract outputs from Stage-1 + anatomy, restores outline roots, and clears any prior skeleton_* so + anchoring writes a fresh result. + """ + from app.services.document_agent.validators import single_shard_plan + + load_stage0_into_coordinator(coordinator, out_dir) + bb = coordinator.blackboard + page_count = int(getattr(anatomy, "page_count", None) or bb.page_count or 0) + bb.toc_result = anatomy.toc_result + bb.toc_hierarchies = list(getattr(anatomy, "toc_hierarchies", None) or []) + bb.shard_plan = getattr(anatomy, "shard_plan", None) or single_shard_plan( + page_count + ) + bb.skeleton_anchor = None + bb.skeleton_nodes = None + bb.pending_skeleton_anchors = [] + restore_pdf_outline_roots_for_anchoring(coordinator, out_dir) + + def load_pipeline_state( state_path: Path, ) -> dict[str, Any]: diff --git a/apps/worker/scripts/page_memory/debug_pm_null_page_react.py b/apps/worker/scripts/page_memory/debug_pm_null_page_react.py index 126474e0..7b3e8c16 100644 --- a/apps/worker/scripts/page_memory/debug_pm_null_page_react.py +++ b/apps/worker/scripts/page_memory/debug_pm_null_page_react.py @@ -24,7 +24,7 @@ _build_debug_coordinator, base_argparser, load_anatomy_cache, - load_stage0_into_coordinator, + load_stage1_into_coordinator_for_anchoring, page_text_cache_path, require_file, resolve_anatomy_cache_path, @@ -46,7 +46,6 @@ def main() -> int: args = parser.parse_args() from app.services.document_agent.structure.toc_anchoring import run_toc_anchoring - from app.services.document_agent.validators import single_shard_plan from shared.core.config import settings pdf_path, filename, out_dir = resolve_paths(args) @@ -56,8 +55,6 @@ def main() -> int: require_file(anatomy_cache, hint="Run Stage 1 first") anatomy = load_anatomy_cache(anatomy_cache, pdf_path, filename) - page_count = int(anatomy.page_count or 0) - hierarchies = list(getattr(anatomy, "toc_hierarchies", None) or []) logger.info("█" * 70) logger.info(" Production null-page locate dump — {}", filename) @@ -74,14 +71,8 @@ def main() -> int: model=None if args.no_vlm else args.model, settings_extra={"skip_toc_anchoring": False}, ) - load_stage0_into_coordinator(coordinator, out_dir) + load_stage1_into_coordinator_for_anchoring(coordinator, out_dir, anatomy) bb = coordinator.blackboard - bb.toc_result = anatomy.toc_result - bb.toc_hierarchies = hierarchies - bb.shard_plan = anatomy.shard_plan or single_shard_plan(page_count) - bb.skeleton_anchor = None - bb.skeleton_nodes = None - bb.pending_skeleton_anchors = [] run_toc_anchoring(coordinator.ctx) diff --git a/apps/worker/scripts/page_memory/debug_pm_stage1_hierarchy.py b/apps/worker/scripts/page_memory/debug_pm_stage1_hierarchy.py index c30c95a8..3a66cc01 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage1_hierarchy.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage1_hierarchy.py @@ -4,8 +4,8 @@ Resumes Stage-0 blackboard (``stage0_state.json`` + ``page_full_text_cache.json``, including asset-probe ``has_asset`` flags) and runs the production TOC segment: - find.toc_anchor_pages → extract.toc_with_boundaries - → persist doc_profile.json + find.toc_anchor_pages → probe.outline → extract.toc_with_boundaries + → persist doc_profile.json + ``_doc_agent/pdf_outline_roots.json`` Requires Stage 0 first: uv run python scripts/page_memory/debug_pm_stage0_bootstrap.py --file ... diff --git a/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py b/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py index 775e25b9..36a448a6 100644 --- a/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py +++ b/apps/worker/scripts/page_memory/debug_pm_stage2_calibration.py @@ -4,7 +4,7 @@ Same PROFILE anchoring path as production PAGE/TEXT: - select primary/pending → calibrate (Agent Phase-1 + Phase-2) → + restore Stage-1 outline (or re-probe) → calibrate / outline route → classify contained/parallel → graft contained → write skeleton_* Also resolves coarse skeletons (C4 resolve-only) into pipeline state so @@ -36,7 +36,7 @@ _serialize_skeletons, base_argparser, load_anatomy_cache, - load_stage0_into_coordinator, + load_stage1_into_coordinator_for_anchoring, page_text_cache_path, pipeline_state_path, record_stage, @@ -125,14 +125,8 @@ def main() -> int: "skip_toc_anchoring": False, }, ) - load_stage0_into_coordinator(coordinator, out_dir) + load_stage1_into_coordinator_for_anchoring(coordinator, out_dir, anatomy) bb = coordinator.blackboard - bb.toc_result = anatomy.toc_result - bb.toc_hierarchies = hierarchies - bb.shard_plan = anatomy.shard_plan or single_shard_plan(page_count) - bb.skeleton_anchor = None - bb.skeleton_nodes = None - bb.pending_skeleton_anchors = [] run_toc_anchoring(coordinator.ctx) diff --git a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py index b73f1aa2..c88eea99 100644 --- a/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py +++ b/apps/worker/tests/contract/test_doc_profile_anatomy_contract.py @@ -84,18 +84,20 @@ def _seed_preprobed_pages( coordinator.blackboard.global_signals["assets_probed"] = True -def test_toc_anchor_text_scan_whole_line_keyword_and_split_repair() -> None: +def test_toc_anchor_text_scan_line_contains_keyword_and_split_repair() -> None: late_lines = [f"body line {idx}" for idx in range(60)] + ["目录"] split_lines = ["Table of", "Con", "tents"] - false_positive_lines = [ + # Containment (not whole-line equality): prefixed titles and body mentions hit. + contains_lines = [ + "General table of contents", "Commentary provides guidance on minimum cement contents in different situations.", "The basic contents of a typical contract document are shown below:", ] late_matches = toc_anchor_tool._find_toc_text_matches(late_lines) # noqa: SLF001 split_matches = toc_anchor_tool._find_toc_text_matches(split_lines) # noqa: SLF001 - false_matches = toc_anchor_tool._find_toc_text_matches( # noqa: SLF001 - false_positive_lines + contains_matches = toc_anchor_tool._find_toc_text_matches( # noqa: SLF001 + contains_lines ) assert late_matches[0]["line_index"] == 60 @@ -103,7 +105,11 @@ def test_toc_anchor_text_scan_whole_line_keyword_and_split_repair() -> None: assert split_matches[0]["match_kind"] == "keyword:table of contents" assert split_matches[0]["line_index"] == 0 assert split_matches[0]["line_end_index"] == 2 - assert false_matches == [] + assert [m["match_kind"] for m in contains_matches] == [ + "keyword:table of contents", + "keyword:contents", + "keyword:contents", + ] def test_toc_extraction_raises_on_pipeline_failure(tmp_path: Path) -> None: diff --git a/apps/worker/tests/contract/test_outline_short_circuit_contract.py b/apps/worker/tests/contract/test_outline_short_circuit_contract.py index 053ee9bb..24d6c93b 100644 --- a/apps/worker/tests/contract/test_outline_short_circuit_contract.py +++ b/apps/worker/tests/contract/test_outline_short_circuit_contract.py @@ -30,7 +30,13 @@ def _outline_roots() -> list[dict[str, Any]]: "level": 2, "page": 15, "children": [], - } + }, + { + "title": "截断章节", + "level": 2, + "page": None, + "children": [], + }, ], } ] @@ -127,6 +133,8 @@ def fake_calibrate(*args: Any, **kwargs: Any) -> Any: assert calibrate_calls["count"] == 0 assert len(judge_calls) == 1 assert judge_calls[0]["toc_pages"] == [2, 3, 4, 5] + # Null outline destinations stay in the judge digest… + assert "截断章节" in str(judge_calls[0].get("outline_digest") or "") assert ctx.blackboard.toc_result is not None assert ctx.blackboard.toc_result.toc_pages == [2, 3, 4, 5] assert ctx.blackboard.toc_result.method == "pdf_outline" @@ -135,6 +143,10 @@ def fake_calibrate(*args: Any, **kwargs: Any) -> Any: assert ctx.blackboard.skeleton_anchor["offset_status"] == "ok" assert ctx.blackboard.toc_hierarchies is not None assert ctx.blackboard.toc_hierarchies[0]["source"] == "pdf_outline" + # …but are omitted from the anchoring hierarchy (no null-page ReAct targets). + anchored_rows = ctx.blackboard.toc_hierarchies[0]["toc_with_level"] + assert [row["heading"] for row in anchored_rows] == ["第一章", "第二章"] + assert all("physical_page" in row for row in anchored_rows) def test_outline_without_toc_pages_adopts_without_judge() -> None: diff --git a/apps/worker/tests/contract/test_probe_outline_contract.py b/apps/worker/tests/contract/test_probe_outline_contract.py index a956a932..e995b824 100644 --- a/apps/worker/tests/contract/test_probe_outline_contract.py +++ b/apps/worker/tests/contract/test_probe_outline_contract.py @@ -1,4 +1,4 @@ -"""Contract tests for probe.outline forest prune and physical pages.""" +"""Contract tests for probe.outline forest build and physical pages.""" from __future__ import annotations @@ -32,13 +32,17 @@ def test_outline_keeps_no_page_parent_with_paged_children() -> None: assert [child["page"] for child in forest[0]["children"]] == [10, 20] -def test_outline_drops_no_page_leaf_and_empty_subtree() -> None: +def test_outline_keeps_no_page_leaf_and_no_page_subtree() -> None: forest = build_outline_forest( [ [1, "Keep", 5], - [1, "DropLeaf", -1], - [1, "DropParent", -1], - [2, "DropChild", 0], + [1, "KeepLeaf", -1], + [1, "KeepParent", -1], + [2, "KeepChild", 0], ] ) - assert [node["title"] for node in forest] == ["Keep"] + assert [node["title"] for node in forest] == ["Keep", "KeepLeaf", "KeepParent"] + assert forest[1]["page"] is None + assert forest[2]["page"] is None + assert [child["title"] for child in forest[2]["children"]] == ["KeepChild"] + assert forest[2]["children"][0]["page"] is None diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index dd0742b4..7b7d6170 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -32,6 +32,8 @@ FRONTEND_URL=https://staging.knowhereto.ai \ API_WEBHOOK_ENDPOINT=https://api-staging.knowhereto.ai/v1/internal/s3-events \ SNS_TOPIC_ARN=arn:aws:sns:us-east-1:107424103509:knowhere-staging-s3-events \ QSTASH_CALLBACK_BASE_URL=https://api-staging.knowhereto.ai/api/v1 \ +API_CPU=256 \ +API_MEMORY=1024 \ WORKER_CPU=2048 \ WORKER_MEMORY=4096 \ python deploy/ecs/render_task_definitions.py --environment staging --output-dir /tmp/knowhere-ecs-rendered @@ -39,6 +41,10 @@ python deploy/ecs/render_task_definitions.py --environment staging --output-dir The output directory is deployment-only and must not be committed. The renderer fails on missing inputs, unresolved placeholders, or either long-lived S3 credential variable. +The staging workflow keeps the API at 256 CPU / 1024 MiB. The production release +workflow sets the API to 512 CPU / 2048 MiB (2 GiB) and the worker to 2048 CPU / +4096 MiB, so production API memory remains at 2 GiB across future releases. + ## Staging workflow prerequisites The staging workflow in `.github/workflows/build-images.yml` expects these GitHub Actions secrets: diff --git a/deploy/ecs/render_task_definitions.py b/deploy/ecs/render_task_definitions.py index 458cb8a7..33e2b9e5 100644 --- a/deploy/ecs/render_task_definitions.py +++ b/deploy/ecs/render_task_definitions.py @@ -32,6 +32,8 @@ "API_WEBHOOK_ENDPOINT", "SNS_TOPIC_ARN", "QSTASH_CALLBACK_BASE_URL", + "API_CPU", + "API_MEMORY", "WORKER_CPU", "WORKER_MEMORY", ) diff --git a/deploy/ecs/task-definition-api.staging.json b/deploy/ecs/task-definition-api.staging.json index 0e1013c9..0b4af2c5 100644 --- a/deploy/ecs/task-definition-api.staging.json +++ b/deploy/ecs/task-definition-api.staging.json @@ -4,8 +4,8 @@ "executionRoleArn": "${EXECUTION_ROLE_ARN}", "networkMode": "awsvpc", "requiresCompatibilities": ["FARGATE"], - "cpu": "256", - "memory": "1024", + "cpu": "${API_CPU}", + "memory": "${API_MEMORY}", "runtimePlatform": { "cpuArchitecture": "X86_64", "operatingSystemFamily": "LINUX" diff --git a/deploy/ecs/test_render_task_definitions.py b/deploy/ecs/test_render_task_definitions.py index b9c1f569..ff7efbee 100644 --- a/deploy/ecs/test_render_task_definitions.py +++ b/deploy/ecs/test_render_task_definitions.py @@ -33,6 +33,8 @@ "API_WEBHOOK_ENDPOINT": "https://api-staging.knowhereto.ai/v1/internal/s3-events", "SNS_TOPIC_ARN": "arn:aws:sns:us-east-1:107424103509:knowhere-staging-s3-events", "QSTASH_CALLBACK_BASE_URL": "https://api-staging.knowhereto.ai/api/v1", + "API_CPU": "256", + "API_MEMORY": "1024", "WORKER_CPU": "2048", "WORKER_MEMORY": "4096", } @@ -155,6 +157,20 @@ def test_staging_worker_preserves_evidence_selected_capacity(tmp_path: Path) -> assert definition["memory"] == "4096" +def test_staging_api_preserves_selected_capacity(tmp_path: Path) -> None: + """API capacity remains at the verified staging size.""" + output_path: Path = tmp_path / "task-definition-api.staging.json" + render_template( + TEMPLATE_DIRECTORY / "task-definition-api.staging.json", + output_path, + RENDER_VARIABLES, + ) + definition: dict[str, object] = json.loads(output_path.read_text(encoding="utf-8")) + + assert definition["cpu"] == "256" + assert definition["memory"] == "1024" + + def test_renderer_rejects_forbidden_s3_credential_variable() -> None: """Task definitions must never inject long-lived S3 credentials.""" definition: dict[str, object] = { diff --git a/packages/shared-python/shared/services/retrieval/nav/__init__.py b/packages/shared-python/shared/services/retrieval/nav/__init__.py index 8f294b4b..ca9233ee 100644 --- a/packages/shared-python/shared/services/retrieval/nav/__init__.py +++ b/packages/shared-python/shared/services/retrieval/nav/__init__.py @@ -4,7 +4,6 @@ ActionKind, NavConfig, NavState, - RegionReport, SubgoalResult, map_mode_enabled, ) @@ -22,7 +21,6 @@ "ActionKind", "NavConfig", "NavState", - "RegionReport", "SubgoalResult", "map_mode_enabled", "run_nav_episode", diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_actions.py b/packages/shared-python/shared/services/retrieval/nav/nav_actions.py index 1c50e245..dc756622 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_actions.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_actions.py @@ -103,7 +103,6 @@ def view_score(view: SectionView) -> float: score=score, metadata={ "map_id": view.map_id, - "n_chunks": view.n_chunks, "highlight": is_hit, "multi": True, }, @@ -186,9 +185,8 @@ def node_actions(sid: str) -> str: hit_tag = format_hit_tag(is_highlight=bool(view.is_highlight)) harvested_tag = format_harvested_tag(getattr(view, "harvested_by", "") or "") map_id = view.map_id or "?" - meta = f"({view.n_chunks} chunks)" lines.append( - f"{indent}[{map_id}] {view.title or view.section_id} {meta}" + f"{indent}[{map_id}] {view.title or view.section_id}" f"{leaf_tag}{hit_tag}{harvested_tag} actions: {node_actions(view.section_id)}" ) if inline_summary and view.summary: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_address.py b/packages/shared-python/shared/services/retrieval/nav/nav_address.py index 39ece29b..0e13d8b4 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_address.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_address.py @@ -138,12 +138,11 @@ def next_dispatch_depth( child_id: str, depth: int, ) -> int: - """Child navigate/harvest depth after DISPATCH into ``child_id``. + """Child harvest depth after DISPATCH into ``child_id``. Namespace → document is depth-neutral (child starts at 0). Namespace → section under a document starts at 1. Once already inside a document or - section scope, always ``depth + 1``. Shared by ``dispatch()`` and harvest - recursion so the two paths agree. + section scope, always ``depth + 1``. """ parent_sid = str(parent_scope or "").strip() if parent_sid: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py index 197af5af..7df6eca3 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py @@ -23,12 +23,6 @@ pack_nav_evidence, unit_score_for_evidence_chunk, ) -from .nav_map_scores import ( - compute_corpus_map_and_unit_scores, - compute_map_and_unit_scores, - select_map_highlights, -) -from .nav_navigate import navigate from .nav_types import ( LegalAction, NavConfig, @@ -379,10 +373,6 @@ def _run_nav_episode_body( mult = float(getattr(cfg, "scope_inline_summary_budget_mult", 0.0) or 0.0) if mult > 0.0 and int(budget_chars) > 0: cfg.scope_inline_summary_char_limit = max(1, int(budget_chars * mult)) - # Depth-0 oversize→DISPATCH threshold defaults to the evidence budget. - if bool(getattr(cfg, "enable_depth0_oversize_to_dispatch", False)): - if int(getattr(cfg, "depth0_oversize_char_limit", 0) or 0) <= 0: - cfg.depth0_oversize_char_limit = max(1, int(budget_chars)) retrieval_t0 = time.perf_counter() if toolspace is not None: ts = toolspace @@ -400,88 +390,63 @@ def _run_nav_episode_body( state = NavState(doc_id=episode_doc, query=query, task_type=task_type) steps: List[AgentStep] = [] - map_started = time.perf_counter() if namespace_mode: section_ids = list(ts.sections_for_doc("")) - state.map_scores, state.unit_scores = compute_corpus_map_and_unit_scores( - ts, doc_ids=corpus_ids, query=query - ) else: section_ids = ts.sections_for_doc(episode_doc) - state.map_scores, state.unit_scores = compute_map_and_unit_scores( - ts, doc_id=episode_doc, query=query, root_ids=section_ids + + from .nav_plan import plan_query + + # Planner is query-only: no pre-lit map. Score/light after the plan exists. + plan_t0 = time.perf_counter() + retrieval_plan = plan_query(state, cfg) + state.retrieval_plan = retrieval_plan + steps.append( + AgentStep( + step_idx=len(steps) + 1, + action="query_plan", + detail=stamp_step_detail({ + "fallback": bool(retrieval_plan.fallback), + "n_subgoals": len(retrieval_plan.subgoals), + "reason": retrieval_plan.reason, + "plan": retrieval_plan.to_dict(), + "seconds": time.perf_counter() - plan_t0, + }, t0=plan_t0), ) - _logger.info( - "retrieval mapnav phase=map_scoring seconds=%.3f documents=%d sections=%d", - time.perf_counter() - map_started, - len(corpus_ids), - len(section_ids), ) - state.highlight_ids = select_map_highlights( - state.unit_scores, k=int(cfg.collect_top_k) + _logger.info( + "retrieval mapnav phase=planner seconds=%.3f subgoals=%d", + time.perf_counter() - plan_t0, + len(retrieval_plan.subgoals), ) - if cfg.is_checklist: - from .nav_plan import plan_query + # Checklist: wave orchestration. Every episode runs plan + harvest + control. + if state.retrieval_plan is None: + from .nav_plan import fallback_plan - plan_t0 = time.perf_counter() - retrieval_plan = plan_query(ts, state, cfg) - state.retrieval_plan = retrieval_plan - steps.append( - AgentStep( - step_idx=len(steps) + 1, - action="query_plan", - detail=stamp_step_detail({ - "fallback": bool(retrieval_plan.fallback), - "n_subgoals": len(retrieval_plan.subgoals), - "reason": retrieval_plan.reason, - "plan": retrieval_plan.to_dict(), - "planning_map_char_limit": int( - getattr(cfg, "planning_map_char_limit", 0) or cfg.map_char_limit - ), - "seconds": time.perf_counter() - plan_t0, - }, t0=plan_t0), - ) - ) - _logger.info( - "retrieval mapnav phase=planner seconds=%.3f subgoals=%d", - time.perf_counter() - plan_t0, - len(retrieval_plan.subgoals), - ) + state.retrieval_plan = fallback_plan(query, reason="missing_plan") + from .nav_orchestrate import execute_plan, seed_episode_map_scores_from_plan - # Checklist: wave orchestration; navigate mode: classic single navigate. - if cfg.is_checklist and state.retrieval_plan is not None: - from .nav_orchestrate import execute_plan + seed_episode_map_scores_from_plan(ts, state, cfg, state.retrieval_plan) - orch_t0 = time.perf_counter() - orch_detail = execute_plan( - ts, state, cfg, steps_out=steps, episode_query=query - ) - orch_detail = dict(orch_detail or {}) - orch_detail["seconds"] = time.perf_counter() - orch_t0 - steps.append( - AgentStep( - step_idx=len(steps) + 1, - action="plan_orchestrate", - detail=stamp_step_detail(orch_detail, t0=orch_t0), - ) - ) - _logger.info( - "retrieval mapnav phase=orchestration seconds=%.3f waves=%d", - time.perf_counter() - orch_t0, - len(orch_detail.get("waves", [])), - ) - else: - navigate( - ts, - state=state, - scope=None, - query=query, - config=cfg, - depth=0, - budget=int(cfg.map_char_limit), - steps_out=steps, + orch_t0 = time.perf_counter() + orch_detail = execute_plan( + ts, state, cfg, steps_out=steps, episode_query=query + ) + orch_detail = dict(orch_detail or {}) + orch_detail["seconds"] = time.perf_counter() - orch_t0 + steps.append( + AgentStep( + step_idx=len(steps) + 1, + action="plan_orchestrate", + detail=stamp_step_detail(orch_detail, t0=orch_t0), ) + ) + _logger.info( + "retrieval mapnav phase=orchestration seconds=%.3f waves=%d", + time.perf_counter() - orch_t0, + len(orch_detail.get("waves", [])), + ) evidence_started = time.perf_counter() fill = pack_nav_evidence( diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_compose.py b/packages/shared-python/shared/services/retrieval/nav/nav_compose.py index 696328b4..3951a70c 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_compose.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_compose.py @@ -111,6 +111,18 @@ def _clip(text: str) -> str: return t[:max_chars].rstrip() return t + # ``path_titles`` is a lightweight title lookup. Prefer it over + # ``get_structure`` because the latter may calculate the complete subtree + # chunk count, which is unnecessary while rendering evidence headers. + path_titles = getattr(ts, "path_titles", None) + if callable(path_titles): + try: + path = str(path_titles(sid, doc_id) or "").strip() + except TypeError: + path = str(path_titles(sid) or "").strip() + if path: + return _clip(path.rsplit(" / ", 1)[-1]) + # Prefer structure title (Knowhere / ProviderToolSpace); never parse ids. try: st = ts.get_structure(sid) @@ -172,25 +184,43 @@ def _child_final_score( return float(own_unit) + float(w_conf) * conf -def _is_header_only_owner(owner: str, owners: set[str], ts: ToolSpace, doc_id: str) -> bool: - """True if owner is a structural ancestor of another collected owner.""" - if not owner or owner not in owners: - return False - owner_doc = _section_doc_id(ts, owner, doc_id) - for other in owners: - if other == owner: - continue - if _section_doc_id(ts, other, doc_id) != owner_doc: +def _parent_id_for_compose(ts: ToolSpace, section_id: str, doc_id: str) -> Optional[str]: + """Resolve a direct parent without materializing a section subtree.""" + parent_fn = getattr(ts, "parent_id", None) + if callable(parent_fn): + try: + parent = parent_fn(section_id) + except Exception: + parent = None + if parent: + return str(parent) + return direct_parent_id(ts, section_id, doc_id) + + +def _header_only_owners( + owners: set[str], ts: ToolSpace, doc_id: str +) -> set[str]: + """Find collected owners that are ancestors of another collected owner. + + Walk each owner's parent chain once instead of comparing every owner with + every other owner. This keeps compose grouping proportional to the number + of owners and hierarchy depth, even when a broad collect returns thousands + of chunks. + """ + headers: set[str] = set() + for owner in owners: + if not owner: continue - cur = other + owner_doc = _section_doc_id(ts, owner, doc_id) + current = owner for _ in range(64): - p = direct_parent_id(ts, cur, owner_doc) - if p is None: + parent = _parent_id_for_compose(ts, current, owner_doc) + if parent is None: break - if p == owner: - return True - cur = p - return False + if parent in owners and _section_doc_id(ts, parent, doc_id) == owner_doc: + headers.add(parent) + current = parent + return headers @dataclass @@ -259,9 +289,7 @@ def _build_groups( items.append((chunk, owner, score)) owners = {owner for _c, owner, _s in items if owner} - header_owners = { - o for o in owners if _is_header_only_owner(o, owners, ts, state.doc_id) - } + header_owners = _header_only_owners(owners, ts, state.doc_id) groups: Dict[Optional[str], _ParentGroup] = {} for chunk, owner, score in items: @@ -292,65 +320,6 @@ def _build_groups( return list(groups.values()) -def _group_summary_text(ts: ToolSpace, section_id: str) -> str: - """Prebuilt section summary for a compose group header (no body heuristics).""" - sid = str(section_id or "").strip() - if not sid: - return "" - try: - from section_summary_store import get_summary - - doc = _section_doc_id(ts, sid, "") - text = str(get_summary(sid, doc_id=doc) or "").strip() - if text: - return text - except Exception: - pass - try: - st = ts.get_structure(sid) or {} - text = str(st.get("summary") or "").strip() - if text: - return text - return str(st.get("preview") or "").strip() - except Exception: - return "" - - -def build_compose_preview( - collected: Sequence[Tuple[Chunk, float]], - ts: ToolSpace, - state: NavState, - config: NavConfig, -) -> Tuple[str, Dict[str, str]]: - """Assemble current pool into [G*] parent-group preview for group_rank. - - One line per group: ``[Gi] §title`` plus that group's header-node summary. - No child expansion, unit scores, or char counts. If the preview exceeds - ``compose_group_rank_max_chars``, return empty so callers skip group_rank. - """ - scored = dedupe_scored(list(collected)) - groups = _build_groups(scored, ts, state, config) - groups.sort(key=lambda g: (-g.group_key, g.doc_order_key)) - - lines: List[str] = [] - g_map: Dict[str, str] = {} - for i, g in enumerate(groups, 1): - gid = f"G{i}" - if g.parent_id: - g_map[gid] = str(g.parent_id) - title = g.parent_title or str(g.parent_id or "").strip() or gid - lines.append(f"[{gid}] §{title}") - summary = _group_summary_text(ts, str(g.parent_id or "")) - if summary: - lines.append(f" {summary}") - - text = "\n".join(lines) - max_chars = max(0, int(getattr(config, "compose_group_rank_max_chars", 10000) or 0)) - if max_chars > 0 and len(text) > max_chars: - return "", {} - return text, g_map - - def _render_group( group: _ParentGroup, selected: Sequence[_ChildItem], @@ -466,27 +435,180 @@ def _render_kept( ) +def _count_power10_boundaries(lower: int, upper: int) -> int: + """Count decimal digit-boundary positions in ``(lower, upper]``.""" + if upper <= lower: + return 0 + boundary = 10 + count = 0 + while boundary <= upper: + if boundary > lower: + count += 1 + boundary *= 10 + return count + + +class _SelectionLength: + """Maintain exact rendered selection length under one-child mutations.""" + + def __init__(self, groups: Sequence[_ParentGroup], kept_ids: Set[str]) -> None: + self._groups = groups + self.kept_ids = kept_ids + self._group_by_node_id: Dict[str, int] = {} + self._plain_lengths: Dict[str, int] = {} + self._indented_lengths: Dict[str, int] = {} + self._counts: List[int] = [] + self._plain_totals: List[int] = [] + self._indented_totals: List[int] = [] + self._nonempty_count = 0 + self.total = 0 + + for group_index, group in enumerate(groups): + plain_total = 0 + indented_total = 0 + count = 0 + for child in group.children: + node_id = str(child.chunk.node_id or "") + if not node_id: + continue + body = _chunk_body(child.chunk) + plain_length = len(body) + indented_length = len( + "\n".join( + (" " + line if line.strip() else line) + for line in body.splitlines() + ) + ) + self._group_by_node_id[node_id] = group_index + self._plain_lengths[node_id] = plain_length + self._indented_lengths[node_id] = indented_length + if node_id in kept_ids: + count += 1 + plain_total += plain_length + indented_total += indented_length + self._counts.append(count) + self._plain_totals.append(plain_total) + self._indented_totals.append(indented_total) + + self._recalculate_total() + + @staticmethod + def _block_length( + group: _ParentGroup, + *, + count: int, + plain_total: int, + indented_total: int, + evidence_index: int, + ) -> int: + if count <= 0: + return 0 + prefix = len(f"[E{evidence_index}]") + if group.parent_title: + prefix += 1 + len(f"[§ {group.parent_title}]") + body_length = indented_total if count >= 2 else plain_total + return prefix + count + body_length + + def _recalculate_total(self) -> None: + total = 0 + evidence_index = 0 + for index, group in enumerate(self._groups): + count = self._counts[index] + if count <= 0: + continue + evidence_index += 1 + total += self._block_length( + group, + count=count, + plain_total=self._plain_totals[index], + indented_total=self._indented_totals[index], + evidence_index=evidence_index, + ) + self._nonempty_count = evidence_index + self.total = total + max(0, evidence_index - 1) * 2 + + def _evidence_index(self, group_index: int) -> int: + return sum(1 for count in self._counts[:group_index] if count > 0) + 1 + + def remove(self, node_id: str) -> None: + group_index = self._group_by_node_id.get(node_id) + if group_index is None or node_id not in self.kept_ids: + return + old_count = self._counts[group_index] + old_index = self._evidence_index(group_index) + old_block = self._block_length( + self._groups[group_index], + count=old_count, + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=old_index, + ) + self.kept_ids.remove(node_id) + self._counts[group_index] -= 1 + self._plain_totals[group_index] -= self._plain_lengths[node_id] + self._indented_totals[group_index] -= self._indented_lengths[node_id] + if old_count > 1: + new_block = self._block_length( + self._groups[group_index], + count=old_count - 1, + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=old_index, + ) + self.total += new_block - old_block + return + + self.total -= old_block + if self._nonempty_count > 1: + self.total -= 2 + self.total -= _count_power10_boundaries(old_index, self._nonempty_count) + self._nonempty_count -= 1 + + def add(self, node_id: str) -> bool: + group_index = self._group_by_node_id.get(node_id) + if group_index is None or node_id in self.kept_ids: + return False + old_count = self._counts[group_index] + evidence_index = self._evidence_index(group_index) + if old_count > 0: + old_block = self._block_length( + self._groups[group_index], + count=old_count, + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=evidence_index, + ) + + self.kept_ids.add(node_id) + self._counts[group_index] += 1 + self._plain_totals[group_index] += self._plain_lengths[node_id] + self._indented_totals[group_index] += self._indented_lengths[node_id] + new_block = self._block_length( + self._groups[group_index], + count=self._counts[group_index], + plain_total=self._plain_totals[group_index], + indented_total=self._indented_totals[group_index], + evidence_index=evidence_index, + ) + if old_count > 0: + self.total += new_block - old_block + else: + old_nonempty_count = self._nonempty_count + self.total += new_block + if old_nonempty_count > 0: + self.total += 2 + self.total += _count_power10_boundaries( + evidence_index, old_nonempty_count + 1 + ) + self._nonempty_count += 1 + return True + + def _selection_chars( groups: Sequence[_ParentGroup], kept_ids: Set[str] ) -> int: - """Exact rendered length of the current full-text selection (no truncation).""" - parts: List[str] = [] - sep = "\n\n" - for g in groups: - entries = [ - c - for c in sorted(g.children, key=lambda x: x.line_key) - if c.chunk.node_id in kept_ids - ] - if not entries: - continue - block = _render_group( - g, entries, evidence_index=len(parts) + 1, indent=len(entries) >= 2 - ) - parts.append(block) - if not parts: - return 0 - return len(sep.join(parts)) + """Exact rendered length of the current full-text selection.""" + return _SelectionLength(groups, kept_ids).total def _refill_to_budget( @@ -519,18 +641,16 @@ def _refill_to_budget( c.line_key, ) ) - used = _selection_chars(groups, kept_ids) + selection = _SelectionLength(groups, kept_ids) for child in dropped: - if used >= budget_chars: + if selection.total >= budget_chars: break - if len(child.chunk.text or "") > budget_chars - used: + if len(child.chunk.text or "") > budget_chars - selection.total: continue - kept_ids.add(child.chunk.node_id) - size = _selection_chars(groups, kept_ids) - if size > budget_chars: - kept_ids.discard(child.chunk.node_id) + if not selection.add(child.chunk.node_id): continue - used = size + if selection.total > budget_chars: + selection.remove(child.chunk.node_id) def _pack_trim( @@ -556,8 +676,10 @@ def _pack_trim( if not kept_ids: return ComposeFillResult([], "", 0, 0, False, scored_flat, dropped_any=False) + selection = _SelectionLength(groups, kept_ids) + def fits() -> bool: - return _selection_chars(groups, kept_ids) <= budget_chars + return selection.total <= budget_chars if fits(): return _render_kept( @@ -575,7 +697,7 @@ def fits() -> bool: for child in candidates: if fits(): break - kept_ids.discard(child.chunk.node_id) + selection.remove(child.chunk.node_id) if fits(): break @@ -591,7 +713,7 @@ def fits() -> bool: break if len(kept_ids) <= 1: break - kept_ids.discard(child.chunk.node_id) + selection.remove(child.chunk.node_id) if fits() or len(kept_ids) <= 1: break diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_harvest.py b/packages/shared-python/shared/services/retrieval/nav/nav_harvest.py index 22ab1de7..dbd27638 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_harvest.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_harvest.py @@ -14,16 +14,16 @@ toward DISPATCH — no separate "scope overflow" special case is needed here. Depends only on the existing map/action primitives (``nav_projection``, -``nav_actions``) plus ``nav_navigate._apply_collect`` for hydration — the same -kernel surface ``navigate()`` already uses. No new ToolSpace capability is -required beyond the 5 documented in docs/audit_plan_nav_overlap.md. +``nav_actions``) plus ``nav_navigate._apply_collect`` for hydration. No new +ToolSpace capability is required beyond the 5 documented in +docs/audit_plan_nav_overlap.md. """ from __future__ import annotations import re from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple from .nav_actions import build_legal_actions, format_actionable_map_observation from .nav_compose import parse_collect_confidence @@ -263,6 +263,7 @@ def harvest( entry_scope: Optional[str], query: str, steps_out: Optional[List[Any]] = None, + allowed_section_ids: Optional[Set[str]] = None, ) -> HarvestResult: """Single-decision-per-node evidence harvest for one subgoal.""" result = HarvestResult(subgoal_id=subgoal.id) @@ -287,6 +288,7 @@ def harvest( depth=initial_depth, steps_out=steps_out, result=result, + allowed_section_ids=allowed_section_ids, ) return result @@ -302,6 +304,7 @@ def _harvest_node( depth: int, steps_out: Optional[List[Any]], result: HarvestResult, + allowed_section_ids: Optional[Set[str]] = None, ) -> None: from .nav_navigate import _apply_collect # late import avoids cycle from .nav_token_budget import stamp_step_detail @@ -320,6 +323,7 @@ def _harvest_node( dismissed_section_ids=state.dismissed_section_ids | subgoal_dismissed, highlight_ids=state.highlight_ids, harvested_section_ids=state.harvested_owner_subgoal if show_harvested else None, + allowed_section_ids=allowed_section_ids, ) actions = build_legal_actions( state, projection, step_idx=0, config=config, depth=depth, ts=ts @@ -467,6 +471,7 @@ def _harvest_node( depth=child_depth, steps_out=steps_out, result=result, + allowed_section_ids=allowed_section_ids, ) # Whole dispatched subtree came up empty: a dead end, not just "not # yet explored" — dismiss it too so a later widen doesn't re-dispatch diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py index 8c2a3f4f..80b9cde0 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_hierarchy.py @@ -45,7 +45,6 @@ class NodeMeta: title: str = "" summary: str = "" has_children: bool = False - n_chunks: int = 0 @runtime_checkable @@ -61,7 +60,7 @@ def children(self, section_id: str) -> Sequence[str]: ... def node_meta(self, section_id: str) -> NodeMeta: - """Title/summary/chunk-count/has_children for one node.""" + """Title/summary/has_children for one node.""" ... def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: @@ -134,7 +133,6 @@ def get_structure(self, section_id: str) -> dict: "preview": meta.title, "summary": str(meta.summary or ""), "n_lines": 1, - "n_chunks": int(meta.n_chunks), "children": [ {"section_id": cid, "preview": self._provider.node_meta(cid).title} for cid in child_ids @@ -332,7 +330,6 @@ def node_meta(self, section_id: str) -> NodeMeta: title=node.title, summary=self._summaries.get(section_id, ""), has_children=bool(node.children), - n_chunks=1, ) def parent_id(self, section_id: str) -> Optional[str]: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 38cef224..bd4e24f7 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -24,6 +24,7 @@ import logging import time from dataclasses import dataclass, field +from hashlib import sha256 from typing import ( Any, Callable, @@ -323,6 +324,9 @@ def load_persisted_score_corpus( for token in tokenize_query_for_ranker(query) ) ) + query_token_hashes = [ + sha256(token.encode("utf-8")).hexdigest() for token in query_tokens + ] cur = self._connection().cursor() try: revision_key = tuple(revisions) @@ -431,20 +435,21 @@ def load_persisted_score_corpus( frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} if unit_rows and query_tokens: stage_started = time.perf_counter() + # Restrict the token scan to this episode's map units instead of + # matching token_hash across the whole table then filtering. + allowed_map_unit_ids = [str(row["map_unit_id"]) for row in unit_rows] cur.execute( - "SELECT units.id, tokens.channel, tokens.token, tokens.frequency " - "FROM document_map_unit_tokens AS tokens " - "JOIN document_map_units AS units ON units.id = tokens.map_unit_id " - f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " - "ON units.document_id = revisions.document_id " - "AND units.job_result_id = revisions.job_result_id " - "WHERE tokens.token = ANY(%s) AND tokens.channel = ANY(%s)", - [*revision_params, list(query_tokens), list(_MAP_SCORE_CHANNELS)], + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE map_unit_id = ANY(%s) " + "AND token_hash = ANY(%s) AND channel = ANY(%s)", + [ + allowed_map_unit_ids, + list(query_token_hashes), + list(_MAP_SCORE_CHANNELS), + ], ) - allowed_map_unit_ids = {str(row["map_unit_id"]) for row in unit_rows} for map_unit_id, channel, token, frequency in cur.fetchall(): - if str(map_unit_id) not in allowed_map_unit_ids: - continue frequencies.setdefault((str(map_unit_id), str(channel)), {})[ str(token) ] = int(frequency) @@ -492,6 +497,11 @@ def load_persisted_score_corpus( ), ) for row in unit_rows + # Only units with at least one query-token frequency can + # score > 0; zero-frequency units are implicit 0 and are + # never materialized for BM25. + if frequencies.get((str(row["map_unit_id"]), "path")) + or frequencies.get((str(row["map_unit_id"]), "content")) ], path_stats=path_stats, content_stats=content_stats, @@ -698,7 +708,6 @@ def node_meta(self, section_id: str) -> NodeMeta: title=row.section_title, summary=row.summary, has_children=bool(self._children.get(section_id)), - n_chunks=len(self.subtree_units(section_id)), ) def relations(self, section_id: str) -> Tuple[Set[str], Set[str]]: @@ -789,9 +798,6 @@ def summaries(self) -> Dict[str, str]: def all_section_ids(self) -> List[str]: return list(self._sections) - def chunk_count(self) -> int: - return len(self._chunk_ids) - class LazyKnowhereProvider(KnowhereProvider): """Hierarchy provider that loads full chunk rows only on first access.""" @@ -1103,19 +1109,10 @@ def node_meta(self, section_id: str) -> NodeMeta: sid = str(section_id or "").strip() if sid in self._docs: provider = self._docs[sid] - count_fn = getattr(provider, "chunk_count", None) - n_chunks = ( - int(count_fn()) - if callable(count_fn) - else sum( - len(provider.self_units(sec)) for sec in provider.all_section_ids() - ) - ) return NodeMeta( title=self._titles.get(sid, sid), summary="", has_children=bool(provider.roots(sid)), - n_chunks=n_chunks, ) owner = self._section_owner.get(sid) if not owner: diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py index 767c1675..4cd0bfa8 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py @@ -7,7 +7,7 @@ Thinking policy (DeepSeek V4 defaults thinking ON if omitted): -- ``action`` (navigate / harvest / refine / verify / score): +- ``action`` (harvest / refine / verify / score): always disabled — short JSON under ``llm_max_tokens`` (often 256). - ``planner`` (plan_query / replan only): episode-bound ``NavConfig.planner_thinking``, else ``NAV_PLANNER_THINKING`` for EXP diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index 2458661b..28c04600 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -330,7 +330,7 @@ def compute_map_scores( query: str, root_ids: Optional[Sequence[str]] = None, ) -> Dict[str, float]: - """Leaf 3-channel scores + parent max-pool (self_only only if interstitial).""" + """Leaf path+content scores + parent max-pool (self_only only if interstitial).""" map_scores, _unit_scores = compute_map_and_unit_scores( ts, doc_id=doc_id, query=query, root_ids=root_ids ) @@ -395,15 +395,25 @@ def compute_corpus_map_and_unit_scores_many( del namespace # Dense scoring is intentionally disabled for the corpus path. + # Tree shape is query-independent; reuse it across the episode's two + # scoring passes (user query + per-subgoal retrieval_query) on the same + # ToolSpace instead of re-walking every document each time. + tree_cache = getattr(ts, "_mapnav_tree_cache", None) + if not isinstance(tree_cache, dict): + tree_cache = {} + setattr(ts, "_mapnav_tree_cache", tree_cache) tree_by_doc: Dict[ str, Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], ] = {} tree_started = time.perf_counter() for doc_id in valid_doc_ids: - root_ids = list(ts.sections_for_doc(doc_id)) - children_map, leaves, titles = _walk_tree(ts, doc_id, root_ids) - tree_by_doc[doc_id] = (children_map, leaves, titles) + cached = tree_cache.get(doc_id) + if cached is None: + root_ids = list(ts.sections_for_doc(doc_id)) + cached = _walk_tree(ts, doc_id, root_ids) + tree_cache[doc_id] = cached + tree_by_doc[doc_id] = cached _logger.info( "retrieval mapnav phase=tree_build seconds=%.3f documents=%d sections=%d", time.perf_counter() - tree_started, @@ -508,8 +518,7 @@ def relight_map_for_query( ) -> Tuple[Dict[str, float], Dict[str, float], List[str]]: """Re-score the whole shared map against ``query``. - Same namespace / single-doc split as the episode-level pass in ``nav_agent``: - an empty ``doc_id`` means the corpus root, where document ids are map nodes + An empty ``doc_id`` means the corpus root, where document ids are map nodes and ``ts.document_ids()`` is already restricted to the episode's corpus. """ doc = str(doc_id or "").strip() diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_navigate.py b/packages/shared-python/shared/services/retrieval/nav/nav_navigate.py index 96d64642..924cd920 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_navigate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_navigate.py @@ -1,26 +1,10 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List from ._compat import ToolSpace -from ._compat import Chunk -from .nav_address import ( - NavLevel, - address_level, - next_dispatch_depth, - owner_document, - uses_document_nodes, -) -from .nav_actions import build_legal_actions, format_actionable_map_observation -from .nav_policy import choose_llm_action, choose_rule_action -from .nav_projection import build_projection -from .nav_types import ( - ActionKind, - LegalAction, - NavConfig, - NavState, - RegionReport, -) +from .nav_address import owner_document +from .nav_types import LegalAction, NavConfig, NavState def _batch_actions(chosen: LegalAction) -> List[LegalAction]: @@ -66,203 +50,6 @@ def _batch_collect_deepest_first( return acts -def _chunk_plain_chars(chunk: Chunk) -> int: - text = (getattr(chunk, "text", None) or "").strip() - if not text: - return 0 - lines = text.splitlines() - if lines and lines[0].strip().startswith("[§"): - text = "\n".join(lines[1:]).strip() - return len(text) - - -def _estimate_branch_chars(ts: ToolSpace, section_id: str, doc_id: str) -> int: - """Evidence-sized estimate of hydrating section_id ∪ descendants.""" - sid = str(section_id or "").strip() - if not sid: - return 0 - materialize = getattr(ts, "_materialize_leaf_path_chunks", None) - if not callable(materialize): - return 0 - try: - pool = list(materialize(sid, doc_id) or []) - except Exception: - return 0 - return sum(_chunk_plain_chars(c) for c in pool) - - -def _section_has_children( - ts: ToolSpace, - section_id: str, - doc_id: str, - projection: Any = None, -) -> bool: - sid = str(section_id or "").strip() - if not sid: - return False - if projection is not None: - for view in list(getattr(projection, "tree_sections", None) or []) + list( - getattr(projection, "visible_sections", None) or [] - ): - if str(getattr(view, "section_id", "") or "") == sid: - return bool(getattr(view, "has_children", False)) - relations = getattr(ts, "section_relation_ids", None) - if callable(relations): - try: - _anc, desc = relations(sid, doc_id) - desc = {str(x).strip() for x in (desc or set()) if str(x).strip()} - desc.discard(sid) - return bool(desc) - except Exception: - pass - materialize = getattr(ts, "_materialize_leaf_path_chunks", None) - if callable(materialize): - try: - pool = list(materialize(sid, doc_id) or []) - return len(pool) > 1 - except Exception: - return False - return False - - -def _split_oversize_collect_actions( - ts: ToolSpace, - state: NavState, - chosen: LegalAction, - config: NavConfig, - projection: Any, -) -> Tuple[List[LegalAction], List[LegalAction], List[Dict[str, Any]]]: - """Split a COLLECT batch into (keep_collect, rewrite_dispatch, rewrite_info).""" - limit = int(getattr(config, "depth0_oversize_char_limit", 0) or 0) - keep: List[LegalAction] = [] - rewrite: List[LegalAction] = [] - info: List[Dict[str, Any]] = [] - for act in _batch_actions(chosen): - sid = str(act.section_id or "").strip() - if not sid: - continue - act_doc = owner_document(ts, sid, "") or str(state.doc_id or "") - chars = _estimate_branch_chars(ts, sid, act_doc) if act_doc else 0 - has_kids = _section_has_children(ts, sid, act_doc or state.doc_id, projection) - if limit > 0 and chars > limit and has_kids: - rewrite.append( - LegalAction( - action_id=str(act.action_id or ""), - kind=ActionKind.DISPATCH, - section_id=sid, - label=str(act.label or ""), - score=float(act.score or 0.0), - metadata=dict(act.metadata or {}), - ) - ) - info.append( - { - "section_id": sid, - "branch_chars": chars, - "limit": limit, - "from_action_id": str(act.action_id or ""), - } - ) - else: - keep.append(act) - return keep, rewrite, info - - -def _estimate_region_chars(projection_text: str) -> int: - return len(projection_text or "") - - -def _fork_nav_state(state: NavState, *, doc_id: Optional[str] = None) -> NavState: - """Copy mutable evidence fields for an isolated child navigate().""" - return NavState( - doc_id=str(doc_id) if doc_id is not None else state.doc_id, - query=state.query, - task_type=state.task_type, - current_scope=state.current_scope, - collected_ids=set(state.collected_ids), - collected=list(state.collected), - map_scores=dict(state.map_scores or {}), - unit_scores=dict(state.unit_scores or {}), - highlight_ids=list(state.highlight_ids), - collected_section_ids=set(state.collected_section_ids), - blocked_collect_section_ids=set(state.blocked_collect_section_ids), - action_history=[], - refusal_events=[], - reports_context="", - investigated_section_ids=set(), - dismissed_section_ids=set(state.dismissed_section_ids), - collect_confidence=dict(state.collect_confidence), - explicit_collect_ids=set(state.explicit_collect_ids), - group_priority=dict(state.group_priority), - retrieval_plan=state.retrieval_plan, - slot_bindings=dict(state.slot_bindings), - satisfied_subgoal_ids=set(state.satisfied_subgoal_ids), - attempted_subgoal_ids=set(state.attempted_subgoal_ids), - focus_subgoal_id=state.focus_subgoal_id, - focus_subgoal_need=state.focus_subgoal_need, - focus_subgoal_contract=state.focus_subgoal_contract, - focus_retrieval_query=state.focus_retrieval_query, - focus_contract_kind=state.focus_contract_kind, - subgoal_results=dict(state.subgoal_results), - replan_count=int(state.replan_count or 0), - harvested_owner_subgoal=dict(state.harvested_owner_subgoal), - subgoal_widen_gaps=dict(state.subgoal_widen_gaps), - subgoal_refined_queries=dict(state.subgoal_refined_queries), - subgoal_dismissed_section_ids={ - k: set(v) for k, v in (state.subgoal_dismissed_section_ids or {}).items() - }, - subgoal_attempt_counts=dict(state.subgoal_attempt_counts), - dropped_subgoal_ids=set(state.dropped_subgoal_ids), - asset_observation_context=str( - getattr(state, "asset_observation_context", "") or "" - ), - ) - - -def _merge_nav_state(parent: NavState, child: NavState) -> None: - """Merge a forked subagent state into the parent (called under lock).""" - for chunk, score in child.collected: - nid = getattr(chunk, "node_id", None) - if nid is None or nid in parent.collected_ids: - continue - parent.collected_ids.add(nid) - parent.collected.append((chunk, float(score))) - parent.collected_section_ids.update(child.collected_section_ids) - parent.blocked_collect_section_ids.update(child.blocked_collect_section_ids) - parent.investigated_section_ids.update(child.investigated_section_ids) - parent.dismissed_section_ids.update(child.dismissed_section_ids) - parent.refusal_events.extend(child.refusal_events) - parent.action_history.extend(child.action_history) - parent.collect_confidence.update(child.collect_confidence) - parent.explicit_collect_ids.update(child.explicit_collect_ids) - parent.group_priority.update(child.group_priority) - child_asset = str(getattr(child, "asset_observation_context", "") or "").strip() - if child_asset: - prev = str(getattr(parent, "asset_observation_context", "") or "").strip() - parent.asset_observation_context = ( - f"{prev}\n\n{child_asset}".strip() if prev else child_asset - ) - parent.slot_bindings.update(child.slot_bindings) - parent.satisfied_subgoal_ids.update(child.satisfied_subgoal_ids) - parent.attempted_subgoal_ids.update(child.attempted_subgoal_ids) - parent.subgoal_results.update(child.subgoal_results) - parent.harvested_owner_subgoal.update(child.harvested_owner_subgoal) - parent.subgoal_widen_gaps.update(child.subgoal_widen_gaps) - parent.subgoal_refined_queries.update(child.subgoal_refined_queries) - for sid, ids in (child.subgoal_dismissed_section_ids or {}).items(): - parent.subgoal_dismissed_section_ids.setdefault(sid, set()).update(ids) - parent.dropped_subgoal_ids.update(child.dropped_subgoal_ids) - for sid, n in child.subgoal_attempt_counts.items(): - parent.subgoal_attempt_counts[sid] = max( - int(parent.subgoal_attempt_counts.get(sid, 0)), int(n) - ) - if child.reports_context: - if parent.reports_context: - parent.reports_context = parent.reports_context + "\n" + child.reports_context - else: - parent.reports_context = child.reports_context - - def _apply_collect( ts: ToolSpace, state: NavState, @@ -320,475 +107,3 @@ def _apply_collect( if sids: detail["section_id"] = sids[0] return detail - - -def _format_region_reports(reports: List[RegionReport]) -> str: - if not reports: - return "" - lines = [f"=== Investigate results ({len(reports)} region(s)) ==="] - for i, rep in enumerate(reports, 1): - scope = rep.scope or "" - status = "skipped" if rep.skipped else "ok" - lines.append(f"[region {i}] {scope} ({status})") - if rep.summary: - lines.append(rep.summary) - if rep.collected_section_ids: - lines.append( - "collected: " + ", ".join(rep.collected_section_ids[:20]) - ) - if rep.reason: - lines.append(f"reason: {rep.reason}") - lines.append("---") - lines.append("=== End Investigate ===") - return "\n".join(lines) - - -def dispatch( - ts: ToolSpace, - state: NavState, - ids: List[str], - *, - query: str, - config: NavConfig, - depth: int, - budget: int, - steps_out: Optional[List[Any]] = None, -) -> List[RegionReport]: - """Run navigate() on each region id (fork/merge state). - - Namespace→document DISPATCH is depth-neutral (document episode starts at - depth 0). Namespace→section DISPATCH starts at depth 1. - """ - scope_now = str(state.current_scope or "").strip() - region_ids = [ - rid - for rid in (str(x).strip() for x in ids if str(x).strip()) - if rid != scope_now # never re-enter the current scope (self-dispatch) - ] - if not region_ids: - return [] - - # Serial DISPATCH (asyncio.gather reserved for Knowhere production). - namespace_parent = uses_document_nodes(ts) and not str(state.doc_id or "").strip() - reports: List[RegionReport] = [] - - for rid in region_ids: - level = address_level(ts, rid) - child_doc = owner_document(ts, rid, "") - child_depth = next_dispatch_depth( - ts, - parent_doc_id=str(state.doc_id or ""), - parent_scope=scope_now, - child_id=rid, - depth=depth, - ) - # Namespace parent: switch episode doc_id to the real document. - enter_doc = False - if namespace_parent and level == NavLevel.DOCUMENT: - enter_doc = True - child_doc = rid - elif namespace_parent and child_doc: - enter_doc = True - if enter_doc and child_doc: - child_state = _fork_nav_state(state, doc_id=str(child_doc)) - else: - child_state = _fork_nav_state(state) - try: - report = navigate( - ts, - state=child_state, - scope=rid, - query=query, - config=config, - depth=child_depth, - budget=budget, - steps_out=None, # parent records dispatch; child history merges via state - ) - except Exception as exc: - from .nav_token_budget import NavTokenLimit, stamp_step_detail - - if isinstance(exc, NavTokenLimit): - report = RegionReport( - scope=rid, - summary="", - reason="token_limit", - skipped=True, - depth=child_depth, - ) - else: - report = RegionReport( - scope=rid, - summary="", - reason=f"dispatch_failed: {exc}", - skipped=True, - depth=child_depth, - ) - _merge_nav_state(state, child_state) - if steps_out is not None: - from ._compat import AgentStep - - for h in child_state.action_history: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action=f"nav_{h.get('kind', 'step')}", - detail=stamp_step_detail(dict(h)), - ) - ) - reports.append(report) - return reports - - -def navigate( - ts: ToolSpace, - *, - state: NavState, - scope: Optional[str], - query: str, - config: NavConfig, - depth: int = 0, - budget: Optional[int] = None, - steps_out: Optional[List[Any]] = None, -) -> RegionReport: - """Recursive observe-act loop: COLLECT / DISPATCH / FINISH. - - When enable_recursive_dispatch is False, only depth==0 may DISPATCH; deeper - regions hard-COLLECT visible nodes or skip on overflow/error. - """ - from ._compat import AgentStep - from .nav_token_budget import stamp_step_detail - - char_budget = int(budget if budget is not None else config.map_char_limit) - prev_scope = state.current_scope - state.current_scope = scope - collected_before = set(state.collected_section_ids) - max_steps = max(1, int(config.navigate_max_steps if depth > 0 else config.max_steps)) - report = RegionReport(scope=scope, depth=depth) - - try: - for step_idx in range(max_steps): - projection = build_projection( - ts, - doc_id=state.doc_id, - query=query, - scope=scope, - config=config, - map_scores=state.map_scores, - collected_section_ids=state.collected_section_ids, - dismissed_section_ids=state.dismissed_section_ids, - highlight_ids=state.highlight_ids, - harvested_section_ids=( - state.harvested_owner_subgoal if config.is_checklist else None - ), - ) - # Experimental non-recursive mode: if a deep region overflows the - # map budget after folding, skip rather than invent hard truncation. - if ( - depth > 0 - and not config.enable_recursive_dispatch - and _estimate_region_chars(projection.text) > char_budget * 2 - and projection.truncated - ): - report.skipped = True - report.reason = "region_overflow_skip" - break - - actions = build_legal_actions( - state, - projection, - step_idx=step_idx, - config=config, - depth=depth, - max_steps=max_steps, - ts=ts, - ) - if not actions: - report.reason = "no_legal_actions" - break - - obs = format_actionable_map_observation( - projection, - actions, - inline_summary=scope is not None, - ) - projection.text = obs - - group_map: Dict[str, str] = {} - assembled_preview = "" - if depth == 0 and state.collected: - from .nav_compose import build_compose_preview, dedupe_scored - - # Empty when preview exceeds compose_group_rank_max_chars → skip rank. - assembled_preview, group_map = build_compose_preview( - dedupe_scored(list(state.collected)), - ts, - state, - config, - ) - - if (config.policy or "").strip().lower() == "llm": - chosen, meta = choose_llm_action( - state, - projection, - actions, - step_idx=step_idx, - config=config, - depth=depth, - max_steps=max_steps, - group_map=group_map or None, - assembled_preview=assembled_preview or None, - ) - else: - chosen = choose_rule_action( - state, projection, actions, step_idx=step_idx, config=config - ) - meta = {"reason": "rule_policy"} - - detail: Dict[str, Any] = { - "action_id": chosen.action_id, - "kind": chosen.kind.value, - "section_id": chosen.section_id, - "scope": scope, - "llm_reason": meta.get("reason"), - "llm_raw": meta.get("raw"), - "depth": depth, - "n_legal_actions": len(actions), - "legal_actions_preview": [a.prompt_line() for a in actions[:16]], - "projection_chars": len(obs), - } - if meta.get("group_rank"): - detail["group_rank"] = meta.get("group_rank") - - if chosen.kind == ActionKind.FINISH: - report.reason = str(meta.get("reason") or "finish") - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="nav_finish", - detail=stamp_step_detail(detail), - ) - ) - state.action_history.append({**detail, "step_idx": step_idx}) - break - - if chosen.kind == ActionKind.COLLECT: - keep_acts = _batch_actions(chosen) - rewrite_acts: List[LegalAction] = [] - rewrite_info: List[Dict[str, Any]] = [] - if depth == 0 and bool( - getattr(config, "enable_depth0_oversize_to_dispatch", False) - ): - keep_acts, rewrite_acts, rewrite_info = _split_oversize_collect_actions( - ts, state, chosen, config, projection - ) - - # Oversized branches first: rewrite COLLECT -> DISPATCH. - if rewrite_acts: - region_ids = [ - str(a.section_id or "").strip() - for a in rewrite_acts - if a.section_id - ] - child_reports = dispatch( - ts, - state, - region_ids, - query=query, - config=config, - depth=depth, - budget=char_budget, - steps_out=steps_out, - ) - for rid in region_ids: - state.investigated_section_ids.add(rid) - block = _format_region_reports(child_reports) - if block: - if state.reports_context: - state.reports_context = ( - state.reports_context + "\n" + block - ) - else: - state.reports_context = block - ddetail = { - **detail, - "kind": "dispatch", - "rewritten_collect_to_dispatch": True, - "rewrite_info": rewrite_info, - "dispatch_regions": region_ids, - "n_child_reports": len(child_reports), - "n_child_skipped": sum(1 for r in child_reports if r.skipped), - "reports_snippet": (block or "")[:2000], - "section_id": region_ids[0] if region_ids else chosen.section_id, - } - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="nav_dispatch", - detail=stamp_step_detail(ddetail), - ) - ) - state.action_history.append({**ddetail, "step_idx": step_idx}) - - # Remaining non-oversize COLLECTs (if any). - if keep_acts: - collect_chosen = keep_acts[0] - base_meta = dict(chosen.metadata or {}) - # Drop the original full batch; rebuild from keep_acts only. - base_meta.pop("batch_actions", None) - if len(keep_acts) > 1: - base_meta["batch_actions"] = keep_acts - collect_chosen.metadata = base_meta - cdetail = _apply_collect(ts, state, collect_chosen, config) - cdetail_full = { - **detail, - **cdetail, - "kind": "collect", - "rewritten_collect_to_dispatch": bool(rewrite_info), - "rewrite_info": rewrite_info or None, - } - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="nav_collect", - detail=stamp_step_detail(cdetail_full), - ) - ) - state.action_history.append({**cdetail_full, "step_idx": step_idx}) - elif not rewrite_acts: - # Empty selection — should not happen; fall back to original. - cdetail = _apply_collect(ts, state, chosen, config) - detail.update(cdetail) - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="nav_collect", - detail=stamp_step_detail(detail), - ) - ) - state.action_history.append({**detail, "step_idx": step_idx}) - continue - - if chosen.kind == ActionKind.DISPATCH: - region_ids = [ - str(a.section_id or "").strip() - for a in _batch_actions(chosen) - if a.section_id - ] - # Non-recursive experiment: deep agents should not see DISPATCH - # (build_legal_actions gates it); still guard here. - if depth > 0 and not config.enable_recursive_dispatch: - detail["skipped_dispatch"] = True - detail["reason"] = "recursive_dispatch_disabled" - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="nav_dispatch_skipped", - detail=stamp_step_detail(detail), - ) - ) - continue - - child_reports = dispatch( - ts, - state, - region_ids, - query=query, - config=config, - depth=depth, - budget=char_budget, - steps_out=steps_out, - ) - for rid in region_ids: - state.investigated_section_ids.add(rid) - block = _format_region_reports(child_reports) - if block: - if state.reports_context: - state.reports_context = state.reports_context + "\n" + block - else: - state.reports_context = block - detail["dispatch_regions"] = region_ids - detail["n_child_reports"] = len(child_reports) - detail["n_child_skipped"] = sum(1 for r in child_reports if r.skipped) - detail["reports_snippet"] = (block or "")[:2000] - if steps_out is not None: - steps_out.append( - AgentStep( - step_idx=len(steps_out) + 1, - action="nav_dispatch", - detail=stamp_step_detail(detail), - ) - ) - state.action_history.append({**detail, "step_idx": step_idx}) - continue - - # Unknown kind — stop. - report.reason = f"unknown_action:{chosen.kind}" - break - else: - report.reason = report.reason or "max_steps" - - except Exception as exc: - from .nav_token_budget import NavTokenLimit - - if isinstance(exc, NavTokenLimit): - report.skipped = True - report.reason = "token_limit" - else: - report.skipped = True - report.reason = f"navigate_error: {exc}" - finally: - state.current_scope = prev_scope - - newly = sorted(state.collected_section_ids - collected_before) - report.collected_section_ids = newly - roots = [ - str(h.get("section_id") or "") - for h in state.action_history - if h.get("kind") == "collect" - and int(h.get("n_added", 0) or 0) > 0 - and h.get("section_id") - ] - report.summary = ( - f"collected {len(newly)} branch node(s); explicit roots={roots[-8:]}" - if newly - else (report.reason or "no new evidence") - ) - return report - - -def sort_collected_by_doc_order( - scored: List[Tuple[Chunk, float]], - ts: ToolSpace, - doc_id: str, -) -> List[Tuple[Chunk, float]]: - """Order evidence by (doc_id, line). Cross-doc when episode doc_id is empty.""" - idx = getattr(ts, "_idx", None) - node_map = getattr(idx, "_node_to_doc_line", {}) if idx is not None else {} - cross_doc = not str(doc_id or "").strip() - - def key(item: Tuple[Chunk, float]) -> Tuple[str, int, int, str]: - chunk, _score = item - cdoc = str(getattr(chunk, "doc_id", "") or "") - line_ids = list(chunk.line_ids or ()) - if line_ids: - ln = min(line_ids) - return (cdoc if cross_doc else "", ln, ln, chunk.node_id) - loc = node_map.get(chunk.node_id) or node_map.get( - str(getattr(chunk, "section_id", "") or "") - ) - if loc and len(loc) >= 2: - loc_doc, loc_line = str(loc[0]), loc[1] - if cross_doc or loc_doc == doc_id: - try: - li = int(loc_line) - return (loc_doc if cross_doc else "", li, li, chunk.node_id) - except Exception: - pass - return (cdoc if cross_doc else "\uffff", 10**9, 10**9, chunk.node_id) - - return sorted(scored, key=key) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_node_filter.py b/packages/shared-python/shared/services/retrieval/nav/nav_node_filter.py new file mode 100644 index 00000000..9de89760 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/nav/nav_node_filter.py @@ -0,0 +1,271 @@ +"""Deterministic WHERE filter over the in-memory map-nav hierarchy. + +Agent-authored predicates run on ``path`` (filename + title chain via +``path_titles``) and ``summary``. Field predicates AND together; terms inside +one field OR together. No top-K and no result truncation for substring +matches. Regex is bounded by pattern length, node count, and compile/search +exceptions. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any, Dict, Iterable, List, Literal, Sequence, Tuple + +MatchKind = Literal["substring", "regex"] +FilterField = Literal["path", "summary"] + +_MAX_REGEX_PATTERN_LEN = 256 +_MAX_REGEX_NODES = 100_000 + + +@dataclass(frozen=True) +class FieldPredicate: + field: FilterField + terms: Tuple[str, ...] + match: MatchKind = "substring" + + +@dataclass(frozen=True) +class NodeFilter: + predicates: Tuple[FieldPredicate, ...] = () + + +@dataclass +class FilterResult: + matched_section_ids: List[str] + matched_doc_ids: List[str] + cardinality: int + truncated: bool = False + failed_predicates: List[str] = field(default_factory=list) + + +def field_predicate( + field: str, + terms: Sequence[str], + match: str = "substring", +) -> FieldPredicate: + key = str(field or "").strip().lower() + if key not in {"path", "summary"}: + raise ValueError(f"unsupported filter field: {field!r}") + kind = str(match or "substring").strip().lower() + if kind not in {"substring", "regex"}: + raise ValueError(f"unsupported filter match: {match!r}") + cleaned = tuple(str(term) for term in terms if str(term)) + return FieldPredicate(field=key, terms=cleaned, match=kind) # type: ignore[arg-type] + + +def node_filter(predicates: Sequence[FieldPredicate]) -> NodeFilter: + return NodeFilter(predicates=tuple(predicates)) + + +def apply_node_filter( + ts: Any, + doc_ids: Sequence[str], + nf: NodeFilter, +) -> FilterResult: + """Walk the named documents and evaluate ``nf`` on every node.""" + wanted = [str(did).strip() for did in doc_ids if str(did).strip()] + compiled, failed = _compile_predicates(nf.predicates) + if failed: + return FilterResult( + matched_section_ids=[], + matched_doc_ids=[], + cardinality=0, + truncated=False, + failed_predicates=failed, + ) + summaries = _load_summaries(ts) + matched_sections: List[str] = [] + matched_docs: List[str] = [] + seen_sections: set[str] = set() + seen_docs: set[str] = set() + visited = 0 + truncated = False + uses_regex = any(pred.match == "regex" for pred in nf.predicates) + + for doc_id in wanted: + for sid, owner_doc, is_doc_node in _iter_doc_nodes(ts, doc_id): + visited += 1 + if uses_regex and visited > _MAX_REGEX_NODES: + truncated = True + break + path_text = _path_text(ts, sid, owner_doc) + summary_text = "" if is_doc_node else str(summaries.get(sid) or "") + if not is_doc_node and not summary_text: + summary_text = _summary_fallback(ts, sid) + values = {"path": path_text, "summary": summary_text} + if not _node_matches(values, compiled): + continue + if is_doc_node: + if owner_doc not in seen_docs: + seen_docs.add(owner_doc) + matched_docs.append(owner_doc) + continue + if sid in seen_sections: + continue + seen_sections.add(sid) + matched_sections.append(sid) + if owner_doc and owner_doc not in seen_docs: + seen_docs.add(owner_doc) + matched_docs.append(owner_doc) + if truncated: + break + + return FilterResult( + matched_section_ids=matched_sections, + matched_doc_ids=matched_docs, + cardinality=len(matched_sections), + truncated=truncated, + failed_predicates=failed, + ) + + +def render_submap_observation( + ts: Any, + result: FilterResult, + *, + doc_ids: Sequence[str] | None = None, +) -> str: + """Hit-count line plus every matched node (path + summary).""" + del doc_ids + header = f"hits={result.cardinality}" + if result.truncated: + header = f"{header} truncated=true" + if result.failed_predicates: + header = f"{header} failed_predicates={len(result.failed_predicates)}" + if result.cardinality == 0: + return header + + summaries = _load_summaries(ts) + lines = [header] + for sid in result.matched_section_ids: + owner = _owner_document(ts, sid) + title = _path_text(ts, sid, owner) or sid + block = [f"{title}"] + summary = str(summaries.get(sid) or "").strip() + if summary: + block.append(f" summary: {summary}") + lines.append("\n".join(block)) + return "\n".join(lines) + + +def _compile_predicates( + predicates: Sequence[FieldPredicate], +) -> Tuple[List[Tuple[FieldPredicate, List[Any]]], List[str]]: + compiled: List[Tuple[FieldPredicate, List[Any]]] = [] + failed: List[str] = [] + for pred in predicates: + if pred.match != "regex": + compiled.append((pred, [])) + continue + patterns: List[Any] = [] + ok = True + for term in pred.terms: + if len(term) > _MAX_REGEX_PATTERN_LEN: + failed.append(f"{pred.field}:regex:too_long") + ok = False + break + try: + patterns.append(re.compile(term, flags=re.IGNORECASE)) + except re.error: + failed.append(f"{pred.field}:regex:invalid") + ok = False + break + if ok: + compiled.append((pred, patterns)) + return compiled, failed + + +def _node_matches( + values: Dict[str, str], + compiled: Sequence[Tuple[FieldPredicate, List[Any]]], +) -> bool: + if not compiled: + return True + for pred, patterns in compiled: + text = values.get(pred.field, "") + if pred.match == "regex": + if not patterns or not any(p.search(text or "") for p in patterns): + return False + continue + haystack = (text or "").lower() + if not pred.terms or not any(term.lower() in haystack for term in pred.terms): + return False + return True + + +def _iter_doc_nodes(ts: Any, doc_id: str) -> Iterable[Tuple[str, str, bool]]: + yield doc_id, doc_id, True + stack = list(_roots(ts, doc_id)) + seen: set[str] = set() + while stack: + sid = stack.pop(0) + if not sid or sid in seen: + continue + seen.add(sid) + yield sid, doc_id, False + stack[0:0] = [str(child) for child in _children(ts, sid) if str(child)] + + +def _roots(ts: Any, doc_id: str) -> List[str]: + fn = getattr(ts, "sections_for_doc", None) + if callable(fn): + return [str(sid) for sid in (fn(doc_id) or ()) if str(sid).strip()] + provider = getattr(ts, "_provider", None) + root_fn = getattr(provider, "roots", None) if provider is not None else None + if callable(root_fn): + return [str(sid) for sid in (root_fn(doc_id) or ()) if str(sid).strip()] + return [] + + +def _children(ts: Any, section_id: str) -> List[str]: + fn = getattr(ts, "_provider", None) + child_fn = getattr(fn, "children", None) if fn is not None else None + if callable(child_fn): + return [str(sid) for sid in (child_fn(section_id) or ()) if str(sid).strip()] + return [] + + +def _path_text(ts: Any, section_id: str, doc_id: str) -> str: + fn = getattr(ts, "path_titles", None) + if callable(fn): + try: + return str(fn(section_id, doc_id) or "") + except TypeError: + return str(fn(section_id) or "") + return "" + + +def _load_summaries(ts: Any) -> Dict[str, str]: + provider = getattr(ts, "_provider", None) + fn = getattr(provider, "summaries", None) if provider is not None else None + if not callable(fn): + return {} + raw = fn() or {} + return { + str(sid): str(summary or "") + for sid, summary in raw.items() + if str(sid).strip() and str(summary or "").strip() + } + + +def _summary_fallback(ts: Any, section_id: str) -> str: + structure_fn = getattr(ts, "get_structure", None) + if callable(structure_fn): + try: + st = structure_fn(section_id) or {} + return str(st.get("summary") or "").strip() + except Exception: + return "" + return "" + + +def _owner_document(ts: Any, section_id: str) -> str: + fn = getattr(ts, "owner_document", None) + if callable(fn): + got = fn(section_id) + if got: + return str(got) + return "" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 52cf8867..4ffae842 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -1,7 +1,7 @@ """M4/M5: wave orchestration over a RetrievalPlan. Execution order = dependency DAG ∩ soft prefer_after. Each subgoal runs its own -harvest/navigate so evidence attribution stays per-subgoal. Slot values are +harvest so evidence attribution stays per-subgoal. Slot values are extracted only when a later subgoal references them; checklist acceptance is owned by ``plan_control``. """ @@ -16,7 +16,6 @@ from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple from .nav_token_budget import stamp_step_detail -from .nav_navigate import navigate from .nav_plan import ( RetrievalPlan, Subgoal, @@ -137,24 +136,92 @@ def _resolve_subgoal_query(state: NavState, subgoal: Subgoal) -> str: return refined or query -def _run_navigate_for_query( +def seed_episode_map_scores_from_plan( ts: Any, state: NavState, config: NavConfig, - *, - query: str, - steps_out: Optional[List[Any]], -) -> None: - navigate( + plan: RetrievalPlan, +) -> int: + """Light the map from executable plan retrieval_query strings; set baseline. + + Scores bound (or refined) retrieval_query values only — queries still holding + ``{{slot}}`` placeholders are skipped and scored later when harvest resolves + them. Empty plan falls back to ``state.query``. Harvest may temporarily + relight per subgoal; ``_relit_map`` restores this baseline for evidence pack. + Returns the number of unique queries scored. + """ + queries: List[str] = [] + seen: Set[str] = set() + for subgoal in plan.subgoals: + refined = str( + (state.subgoal_refined_queries or {}).get(subgoal.id) or "" + ).strip() + if refined: + q = refined + else: + bound = bind_slots(subgoal.retrieval_query, state.slot_bindings) + # Skip not-yet-bound slot queries — scoring a stripped skeleton + # cannot cache-hit the later bound harvest string. + if unbound_slots(bound): + continue + q = str(bound or "").strip() + if q and q not in seen: + seen.add(q) + queries.append(q) + if not queries: + q = str(state.query or "").strip() + if q: + queries = [q] + if not queries: + state.map_scores = {} + state.unit_scores = {} + state.highlight_ids = [] + state.relit_map_cache = {} + return 0 + + # Replace prior episode lighting entirely (including after replan). + state.relit_map_cache = {} + + map_started = time.perf_counter() + from .nav_map_scores import relight_maps_for_queries + + prepared = relight_maps_for_queries( ts, - state=state, - scope=None, - query=query, - config=config, - depth=0, - budget=int(config.map_char_limit), - steps_out=steps_out, + doc_id=state.doc_id, + queries=queries, + top_k=int(config.collect_top_k), + ) + for query, triple in prepared.items(): + scores, units, highlights = triple + if scores: + state.relit_map_cache[query] = ( + dict(scores), + dict(units), + list(highlights), + ) + + baseline = None + for query in queries: + triple = prepared.get(query) + if triple and triple[0]: + baseline = triple + break + if baseline is not None: + scores, units, highlights = baseline + state.map_scores = dict(scores) + state.unit_scores = dict(units) + state.highlight_ids = list(highlights) + else: + state.map_scores = {} + state.unit_scores = {} + state.highlight_ids = [] + + _logger.info( + "retrieval mapnav phase=map_scoring seconds=%.3f queries=%d", + time.perf_counter() - map_started, + len(queries), ) + return len(queries) def _wave_subgoal_result( @@ -195,11 +262,10 @@ def _relit_map( ) -> Iterator[None]: """Score the shared map against the harvest ``query`` for one call. - Episode-level ``state.map_scores`` is computed from the original user query. - Checklist harvests run under a per-subgoal ``retrieval_query``, so the map - must be re-scored against that string — otherwise the ranking disagrees with - the query the policy is told to pursue. Scoring failures degrade to the - episode lighting. + Episode-level ``state.map_scores`` is the post-plan baseline (lit from + plan ``retrieval_query`` strings). Checklist harvests may run under a + different per-subgoal string, so the map is re-scored for that call and + restored to the episode baseline afterward for evidence packing. """ relit = prepared q = (query or "").strip() @@ -216,15 +282,19 @@ def _relit_map( top_k=int(config.collect_top_k), ) if scores: - relit = (scores, units, highlights) + relit = (dict(scores), dict(units), list(highlights)) state.relit_map_cache[q] = relit except Exception: relit = None if relit is None: yield return + scores, units, highlights = relit saved = (state.map_scores, state.unit_scores, state.highlight_ids) - state.map_scores, state.unit_scores, state.highlight_ids = relit + # Copy so harvest readers cannot mutate the cached baseline triple. + state.map_scores = dict(scores) + state.unit_scores = dict(units) + state.highlight_ids = list(highlights) try: yield finally: @@ -249,8 +319,6 @@ def _execute_subgoal_harvest_once( Retry / widen / drop / replan authority belongs to ``plan_control`` across waves (see ``nav_control.plan_control``), not to this single call. """ - from .nav_harvest import harvest - rq = retrieval_query or _resolve_subgoal_query(state, subgoal) refined = str((state.subgoal_refined_queries or {}).get(subgoal.id) or "").strip() _set_focus(state, subgoal, rq) @@ -266,12 +334,11 @@ def _execute_subgoal_harvest_once( query=rq, prepared=prepared_relight, ): - harvest_result = harvest( + harvest_result = _harvest_after_node_filter( ts, state, config, subgoal=subgoal, - entry_scope=None, query=rq, steps_out=steps_out, ) @@ -303,6 +370,126 @@ def _execute_subgoal_harvest_once( } +def _harvest_after_node_filter( + ts: Any, + state: NavState, + config: NavConfig, + *, + subgoal: Subgoal, + query: str, + steps_out: Optional[List[Any]], +) -> Any: + from .nav_harvest import harvest + + if not bool(getattr(config, "enable_node_filter", False)) or not bool( + getattr(subgoal, "use_node_filter", False) + ): + return harvest( + ts, + state, + config, + subgoal=subgoal, + entry_scope=None, + query=query, + steps_out=steps_out, + ) + + from .nav_scope_filter import run_scope_filter + + doc_ids = list(ts.document_ids() or ()) + if not doc_ids and state.doc_id: + doc_ids = [str(state.doc_id)] + seed = None + raw_seed = list(getattr(subgoal, "node_filter_predicates", None) or []) + if raw_seed: + from .nav_scope_filter import parse_node_filter + + seed = parse_node_filter({"predicates": raw_seed}) + outcome = run_scope_filter( + ts, + config, + query=query, + doc_ids=doc_ids, + seed_filter=seed, + steps_out=steps_out, + ) + if outcome.decision == "collect_all": + collected = _collect_filtered_sections( + ts, state, config, subgoal=subgoal, section_ids=outcome.settled_section_ids + ) + if collected is not None: + return collected + if outcome.decision == "scoped_harvest" and outcome.settled_section_ids: + return harvest( + ts, + state, + config, + subgoal=subgoal, + entry_scope=None, + query=query, + steps_out=steps_out, + allowed_section_ids=set(outcome.settled_section_ids), + ) + return harvest( + ts, + state, + config, + subgoal=subgoal, + entry_scope=None, + query=query, + steps_out=steps_out, + ) + + +def _collect_filtered_sections( + ts: Any, + state: NavState, + config: NavConfig, + *, + subgoal: Subgoal, + section_ids: Sequence[str], +) -> Any: + from .nav_address import is_dispatch_only_node + from .nav_harvest import HarvestResult + from .nav_knowhere import is_root_section + from .nav_navigate import _apply_collect + from .nav_types import ActionKind, LegalAction + + collectable = [ + sid + for sid in section_ids + if str(sid).strip() + and not is_dispatch_only_node(ts, sid) + and not is_root_section(ts, sid) + ] + if not collectable: + return None + actions = [ + LegalAction( + action_id=f"C{i}", + kind=ActionKind.COLLECT, + section_id=sid, + metadata={"multi": True}, + ) + for i, sid in enumerate(collectable, start=1) + ] + primary = actions[0] + primary.metadata = dict(primary.metadata or {}) + primary.metadata["batch_actions"] = actions + primary.metadata["confidence_by_section"] = {sid: 1.0 for sid in collectable} + detail = _apply_collect(ts, state, primary, config) + new_roots = list(detail.get("collect_section_ids") or []) + if bool(config.is_checklist): + for sid in new_roots: + state.harvested_owner_subgoal[sid] = subgoal.id + return HarvestResult( + subgoal_id=subgoal.id, + new_section_ids=new_roots, + n_policy_calls=0, + reason="node_filter_collect_all", + ) + + def _apply_plan_control( ts: Any, state: NavState, @@ -418,13 +605,13 @@ def execute_plan( plan = state.retrieval_plan if plan is None or not getattr(plan, "subgoals", None): - _run_navigate_for_query( - ts, state, config, query=episode_query or state.query, steps_out=steps_out - ) - return {"fallback_navigate": True} + # Retired: no multi-step navigate fallback. Planning always yields a + # subgoal (fallback_plan on failure); an empty plan here is a bug. + _logger.warning("execute_plan called with empty retrieval_plan; skipping wave") + return {"waves": [], "results": {}} max_waves = int(getattr(config, "max_waves", 0) or 0) - # Checklist mode always harvests + plan_controls (navigate-per-subgoal retired). + # Checklist mode always harvests + plan_controls. wave_idx = 0 summary: Dict[str, Any] = {"waves": [], "results": {}} episode_done = False @@ -452,23 +639,45 @@ def execute_plan( by_id = {s.id: s for s in plan.subgoals} outputs: List[Dict[str, Any]] = [] query_by_subgoal = { - sid: _resolve_subgoal_query(state, by_id[sid]) for sid in ready + sid: str(_resolve_subgoal_query(state, by_id[sid]) or "").strip() + for sid in ready } prepared_relights: Dict[ str, Tuple[Dict[str, float], Dict[str, float], List[str]], ] = {} - try: - from .nav_map_scores import relight_maps_for_queries + missing_queries: List[str] = [] + for query in query_by_subgoal.values(): + if not query: + continue + cached = state.relit_map_cache.get(query) + if cached is not None: + prepared_relights[query] = cached + elif query not in missing_queries: + missing_queries.append(query) + if missing_queries: + try: + from .nav_map_scores import relight_maps_for_queries - prepared_relights = relight_maps_for_queries( - ts, - doc_id=state.doc_id, - queries=list(query_by_subgoal.values()), - top_k=int(config.collect_top_k), - ) - except Exception: - prepared_relights = {} + scored = relight_maps_for_queries( + ts, + doc_id=state.doc_id, + queries=missing_queries, + top_k=int(config.collect_top_k), + ) + for query, triple in scored.items(): + scores, units, highlights = triple + prepared = (dict(scores), dict(units), list(highlights)) + prepared_relights[query] = prepared + if scores: + state.relit_map_cache[query] = prepared + except Exception: + _logger.exception( + "Failed to precompute relight maps; continuing without prepared relights " + "(doc_id=%s, query_count=%d)", + state.doc_id, + len(missing_queries), + ) def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) -> Dict[str, Any]: query = query_by_subgoal[sid] @@ -559,7 +768,7 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) if cap > 0 and int(state.replan_count) < cap: state.replan_count += 1 t0 = time.perf_counter() - new_plan = plan_query(ts, state, config) + new_plan = plan_query(state, config) state.retrieval_plan = new_plan plan = new_plan # A regenerated plan gets fresh subgoal ids (s1, s2, ... again), @@ -567,7 +776,9 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) # qualified "sX.slot" bindings) cannot be safely carried over — # those ids now mean something else. What IS safe and worth # keeping is unqualified slot bindings (plain fact values) and - # every chunk already in state.collected. + # every chunk already in state.collected. Clear before re-seeding + # so stale refined queries / qualified bindings cannot leak into + # the new plan's retrieval_query scoring. state.satisfied_subgoal_ids = set() state.attempted_subgoal_ids = set() state.dropped_subgoal_ids = set() @@ -579,6 +790,7 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) state.slot_bindings = { k: v for k, v in state.slot_bindings.items() if "." not in k } + seed_episode_map_scores_from_plan(ts, state, config, new_plan) if steps_out is not None: steps_out.append( AgentStep( diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_plan.py b/packages/shared-python/shared/services/retrieval/nav/nav_plan.py index 12fe13c8..cf7aef87 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_plan.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_plan.py @@ -1,7 +1,8 @@ """Structure-conditioned query planning (M2). -Looks at a planning map observation and emits a coverage checklist plus an -auditable RetrievalPlan over one shared search space. +``plan_query`` is query-only: it emits a coverage checklist plus an auditable +RetrievalPlan from the user question (no pre-lit map). ``refine_subgoal_query`` +may still read a folded planning map after harvest has lit scores. """ from __future__ import annotations @@ -55,6 +56,8 @@ class Subgoal: prefer_after: List[str] = field(default_factory=list) contract: Contract = field(default_factory=Contract) produces: List[str] = field(default_factory=list) + use_node_filter: bool = False + node_filter_predicates: List[Dict[str, Any]] = field(default_factory=list) @dataclass @@ -112,8 +115,8 @@ def unbound_slots(text: str) -> List[str]: def extract_plan_json(text: str) -> Optional[dict]: """Parse a (possibly nested) JSON object from model output. - Unlike the navigate-action helper, this uses brace balancing so nested - plan objects are not truncated to the first inner ``{...}``. + Uses brace balancing so nested plan objects are not truncated to the first + inner ``{...}``. """ s = (text or "").strip() if not s: @@ -159,6 +162,39 @@ def extract_plan_json(text: str) -> Optional[dict]: return None +def _as_bool(raw: Any) -> bool: + if isinstance(raw, bool): + return raw + return str(raw or "").strip().lower() in {"1", "true", "yes", "on"} + + +def _parse_node_filter_predicates(raw: Any) -> List[Dict[str, Any]]: + if isinstance(raw, dict): + raw = raw.get("predicates") + if not isinstance(raw, list): + return [] + out: List[Dict[str, Any]] = [] + for item in raw: + if not isinstance(item, dict): + continue + field = str(item.get("field") or "").strip().lower() + if field not in {"path", "summary"}: + continue + terms = item.get("terms") or [] + if isinstance(terms, str): + terms = [terms] + if not isinstance(terms, list): + continue + cleaned = [str(t) for t in terms if str(t)] + if not cleaned: + continue + match = str(item.get("match") or "substring").strip().lower() + if match not in {"substring", "regex"}: + match = "substring" + out.append({"field": field, "terms": cleaned, "match": match}) + return out + + def _as_str_list(raw: Any) -> List[str]: if raw is None: return [] @@ -249,9 +285,9 @@ def language_reference_text( ) -> str: """Script reference for retrieval_query checks. - Uses the user query plus visible map *titles* only. The actionable map - observation is intentionally excluded — its English chrome (collect=/dispatch= - /[Hit]/ would falsely dominate script detection. + Planner is query-only (projection=None). Refine still sees the folded map; + its English chrome (collect=/dispatch=/[Hit]) is excluded — only visible + map *titles* are used so they do not dominate script detection. """ parts: List[str] = [] q = (query or "").strip() @@ -385,7 +421,6 @@ def parse_retrieval_plan( obj: dict, *, query: str, - projection: Optional[Projection] = None, ) -> RetrievalPlan: """Parse LLM JSON into a RetrievalPlan; invalid refs are dropped.""" rows = obj.get("subgoals") or obj.get("goals") or obj.get("steps") or [] @@ -421,6 +456,10 @@ def parse_retrieval_plan( prefer_after=_as_str_list(row.get("prefer_after")), contract=_parse_contract(row.get("contract")), produces=produces, + use_node_filter=_as_bool(row.get("use_node_filter")), + node_filter_predicates=_parse_node_filter_predicates( + row.get("node_filter") or row.get("node_filter_predicates") + ), ) ) @@ -581,38 +620,40 @@ def _planner_system_prompt(*, max_subgoals: int) -> str: if max_subgoals > 0: cap = f" Prefer at most {max_subgoals} subgoals." return ( - "You are a retrieval planner for a hierarchical document map.\n" - "You see a folded title map of the corpus/document. Nodes may carry " - "collect=C* and dispatch=D* action ids; [Hit] marks hybrid retrieval beacons.\n" - "Your job is to emit a coverage checklist plus a retrieval plan over ONE " - "shared search space — do not partition the corpus with per-subgoal " - "scopes or map anchors.\n\n" + "You are a retrieval planner. You receive only the user query (no " + "document map). Emit a coverage checklist plus a retrieval plan over " + "ONE shared search space — do not invent per-subgoal corpus partitions.\n\n" "Rules:\n" "1. coverage_checklist lists the facts that episode evidence must cover " "(short, concrete facts in the query's language)." f"{cap}\n" - "2. Default to a SINGLE subgoal over the whole map. Only add more " - "subgoals for a hard data dependency ({{s1.slot}} in a later query) or " - "for clearly independent cross-entity comparisons. Do not split merely " - "to list checklist items.\n" + "2. Default to a SINGLE subgoal. Only add more subgoals for a hard data " + "dependency ({{s1.slot}} in a later query) or for clearly independent " + "cross-entity comparisons. Do not split merely to list checklist items.\n" "3. Each subgoal produces at most ONE slot name in produces " "(enumeration = one list-valued slot).\n" "4. retrieval_query is a SHORT KEYWORD QUERY for THIS subgoal only " "(space-separated entity/role/topic tokens, e.g. \"王仁坤 总工程师 设计成果\"). " "Split or adapt the user question into compact lexical terms — do NOT " - "emit a full natural-language question or long prose. It is scored by " - "lexical/hybrid retrieval and lights the map, so keep entity names and " - "role terms; drop filler words. Same language/script as the map titles " - "and user query — do not translate section terms into another script.\n" + "emit a full natural-language question or long prose. Downstream lexical " + "retrieval scores and lights the map from these tokens, so keep entity " + "names and role terms; drop filler words. Same language/script as the " + "user query — do not translate terms into another script.\n" "5. If a later retrieval_query needs a value from an earlier subgoal, " "write it as {{s1.slot}} (not prose). That implies depends_on.\n" "6. depends_on = hard data dependency. prefer_after = soft ordering only.\n" "7. All subgoals share one search space — do not invent per-subgoal scopes.\n" "8. relations only for parent-child or sibling (omit unrelated pairs).\n" - "9. map_coverage: sufficient | partial | insufficient — whether the planning " - "map shows enough structure to ground this plan.\n" - "10. reason must be English, under 40 words. Document titles stay original " - "language.\n\n" + "9. map_coverage: sufficient | partial | insufficient — whether the user " + "query alone is enough to write executable retrieval_query terms " + "(use sufficient unless the query is empty or unusable).\n" + "10. reason must be English, under 40 words.\n" + "11. Set use_node_filter=true when the subgoal enumerates or compares " + "named facets you can write as path predicates (filenames, section " + "titles). Keep it false for vague semantic needs; " + "retrieval_query remains the fuzzy-leg fallback. Optional node_filter " + "may seed predicates: " + '[{\"field\":\"path\",\"terms\":[\"...\"],\"match\":\"substring|regex\"}].\n\n' "Return ONLY one JSON object:\n" "{\n" ' "reason": "...",\n' @@ -630,7 +671,9 @@ def _planner_system_prompt(*, max_subgoals: int) -> str: ' "prefer_after": [],\n' ' "produces": ["entity"],\n' ' "contract": {"kind": "single_fact|enumeration|span|comparison|existence", ' - '"cardinality": null}\n' + '"cardinality": null},\n' + ' "use_node_filter": false,\n' + ' "node_filter": []\n' " },\n" " {\n" ' "id": "s2",\n' @@ -647,7 +690,6 @@ def _planner_system_prompt(*, max_subgoals: int) -> str: def _language_repair_user( - observation: str, state: NavState, bad_plan: RetrievalPlan, *, @@ -661,26 +703,21 @@ def _language_repair_user( return ( f"User query: {state.query}\n" f"Task type: {state.task_type}\n\n" - f"=== Planning Map ===\n{observation}\n=== End Planning Map ===\n\n" "Your previous plan had retrieval_query values that do not match the " - "language/script of the map titles and user query. Those queries will " - "fail lexical retrieval. Rewrite the FULL plan JSON. Keep structure, but " - "rewrite every mismatched retrieval_query as a short keyword query " - "in the map's own language and terms (space-separated tokens, not a " + "language/script of the user query. Those queries will fail lexical " + "retrieval. Rewrite the FULL plan JSON. Keep structure, but rewrite " + "every mismatched retrieval_query as a short keyword query in the " + "query's own language and terms (space-separated tokens, not a " "full-sentence question).\n" f"Mismatched retrieval_query lines:\n{bad_block}\n" ) def plan_query( - ts: Any, state: NavState, config: NavConfig, - *, - observation: Optional[str] = None, - projection: Optional[Projection] = None, ) -> RetrievalPlan: - """LLM plan over the planning map. Falls back to a single subgoal on failure.""" + """LLM plan from the user query only (no pre-lit map). Falls back to one subgoal.""" from .nav_llm import ( # type: ignore nav_chat, planner_output_max_tokens, @@ -692,9 +729,6 @@ def plan_query( if nav_token_budget_exhausted(): return fallback_plan(state.query, reason="token_limit") - if projection is None or observation is None: - projection, observation = build_planning_observation(ts, state, config) - model = resolve_nav_model( model=config.planner_model, model_env="NAV_PLANNER_MODEL", @@ -707,11 +741,10 @@ def plan_query( timeout_s = 300.0 if thinking_mode == "enabled" else 90.0 max_subgoals = int(getattr(config, "planner_max_subgoals", 0) or 0) system = _planner_system_prompt(max_subgoals=max_subgoals) - reference = language_reference_text(query=state.query, projection=projection) + reference = language_reference_text(query=state.query) user = ( f"User query: {state.query}\n" f"Task type: {state.task_type}\n\n" - f"=== Planning Map ===\n{observation}\n=== End Planning Map ===\n\n" "Return the retrieval plan JSON." ) max_tokens = planner_output_max_tokens( @@ -757,9 +790,7 @@ def plan_query( last_err = "empty_content" continue obj = extract_plan_json(last_raw) or {} - plan = parse_retrieval_plan( - obj, query=state.query, projection=projection - ) + plan = parse_retrieval_plan(obj, query=state.query) plan.raw = last_raw[:2000] ok, why = validate_retrieval_plan(plan) if not ok: @@ -768,9 +799,7 @@ def plan_query( if plan_has_language_mismatch(plan, reference) and not language_repair_used: language_repair_used = True last_err = "language_mismatch" - user = _language_repair_user( - observation, state, plan, reference=reference - ) + user = _language_repair_user(state, plan, reference=reference) continue if plan_has_language_mismatch(plan, reference): last_err = "language_mismatch" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_policy.py b/packages/shared-python/shared/services/retrieval/nav/nav_policy.py index 0bce11d5..666b7ad8 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_policy.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_policy.py @@ -2,54 +2,7 @@ import json import re -import time -from typing import Any, List, Optional - -from .nav_actions import action_by_id, actions_by_ids -from .nav_types import ActionKind, LegalAction, NavConfig, NavState, Projection - - -def choose_rule_action( - state: NavState, - projection: Projection, - actions: List[LegalAction], - *, - step_idx: int, - config: NavConfig, -) -> LegalAction: - """Deterministic pick for ``policy=rule`` (and rare last-resort paths). - - Illegal LLM action_ids do **not** use this COLLECT-first path; they FINISH - the scope and record ``refusal_events`` (see ``choose_llm_action``). - """ - del state, projection, step_idx, config - - def first(kind: ActionKind) -> Optional[LegalAction]: - for a in actions: - if a.kind == kind: - return a - return None - - act = first(ActionKind.COLLECT) or first(ActionKind.DISPATCH) - if act: - return act - return first(ActionKind.FINISH) or actions[-1] - - -def _finish_or_rule( - state: NavState, - projection: Projection, - actions: List[LegalAction], - *, - step_idx: int, - config: NavConfig, -) -> LegalAction: - finish = next((a for a in actions if a.kind == ActionKind.FINISH), None) - if finish is not None: - return finish - return choose_rule_action( - state, projection, actions, step_idx=step_idx, config=config - ) +from typing import Optional def _extract_json_obj(text: str) -> Optional[dict]: @@ -70,426 +23,3 @@ def _extract_json_obj(text: str) -> Optional[dict]: return obj if isinstance(obj, dict) else None except Exception: return None - - -def _extract_action_id_fallback(text: str) -> str: - s = (text or "").strip() - for key in ("action_id", "id"): - m = re.search(rf'"{key}"\s*:\s*"([^"]+)"', s, flags=re.I) - if m: - return str(m.group(1) or "").strip().upper() - return "" - - -def _normalize_id_list(obj: dict, primary_aid: str) -> List[str]: - """Union of action_id and optional ids into one ordered selected set.""" - ids_raw = obj.get("ids") or obj.get("action_ids") or obj.get("action_args", {}) - ids: List[str] = [] - if isinstance(ids_raw, dict): - ids_raw = ids_raw.get("ids") or [] - if isinstance(ids_raw, list): - ids = [str(x).strip().upper() for x in ids_raw if str(x).strip()] - if primary_aid and primary_aid not in ids: - ids = [primary_aid] + ids - out: List[str] = [] - seen: set[str] = set() - for i in ids: - if i not in seen: - out.append(i) - seen.add(i) - return out - - -def _collect_roots_from_history(state: NavState) -> List[str]: - """Sections the agent explicitly COLLECT'd (not descendant auto-marks).""" - roots: List[str] = [] - seen: set[str] = set() - for h in state.action_history: - if h.get("kind") != "collect": - continue - if int(h.get("n_added", 0) or 0) <= 0: - continue - sid = str(h.get("section_id") or "").strip() - if sid and sid not in seen: - roots.append(sid) - seen.add(sid) - return roots - - -def _format_agent_state( - state: NavState, - step_idx: int, - config: NavConfig, - *, - max_steps: Optional[int] = None, -) -> str: - """Agent state block shown before the actionable observation.""" - episode_steps = int(max_steps if max_steps is not None else config.max_steps) - lines = ["=== Agent State ==="] - lines.append(f"Current scope: {state.current_scope or 'document-root'}") - lines.append(f"Step: {step_idx + 1} / {episode_steps}") - lines.append( - "Observation mode: folded hierarchy map " - "(title-only at document-root; summaries inline inside dispatched regions)" - ) - - roots = _collect_roots_from_history(state) - if roots: - lines.append(f"Evidence collected: {len(roots)} section(s)") - for sid in roots: - lines.append(f' - "{sid}"') - else: - lines.append("Evidence collected: none") - - investigated = sorted(state.investigated_section_ids) - if investigated: - lines.append(f"Regions investigated (subagent reports below): {len(investigated)}") - for sid in investigated[:20]: - lines.append(f' - "{sid}"') - if len(investigated) > 20: - lines.append(f" - ... (+{len(investigated) - 20} more)") - - remaining = episode_steps - step_idx - 1 - if remaining <= 2: - lines.append( - f"Only {remaining} step(s) remaining. Prefer COLLECT or FINISH if evidence is sufficient." - ) - - lines.append("=== End Agent State ===") - return "\n".join(lines) - - -def _system_prompt( - *, - depth: int = 0, - dispatch_available: bool = True, - has_preview: bool = False, -) -> str: - """Observe-act navigate prompt (COLLECT / DISPATCH / FINISH). - - The prompt is state-adaptive: when no DISPATCH action is legal at this layer - (e.g. recursion off and depth>0), all DISPATCH semantics/examples/preferences - are removed and the model is told this layer is COLLECT/FINISH only. This - prevents the model from emitting illegal D* that would finish the scope. - - When has_preview is True (depth-0 with assembled evidence groups), FINISH must - include a relative group_rank over [G*] ids. - - Style follows KNOWHERE collector rules (action IDs on node lines, English - reason, no invented targets). Asset SEARCH is via harvest search_assets. - """ - role = ( - "You are a document navigation agent running an observe-act loop." - if depth == 0 - else "You are a region subagent investigating one assigned document subtree." - ) - - ids_hint = ( - '"ids": ["C1","C3",...] or ["D1","D2",...]' - if dispatch_available - else '"ids": ["C1","C3",...]' - ) - - action_semantics = [ - "Action semantics:", - " - collect=C*: add each selected section to evidence. A parent section " - "hydrates its full subtree; a leaf adds only that section. " - "For every selected collect id, provide confidence in [0,1] " - "(object map keyed by action id, or a single scalar for one id).", - ] - if dispatch_available: - action_semantics.append( - " - dispatch=D*: hand the listed region(s) to a child subagent " - "explorers; you receive their reports without moving your own viewpoint." - ) - action_semantics.append(" - finish=F*: end navigation for this scope / document.") - - scope_rule = ( - " - This layer is COLLECT/FINISH only: there is NO dispatch action here. " - "Do not output any D* id; COLLECT the relevant sections directly.\n" - if not dispatch_available - else "" - ) - prefer_rule = ( - " - Prefer DISPATCH for large internal sections when available; prefer " - "COLLECT for leaves or small clearly-relevant sections.\n" - if dispatch_available - else " - COLLECT the sections relevant to the query; use FINISH when done.\n" - ) - - preview_rule = "" - if has_preview: - preview_rule = ( - " - When Assembled Evidence ([G*] groups) is shown, FINISH MUST include " - '"group_rank": an ordered list of those G* ids, most relevant to the ' - "query first (relative ranking, not absolute scores).\n" - " - For list/coverage queries, FINISH only if the assembled groups already " - "cover ALL required items; otherwise COLLECT the missing ones first.\n" - ) - - examples = [ - "Return ONLY one JSON object, e.g.:", - '{"action_id":"C1","confidence":0.8,"reason":"short reason"}', - 'Batch: {"action_id":"C1","ids":["C1","C3"],' - '"confidence":{"C1":0.7,"C3":0.9},"reason":"..."}', - ] - if dispatch_available: - examples.append( - 'Batch: {"action_id":"D1","ids":["D1","D2"],"reason":"..."}' - ) - if has_preview: - examples.append( - 'Finish with rank: {"action_id":"F1","group_rank":["G2","G1"],"reason":"..."}' - ) - - return ( - f"{role}\n\n" - "The observation is a folded hierarchy map. Each visible node lists only the " - "action IDs currently legal for that node. Collected branches are removed from " - "the map. At document-root the map is title-only; inside a dispatched region, " - "node summaries are inlined.\n\n" - "=== Rules ===\n\n" - f"Select one or more action IDs of the same kind. Put them in action_id and " - f"optional {ids_hint}; the final selection is their union. " - "Hydration is decided by hierarchy after selection " - "(parent COLLECT = full subtree; leaf COLLECT = that section only).\n\n" - + "\n".join(action_semantics) - + "\n\n" - " - Use only action IDs shown on a node line or under Global actions. " - "Never invent IDs or write raw section paths as targets.\n" - + scope_rule - + prefer_rule - + preview_rule - + " - Do NOT re-collect a section already listed under Evidence collected.\n" - " - FINISH when this scope is done: evidence is sufficient, or this region is " - "irrelevant / exhausted (especially as a subagent). " - "The system will not infer missing evidence for you.\n" - " - When steps remaining <= 2, prioritize COLLECT or FINISH.\n\n" - "=== End Rules ===\n\n" - + "\n".join(examples) - + "\n" - "Do not include any explanation outside the JSON.\n\n" - "IMPORTANT:\n" - "1. All agent-generated text (reason) MUST be in English.\n" - "2. Document content and section titles MUST remain in their original language.\n" - "3. Keep reason under 25 words.\n" - "4. COLLECT must include confidence for each selected collect id.\n" - ) - - -def choose_llm_action( - state: NavState, - projection: Projection, - actions: List[LegalAction], - *, - step_idx: int, - config: NavConfig, - depth: int = 0, - max_steps: Optional[int] = None, - group_map: Optional[dict[str, str]] = None, - assembled_preview: Optional[str] = None, -) -> tuple[LegalAction, dict]: - from .nav_llm import nav_chat, resolve_nav_model - from .nav_token_budget import NavTokenLimit, nav_token_budget_exhausted - - def _token_limit_finish() -> tuple[LegalAction, dict]: - finish = _finish_or_rule( - state, projection, actions, step_idx=step_idx, config=config - ) - return finish, { - "reason": "token_limit", - "stop_reason": "token_limit", - "depth": depth, - } - - if nav_token_budget_exhausted(): - return _token_limit_finish() - - model = resolve_nav_model( - model=(config.subagent_model if depth > 0 else config.llm_model), - model_env=("NAV_SUBAGENT_MODEL" if depth > 0 else "NAV_LLM_MODEL"), - fallback_envs=("NAV_LLM_MODEL",), - ) - agent_state = _format_agent_state( - state, step_idx, config, max_steps=max_steps - ) - has_preview = bool(assembled_preview and group_map) - dispatch_available = any(a.kind == ActionKind.DISPATCH for a in actions) - system = _system_prompt( - depth=depth, - dispatch_available=dispatch_available, - has_preview=has_preview, - ) - - reports_block = "" - if state.reports_context: - reports_block = ( - f"\n=== Subagent Reports ===\n{state.reports_context}\n" - f"=== End Subagent Reports ===\n" - ) - preview_block = "" - if has_preview: - preview_block = ( - f"\n=== Assembled Evidence (rank these on FINISH) ===\n" - f"{assembled_preview}\n" - f"=== End Assembled Evidence ===\n" - ) - - focus_need = str(getattr(state, "focus_subgoal_need", "") or "").strip() - focus_rq = str(getattr(state, "focus_retrieval_query", "") or "").strip() - effective_query = focus_rq or state.query - user = ( - f"User query: {effective_query}\n" - f"Task type: {state.task_type}\n\n" - f"{agent_state}\n\n" - f"=== Actionable Observation ===\n" - f"{projection.text}\n" - f"=== End Actionable Observation ===\n" - f"{reports_block}" - f"{preview_block}\n" - 'Return: {"action_id":"...","reason":"..."}' - ) - if focus_need: - focus_id = str(getattr(state, "focus_subgoal_id", "") or "").strip() - focus_contract = str(getattr(state, "focus_subgoal_contract", "") or "").strip() - episode_line = "" - if focus_rq and focus_rq.strip() != str(state.query or "").strip(): - episode_line = f"episode_query: {state.query}\n" - focus_block = ( - f"\n=== Current Subgoal (soft focus; action space unchanged) ===\n" - f"id: {focus_id or '-'}\n" - f"{episode_line}" - f"need: {focus_need}\n" - f"retrieval_query: {focus_rq or focus_need}\n" - f"contract: {focus_contract or '-'}\n" - f"Prefer evidence that serves this need, but you may still collect any " - f"visible useful node.\n" - f"=== End Current Subgoal ===\n" - ) - user = ( - f"User query: {effective_query}\n" - f"Task type: {state.task_type}\n\n" - f"{agent_state}\n" - f"{focus_block}\n" - f"=== Actionable Observation ===\n" - f"{projection.text}\n" - f"=== End Actionable Observation ===\n" - f"{reports_block}" - f"{preview_block}\n" - 'Return: {"action_id":"...","reason":"..."}' - ) - - purpose = "nav_navigate_v1" if depth == 0 else "nav_subagent_v1" - last_error: Optional[Exception] = None - for attempt in range(3): - try: - cached = nav_chat( - purpose=purpose, - model=model, - messages=[ - {"role": "system", "content": system}, - {"role": "user", "content": user}, - ], - temperature=float(config.llm_temperature), - max_tokens=int(config.llm_max_tokens), - response_format={"type": "json_object"}, - context="Nav Agent", - usage_tag="nav", - ) - text = str(cached.get("content") or "").strip() - obj = _extract_json_obj(text) or {} - aid = str(obj.get("action_id") or obj.get("id") or "").strip().upper() - if not aid: - aid = _extract_action_id_fallback(text) - primary = action_by_id(actions, aid) - if primary is not None: - meta: dict[str, Any] = { - "model": model, - "reason": str(obj.get("reason") or "")[:300], - "raw": text[:500], - "depth": depth, - } - if primary.kind in {ActionKind.COLLECT, ActionKind.DISPATCH}: - id_list = _normalize_id_list(obj, primary.action_id.upper()) - selected = actions_by_ids(actions, id_list) - selected = [a for a in selected if a.kind == primary.kind] - if not selected: - selected = [primary] - meta["selected_ids"] = [a.action_id for a in selected] - meta["selected_section_ids"] = [a.section_id for a in selected] - primary.metadata = dict(primary.metadata or {}) - primary.metadata["batch_actions"] = selected - if primary.kind == ActionKind.COLLECT: - from .nav_compose import parse_collect_confidence - - conf_by_aid = parse_collect_confidence(obj, selected) - conf_by_sid = { - str(a.section_id): float(conf_by_aid.get(a.action_id.upper(), 0.0)) - for a in selected - if a.section_id - } - meta["confidence_by_action"] = conf_by_aid - meta["confidence_by_section"] = conf_by_sid - primary.metadata["confidence_by_section"] = conf_by_sid - if primary.kind == ActionKind.FINISH and group_map: - rank_raw = obj.get("group_rank") or obj.get("groups") or [] - if isinstance(rank_raw, list) and rank_raw: - n = len(rank_raw) - applied: List[str] = [] - for i, g in enumerate(rank_raw): - gid = str(g).strip().upper() - pid = group_map.get(gid) - if pid: - state.group_priority[pid] = float(n - i) - applied.append(gid) - if applied: - meta["group_rank"] = applied - return primary, meta - last_error = RuntimeError( - "Nav Agent LLM 返回了非法 action_id=" - f"{aid!r}；合法选项={[a.action_id for a in actions]!r}；raw={text[:500]!r}" - ) - if attempt < 2: - time.sleep(min(2.0, 0.4 * (attempt + 1))) - continue - # Do not silently COLLECT the first tree row — finish this scope. - fallback = _finish_or_rule( - state, projection, actions, step_idx=step_idx, config=config - ) - state.refusal_events.append( - { - "tool": "policy", - "status": "illegal_action", - "message": ( - f"illegal action_id={aid!r} after retries; " - f"finishing scope with {fallback.action_id}" - ), - "illegal_action_id": aid, - "fallback_action_id": fallback.action_id, - "fallback_kind": fallback.kind.value, - "depth": depth, - "step_idx": step_idx, - } - ) - return fallback, { - "model": model, - "reason": "illegal_action_finish", - "raw": text[:500], - "illegal_action_id": aid, - "fallback_action_id": fallback.action_id, - "depth": depth, - } - except NavTokenLimit: - return _token_limit_finish() - except RuntimeError as exc: - last_error = exc - if attempt < 2: - time.sleep(min(2.0, 0.4 * (attempt + 1))) - continue - raise - except Exception as exc: - last_error = exc - time.sleep(min(2.0, 0.4 * (attempt + 1))) - raise RuntimeError( - f"Nav Agent LLM 调用失败（model={model!r}，step={step_idx}）：{last_error}" - ) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_projection.py b/packages/shared-python/shared/services/retrieval/nav/nav_projection.py index 452c3fdb..b4ffd579 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_projection.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_projection.py @@ -76,7 +76,6 @@ def _section_view_from_structure( preview=preview, score=_lexical_score(query, f"{section_id} {preview}"), n_lines=int(st.get("n_lines") or 0), - n_chunks=int(st.get("n_chunks") or 0), has_children=bool(children), depth_from_scope=depth_from_scope, title=preview[:80] if preview else section_id, @@ -127,7 +126,6 @@ class _MapNode: title: str score: float n_lines: int - n_chunks: int has_children: bool children: List["_MapNode"] = field(default_factory=list) n_descendants: int = 0 @@ -317,11 +315,13 @@ def _build_map_tree( collected_section_ids: Optional[Set[str]] = None, dismissed_section_ids: Optional[Set[str]] = None, harvested_section_ids: Optional[Dict[str, str]] = None, + keep_ids: Optional[Set[str]] = None, ) -> List[_MapNode]: roots: List[_MapNode] = [] seen: Set[str] = set() node_count = 0 harvested = dict(harvested_section_ids or {}) + keep = set(keep_ids) if keep_ids is not None else None # collected = branch done (sid ∪ descendants already marked by caller). # Harvested roots stay visible as a collapsed single line (fix-map- # visibility); their descendants are still fully removed like any other @@ -351,7 +351,6 @@ def make_node(section_id: str, depth: int, parent_id: Optional[str]) -> Optional title=title, score=score, n_lines=int(st.get("n_lines") or 0), - n_chunks=int(st.get("n_chunks") or 0), has_children=False, parent_id=parent_id, ) @@ -366,6 +365,8 @@ def append_visible_descendants( """Attach section_id if visible. collected/dismissed drop node + subtree.""" if not section_id or section_id in gone: return + if keep is not None and section_id not in keep: + return node = make_node(section_id, depth, parent_id) if node is None: return @@ -457,10 +458,7 @@ def render(node: _MapNode) -> None: leaf_tag = " [Leaf]" if not node.has_children else "" hit_tag = format_hit_tag(is_highlight=is_hit) harvested_tag = format_harvested_tag(node.harvested_by) - line = ( - f"{indent}[{map_id}] {node.title} ({node.n_chunks} chunks)" - f"{leaf_tag}{hit_tag}{harvested_tag}" - ) + line = f"{indent}[{map_id}] {node.title}{leaf_tag}{hit_tag}{harvested_tag}" lines.append(line) summary = "" if inline_summary: @@ -480,7 +478,6 @@ def render(node: _MapNode) -> None: preview="", score=node.score, n_lines=node.n_lines, - n_chunks=node.n_chunks, has_children=node.has_children, depth_from_scope=node.depth, map_id=map_id, @@ -501,6 +498,30 @@ def render(node: _MapNode) -> None: return "\n".join(lines), visible, id_to_section, any_hidden +def _expand_allowed_with_ancestors( + ts: Any, + allowed_section_ids: Optional[Set[str]], +) -> Optional[Set[str]]: + if allowed_section_ids is None: + return None + keep = {str(sid).strip() for sid in allowed_section_ids if str(sid).strip()} + parent_fn = getattr(ts, "parent_id", None) + if not callable(parent_fn): + return keep + extra: Set[str] = set() + for sid in keep: + cur = sid + seen: Set[str] = set() + while cur and cur not in seen: + seen.add(cur) + parent = parent_fn(cur) + if not parent: + break + extra.add(str(parent)) + cur = str(parent) + return keep | extra + + def _fallback_highlights_from_tree(roots: List[_MapNode], k: int) -> List[str]: leaves = [n for n in _flatten(roots, include_hidden=True) if not n.has_children] leaves.sort(key=lambda n: (-n.score, n.section_id)) @@ -520,6 +541,7 @@ def build_map( highlight_ids: Optional[List[str]] = None, extra_hidden_ids: Optional[Set[str]] = None, harvested_section_ids: Optional[Dict[str, str]] = None, + allowed_section_ids: Optional[Set[str]] = None, ) -> Projection: """Full-depth title map with score-ordered budget hiding (+ optional inline summary).""" scores = dict(map_scores or {}) @@ -527,6 +549,7 @@ def build_map( root_ids = [scope] else: root_ids = _top_sections(ts, doc_id) + keep_ids = _expand_allowed_with_ancestors(ts, allowed_section_ids) roots = _build_map_tree( ts, root_ids=root_ids, @@ -535,6 +558,7 @@ def build_map( collected_section_ids=collected_section_ids, dismissed_section_ids=dismissed_section_ids, harvested_section_ids=harvested_section_ids, + keep_ids=keep_ids, ) hits = [str(x).strip() for x in (highlight_ids or []) if str(x).strip()] if not hits: @@ -606,6 +630,7 @@ def build_projection( highlight_ids: Optional[List[str]] = None, extra_hidden_ids: Optional[Set[str]] = None, harvested_section_ids: Optional[Dict[str, str]] = None, + allowed_section_ids: Optional[Set[str]] = None, ) -> Projection: if map_mode_enabled(config): return build_map( @@ -620,6 +645,7 @@ def build_projection( highlight_ids=highlight_ids, extra_hidden_ids=extra_hidden_ids, harvested_section_ids=harvested_section_ids, + allowed_section_ids=allowed_section_ids, ) # Minimal non-map fallback (legacy shallow projection) — kept for ablation only. @@ -670,7 +696,7 @@ def add_line(text: str) -> None: leaf_tag = " [Leaf]" if not view.has_children else "" title = view.preview[:80] if view.preview else view.section_id add_line( - f"{indent}[{view.section_id}] {title} ({view.n_chunks} chunks){leaf_tag}" + f"{indent}[{view.section_id}] {title}{leaf_tag}" ) if view.preview: add_line(f"{indent} Preview: \"{view.preview[:80]}\"") diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_scope_filter.py b/packages/shared-python/shared/services/retrieval/nav/nav_scope_filter.py new file mode 100644 index 00000000..c0be92e9 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/nav/nav_scope_filter.py @@ -0,0 +1,441 @@ +"""Bounded WHERE self-correction loop (pre-harvest). + +The policy writes or revises a ``NodeFilter``. Apply is deterministic. The +loop stops at ``filter_max_rounds``, token exhaustion, explicit fallback, or +a done-in-band settle. Decision is cardinality-driven; the agent may override. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any, Dict, List, Literal, Optional, Sequence + +from .nav_node_filter import ( + FilterResult, + NodeFilter, + apply_node_filter, + field_predicate, + node_filter, + render_submap_observation, +) +from .nav_types import NavConfig + +ScopeDecision = Literal["collect_all", "scoped_harvest", "fallback"] +_SCOPE_FILTER_PURPOSE = "nav_scope_filter_v1" +_DECISIONS = {"collect_all", "scoped_harvest", "fallback"} + + +@dataclass +class ScopeFilterOutcome: + decision: ScopeDecision + settled_section_ids: List[str] = field(default_factory=list) + settled_doc_ids: List[str] = field(default_factory=list) + rounds: int = 0 + last_result: Optional[FilterResult] = None + reason: str = "" + + +def parse_node_filter(obj: Dict[str, Any] | None) -> Optional[NodeFilter]: + if not isinstance(obj, dict): + return None + raw = obj.get("predicates") + if raw is None: + raw = obj.get("filter") + if not isinstance(raw, list): + return None + preds = [] + for item in raw: + if not isinstance(item, dict): + continue + terms = item.get("terms") or [] + if isinstance(terms, str): + terms = [terms] + if not isinstance(terms, list): + continue + try: + preds.append( + field_predicate( + str(item.get("field") or ""), + [str(t) for t in terms], + str(item.get("match") or "substring"), + ) + ) + except ValueError: + continue + if not preds: + return None + return node_filter(preds) + + +def run_scope_filter( + ts: Any, + config: NavConfig, + *, + query: str, + doc_ids: Sequence[str], + map_observation: str = "", + seed_filter: Optional[NodeFilter] = None, + steps_out: Optional[List[Any]] = None, +) -> ScopeFilterOutcome: + """Apply → observe → revise until settle, fallback, or the round cap.""" + if not bool(getattr(config, "enable_node_filter", False)): + return ScopeFilterOutcome(decision="fallback", reason="disabled") + + from .nav_token_budget import nav_token_budget_exhausted, stamp_step_detail + + max_rounds = max(1, int(getattr(config, "filter_max_rounds", 3) or 3)) + min_hits = max(0, int(getattr(config, "filter_min_hits", 1) or 0)) + max_hits = max(min_hits, int(getattr(config, "filter_max_hits", 40) or 0)) + wanted = [str(did).strip() for did in doc_ids if str(did).strip()] + map_text = str(map_observation or "").strip() or _compact_map(ts, wanted) + current = seed_filter + last_result: Optional[FilterResult] = None + last_obs = "" + last_decision: Optional[ScopeDecision] = None + last_reason = "" + + if current is None: + action = _scope_filter_policy_call( + config, + query=query, + full_map=map_text, + last_filter=None, + last_result=None, + last_observation="", + round_idx=0, + max_rounds=max_rounds, + include_full_map=True, + ) + current = action.get("filter") + last_decision = action.get("decision") + last_reason = str(action.get("reason") or "") + if action.get("kind") == "fallback" or current is None: + return ScopeFilterOutcome( + decision="fallback", + reason=last_reason or "policy_fallback", + ) + + for round_idx in range(1, max_rounds + 1): + if nav_token_budget_exhausted(): + return ScopeFilterOutcome( + decision="fallback", + settled_section_ids=list(last_result.matched_section_ids) + if last_result + else [], + settled_doc_ids=list(last_result.matched_doc_ids) if last_result else [], + rounds=round_idx - 1, + last_result=last_result, + reason="token_limit", + ) + assert current is not None + result = apply_node_filter(ts, wanted, current) + last_result = result + last_obs = render_submap_observation(ts, result, doc_ids=wanted) + in_band = min_hits <= result.cardinality <= max_hits + if steps_out is not None: + from ._compat import AgentStep + + steps_out.append( + AgentStep( + step_idx=len(steps_out) + 1, + action="node_filter", + detail=stamp_step_detail( + { + "round": round_idx, + "predicates": _filter_payload(current), + "fields": sorted( + {p.field for p in current.predicates} + ), + "cardinality": result.cardinality, + "truncated": result.truncated, + "failed_predicates": list(result.failed_predicates), + "matched_section_ids": list(result.matched_section_ids), + "action": "", + "decision": "", + "reason": "", + } + ), + ) + ) + + is_last = round_idx >= max_rounds + if is_last: + if steps_out: + steps_out[-1].detail["action"] = "max_rounds" + return _settle( + result, + in_band=in_band, + agent_decision=last_decision, + min_hits=min_hits, + rounds=round_idx, + reason=("max_rounds" if in_band else "max_rounds_out_of_band"), + steps_out=steps_out, + ) + + action = _scope_filter_policy_call( + config, + query=query, + full_map=map_text, + last_filter=current, + last_result=result, + last_observation=last_obs, + round_idx=round_idx, + max_rounds=max_rounds, + include_full_map=False, + ) + if action.get("kind") == "widen": + # Agent judged the sub-map too narrow: one re-look with the full + # map, same round. A second widen settles on whatever it returns. + action = _scope_filter_policy_call( + config, + query=query, + full_map=map_text, + last_filter=current, + last_result=result, + last_observation=last_obs, + round_idx=round_idx, + max_rounds=max_rounds, + include_full_map=True, + ) + if action.get("kind") == "widen": + if action.get("filter") is not None: + action["kind"] = "filter" + else: + action["kind"] = "fallback" + action["reason"] = action.get("reason") or "widen_without_filter" + last_decision = action.get("decision") + last_reason = str(action.get("reason") or "") + kind = str(action.get("kind") or "") + if steps_out: + steps_out[-1].detail["action"] = kind + steps_out[-1].detail["reason"] = last_reason + if kind == "fallback": + if steps_out: + steps_out[-1].detail["decision"] = "fallback" + return ScopeFilterOutcome( + decision="fallback", + settled_section_ids=list(result.matched_section_ids), + settled_doc_ids=list(result.matched_doc_ids), + rounds=round_idx, + last_result=result, + reason=last_reason or "policy_fallback", + ) + if kind == "done": + return _settle( + result, + in_band=in_band, + agent_decision=last_decision, + min_hits=min_hits, + rounds=round_idx, + reason=last_reason or "done", + steps_out=steps_out, + ) + nxt = action.get("filter") + if nxt is not None: + current = nxt + + assert last_result is not None + in_band = min_hits <= last_result.cardinality <= max_hits + return _settle( + last_result, + in_band=in_band, + agent_decision=last_decision, + min_hits=min_hits, + rounds=max_rounds, + reason=last_reason or "max_rounds", + steps_out=steps_out, + ) + + +def _settle( + result: FilterResult, + *, + in_band: bool, + agent_decision: Optional[ScopeDecision], + min_hits: int, + rounds: int, + reason: str, + steps_out: Optional[List[Any]], +) -> ScopeFilterOutcome: + if not in_band or result.cardinality <= 0: + decision: ScopeDecision = "fallback" + settle_reason = reason or "out_of_band" + elif agent_decision in _DECISIONS: + decision = agent_decision + settle_reason = reason or "agent" + if decision == "fallback": + settle_reason = reason or "agent_fallback" + elif result.cardinality <= min_hits: + decision = "collect_all" + settle_reason = reason or "small_cardinality" + else: + decision = "scoped_harvest" + settle_reason = reason or "medium_cardinality" + if steps_out: + steps_out[-1].detail["decision"] = decision + steps_out[-1].detail["reason"] = settle_reason + return ScopeFilterOutcome( + decision=decision, + settled_section_ids=list(result.matched_section_ids), + settled_doc_ids=list(result.matched_doc_ids), + rounds=rounds, + last_result=result, + reason=settle_reason, + ) + + +def _scope_filter_policy_call( + config: NavConfig, + *, + query: str, + full_map: str, + last_filter: Optional[NodeFilter], + last_result: Optional[FilterResult], + last_observation: str, + round_idx: int, + max_rounds: int, + include_full_map: bool, +) -> Dict[str, Any]: + from .nav_llm import nav_chat, resolve_nav_model + from .nav_policy import _extract_json_obj + from .nav_token_budget import NavTokenLimit, nav_token_budget_exhausted + + if nav_token_budget_exhausted(): + return {"kind": "fallback", "reason": "token_limit"} + + model = resolve_nav_model( + model=config.llm_model, + model_env="NAV_LLM_MODEL", + fallback_envs=("NAV_LLM_MODEL",), + ) + user = f"Query: {query}\nRound: {round_idx}/{max_rounds}\n" + if last_filter is not None: + user += f"Last filter: {json.dumps(_filter_payload(last_filter), ensure_ascii=False)}\n" + if last_result is not None: + user += f"Last hits: {last_result.cardinality}\n" + if last_observation: + user += f"=== Sub-map (hits of last filter) ===\n{last_observation}\n" + if include_full_map: + user += f"=== Full map ===\n{full_map}\n=== End Full Map ===\n" + try: + cached = nav_chat( + purpose=_SCOPE_FILTER_PURPOSE, + model=model, + messages=[ + {"role": "system", "content": _scope_filter_system_prompt(seed=last_filter is None)}, + {"role": "user", "content": user}, + ], + temperature=float(config.llm_temperature), + max_tokens=max(256, int(config.llm_max_tokens or 256)), + response_format={"type": "json_object"}, + context="Nav Scope Filter", + usage_tag="nav_scope_filter", + ) + except NavTokenLimit: + return {"kind": "fallback", "reason": "token_limit"} + + text = str(cached.get("content") or "").strip() + obj = _extract_json_obj(text) or {} + kind = str(obj.get("action") or obj.get("kind") or "filter").strip().lower() + if kind not in {"filter", "done", "widen", "fallback"}: + kind = "filter" + decision_raw = str(obj.get("decision") or "").strip().lower() + decision: Optional[ScopeDecision] = ( + decision_raw if decision_raw in _DECISIONS else None + ) + parsed = parse_node_filter(obj) + return { + "kind": kind, + "filter": parsed, + "decision": decision, + "reason": str(obj.get("reason") or "")[:300], + "raw": text, + } + + +def _scope_filter_system_prompt(*, seed: bool) -> str: + field_schema = ( + '"path"' if seed else '"path|summary"' + ) + field_rule = ( + "Match on section paths only (filenames and title chains)." + if seed + else "Fields AND together; terms inside one field OR together." + ) + base = ( + "You write a WHERE filter over document sections. Return json.\n" + "Schema: {\"action\":\"filter|done|widen|fallback\",\"predicates\":" + f"[{{\"field\":{field_schema},\"terms\":[\"...\"],\"match\":\"substring|regex\"}}]," + "\"decision\":\"collect_all|scoped_harvest|fallback\",\"reason\":\"...\"}\n" + f"{field_rule}\n" + "Write terms as natural-language words from the query or the map " + "(entities, topics, aliases); do not rely on section numbering alone.\n" + ) + if seed: + return base + ( + "Write the first filter for the query against the full map. " + "action=filter returns it; action=fallback only when no path " + "predicate can isolate the target sections." + ) + return base + ( + "You see the sub-map hit by the last filter; judge by its content:\n" + "- sub-map covers the query -> action=done\n" + "- sub-map has off-topic nodes -> action=filter with a narrower " + "revision of the last filter\n" + "- sub-map looks too narrow or misses parts of the query -> " + "action=widen to see the full map once, then revise\n" + "Revise the last filter, never restart from the query. " + "action=fallback only when no predicate can isolate the target sections." + ) + + +def _filter_payload(nf: NodeFilter) -> List[Dict[str, Any]]: + return [ + {"field": p.field, "terms": list(p.terms), "match": p.match} + for p in nf.predicates + ] + + +def _compact_map(ts: Any, doc_ids: Sequence[str]) -> str: + path_fn = getattr(ts, "path_titles", None) + structure_fn = getattr(ts, "get_structure", None) + lines: List[str] = [] + + def path_of(sid: str, doc_id: str) -> str: + if not callable(path_fn): + return sid + try: + return str(path_fn(sid, doc_id) or sid) + except TypeError: + return str(path_fn(sid) or sid) + + def summary_of(sid: str) -> str: + if not callable(structure_fn): + return "" + try: + st = structure_fn(sid) or {} + return str(st.get("summary") or "").strip() + except Exception: + return "" + + provider = getattr(ts, "_provider", None) + child_fn = getattr(provider, "children", None) if provider is not None else None + root_fn = getattr(ts, "sections_for_doc", None) + + for doc_id in doc_ids: + lines.append(path_of(doc_id, doc_id)) + stack = [str(s) for s in (root_fn(doc_id) if callable(root_fn) else []) if str(s)] + seen: set[str] = set() + while stack: + sid = stack.pop(0) + if not sid or sid in seen: + continue + seen.add(sid) + path = path_of(sid, doc_id) + summary = summary_of(sid) + line = path if not summary else f"{path} | {summary}" + lines.append(line) + kids = [str(c) for c in (child_fn(sid) if callable(child_fn) else []) if str(c)] + stack[0:0] = kids + return "\n".join(lines) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_types.py b/packages/shared-python/shared/services/retrieval/nav/nav_types.py index 3c6b5241..aebee142 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_types.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_types.py @@ -5,7 +5,7 @@ from enum import Enum from typing import Any, Dict, List, Literal, Optional, Tuple -NavMode = Literal["navigate", "checklist"] +NavMode = Literal["checklist"] class ActionKind(str, Enum): @@ -58,7 +58,6 @@ class NavConfig: # Recursive dispatch. enable_recursive_dispatch: bool = True max_dispatch_depth: int = 3 - navigate_max_steps: int = 8 subagent_model: str = "" # Scoped maps whose estimated (with-summary) size exceeds this threshold drop # inline summaries (title-only), nudging the agent to DISPATCH deeper rather @@ -69,16 +68,8 @@ class NavConfig: # COMPOSE child score = own_unit + compose_confidence_weight * collect_confidence # (see nav_compose._child_final_score); drives group_key / within-group rank. compose_confidence_weight: float = 0.5 - # Depth-0 group_rank preview budget: over this many chars, skip Assembled - # Evidence / group_rank entirely (title+summary per parent group). - compose_group_rank_max_chars: int = 10000 - # Depth-0 hard rewrite: after agent chooses COLLECT, if branch text length - # exceeds the limit and the node has children, rewrite that sid to DISPATCH. - enable_depth0_oversize_to_dispatch: bool = False - # 0 = use episode evidence budget_chars (set in run_nav_episode). - depth0_oversize_char_limit: int = 0 - # Product mode: navigate = classic map loop; checklist = plan+harvest+control. - mode: NavMode = "navigate" + # Product mode: checklist = plan+harvest+control. + mode: NavMode = "checklist" # Display budget for the one-shot planning map; executor still uses map_char_limit. # 0 = reuse map_char_limit. planning_map_char_limit: int = 10000 @@ -90,12 +81,12 @@ class NavConfig: planner_think_max_tokens: int = 0 # 0 → RETRIEVAL_NAV_TOKEN_LIMIT env / default 100000. token_limit: int = 0 - # Separate from navigate llm_max_tokens: plan JSON is larger. + # Plan JSON is larger than a single harvest action; give the planner more room. planner_llm_max_tokens: int = 1024 - # Harvest multi-id JSON (collect_ids + per-id confidence) needs more than - # navigate's 256 — capped completion truncates mid-object and parses as empty. + # Harvest multi-id JSON (collect_ids + per-id confidence) needs more than a + # bare 256 — capped completion truncates mid-object and parses as empty. harvest_llm_max_tokens: int = 1024 - # Checklist: navigate/harvest cycles per subgoal before drop. Min 1. + # Checklist: harvest cycles per subgoal before drop. Min 1. subgoal_max_attempts: int = 2 # Checklist: 0 = never replan; otherwise hard cap on structural replans. # Default 1 so shared-checklist probe/harness match without silent overrides. @@ -107,6 +98,11 @@ class NavConfig: # Retired: plan_control now shows full prebuilt section summaries (already # head/tail clipped at summary-build time), not a raw-evidence char cut. plan_control_digest_chars: int = 600 + # WHERE node filter (pre-harvest). Off until orchestrate enables a subgoal. + enable_node_filter: bool = False + filter_max_rounds: int = 3 + filter_min_hits: int = 1 + filter_max_hits: int = 40 @property def is_checklist(self) -> bool: @@ -123,7 +119,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "NavConfig": flat["tight_remaining_steps"] = int( budget_modes.get("tight_remaining_steps", cls.tight_remaining_steps) ) - # Retired product flags (collapsed into mode=navigate|checklist). + # Retired product flags (dropped after the navigate loop was removed). for dead in ( "expand_top_k", "map_peek_top_k", @@ -157,10 +153,14 @@ def from_dict(cls, data: Dict[str, Any]) -> "NavConfig": "llm_model_env", "planner_model_env", "subagent_model_env", + # Retired with the multi-step navigate loop. + "navigate_max_steps", + "enable_depth0_oversize_to_dispatch", + "depth0_oversize_char_limit", + "compose_group_rank_max_chars", ): flat.pop(dead, None) - raw_mode = str(flat.get("mode") or "navigate").strip().lower() - flat["mode"] = "checklist" if raw_mode == "checklist" else "navigate" + flat["mode"] = "checklist" allowed = {f.name for f in cls.__dataclass_fields__.values()} cfg = cls(**{k: v for k, v in flat.items() if k in allowed}) if cfg.map_mode and cfg.llm_max_tokens < 256: @@ -175,7 +175,6 @@ class SectionView: preview: str score: float = 0.0 n_lines: int = 0 - n_chunks: int = 0 has_children: bool = False depth_from_scope: int = 0 map_id: str = "" @@ -225,19 +224,6 @@ def prompt_line(self) -> str: return " | ".join(bits) -@dataclass -class RegionReport: - """Result of one navigate(scope, ...) call (top-level or dispatched subagent).""" - - scope: Optional[str] - collected_section_ids: List[str] = field(default_factory=list) - suggestions: List[str] = field(default_factory=list) - summary: str = "" - reason: str = "" - skipped: bool = False - depth: int = 0 - - @dataclass class SubgoalResult: """Typed outcome of one subgoal execution (M5).""" @@ -258,7 +244,7 @@ class NavState: doc_id: str query: str task_type: str = "unknown" - # Working scope for the *current* navigate call (set by navigate(), not a stack). + # Working scope for the current harvest level (set by harvest(), not a stack). current_scope: Optional[str] = None collected_ids: set[str] = field(default_factory=set) collected: List[Tuple[Any, float]] = field(default_factory=list) @@ -270,9 +256,6 @@ class NavState: blocked_collect_section_ids: set[str] = field(default_factory=set) action_history: List[Dict[str, Any]] = field(default_factory=list) refusal_events: List[Dict[str, Any]] = field(default_factory=list) - # Subagent / investigate reports shown to the parent agent. - reports_context: str = "" - investigated_section_ids: set[str] = field(default_factory=set) dismissed_section_ids: set[str] = field(default_factory=set) # Explicit COLLECT confidence by section_id; hydration-only descendants stay 0. collect_confidence: Dict[str, float] = field(default_factory=dict) diff --git a/packages/shared-python/shared/services/retrieval/nav_config.py b/packages/shared-python/shared/services/retrieval/nav_config.py index 7005868e..be1d3dd8 100644 --- a/packages/shared-python/shared/services/retrieval/nav_config.py +++ b/packages/shared-python/shared/services/retrieval/nav_config.py @@ -45,14 +45,10 @@ "map_children_limit": 10000, "enable_recursive_dispatch": True, "max_dispatch_depth": 3, - "navigate_max_steps": 8, "subagent_model": MAPNAV_MODEL, "scope_inline_summary_char_limit": 1500, "scope_inline_summary_budget_mult": 3.0, "compose_confidence_weight": 0.5, - "compose_group_rank_max_chars": 10000, - "enable_depth0_oversize_to_dispatch": True, - "depth0_oversize_char_limit": 500, "mode": "checklist", "planning_map_char_limit": 10000, "planner_max_subgoals": 0, @@ -65,6 +61,10 @@ "max_waves": 0, "max_harvest_depth": 3, "plan_control_digest_chars": 600, + "enable_node_filter": False, + "filter_max_rounds": 3, + "filter_min_hits": 1, + "filter_max_hits": 40, } diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index ab76977f..95436458 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -19,6 +19,7 @@ import time from collections.abc import Mapping from dataclasses import dataclass, field +from hashlib import sha256 from typing import Any from loguru import logger @@ -191,6 +192,9 @@ async def map_unit_discovery( query_tokens = tokenize_query_for_ranker(query) if not query_tokens: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + query_token_hashes = [ + sha256(token.encode("utf-8")).hexdigest() for token in query_tokens + ] revision_join, revision_clause, revision_params = _build_revision_scope( revision_pins @@ -228,30 +232,34 @@ async def map_unit_discovery( if not unit_rows: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) - frequency_result = await db.execute( - text( - ( - cte - if signal_paths - else _SCOPED_UNIT_IDS_CTE.format( - revision_join=revision_join, - revision_clause=revision_clause, - exclude_clause=exclude_clause, - type_clause=type_clause, - ) - ) - + """ - SELECT tokens.map_unit_id, tokens.channel, tokens.token, tokens.frequency - FROM document_map_unit_tokens AS tokens - JOIN scoped_units ON scoped_units.map_unit_id = tokens.map_unit_id - WHERE tokens.channel = ANY(:channels) - AND tokens.token = ANY(:tokens) + frequency_scope_cte = cte if signal_paths else _SCOPED_UNIT_IDS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, + exclude_clause=exclude_clause, + type_clause=type_clause, + ) + frequency_query = text( + "WITH matching_tokens AS MATERIALIZED (" + "SELECT map_unit_id, channel, token, frequency " + "FROM document_map_unit_tokens " + "WHERE channel = ANY(:channels) " + "AND token_hash = ANY(:token_hashes)" + "), " + + frequency_scope_cte.lstrip().removeprefix("WITH ") + + """ + SELECT matching_tokens.map_unit_id, matching_tokens.channel, + matching_tokens.token, matching_tokens.frequency + FROM matching_tokens + JOIN scoped_units + ON scoped_units.map_unit_id = matching_tokens.map_unit_id """ - ), + ) + frequency_result = await db.execute( + frequency_query, { **params, "channels": list(_MAP_SCORE_CHANNELS), - "tokens": query_tokens, + "token_hashes": query_token_hashes, }, ) frequencies: dict[tuple[str, str], dict[str, int]] = {} diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py index 533251e5..f2e24b56 100644 --- a/packages/shared-python/shared/services/retrieval/settings.py +++ b/packages/shared-python/shared/services/retrieval/settings.py @@ -2,7 +2,6 @@ CHANNEL_WEIGHT_PATH = 1.0 CHANNEL_WEIGHT_CONTENT = 2.0 -CHANNEL_WEIGHT_TERM = 1.5 INTERNAL_RECALL_K_MULTIPLIER = 2 RRF_K = 60 DEFAULT_TOP_K = 10 diff --git a/packages/shared-python/shared/services/retrieval/trace/mapnav.py b/packages/shared-python/shared/services/retrieval/trace/mapnav.py index 3f243cd7..6a2e8324 100644 --- a/packages/shared-python/shared/services/retrieval/trace/mapnav.py +++ b/packages/shared-python/shared/services/retrieval/trace/mapnav.py @@ -94,9 +94,6 @@ def _map_one( parent_step_index=parent_step_index, scope=scope, observation={ - "projection_chars": detail.get("projection_chars"), - "hit_section_ids": detail.get("hit_section_ids") or [], - "planning_map_char_limit": detail.get("planning_map_char_limit"), "llm_raw": _clip_raw( plan_payload.get("raw") or detail.get("llm_raw") or detail.get("raw") ), @@ -134,7 +131,6 @@ def _map_one( "attempted": detail.get("attempted"), "dropped": detail.get("dropped"), "waves": detail.get("waves"), - "fallback_navigate": detail.get("fallback_navigate"), }, budget=budget, elapsed_ms=elapsed, @@ -200,68 +196,58 @@ def _map_one( elapsed_ms=elapsed, ) - if action == "search_assets": + if action == "node_filter": + step_action = str(detail.get("action") or detail.get("decision") or "filter") return DecisionTraceStep( step_index=step_index, - agent="asset_inspector", - phase="asset_search", + agent="navigator", + phase="node_filter", parent_step_index=parent_step_index, scope=scope, observation={ - "candidates": detail.get("candidates") - or detail.get("requests") - or [], + "predicates": detail.get("predicates") or [], + "fields": detail.get("fields") or [], + "cardinality": detail.get("cardinality"), + "truncated": detail.get("truncated"), + "failed_predicates": detail.get("failed_predicates") or [], + "matched_section_ids": detail.get("matched_section_ids") or [], + "round": detail.get("round"), }, decision={ - "action": "search_assets", - "asset_ids": detail.get("asset_ids") or [], + "action": step_action, + "reason": detail.get("reason") or "", }, result={ - "status": "ok", - "n_added": detail.get("n_added"), - "subgoal_id": detail.get("subgoal_id"), - "requests": detail.get("requests"), + "status": "fallback" if step_action == "fallback" else "ok", + "cardinality": detail.get("cardinality"), + "decision": detail.get("decision") or "", + "reason": detail.get("reason") or "", }, budget=budget, elapsed_ms=elapsed, ) - if action in { - "nav_collect", - "nav_dispatch", - "nav_finish", - "nav_dispatch_skipped", - } or action.startswith("nav_"): - kind = { - "nav_collect": "COLLECT", - "nav_dispatch": "DISPATCH", - "nav_finish": "FINISH", - "nav_dispatch_skipped": "SKIP", - }.get(action) - if kind is None and action.startswith("nav_"): - kind = action[4:].upper() or "STEP" + if action == "search_assets": return DecisionTraceStep( step_index=step_index, - agent="navigator", - phase="navigate", + agent="asset_inspector", + phase="asset_search", parent_step_index=parent_step_index, scope=scope, observation={ - "legal_actions_preview": detail.get("legal_actions_preview") or [], - "projection_chars": detail.get("projection_chars"), - "n_legal_actions": detail.get("n_legal_actions"), - "llm_raw": _clip_raw(detail.get("llm_raw") or detail.get("raw")), + "candidates": detail.get("candidates") + or detail.get("requests") + or [], }, decision={ - "action": kind, - "action_id": detail.get("action_id"), - "ids": detail.get("ids") or detail.get("section_ids"), - "reason": detail.get("reason") or "", + "action": "search_assets", + "asset_ids": detail.get("asset_ids") or [], }, result={ "status": "ok", - "collect_section_ids": detail.get("collect_section_ids") or [], - "kind": detail.get("kind") or kind, + "n_added": detail.get("n_added"), + "subgoal_id": detail.get("subgoal_id"), + "requests": detail.get("requests"), }, budget=budget, elapsed_ms=elapsed, @@ -369,7 +355,7 @@ def build_decision_trace( steps_in: Sequence[Any] = list(getattr(episode, "steps", None) or ()) out: list[DecisionTraceStep] = [] harvest_parent_by_depth: dict[int, int] = {} - layer_counts = {"planner": 0, "harvest": 0, "control": 0, "navigate": 0} + layer_counts = {"planner": 0, "harvest": 0, "control": 0} for raw_step in steps_in: action = str(getattr(raw_step, "action", "") or "").strip() @@ -401,12 +387,10 @@ def build_decision_trace( harvest_parent_by_depth[depth] = mapped.step_index if mapped.phase == "plan": layer_counts["planner"] += 1 - elif mapped.phase in {"harvest", "plan_wave", "asset_search"}: + elif mapped.phase in {"harvest", "plan_wave", "asset_search", "node_filter"}: layer_counts["harvest"] += 1 elif mapped.phase == "plan_control": layer_counts["control"] += 1 - elif mapped.phase == "navigate": - layer_counts["navigate"] += 1 stop = str(getattr(episode, "stop_reason", "") or "completed") last_budget = out[-1].budget if out and out[-1].budget else { diff --git a/packages/shared-python/shared/tests/test_nav_bridge_config.py b/packages/shared-python/shared/tests/test_nav_bridge_config.py index ceac134c..bfd42c33 100644 --- a/packages/shared-python/shared/tests/test_nav_bridge_config.py +++ b/packages/shared-python/shared/tests/test_nav_bridge_config.py @@ -30,8 +30,6 @@ def test_build_nav_config_is_checklist_map_trim_stack() -> None: assert cfg.is_checklist assert cfg.map_mode is True assert cfg.policy == "llm" - assert cfg.enable_depth0_oversize_to_dispatch is True - assert cfg.depth0_oversize_char_limit == 500 assert cfg.subgoal_max_attempts == 2 assert cfg.max_replans == 1 assert cfg.max_waves == 0 diff --git a/packages/shared-python/shared/tests/test_nav_node_filter.py b/packages/shared-python/shared/tests/test_nav_node_filter.py new file mode 100644 index 00000000..7e059d6c --- /dev/null +++ b/packages/shared-python/shared/tests/test_nav_node_filter.py @@ -0,0 +1,169 @@ +"""Unit tests for deterministic WHERE node filter (in-memory hierarchy).""" + +from __future__ import annotations + +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + NamespaceKnowhereProvider, + SectionRow, +) +from shared.services.retrieval.nav.nav_node_filter import ( + apply_node_filter, + field_predicate, + node_filter, + render_submap_observation, +) + + +def _section( + section_id: str, + parent: str | None, + path: str, + title: str, + *, + level: int, + summary: str = "", + order: int = 0, +) -> SectionRow: + return SectionRow( + section_id=section_id, + parent_section_id=parent, + section_path=path, + section_title=title, + section_level=level, + summary=summary, + sort_order=order, + ) + + +def _namespace_ts() -> ProviderToolSpace: + apple = KnowhereProvider( + doc_id="doc_apple", + sections=[ + _section("sec_root_a", None, "Root", "Root", level=0, order=0), + _section( + "sec_q3", + "sec_root_a", + "Q3 Results", + "Q3 Results", + level=1, + summary="Apple quarterly profit and revenue", + order=1, + ), + _section( + "sec_hw", + "sec_root_a", + "Hardware", + "Hardware", + level=1, + summary="iPhone unit sales", + order=2, + ), + ], + units=(), + ) + orange = KnowhereProvider( + doc_id="doc_orange", + sections=[ + _section("sec_root_o", None, "Root", "Root", level=0, order=0), + _section( + "sec_crop", + "sec_root_o", + "Crop Report", + "Crop Report", + level=1, + summary="orange harvest yield", + order=1, + ), + ], + units=(), + ) + provider = NamespaceKnowhereProvider( + [apple, orange], + titles={ + "doc_apple": "AAPL 10-K.pdf", + "doc_orange": "Citrus Outlook.docx", + }, + ) + return ProviderToolSpace(provider) + + +def test_path_filter_returns_complete_set_across_documents() -> None: + ts = _namespace_ts() + result = apply_node_filter( + ts, + ["doc_apple", "doc_orange"], + node_filter([field_predicate("path", ["AAPL", "Apple"])]), + ) + + assert result.truncated is False + assert result.failed_predicates == [] + assert result.matched_doc_ids == ["doc_apple"] + assert "sec_q3" in result.matched_section_ids + assert "sec_hw" in result.matched_section_ids + assert "sec_crop" not in result.matched_section_ids + assert result.cardinality == len(result.matched_section_ids) + assert result.cardinality >= 2 + + +def test_summary_filter_and_path_and_summary() -> None: + ts = _namespace_ts() + summary_only = apply_node_filter( + ts, + ["doc_apple", "doc_orange"], + node_filter([field_predicate("summary", ["profit"])]), + ) + assert summary_only.matched_section_ids == ["sec_q3"] + + both = apply_node_filter( + ts, + ["doc_apple", "doc_orange"], + node_filter( + [ + field_predicate("path", ["AAPL", "Apple"]), + field_predicate("summary", ["profit"]), + ] + ), + ) + assert both.matched_section_ids == ["sec_q3"] + assert both.matched_doc_ids == ["doc_apple"] + + +def test_zero_hits_and_invalid_regex_recorded() -> None: + ts = _namespace_ts() + empty = apply_node_filter( + ts, + ["doc_apple", "doc_orange"], + node_filter([field_predicate("path", ["zzz-not-present"])]), + ) + assert empty.cardinality == 0 + assert empty.matched_section_ids == [] + assert empty.matched_doc_ids == [] + + bad = apply_node_filter( + ts, + ["doc_apple", "doc_orange"], + node_filter([field_predicate("path", ["(unclosed"], match="regex")]), + ) + assert bad.cardinality == 0 + assert bad.matched_section_ids == [] + assert bad.failed_predicates == ["path:regex:invalid"] + + +def test_regex_or_terms_and_preview_budget() -> None: + ts = _namespace_ts() + result = apply_node_filter( + ts, + ["doc_apple", "doc_orange"], + node_filter( + [field_predicate("summary", ["profit|yield"], match="regex")] + ), + ) + assert set(result.matched_section_ids) == {"sec_q3", "sec_crop"} + assert result.cardinality == 2 + + preview = render_submap_observation(ts, result) + assert preview.startswith("hits=2") + assert "sec_q3" not in preview # paths shown, not ids + assert "Q3 Results" in preview diff --git a/packages/shared-python/shared/tests/test_nav_node_filter_wire.py b/packages/shared-python/shared/tests/test_nav_node_filter_wire.py new file mode 100644 index 00000000..c3d0d7eb --- /dev/null +++ b/packages/shared-python/shared/tests/test_nav_node_filter_wire.py @@ -0,0 +1,318 @@ +"""Orchestrate pre-pass + allowed-scope projection tests.""" + +from __future__ import annotations + +from typing import Any + +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + NamespaceKnowhereProvider, + SectionRow, +) +from shared.services.retrieval.nav.nav_orchestrate import _execute_subgoal_harvest_once +from shared.services.retrieval.nav.nav_plan import RetrievalPlan, Subgoal +from shared.services.retrieval.nav.nav_projection import build_projection +from shared.services.retrieval.nav.nav_scope_filter import ScopeFilterOutcome +from shared.services.retrieval.nav.nav_types import NavConfig, NavState + + +def _section( + section_id: str, + parent: str | None, + path: str, + title: str, + *, + level: int, + summary: str = "", + order: int = 0, +) -> SectionRow: + return SectionRow( + section_id=section_id, + parent_section_id=parent, + section_path=path, + section_title=title, + section_level=level, + summary=summary, + sort_order=order, + ) + + +def _ts() -> ProviderToolSpace: + apple = KnowhereProvider( + doc_id="doc_apple", + sections=[ + _section("sec_root_a", None, "Root", "Root", level=0, order=0), + _section( + "sec_q3", + "sec_root_a", + "Q3 Results", + "Q3 Results", + level=1, + summary="profit", + order=1, + ), + _section( + "sec_hw", + "sec_root_a", + "Hardware", + "Hardware", + level=1, + summary="iphone", + order=2, + ), + ], + units=(), + ) + return ProviderToolSpace( + NamespaceKnowhereProvider([apple], titles={"doc_apple": "AAPL 10-K.pdf"}) + ) + + +def _cfg(**kwargs: Any) -> NavConfig: + data = { + "enable_node_filter": True, + "mode": "checklist", + "map_mode": True, + "llm_model": "test-model", + } + data.update(kwargs) + return NavConfig.from_dict(data) + + +def _subgoal(*, use_filter: bool = True) -> Subgoal: + return Subgoal( + id="s1", + need="apple profit", + retrieval_query="apple profit", + use_node_filter=use_filter, + ) + + +def test_collect_all_skips_harvest(monkeypatch: Any) -> None: + harvest_calls: list[dict[str, Any]] = [] + collect_calls: list[list[str]] = [] + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", + lambda *args, **kwargs: ScopeFilterOutcome( + decision="collect_all", + settled_section_ids=["sec_q3"], + settled_doc_ids=["doc_apple"], + rounds=1, + reason="small", + ), + ) + + def fake_harvest(*args: Any, **kwargs: Any) -> Any: + harvest_calls.append({"args": args, "kwargs": kwargs}) + raise AssertionError("harvest should not run") + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_harvest.harvest", + fake_harvest, + ) + + def fake_collect(ts: Any, state: NavState, chosen: Any, config: Any) -> dict[str, Any]: + del ts, config + batch = list((chosen.metadata or {}).get("batch_actions") or [chosen]) + sids = [str(a.section_id) for a in batch] + collect_calls.append(sids) + state.explicit_collect_ids.update(sids) + state.collected_section_ids.update(sids) + return {"collect_section_ids": sids} + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_navigate._apply_collect", + fake_collect, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._relit_map", + lambda *args, **kwargs: _nullcontext(), + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", + lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), + ) + + out = _execute_subgoal_harvest_once( + _ts(), + NavState(doc_id="doc_apple", query="apple profit"), + _cfg(), + RetrievalPlan(subgoals=[_subgoal()]), + _subgoal(), + steps_out=[], + ) + assert harvest_calls == [] + assert collect_calls == [["sec_q3"]] + assert out["harvest"]["reason"] == "node_filter_collect_all" + + +def test_scoped_harvest_passes_allowed_ids(monkeypatch: Any) -> None: + seen: dict[str, Any] = {} + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", + lambda *args, **kwargs: ScopeFilterOutcome( + decision="scoped_harvest", + settled_section_ids=["sec_q3"], + settled_doc_ids=["doc_apple"], + rounds=1, + reason="medium", + ), + ) + + def fake_harvest(*args: Any, **kwargs: Any) -> Any: + seen["allowed"] = kwargs.get("allowed_section_ids") + return type( + "HR", + (), + { + "n_policy_calls": 1, + "visited_section_ids": [], + "max_depth_hit": False, + "reason": "harvested", + }, + )() + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_harvest.harvest", + fake_harvest, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._relit_map", + lambda *args, **kwargs: _nullcontext(), + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", + lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), + ) + + _execute_subgoal_harvest_once( + _ts(), + NavState(doc_id="doc_apple", query="apple profit"), + _cfg(), + RetrievalPlan(subgoals=[_subgoal()]), + _subgoal(), + steps_out=[], + ) + assert seen["allowed"] == {"sec_q3"} + + +def test_fallback_matches_today(monkeypatch: Any) -> None: + seen: dict[str, Any] = {} + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", + lambda *args, **kwargs: ScopeFilterOutcome( + decision="fallback", + reason="policy_fallback", + ), + ) + + def fake_harvest(*args: Any, **kwargs: Any) -> Any: + seen["allowed"] = kwargs.get("allowed_section_ids") + return type( + "HR", + (), + { + "n_policy_calls": 1, + "visited_section_ids": [], + "max_depth_hit": False, + "reason": "", + }, + )() + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_harvest.harvest", + fake_harvest, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._relit_map", + lambda *args, **kwargs: _nullcontext(), + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", + lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), + ) + + _execute_subgoal_harvest_once( + _ts(), + NavState(doc_id="doc_apple", query="apple profit"), + _cfg(), + RetrievalPlan(subgoals=[_subgoal()]), + _subgoal(), + steps_out=[], + ) + assert seen["allowed"] is None + + +def test_flag_off_skips_pre_pass(monkeypatch: Any) -> None: + called = {"n": 0} + + def boom(*args: Any, **kwargs: Any) -> Any: + called["n"] += 1 + raise AssertionError("scope filter should not run") + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_scope_filter.run_scope_filter", + boom, + ) + + def fake_harvest(*args: Any, **kwargs: Any) -> Any: + return type( + "HR", + (), + { + "n_policy_calls": 1, + "visited_section_ids": [], + "max_depth_hit": False, + "reason": "", + }, + )() + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_harvest.harvest", + fake_harvest, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._relit_map", + lambda *args, **kwargs: _nullcontext(), + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._wave_subgoal_result", + lambda *args, **kwargs: type("Sig", (), {"chars_used": 0})(), + ) + _execute_subgoal_harvest_once( + _ts(), + NavState(doc_id="doc_apple", query="apple profit"), + _cfg(), + RetrievalPlan(subgoals=[_subgoal(use_filter=False)]), + _subgoal(use_filter=False), + steps_out=[], + ) + assert called["n"] == 0 + + +def test_projection_keeps_allowed_and_ancestors() -> None: + ts = _ts() + cfg = _cfg() + projection = build_projection( + ts, + doc_id="", + query="profit", + scope=None, + config=cfg, + allowed_section_ids={"sec_q3"}, + ) + ids = {v.section_id for v in projection.tree_sections} + assert "sec_q3" in ids + assert "sec_hw" not in ids + assert "sec_root_a" in ids or "doc_apple" in ids + + +class _nullcontext: + def __enter__(self) -> None: + return None + + def __exit__(self, *args: Any) -> None: + return None diff --git a/packages/shared-python/shared/tests/test_nav_plan_node_filter.py b/packages/shared-python/shared/tests/test_nav_plan_node_filter.py new file mode 100644 index 00000000..9099034e --- /dev/null +++ b/packages/shared-python/shared/tests/test_nav_plan_node_filter.py @@ -0,0 +1,66 @@ +"""Planner parse compat for use_node_filter.""" + +from __future__ import annotations + +from shared.services.retrieval.nav.nav_plan import ( + _planner_system_prompt, + parse_retrieval_plan, +) + + +def test_parse_missing_filter_flag_defaults_false() -> None: + plan = parse_retrieval_plan( + { + "reason": "old plan", + "map_coverage": "sufficient", + "subgoals": [ + { + "id": "s1", + "need": "profit", + "retrieval_query": "apple profit", + } + ], + }, + query="apple profit", + ) + assert len(plan.subgoals) == 1 + assert plan.subgoals[0].use_node_filter is False + assert plan.subgoals[0].node_filter_predicates == [] + + +def test_parse_use_node_filter_and_seed() -> None: + plan = parse_retrieval_plan( + { + "reason": "enum tickers", + "map_coverage": "sufficient", + "subgoals": [ + { + "id": "s1", + "need": "apple profit", + "retrieval_query": "AAPL profit", + "use_node_filter": True, + "node_filter": [ + { + "field": "path", + "terms": ["AAPL", "Apple"], + "match": "substring", + } + ], + } + ], + }, + query="apple profit", + ) + sg = plan.subgoals[0] + assert sg.use_node_filter is True + assert sg.node_filter_predicates == [ + {"field": "path", "terms": ["AAPL", "Apple"], "match": "substring"} + ] + + +def test_planner_prompt_mentions_where_vs_fuzzy() -> None: + text = _planner_system_prompt(max_subgoals=0) + assert "use_node_filter" in text + assert "fallback" in text + assert "only the user query" in text.lower() + assert "folded" not in text.lower() diff --git a/packages/shared-python/shared/tests/test_nav_plan_query_only.py b/packages/shared-python/shared/tests/test_nav_plan_query_only.py new file mode 100644 index 00000000..d94ce9f2 --- /dev/null +++ b/packages/shared-python/shared/tests/test_nav_plan_query_only.py @@ -0,0 +1,253 @@ +"""Planner is query-only; map lighting happens after the plan exists.""" + +from __future__ import annotations + +from typing import Any + +from shared.services.retrieval.nav.nav_orchestrate import seed_episode_map_scores_from_plan +from shared.services.retrieval.nav.nav_plan import ( + RetrievalPlan, + Subgoal, + _planner_system_prompt, + fallback_plan, +) +from shared.services.retrieval.nav.nav_types import NavConfig, NavState + + +def test_planner_system_prompt_is_query_only() -> None: + text = _planner_system_prompt(max_subgoals=3) + assert "only the user query" in text.lower() + assert "no document map" in text.lower() + assert "folded" not in text.lower() + assert "Hit" not in text + assert "collect=C" not in text + assert "whether the user query alone is enough" in text.lower() + + +def test_seed_episode_scores_plan_retrieval_queries(monkeypatch: Any) -> None: + seen: dict[str, Any] = {} + + def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): + seen["queries"] = list(queries) + seen["top_k"] = top_k + out = {} + for i, q in enumerate(queries): + out[q] = ({f"n{i}": 1.0}, {f"u{i}": 2.0}, [f"u{i}"]) + return out + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", + fake_relight, + ) + state = NavState(doc_id="", query="user original", task_type="unknown") + state.relit_map_cache = {"stale": ({}, {}, [])} + plan = RetrievalPlan( + subgoals=[ + Subgoal(id="s1", need="a", retrieval_query="alpha terms"), + Subgoal(id="s2", need="b", retrieval_query="beta terms"), + Subgoal(id="s3", need="c", retrieval_query="alpha terms"), + ] + ) + n = seed_episode_map_scores_from_plan( + object(), state, NavConfig(collect_top_k=7), plan + ) + assert n == 2 + assert seen["queries"] == ["alpha terms", "beta terms"] + assert seen["top_k"] == 7 + assert "stale" not in state.relit_map_cache + assert state.map_scores == {"n0": 1.0} + assert state.unit_scores == {"u0": 2.0} + assert state.highlight_ids == ["u0"] + assert "alpha terms" in state.relit_map_cache + assert "beta terms" in state.relit_map_cache + + +def test_seed_episode_fallback_uses_user_query(monkeypatch: Any) -> None: + seen: list[str] = [] + + def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): + seen.extend(queries) + return {queries[0]: ({"n": 1.0}, {"u": 1.0}, ["u"])} + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", + fake_relight, + ) + state = NavState(doc_id="doc1", query="心血管", task_type="unknown") + plan = RetrievalPlan( + subgoals=[Subgoal(id="s1", need="should not score", retrieval_query="")] + ) + n = seed_episode_map_scores_from_plan( + object(), state, NavConfig(collect_top_k=5), plan + ) + assert n == 1 + assert seen == ["心血管"] + assert state.map_scores == {"n": 1.0} + + +def test_seed_skips_unbound_slot_queries(monkeypatch: Any) -> None: + seen: list[str] = [] + + def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): + seen.extend(queries) + return {q: ({f"n:{q}": 1.0}, {f"u:{q}": 1.0}, [f"u:{q}"]) for q in queries} + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", + fake_relight, + ) + state = NavState(doc_id="", query="user original", task_type="unknown") + plan = RetrievalPlan( + subgoals=[ + Subgoal(id="s1", need="a", retrieval_query="alpha"), + Subgoal( + id="s2", + need="b", + retrieval_query="beta {{s1.entity}}", + depends_on=["s1"], + ), + ] + ) + n = seed_episode_map_scores_from_plan(object(), state, NavConfig(), plan) + assert n == 1 + assert seen == ["alpha"] + assert "alpha" in state.relit_map_cache + assert not any("beta" in q for q in state.relit_map_cache) + + +def test_wave_uses_seeded_cache_and_scores_missing_once(monkeypatch: Any) -> None: + score_calls: list[list[str]] = [] + + def fake_relight(ts: Any, *, doc_id: str, queries: list[str], top_k: int): + score_calls.append(list(queries)) + return { + q: ({f"n:{q}": 1.0}, {f"u:{q}": 1.0}, [f"u:{q}"]) + for q in queries + } + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_map_scores.relight_maps_for_queries", + fake_relight, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._apply_plan_control", + lambda *args, **kwargs: {"done": True}, + ) + + harvest_queries: list[str] = [] + + def fake_harvest(ts: Any, state: NavState, config: NavConfig, **kwargs: Any): + harvest_queries.append(str(kwargs.get("query") or "")) + return type( + "HR", + (), + { + "n_policy_calls": 0, + "visited_section_ids": [], + "max_depth_hit": False, + "reason": "test", + }, + )() + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate._harvest_after_node_filter", + fake_harvest, + ) + + state = NavState(doc_id="", query="user original", task_type="unknown") + plan = RetrievalPlan( + subgoals=[ + Subgoal(id="s1", need="a", retrieval_query="alpha"), + Subgoal(id="s2", need="b", retrieval_query="beta"), + ] + ) + state.retrieval_plan = plan + seed_episode_map_scores_from_plan(object(), state, NavConfig(), plan) + assert score_calls == [["alpha", "beta"]] + + from shared.services.retrieval.nav.nav_orchestrate import execute_plan + + execute_plan(object(), state, NavConfig(), episode_query="user original") + # Wave relight must hit the seeded cache; no second full-corpus scoring. + assert score_calls == [["alpha", "beta"]] + assert harvest_queries == ["alpha", "beta"] + + +def test_run_nav_episode_scores_after_plan(monkeypatch: Any) -> None: + order: list[str] = [] + + class _TS: + def sections_for_doc(self, _doc_id: str) -> list[str]: + return ["sec_root"] + + monkeypatch.setattr( + "shared.services.retrieval.nav._compat.load_llm_env", + lambda: None, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav._compat.require_llm_env", + lambda *_args, **_kwargs: None, + ) + + def fake_plan_query(state: NavState, config: NavConfig, **_kwargs: Any): + order.append("plan") + assert state.map_scores == {} + assert state.unit_scores == {} + return fallback_plan(state.query, reason="test_plan") + + def fake_seed(ts: Any, state: NavState, config: NavConfig, plan: RetrievalPlan): + order.append("score") + assert plan.subgoals + state.map_scores = {"sec_root": 1.0} + state.unit_scores = {"u1": 1.0} + state.highlight_ids = ["u1"] + return 1 + + def fake_execute(ts: Any, state: NavState, config: NavConfig, **_kwargs: Any): + order.append("orch") + assert state.map_scores == {"sec_root": 1.0} + return {"waves": []} + + class _Fill: + scored_chunks = [] + kept_chunks = [] + evidence_text = "" + evidence_chars_actual = 0 + n_chunks_kept = 0 + truncated_last = False + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_plan.plan_query", + fake_plan_query, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate.seed_episode_map_scores_from_plan", + fake_seed, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_orchestrate.execute_plan", + fake_execute, + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_agent.pack_nav_evidence", + lambda *_args, **_kwargs: _Fill(), + ) + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_agent.uses_document_nodes", + lambda _ts: False, + ) + + from shared.services.retrieval.nav.nav_agent import _run_nav_episode_body + + result = _run_nav_episode_body( + None, + "心血管", + doc_id="doc1", + budget_chars=2000, + compose_answer=False, + policy="llm", + config=NavConfig(policy="llm"), + toolspace=_TS(), + ) + assert order == ["plan", "score", "orch"] + assert result.section_ids == ["sec_root"] diff --git a/packages/shared-python/shared/tests/test_nav_scope_filter.py b/packages/shared-python/shared/tests/test_nav_scope_filter.py new file mode 100644 index 00000000..b762e085 --- /dev/null +++ b/packages/shared-python/shared/tests/test_nav_scope_filter.py @@ -0,0 +1,275 @@ +"""Bounded WHERE scope-filter loop tests (mocked policy LLM).""" + +from __future__ import annotations + +import json +from typing import Any, List + +from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace +from shared.services.retrieval.nav.nav_knowhere import ( + KnowhereProvider, + NamespaceKnowhereProvider, + SectionRow, +) +from shared.services.retrieval.nav.nav_node_filter import field_predicate, node_filter +from shared.services.retrieval.nav.nav_scope_filter import run_scope_filter +from shared.services.retrieval.nav.nav_types import NavConfig + + +def _section( + section_id: str, + parent: str | None, + path: str, + title: str, + *, + level: int, + summary: str = "", + order: int = 0, +) -> SectionRow: + return SectionRow( + section_id=section_id, + parent_section_id=parent, + section_path=path, + section_title=title, + section_level=level, + summary=summary, + sort_order=order, + ) + + +def _ts() -> ProviderToolSpace: + apple = KnowhereProvider( + doc_id="doc_apple", + sections=[ + _section("sec_root_a", None, "Root", "Root", level=0, order=0), + _section( + "sec_q3", + "sec_root_a", + "Q3 Results", + "Q3 Results", + level=1, + summary="Apple quarterly profit", + order=1, + ), + ], + units=(), + ) + filler = KnowhereProvider( + doc_id="doc_other", + sections=[ + _section("sec_root_o", None, "Root", "Root", level=0, order=0), + _section( + "sec_misc", + "sec_root_o", + "Notes", + "Notes", + level=1, + summary="unrelated notes", + order=1, + ), + ], + units=(), + ) + return ProviderToolSpace( + NamespaceKnowhereProvider( + [apple, filler], + titles={"doc_apple": "AAPL 10-K.pdf", "doc_other": "Misc.docx"}, + ) + ) + + +def _cfg(**kwargs: Any) -> NavConfig: + data = { + "enable_node_filter": True, + "filter_max_rounds": 3, + "filter_min_hits": 1, + "filter_max_hits": 40, + "llm_model": "test-model", + "llm_max_tokens": 256, + } + data.update(kwargs) + return NavConfig.from_dict(data) + + +def _install_script(monkeypatch: Any, replies: List[dict[str, Any]]) -> None: + queue = list(replies) + + def fake_nav_chat(**kwargs: Any) -> dict[str, Any]: + del kwargs + obj = queue.pop(0) if queue else {"action": "fallback", "reason": "empty"} + return {"content": json.dumps(obj)} + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_llm.nav_chat", + fake_nav_chat, + ) + + +def test_widen_relooks_full_map(monkeypatch: Any) -> None: + seen_users: List[str] = [] + queue: List[dict[str, Any]] = [ + {"action": "widen", "reason": "sub-map too narrow"}, + {"action": "done", "decision": "collect_all", "reason": "ok"}, + ] + + def fake_nav_chat(**kwargs: Any) -> dict[str, Any]: + messages = kwargs.get("messages") or [] + seen_users.append(str(messages[-1].get("content") or "")) + obj = queue.pop(0) if queue else {"action": "fallback", "reason": "empty"} + return {"content": json.dumps(obj)} + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_llm.nav_chat", + fake_nav_chat, + ) + out = run_scope_filter( + _ts(), + _cfg(), + query="apple q3 profit", + doc_ids=["doc_apple", "doc_other"], + seed_filter=node_filter([field_predicate("path", ["Q3"])]), + ) + assert out.decision == "collect_all" + assert len(seen_users) == 2 + assert "Full map" not in seen_users[0] + assert "Sub-map" in seen_users[0] + assert "Full map" in seen_users[1] + assert "Last filter" in seen_users[1] + + +def test_zero_hits_widen_then_done(monkeypatch: Any) -> None: + _install_script( + monkeypatch, + [ + { + "action": "filter", + "predicates": [ + {"field": "path", "terms": ["AAPL"], "match": "substring"} + ], + }, + {"action": "done", "decision": "collect_all", "reason": "ok"}, + ], + ) + out = run_scope_filter( + _ts(), + _cfg(), + query="apple q3 profit", + doc_ids=["doc_apple", "doc_other"], + seed_filter=node_filter([field_predicate("path", ["zzz-missing"])]), + ) + assert out.decision == "collect_all" + assert "sec_q3" in out.settled_section_ids + assert out.rounds == 2 + + +def test_zero_hits_policy_fallback(monkeypatch: Any) -> None: + _install_script( + monkeypatch, + [{"action": "fallback", "reason": "cannot widen"}], + ) + out = run_scope_filter( + _ts(), + _cfg(), + query="missing topic", + doc_ids=["doc_apple", "doc_other"], + seed_filter=node_filter([field_predicate("path", ["zzz-missing"])]), + ) + assert out.decision == "fallback" + assert out.reason == "cannot widen" + assert out.rounds == 1 + + +def test_too_many_hits_tighten(monkeypatch: Any) -> None: + _install_script( + monkeypatch, + [ + { + "action": "filter", + "predicates": [ + {"field": "summary", "terms": ["profit"], "match": "substring"} + ], + }, + {"action": "done", "decision": "collect_all", "reason": "tight"}, + ], + ) + out = run_scope_filter( + _ts(), + _cfg(filter_max_hits=1), + query="everything", + doc_ids=["doc_apple", "doc_other"], + seed_filter=node_filter([field_predicate("path", ["Root"])]), + ) + assert out.decision == "collect_all" + assert out.settled_section_ids == ["sec_q3"] + assert out.rounds == 2 + + +def test_max_rounds_hard_stop(monkeypatch: Any) -> None: + _install_script( + monkeypatch, + [ + { + "action": "filter", + "predicates": [ + {"field": "path", "terms": ["zzz"], "match": "substring"} + ], + }, + { + "action": "filter", + "predicates": [ + {"field": "path", "terms": ["still-missing"], "match": "substring"} + ], + }, + ], + ) + out = run_scope_filter( + _ts(), + _cfg(filter_max_rounds=2), + query="no hits", + doc_ids=["doc_apple"], + seed_filter=node_filter([field_predicate("path", ["absent"])]), + ) + assert out.decision == "fallback" + assert out.rounds == 2 + assert out.reason == "max_rounds_out_of_band" + + +def test_cardinality_drives_scoped_harvest(monkeypatch: Any) -> None: + _install_script( + monkeypatch, + [{"action": "done", "reason": "keep"}], + ) + out = run_scope_filter( + _ts(), + _cfg(filter_min_hits=1, filter_max_hits=40), + query="apple", + doc_ids=["doc_apple", "doc_other"], + seed_filter=node_filter([field_predicate("path", ["AAPL"])]), + ) + # filename + Root + Q3 → more than min_hits → scoped_harvest + assert out.decision == "scoped_harvest" + assert out.rounds == 1 + assert "sec_q3" in out.settled_section_ids + + +def test_disabled_skips_policy(monkeypatch: Any) -> None: + called = {"n": 0} + + def boom(**kwargs: Any) -> dict[str, Any]: + called["n"] += 1 + raise AssertionError("policy should not run") + + monkeypatch.setattr( + "shared.services.retrieval.nav.nav_llm.nav_chat", + boom, + ) + out = run_scope_filter( + _ts(), + _cfg(enable_node_filter=False), + query="apple", + doc_ids=["doc_apple"], + seed_filter=node_filter([field_predicate("path", ["AAPL"])]), + ) + assert out.decision == "fallback" + assert out.reason == "disabled" + assert called["n"] == 0 diff --git a/packages/shared-python/shared/tests/test_nav_trace_map.py b/packages/shared-python/shared/tests/test_nav_trace_map.py index 6c022ce6..e25f6cfe 100644 --- a/packages/shared-python/shared/tests/test_nav_trace_map.py +++ b/packages/shared-python/shared/tests/test_nav_trace_map.py @@ -124,6 +124,43 @@ def test_build_decision_trace_maps_three_layers_and_terminal() -> None: assert steps[-1].result["layer_llm_steps"]["control"] >= 1 +def test_node_filter_steps_map_and_count_tokens() -> None: + episode = SimpleNamespace( + stop_reason="completed", + evidence_chars_actual=0, + steps=[ + _step( + "node_filter", + { + "round": 1, + "predicates": [ + {"field": "path", "terms": ["AAPL"], "match": "substring"} + ], + "fields": ["path"], + "cardinality": 2, + "action": "done", + "decision": "collect_all", + "reason": "small_cardinality", + "matched_section_ids": ["sec_q3"], + "token_limit": 100000, + "tokens_used_total": 80, + "tokens_used_delta": 80, + "elapsed_ms": 12, + }, + ) + ], + ) + steps = build_decision_trace(episode, evidence_char_budget=12000, n_refs=0) + assert steps[0].phase == "node_filter" + assert steps[0].observation["cardinality"] == 2 + assert steps[0].observation["fields"] == ["path"] + assert steps[0].decision["action"] == "done" + assert steps[0].result["decision"] == "collect_all" + assert steps[0].budget["tokens_used_delta"] == 80 + assert steps[0].budget["token_limit"] == 100000 + assert steps[-1].result["layer_llm_steps"]["harvest"] >= 1 + + def test_episode_helpers() -> None: episode = SimpleNamespace( steps=[