Fix live multi-turn: pass the resolved drug, stop did-you-mean garbage

apps/ai-service/rag/service.py

@@ -3,15 +3,35 @@ from __future__ import annotations
 from dataclasses import dataclass
 
 from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit
-from .ports import ParentStore, QueryEmbeddingUnavailable, Retriever
+from .ports import (
+    ParentStore,
+    QueryEmbeddingUnavailable,
+    Reranker,
+    RerankUnavailable,
+    Retriever,
+)
 from .sections import SectionResolver
 
 
+# The sections that introduce a drug: what it is, its class, its main use, its
+# mechanism — in book order. A bare drug name is answered from these, not from
+# the dosage-forms table that happens to sit near the top of the monograph.
+INTRO_SECTIONS = (
+    "ten_chung_quoc_te",
+    "loai_thuoc",
+    "chi_dinh",
+    "duoc_ly_va_co_che_tac_dung",
+)
+
+
 @dataclass(frozen=True)
 class EvidencePolicy:
     minimum_score: float = 0.12
     candidate_limit: int = 5
     evidence_limit: int = 3
+    # 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
 
 
 class RetrievalService:
@@ -30,11 +50,13 @@ class RetrievalService:
         parent_store: ParentStore,
         policy: EvidencePolicy | None = None,
         section_resolver: SectionResolver | None = None,
+        reranker: Reranker | None = None,
     ) -> None:
         self._retriever = retriever
         self._parent_store = parent_store
         self._policy = policy or EvidencePolicy()
         self._section_resolver = section_resolver
+        self._reranker = reranker
 
     def retrieve(self, query: str, drug_id: str) -> RetrievalResult:
         if not query.strip() or not drug_id.strip():
@@ -46,12 +68,23 @@ class RetrievalService:
             # truncated list of contraindications reads as a complete one.
             return self._decide(self._hydrate(section_hits, limit=None))
 
-        # Drug resolved but no attribute named ("PARACETAMOL"): show the whole
-        # monograph, in book order, rather than dead-ending on "specify an
-        # attribute". A drug reference answers a drug name with the drug.
+        # Drug resolved but no attribute named. A *bare* drug name ("PARACETAMOL")
+        # shows the whole monograph in book order. A *free-form question* about
+        # the drug ("sốt cao uống được không?") would otherwise dump all ~29
+        # sections at the model; rerank keeps only the sections that answer it.
         overview_hits = self._drug_overview(drug_id)
         if overview_hits is not None:
-            return self._decide(self._hydrate(overview_hits, limit=None))
+            if self._is_question(query):
+                overview_hits = self._rerank(query, overview_hits)
+                return self._decide(self._hydrate(overview_hits, limit=None))
+            # A bare drug name is not a question — introduce the drug from its
+            # identity sections (what it is, its class, its main indication),
+            # not the whole monograph starting with the dosage-forms table.
+            intro = [h for h in overview_hits if h.document.section_key in INTRO_SECTIONS]
+            return self._decide(
+                self._hydrate(intro or overview_hits, limit=None),
+                is_drug_overview=True,
+            )
 
         try:
             hits = self._retriever.search(
@@ -69,7 +102,33 @@ class RetrievalService:
             )
         if not hits or hits[0].score < self._policy.minimum_score:
             return RetrievalResult(EvidenceDecision.ABSTAIN, "insufficient_retrieval_score")
-        return self._decide(self._hydrate(hits))
+        return self._decide(self._hydrate(self._rerank(query, hits)))
+
+    @staticmethod
+    def _is_question(query: str) -> bool:
+        """A bare drug name (one or two tokens) wants the whole monograph; more
+        than that is a question whose overview should be reranked to the point."""
+        return len(query.split()) > 2
+
+    def _rerank(self, query: str, hits: list[SearchHit]) -> list[SearchHit]:
+        """Reorder hits by cross-encoder relevance, keep the top-k.
+
+        Fail-open: no reranker configured, or the provider is unreachable, and
+        the original order is returned unchanged — an ordering aid must never be
+        able to lose an answer. The section route never reaches this.
+        """
+        if self._reranker is None or len(hits) <= 1:
+            return hits
+        try:
+            order = self._reranker.rerank(
+                query,
+                [hit.document.text for hit in hits],
+                top_n=self._policy.rerank_top_k,
+            )
+        except RerankUnavailable:
+            return hits
+        reranked = [hits[index] for index in order if 0 <= index < len(hits)]
+        return reranked[: self._policy.rerank_top_k] or hits
 
     def _drug_overview(self, drug_id: str) -> list[SearchHit] | None:
         """Every prose section of the drug, or None if the store cannot scroll."""
@@ -96,14 +155,21 @@ class RetrievalService:
         hits = find_by_section(drug_id, match.section_key)
         return hits or None
 
-    def _decide(self, evidence: tuple[Evidence, ...]) -> RetrievalResult:
+    def _decide(
+        self, evidence: tuple[Evidence, ...], is_drug_overview: bool = False
+    ) -> RetrievalResult:
         if not evidence:
             return RetrievalResult(EvidenceDecision.ABSTAIN, "parent_hydration_failed")
         if any(not item.source_refs for item in evidence):
             return RetrievalResult(EvidenceDecision.ABSTAIN, "missing_provenance")
         if any(item.requires_visual_check for item in evidence):
             return RetrievalResult(EvidenceDecision.VERIFY_PDF, "visual_verification_required", evidence)
-        return RetrievalResult(EvidenceDecision.ANSWERABLE, "grounded_evidence_available", evidence)
+        return RetrievalResult(
+            EvidenceDecision.ANSWERABLE,
+            "grounded_evidence_available",
+            evidence,
+            is_drug_overview=is_drug_overview,
+        )
 
     def _hydrate(
         self, hits: list[SearchHit], limit: int | None = -1