refactor(retrieval): plan query-only then light from retrieval_query

apps/api/tests/contract/test_retrieval_map_unit_index_contract.py

@@ -59,7 +59,7 @@ def execute(self, statement: str, parameters: object = None) -> None:
                 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 "matching_tokens AS MATERIALIZED" in statement:
+            elif "FROM document_map_unit_tokens" in statement:
                 self.rows = [("unit-frequency", "path", "retrieval", 1)]
             else:
                 self.rows = []
@@ -101,14 +101,15 @@ def close(self) -> None:
     frequency_executions = [
         (statement, parameters)
         for statement, parameters in executions
-        if "matching_tokens AS MATERIALIZED" in statement
+        if "FROM document_map_unit_tokens" in statement
     ]
     assert len(frequency_executions) == 1
     statement, parameters = frequency_executions[0]
-    assert "FROM matching_tokens" in statement
+    assert "map_unit_id = ANY" in statement
     assert "token_hash = ANY" in statement
     assert isinstance(parameters, list)
-    assert parameters[0] == [
+    assert parameters[0] == ["unit-frequency"]
+    assert parameters[1] == [
         "6e51d6a3d90b6a3243d38e6da6b3f31f49867c1360beba83da8ca9630f9672c7"
     ]
 

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:

packages/shared-python/shared/services/retrieval/nav/nav_agent.py

@@ -23,11 +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_types import (
     LegalAction,
     NavConfig,
@@ -395,31 +390,16 @@ 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
-        )
-    _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)
-    )
 
     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(ts, state, cfg)
+    retrieval_plan = plan_query(state, cfg)
     state.retrieval_plan = retrieval_plan
     steps.append(
         AgentStep(
@@ -430,9 +410,6 @@ def _run_nav_episode_body(
                 "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),
         )
@@ -448,7 +425,9 @@ def _run_nav_episode_body(
         from .nav_plan import fallback_plan
 
         state.retrieval_plan = fallback_plan(query, reason="missing_plan")
-    from .nav_orchestrate import execute_plan
+    from .nav_orchestrate import execute_plan, seed_episode_map_scores_from_plan
+
+    seed_episode_map_scores_from_plan(ts, state, cfg, state.retrieval_plan)
 
     orch_t0 = time.perf_counter()
     orch_detail = execute_plan(

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]:

packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py

@@ -435,26 +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(
-                    "WITH matching_tokens AS MATERIALIZED ("
                     "SELECT map_unit_id, channel, token, frequency "
                     "FROM document_map_unit_tokens "
-                    "WHERE token_hash = ANY(%s) AND channel = ANY(%s)"
-                    ") "
-                    "SELECT matching_tokens.map_unit_id, matching_tokens.channel, "
-                    "matching_tokens.token, matching_tokens.frequency "
-                    "FROM matching_tokens "
-                    "JOIN document_map_units AS units "
-                    "ON units.id = matching_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",
-                    [list(query_token_hashes), list(_MAP_SCORE_CHANNELS), *revision_params],
+                    "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)
@@ -502,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,
@@ -708,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]]:
@@ -799,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."""
@@ -1113,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:

packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py

@@ -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()