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

apps/ai-service/rag/answer.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import re
 from dataclasses import dataclass, replace
 
 from . import grounding, metrics as metric_names
@@ -57,14 +58,24 @@ class GroundedAnswerService:
         query: str,
         subject_scope: SubjectScope,
         intent: QueryIntent,
+        drug_id: str | None = None,
     ) -> GroundedAnswer:
-        result = self._routing.retrieve(query, subject_scope, intent)
+        # When the caller already resolved the drug (e.g. the conversational
+        # layer, incl. an inherited follow-up), retrieve for it directly instead
+        # of re-resolving from the turn text — re-resolution from a rewritten
+        # turn is what abstained good follow-ups as "ambiguous".
+        if drug_id is not None:
+            result = self._routing.retrieve_for_drug(
+                query, drug_id, subject_scope, intent
+            )
+        else:
+            result = self._routing.retrieve(query, subject_scope, intent)
         if result.decision == EvidenceDecision.ABSTAIN:
             self._metrics.increment(metric_names.ABSTENTION, reason=result.reason)
             return GroundedAnswer(result, None)
 
-        citations = self._citations(result)
-        if citations is None:
+        indexed = self._indexed_citations(result)
+        if indexed is None:
             return GroundedAnswer(
                 replace(
                     result,
@@ -74,6 +85,7 @@ class GroundedAnswerService:
                 ),
                 None,
             )
+        all_citations = tuple(citation for _, citation in indexed)
         if result.decision == EvidenceDecision.VERIFY_PDF:
             # Never generated over. A quarantined table or formula is exactly
             # the evidence whose numbers were not reliably reconstructed, so
@@ -82,7 +94,7 @@ class GroundedAnswerService:
                 result,
                 "Nguồn có bảng hoặc công thức cần đối chiếu trực tiếp với ảnh PDF; "
                 "không tự động trích số liệu.",
-                citations,
+                all_citations,
             )
 
         evidence_texts = tuple(item.text for item in result.evidence)
@@ -90,7 +102,12 @@ class GroundedAnswerService:
             f"{text} [{index}]" for index, text in enumerate(evidence_texts, start=1)
         )
 
-        generated = self._generate(query, evidence_texts)
+        generated = self._generate(query, evidence_texts, intro=result.is_drug_overview)
+        answer_text = extractive if generated is None else generated
+        # Show only the sources the answer actually cited, not every chunk that
+        # was retrieved — a paragraph that cites [4] must not drag 13 citation
+        # chips onto the screen. Falls back to all when the text cites nothing.
+        citations = self._cited_only(indexed, answer_text) or all_citations
         if generated is None:
             self._metrics.increment(metric_names.ANSWER_EXTRACTIVE)
             return GroundedAnswer(result, extractive, citations)
@@ -98,12 +115,14 @@ class GroundedAnswerService:
         self._metrics.increment(metric_names.GENERATION_SERVED)
         return GroundedAnswer(result, generated, citations, generated=True)
 
-    def _generate(self, query: str, evidence_texts: tuple[str, ...]) -> str | None:
+    def _generate(
+        self, query: str, evidence_texts: tuple[str, ...], intro: bool = False
+    ) -> str | None:
         """A verified generation, or None to fall back to the source text."""
         if self._generator is None or not evidence_texts:
             return None
 
-        request = build_request(query, evidence_texts)
+        request = build_request(query, evidence_texts, intro=intro)
         try:
             raw = self._generator.generate(request.system, request.user, request.schema)
         except AnswerGenerationUnavailable:
@@ -144,9 +163,21 @@ class GroundedAnswerService:
         return answer
 
     @staticmethod
-    def _citations(result: RetrievalResult) -> tuple[Citation, ...] | None:
-        citations = []
-        for evidence in result.evidence:
+    def _cited_only(
+        indexed: list[tuple[int, Citation]], answer_text: str
+    ) -> tuple[Citation, ...]:
+        """Keep citations whose 1-based evidence marker [n] appears in the text."""
+        used = {int(m) for m in re.findall(r"\[(\d+)\]", answer_text)}
+        return tuple(citation for index, citation in indexed if index in used)
+
+    @staticmethod
+    def _indexed_citations(
+        result: RetrievalResult,
+    ) -> list[tuple[int, Citation]] | None:
+        """Citations tagged with the 1-based evidence index the prompt gives them,
+        so the response can show only the ones the answer cited."""
+        citations: list[tuple[int, Citation]] = []
+        for index, evidence in enumerate(result.evidence, start=1):
             if not evidence.source_refs:
                 return None
             for source in evidence.source_refs:
@@ -157,7 +188,7 @@ class GroundedAnswerService:
                     start = end = source.printed_page
                 else:
                     return None
-                citations.append(Citation(
+                citations.append((index, Citation(
                     chunk_id=evidence.matched_doc_id,
                     printed_page_start=int(start),
                     printed_page_end=int(end),
@@ -169,5 +200,5 @@ class GroundedAnswerService:
                     # real crop path wins; otherwise the block id plus the
                     # structured page/bbox fields is enough to render later.
                     attachment=source.source_crop or source.block_id,
-                ))
-        return tuple(citations)
+                )))
+        return citations