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

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)