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

apps/ai-service/tests/test_retrieval_service.py

@@ -59,10 +59,19 @@ def table_service(*, visual: bool = False) -> RetrievalService:
 class _OverviewRetriever:
     """A fake with `find_by_drug`/`find_by_section` (the Qdrant adapter's
     shape) — `InMemoryLexicalRetriever` doesn't implement either, so
-    `retrieve_framed`'s overview path is otherwise untestable."""
+    `retrieve_framed`'s overview path is otherwise untestable.
 
-    def __init__(self, documents: list[RetrievalDocument]) -> None:
+    `search_lexical` is scripted per test (`lexical_hits`), not real text
+    matching — this file is about `retrieve_framed`'s own wiring, not the
+    scorer (see `test_qdrant_adapter.py`/`test_section_routing.py` for
+    that)."""
+
+    def __init__(
+        self, documents: list[RetrievalDocument], lexical_hits: list[SearchHit] = ()
+    ) -> None:
         self._documents = documents
+        self._lexical_hits = lexical_hits
+        self.lexical_calls: list[tuple[str, str]] = []
 
     def find_by_drug(self, drug_id: str) -> list[SearchHit]:
         return [
@@ -80,6 +89,10 @@ class _OverviewRetriever:
     def search(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:
         return []
 
+    def search_lexical(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:  # noqa: ARG002
+        self.lexical_calls.append((query, drug_id))
+        return list(self._lexical_hits)
+
 
 _MONOGRAPH_SECTIONS = (
     "ten_chung_quoc_te", "ma_atc", "loai_thuoc", "dang_thuoc_va_ham_luong",
@@ -89,7 +102,7 @@ _MONOGRAPH_SECTIONS = (
 )
 
 
-def _monograph_service() -> RetrievalService:
+def _monograph_service(max_context_tokens: int = 6000) -> RetrievalService:
     documents = [
         RetrievalDocument(
             doc_id=f"paracetamol::{section}::0", drug_id="paracetamol",
@@ -100,7 +113,7 @@ def _monograph_service() -> RetrievalService:
     ]
     return RetrievalService(
         _OverviewRetriever(documents), InMemoryParentStore([]),
-        EvidencePolicy(evidence_limit=3),
+        EvidencePolicy(evidence_limit=3, max_context_tokens=max_context_tokens),
     )
 
 
@@ -123,13 +136,65 @@ def test_retrieve_framed_overview_answers_from_intro_sections_only():
 
 def test_retrieve_framed_question_without_section_is_capped_even_without_rerank():
     # No reranker configured: `_rerank` fails open and returns everything
-    # unfiltered. Hydration must still bound it — an ordering aid failing
-    # open must not also remove the size cap.
-    result = _monograph_service().retrieve_framed(
+    # unfiltered. Hydration must still bound it by TOKEN budget (2026-08-10:
+    # was a flat evidence_limit count, now pack_evidence) — an ordering aid
+    # failing open must not also remove the size cap.
+    result = _monograph_service(max_context_tokens=50).retrieve_framed(
         "paracetamol", None, "thuốc này có tác dụng phụ gì", is_overview=False
     )
     assert result.decision == EvidenceDecision.ANSWERABLE
-    assert len(result.evidence) <= 3
+    assert len(result.evidence) < len(_MONOGRAPH_SECTIONS)
+
+
+def test_retrieve_framed_packs_overview_by_token_budget_not_flat_count():
+    """2026-08-10 pipeline audit priority #1: a tight token budget can admit
+    FEWER than the old flat evidence_limit=3 when blocks are long, and a
+    generous one can admit MORE when blocks are short — proving this is
+    genuinely token-driven, not a renamed count cap."""
+    tiny_budget_result = _monograph_service(max_context_tokens=20).retrieve_framed(
+        "paracetamol", None, "thuốc này có tác dụng phụ gì", is_overview=False
+    )
+    generous_budget_result = _monograph_service(max_context_tokens=6000).retrieve_framed(
+        "paracetamol", None, "thuốc này có tác dụng phụ gì", is_overview=False
+    )
+    assert len(tiny_budget_result.evidence) < 3
+    assert len(generous_budget_result.evidence) > 3
+
+
+def test_retrieve_framed_pools_lexically_strong_neighbour_section():
+    """The LIVE agent path (`RagAgent` -> understanding -> `retrieve_framed`)
+    must get the same neighbour-pooling `retrieve()` does — found live
+    2026-08-10 that the first version of this fix only wired into
+    `retrieve()`, which the real HTTP request path does not call at all;
+    `retrieve_framed` has its own separate `if section_key:` branch."""
+    documents = [
+        RetrievalDocument(
+            doc_id=f"aspirin::{section}::0", drug_id="aspirin",
+            kind="prose", section_key=section,
+            text=f"Nội dung mục {section}.", source_refs=(SOURCE,),
+        )
+        for section in ("than_trong", "chong_chi_dinh")
+    ]
+    chong_chi_dinh_hit = SearchHit(
+        document=next(d for d in documents if d.section_key == "chong_chi_dinh"),
+        score=1.0,
+    )
+    retriever = _OverviewRetriever(documents, lexical_hits=[
+        SearchHit(document=chong_chi_dinh_hit.document, score=7.0),
+    ])
+    service = RetrievalService(retriever, InMemoryParentStore([]), EvidencePolicy())
+
+    result = service.retrieve_framed(
+        "aspirin", "than_trong",
+        "Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?",
+    )
+
+    assert result.decision == EvidenceDecision.ANSWERABLE
+    returned_sections = {e.matched_doc_id.split("::")[1] for e in result.evidence}
+    assert returned_sections == {"than_trong", "chong_chi_dinh"}
+    assert retriever.lexical_calls == [
+        ("Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?", "aspirin"),
+    ]
 
 
 def test_row_hit_hydrates_complete_parent_and_keeps_citation():