Wire token-budget packing into the overview/rerank fallback path

apps/ai-service/rag/service.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
+from .context import pack_evidence
 from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit
 from .ports import (
     ParentStore,
@@ -32,6 +33,13 @@ class EvidencePolicy:
     # A free-form question about a resolved drug otherwise hands the LLM the
     # entire monograph; rerank trims it to the sections that actually answer.
     rerank_top_k: int = 6
+    # Token budget for the overview/rerank fallback's evidence pool (2026-08-10
+    # pipeline audit, priority #1). Replaces a flat evidence_limit COUNT: 3
+    # short chunks wastes budget a real model has, 3 long ones can silently
+    # exceed it. Never applied to the deterministic section route — a whole
+    # section is the answer there, and a truncated contraindication list
+    # reads as a complete one (see `_section_hits`'s own comment).
+    max_context_tokens: int = 6000
     # symptom_to_drug: a common symptom can match far more drugs than is
     # useful to show in one answer.
     indication_candidate_limit: int = 8
@@ -94,6 +102,9 @@ class RetrievalService:
             if find_by_section is not None:
                 hits = find_by_section(drug_id, section_key)
                 if hits:
+                    hits = hits + self._pooled_neighbour_hits(
+                        query, drug_id, section_key, find_by_section
+                    )
                     return self._decide(self._hydrate(hits, limit=None))
 
         overview_hits = self._drug_overview(drug_id)
@@ -110,9 +121,12 @@ class RetrievalService:
         # Capped even when rerank is disabled/unavailable and fails open to
         # the unfiltered list — an ordering aid must never remove the size
         # bound too, or the same 29-section explosion returns through here.
-        return self._decide(
-            self._hydrate(overview_hits, limit=self._policy.evidence_limit)
-        )
+        # Token budget (not a flat count): rerank already put the best match
+        # first, so packing in that order keeps as much of it as a real
+        # model's context can hold instead of an arbitrary fixed count.
+        hydrated = self._hydrate(overview_hits, limit=None)
+        packed = pack_evidence(hydrated, max_tokens=self._policy.max_context_tokens)
+        return self._decide(packed.evidence)
 
     def retrieve(self, query: str, drug_id: str) -> RetrievalResult:
         if not query.strip() or not drug_id.strip():
@@ -235,6 +249,28 @@ class RetrievalService:
         hits = find_by_drug(drug_id)
         return hits or None
 
+    # Bounded: a query naming a specific condition can legitimately need one
+    # neighbouring section (than_trong -> chong_chi_dinh was the case found
+    # live); more than this starts pooling tangential sections rather than
+    # the one true answer, which is precision loss dressed as recall.
+    _MAX_LEXICAL_POOLED_SECTIONS = 2
+    # Measured live 2026-08-10 on the real corpus: a genuine neighbour match
+    # (chong_chi_dinh, the true positive) scored 7 matched terms; the same
+    # query's false-positive attractor scored 6 — close enough that no
+    # threshold alone separates them (see the exclusion below instead). 5
+    # keeps clearly-incidental overlap (3-4, seen on unrelated sections in
+    # the same measurement) out while still admitting real matches.
+    _LEXICAL_POOL_MIN_SCORE = 5.0
+    # `duoc_ly_va_co_che_tac_dung` is the corpus's documented false-positive
+    # attractor (see `sections.py`'s own module docstring: it's the largest,
+    # most generic section and "sits close to any question about the drug")
+    # — true for embedding similarity there, and measured true for lexical
+    # overlap here too: it scored a close second (6) right behind the real
+    # answer (7) on the exact query that motivated this pooling mechanism.
+    # Excluded from pooling outright rather than trusting a score margin
+    # that isn't reliably wide enough on its own.
+    _LEXICAL_POOL_EXCLUDED_SECTIONS = frozenset({"duoc_ly_va_co_che_tac_dung"})
+
     def _section_hits(self, query: str, drug_id: str) -> list[SearchHit] | None:
         """Hits for an explicitly named section, or None to fall back.
 
@@ -250,19 +286,43 @@ class RetrievalService:
         if match is None:
             return None
         hits = find_by_section(drug_id, match.section_key)
-        if match.section_key == "than_trong":
-            # A "thận trọng" question about a specific condition sometimes has
-            # its real answer filed under "chống chỉ định" instead — found live
-            # 2026-08-10: Aspirin's own "thận trọng" text never says "loét dạ
-            # dày", the fact only exists in its "chống chỉ định" text ("loét
-            # dạ dày hoặc tá tràng đang hoạt động"). The two are the closest
-            # pair of "is this safe for my patient" categories the book has,
-            # and chống chỉ định text is short — pooling it costs nothing on
-            # a drug where than_trong already answers, and prevents a false
-            # "not in this source" clarify/abstain on one where it doesn't.
-            hits = hits + find_by_section(drug_id, "chong_chi_dinh")
+        hits = hits + self._pooled_neighbour_hits(query, drug_id, match.section_key, find_by_section)
         return hits or None
 
+    def _pooled_neighbour_hits(
+        self, query: str, drug_id: str, resolved_section: str, find_by_section
+    ) -> list[SearchHit]:
+        """Other sections of the SAME drug whose text lexically matches the
+        query strongly enough to suggest the resolved section alone may not
+        answer it — found live 2026-08-10: a "thận trọng" question about a
+        specific condition (loét dạ dày) had its real answer filed only
+        under "chống chỉ định" instead, a category the deterministic
+        keyword route never considers once "thận trọng" itself matched.
+        `search_lexical` generalizes the one hardcoded pairing this started
+        as into a query-driven check across every section, still bounded and
+        still whole-section (never a partial, out-of-context fragment).
+        """
+        search_lexical = getattr(self._retriever, "search_lexical", None)
+        if search_lexical is None:
+            return []
+        lexical_hits = search_lexical(query, drug_id, limit=20)
+        pooled: list[SearchHit] = []
+        ineligible = {resolved_section} | self._LEXICAL_POOL_EXCLUDED_SECTIONS
+        pooled_sections: set[str] = set()
+        for hit in lexical_hits:
+            section = hit.document.section_key
+            if (
+                section in ineligible
+                or section in pooled_sections
+                or hit.score < self._LEXICAL_POOL_MIN_SCORE
+            ):
+                continue
+            pooled_sections.add(section)
+            pooled.extend(find_by_section(drug_id, section))
+            if len(pooled_sections) >= self._MAX_LEXICAL_POOLED_SECTIONS:
+                break
+        return pooled
+
     def decide(self, evidence: tuple[Evidence, ...]) -> RetrievalResult:
         """Public entry point for a caller that assembles its own evidence
         pool across several `retrieve_framed` calls — e.g. `RagAgent`'s