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

apps/ai-service/tests/test_fusion.py

@@ -0,0 +1,57 @@
+import pytest
+
+from rag.fusion import reciprocal_rank_fusion
+from rag.models import RetrievalDocument, SearchHit, SourceRef
+
+
+SOURCE = SourceRef(physical_page=1, precision="page")
+
+
+def _hit(doc_id: str, score: float = 1.0) -> SearchHit:
+    return SearchHit(
+        RetrievalDocument(
+            doc_id=doc_id,
+            drug_id="drug",
+            kind="prose",
+            text=doc_id,
+            section_key="section",
+            source_refs=(SOURCE,),
+        ),
+        score,
+    )
+
+
+def test_rrf_promotes_candidate_supported_by_dense_and_lexical_rankings() -> None:
+    fused = reciprocal_rank_fusion(
+        [[_hit("dense-only"), _hit("shared")], [_hit("shared"), _hit("lexical-only")]],
+        rank_constant=60,
+    )
+    assert [hit.document.doc_id for hit in fused] == [
+        "shared", "dense-only", "lexical-only",
+    ]
+
+
+def test_rrf_deduplicates_a_document_within_one_ranking() -> None:
+    fused = reciprocal_rank_fusion(
+        [[_hit("duplicate"), _hit("duplicate")], [_hit("other")]],
+        rank_constant=10,
+    )
+    duplicate = next(hit for hit in fused if hit.document.doc_id == "duplicate")
+    assert duplicate.score == pytest.approx(1 / 11)
+
+
+def test_rrf_limit_and_ties_are_deterministic() -> None:
+    fused = reciprocal_rank_fusion(
+        [[_hit("first")], [_hit("second")]], rank_constant=60, limit=1,
+    )
+    assert [hit.document.doc_id for hit in fused] == ["first"]
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [({"rank_constant": 0}, "rank_constant"), ({"limit": -1}, "limit")],
+)
+def test_rrf_rejects_invalid_configuration(kwargs, message: str) -> None:
+    with pytest.raises(ValueError, match=message):
+        reciprocal_rank_fusion([], **kwargs)
+

apps/ai-service/tests/test_qdrant_adapter.py

@@ -159,3 +159,51 @@ def test_find_by_indication_respects_the_limit():
     hits = retriever.find_by_indication("đau", limit=2)
 
     assert len(hits) == 2
+
+
+def _section_payload(drug_id: str, section_key: str, text: str) -> dict:
+    return {
+        "chunk_id": f"{drug_id}__{section_key}__0", "drug_id": drug_id,
+        "drug_name": drug_id.upper(), "section_key": section_key,
+        "chunk_kind": "prose", "text": text,
+        "heading_physical_page": 100, "printed_page_range": [101, 101],
+    }
+
+
+def test_search_lexical_ranks_by_distinct_matched_term_count():
+    client = _FakeScrollClient([
+        _section_payload("aspirin", "than_trong", "Thận trọng với suy thận."),
+        _section_payload(
+            "aspirin", "chong_chi_dinh",
+            "Không dùng cho người có loét dạ dày tá tràng đang hoạt động.",
+        ),
+    ])
+    retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
+
+    hits = retriever.search_lexical(
+        "Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày", "aspirin", limit=5
+    )
+
+    assert [h.document.section_key for h in hits] == ["chong_chi_dinh", "than_trong"]
+    assert hits[0].score > hits[1].score
+
+
+def test_search_lexical_drops_stopword_only_queries():
+    client = _FakeScrollClient([
+        _section_payload("aspirin", "than_trong", "Thận trọng với suy thận."),
+    ])
+    retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
+
+    assert retriever.search_lexical("là gì và của", "aspirin", limit=5) == []
+
+
+def test_search_lexical_excludes_non_matching_sections():
+    client = _FakeScrollClient([
+        _section_payload("aspirin", "than_trong", "Thận trọng với suy thận."),
+        _section_payload("aspirin", "chi_dinh", "Giảm đau hạ sốt."),
+    ])
+    retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
+
+    hits = retriever.search_lexical("loét dạ dày", "aspirin", limit=5)
+
+    assert hits == []

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

apps/ai-service/tests/test_section_routing.py

@@ -56,6 +56,22 @@ class SectionAwareRetriever:
         ]
 
 
+class LexicalAwareRetriever(SectionAwareRetriever):
+    """Adds `search_lexical`, scripted per test rather than doing real
+    text matching — this suite is about `_pooled_neighbour_hits`' bounding
+    and threshold logic, not the lexical scorer itself (see
+    `adapters/qdrant.py`'s own coverage for that)."""
+
+    def __init__(self, docs: list[RetrievalDocument], lexical_hits: list[SearchHit]) -> None:
+        super().__init__(docs)
+        self._lexical_hits = lexical_hits
+        self.lexical_calls: list[tuple[str, str]] = []
+
+    def search_lexical(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:  # noqa: ARG002
+        self.lexical_calls.append((query, drug_id))
+        return self._lexical_hits
+
+
 class SimilarityOnlyRetriever:
     def __init__(self, docs: list[RetrievalDocument]) -> None:
         self._docs = docs
@@ -78,7 +94,8 @@ CONTRA = [
 INDICATION = [_doc("i1", "chi_dinh", "Giảm đau, hạ sốt, chống viêm.")]
 PHARMACOLOGY = [_doc("p1", "duoc_ly_va_co_che_tac_dung", "Ức chế cyclooxygenase.")]
 PRECAUTION = [_doc("t1", "than_trong", "Thận trọng với người suy thận.")]
-ALL_DOCS = CONTRA + INDICATION + PHARMACOLOGY + PRECAUTION
+INTERACTION = [_doc("x1", "tuong_tac_thuoc", "Tương tác với thuốc chống đông máu.")]
+ALL_DOCS = CONTRA + INDICATION + PHARMACOLOGY + PRECAUTION + INTERACTION
 
 
 def _service(retriever, resolver: SectionResolver | None) -> RetrievalService:
@@ -208,22 +225,89 @@ class TestSectionRouting:
         assert retriever.section_calls == []
         assert retriever.search_calls
 
-    def test_than_trong_also_pools_chong_chi_dinh(self) -> None:
+    def test_lexically_strong_neighbour_section_is_pooled_in(self) -> None:
         """A precaution some drug's own "thận trọng" text never mentions can
         still be filed under "chống chỉ định" (found live: Aspirin + loét dạ
-        dày). Pooling both keeps that answerable instead of a false "not in
-        this source" clarify/abstain."""
-        retriever = SectionAwareRetriever(ALL_DOCS)
+        dày). A neighbour section with real term overlap (score above
+        threshold) gets pooled in whole, not just the matching fragment."""
+        retriever = LexicalAwareRetriever(
+            ALL_DOCS, lexical_hits=[SearchHit(CONTRA[1], 6.0)],  # c2: "chong_chi_dinh"
+        )
         result = _service(retriever, SectionResolver()).retrieve(
-            "Thận trọng khi dùng aspirin là gì?", "aspirin"
+            "Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?", "aspirin"
         )
         assert retriever.section_calls == [
             ("aspirin", "than_trong"), ("aspirin", "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"),
+        ]
+        # Whole section pooled (c1-c5), not just the one matching chunk (c2).
         assert {item.evidence_id for item in result.evidence} == {
             "t1", "c1", "c2", "c3", "c4", "c5",
         }
 
+    def test_weak_lexical_match_is_not_pooled(self) -> None:
+        """One incidental token overlap must not drag in a whole tangential
+        section — only real term overlap (>= threshold) counts."""
+        retriever = LexicalAwareRetriever(
+            ALL_DOCS, lexical_hits=[SearchHit(CONTRA[1], 4.0)],
+        )
+        result = _service(retriever, SectionResolver()).retrieve(
+            "Thận trọng khi dùng aspirin là gì?", "aspirin"
+        )
+        assert retriever.section_calls == [("aspirin", "than_trong")]
+        assert {item.evidence_id for item in result.evidence} == {"t1"}
+
+    def test_lexical_pooling_is_bounded(self) -> None:
+        """At most 2 neighbour sections pool in, even if more score above
+        threshold — unbounded pooling is precision loss dressed as recall."""
+        retriever = LexicalAwareRetriever(
+            ALL_DOCS,
+            lexical_hits=[
+                SearchHit(CONTRA[0], 7.0),      # chong_chi_dinh
+                SearchHit(INDICATION[0], 6.0),  # chi_dinh
+                SearchHit(INTERACTION[0], 5.0), # tuong_tac_thuoc — 3rd, over bound
+            ],
+        )
+        result = _service(retriever, SectionResolver()).retrieve(
+            "Thận trọng khi dùng aspirin là gì?", "aspirin"
+        )
+        assert retriever.section_calls == [
+            ("aspirin", "than_trong"),
+            ("aspirin", "chong_chi_dinh"),
+            ("aspirin", "chi_dinh"),
+        ]
+        # tuong_tac_thuoc (3rd-ranked lexical hit) never pooled — the bound
+        # (2 neighbours) already used by chong_chi_dinh + chi_dinh.
+        assert "x1" not in {item.evidence_id for item in result.evidence}
+
+    def test_generic_pharmacology_section_is_never_pooled(self) -> None:
+        """`duoc_ly_va_co_che_tac_dung` is the corpus's documented
+        false-positive attractor (see `sections.py`) — excluded outright,
+        even with a lexical score that would otherwise clear the threshold
+        (measured live: it scored a close second right behind the real
+        answer on the exact query that motivated this whole mechanism)."""
+        retriever = LexicalAwareRetriever(
+            ALL_DOCS, lexical_hits=[SearchHit(PHARMACOLOGY[0], 9.0)],
+        )
+        result = _service(retriever, SectionResolver()).retrieve(
+            "Thận trọng khi dùng aspirin là gì?", "aspirin"
+        )
+        assert retriever.section_calls == [("aspirin", "than_trong")]
+        assert "p1" not in {item.evidence_id for item in result.evidence}
+
+    def test_retriever_without_lexical_search_still_works(self) -> None:
+        """No `search_lexical` on the retriever (interface segregation, same
+        pattern as `find_by_section`/`find_by_drug`) — pooling is skipped,
+        not an error."""
+        retriever = SectionAwareRetriever(ALL_DOCS)
+        result = _service(retriever, SectionResolver()).retrieve(
+            "Thận trọng khi dùng aspirin là gì?", "aspirin"
+        )
+        assert retriever.section_calls == [("aspirin", "than_trong")]
+        assert {item.evidence_id for item in result.evidence} == {"t1"}
+
     def test_named_but_empty_section_falls_back(self) -> None:
         """A drug with no such section must not abstain — similarity still tries."""
         retriever = SectionAwareRetriever(INDICATION + PHARMACOLOGY)