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

apps/ai-service/rag/artifacts.py

@@ -42,6 +42,9 @@ def load_documents(path: Path) -> list[RetrievalDocument]:
             parent_id=raw.get("parent_id"),
             requires_visual_check=raw.get("requires_visual_check", False),
             drug_name=raw.get("drug_name"),
+            part_index=raw.get("part_index"),
+            part_count=raw.get("part_count"),
+            context_labels=tuple(raw.get("context_labels") or ()),
         ))
     return documents
 

apps/ai-service/rag/context.py

@@ -0,0 +1,61 @@
+"""Evidence-safe context packing for generation."""
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+from .models import Evidence
+
+
+TokenCounter = Callable[[str], int]
+
+
+def conservative_token_count(text: str) -> int:
+    """Dependency-free estimate, conservative for Vietnamese text."""
+    return (len(text) + 2) // 3 if text else 0
+
+
+@dataclass(frozen=True)
+class PackedContext:
+    evidence: tuple[Evidence, ...]
+    omitted_evidence_ids: tuple[str, ...]
+    estimated_tokens: int
+
+
+def pack_evidence(
+    evidence: Sequence[Evidence],
+    *,
+    max_tokens: int,
+    max_items: int | None = None,
+    token_counter: TokenCounter = conservative_token_count,
+    block_overhead_tokens: int = 4,
+) -> PackedContext:
+    """Pack whole blocks in retrieval order; never truncate clinical text."""
+    if max_tokens < 1:
+        raise ValueError("max_tokens must be positive")
+    if max_items is not None and max_items < 1:
+        raise ValueError("max_items must be positive or None")
+    if block_overhead_tokens < 0:
+        raise ValueError("block_overhead_tokens must be non-negative")
+
+    selected: list[Evidence] = []
+    omitted: list[str] = []
+    seen: set[str] = set()
+    used = 0
+    for item in evidence:
+        if item.evidence_id in seen:
+            continue
+        seen.add(item.evidence_id)
+        count = token_counter(item.text)
+        if count < 0:
+            raise ValueError("token_counter must return a non-negative value")
+        cost = count + block_overhead_tokens
+        if (
+            (max_items is not None and len(selected) >= max_items)
+            or used + cost > max_tokens
+        ):
+            omitted.append(item.evidence_id)
+            continue
+        selected.append(item)
+        used += cost
+    return PackedContext(tuple(selected), tuple(omitted), used)

apps/ai-service/rag/expansion.py

@@ -0,0 +1,62 @@
+"""Bounded sibling expansion for split clinical sections."""
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from .models import RetrievalDocument, SearchHit
+
+
+def _compatible(left: RetrievalDocument, right: RetrievalDocument) -> bool:
+    """Prevent expansion across two explicitly different clinical scopes."""
+    if left.drug_id != right.drug_id or left.section_key != right.section_key:
+        return False
+    if left.context_labels and right.context_labels:
+        return bool(set(left.context_labels) & set(right.context_labels))
+    return not left.context_labels and not right.context_labels
+
+
+def expand_siblings(
+    seeds: Sequence[SearchHit],
+    documents: Sequence[RetrievalDocument],
+    *,
+    window: int = 1,
+    limit: int | None = None,
+) -> list[SearchHit]:
+    """Add adjacent parts without crossing drug, section, or scope labels.
+
+    Each seed remains first in its group. Neighbours retain the seed score so
+    later context packing treats them as context for that match, not as a new
+    independently scored retrieval result.
+    """
+    if window < 0:
+        raise ValueError("window must be non-negative")
+    if limit is not None and limit < 0:
+        raise ValueError("limit must be non-negative or None")
+
+    by_position = {
+        (doc.drug_id, doc.section_key, doc.part_index): doc
+        for doc in documents
+        if doc.part_index is not None
+    }
+    output: list[SearchHit] = []
+    seen: set[str] = set()
+
+    for seed in seeds:
+        doc = seed.document
+        group = [doc]
+        if doc.part_index is not None and window:
+            group = []
+            for index in range(doc.part_index - window, doc.part_index + window + 1):
+                sibling = by_position.get((doc.drug_id, doc.section_key, index))
+                if sibling is not None and _compatible(doc, sibling):
+                    group.append(sibling)
+            group.sort(key=lambda item: item.part_index if item.part_index is not None else 0)
+
+        for item in group:
+            if item.doc_id in seen:
+                continue
+            seen.add(item.doc_id)
+            output.append(SearchHit(item, seed.score))
+            if limit is not None and len(output) >= limit:
+                return output
+    return output

apps/ai-service/rag/fusion.py

@@ -0,0 +1,55 @@
+"""Rank fusion primitives for hybrid and multi-query retrieval.
+
+This module is deliberately store-agnostic: dense, lexical, and rewritten-query
+retrievers only need to return ranked ``SearchHit`` lists. Keeping fusion pure
+makes its ordering deterministic and lets the live service add parallel I/O
+without coupling domain code to Qdrant or PostgreSQL.
+"""
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from .models import SearchHit
+
+
+def reciprocal_rank_fusion(
+    rankings: Sequence[Sequence[SearchHit]],
+    *,
+    rank_constant: int = 60,
+    limit: int | None = None,
+) -> list[SearchHit]:
+    """Fuse ranked candidate lists using reciprocal-rank fusion (RRF).
+
+    A document contributes at most once per input ranking. The returned score
+    is the accumulated RRF score, not a provider-specific similarity score, so
+    dense and lexical results remain comparable without score normalization.
+    Ties are stable by first appearance, which keeps results reproducible.
+    """
+    if rank_constant < 1:
+        raise ValueError("rank_constant must be positive")
+    if limit is not None and limit < 0:
+        raise ValueError("limit must be non-negative or None")
+
+    scores: dict[str, float] = {}
+    documents = {}
+    first_seen: dict[str, int] = {}
+    seen_order = 0
+
+    for ranking in rankings:
+        seen_in_ranking: set[str] = set()
+        for rank, hit in enumerate(ranking, start=1):
+            doc_id = hit.document.doc_id
+            if doc_id in seen_in_ranking:
+                continue
+            seen_in_ranking.add(doc_id)
+            if doc_id not in documents:
+                documents[doc_id] = hit.document
+                first_seen[doc_id] = seen_order
+                seen_order += 1
+            scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (rank_constant + rank)
+
+    ordered_ids = sorted(scores, key=lambda doc_id: (-scores[doc_id], first_seen[doc_id]))
+    if limit is not None:
+        ordered_ids = ordered_ids[:limit]
+    return [SearchHit(document=documents[doc_id], score=scores[doc_id]) for doc_id in ordered_ids]
+

apps/ai-service/rag/models.py

@@ -45,6 +45,9 @@ class RetrievalDocument:
     parent_id: str | None = None
     requires_visual_check: bool = False
     drug_name: str | None = None
+    part_index: int | None = None
+    part_count: int | None = None
+    context_labels: tuple[str, ...] = field(default_factory=tuple)
 
 
 @dataclass(frozen=True)

apps/ai-service/rag/prompt.py

@@ -103,6 +103,13 @@ Quy tắc:
   tuổi và cân nặng). "người lớn" thường ĐỦ cho liều người lớn tiêu chuẩn.
 - Câu hỏi KHÔNG về liều (chống chỉ định, tương tác, tác dụng phụ, giới thiệu
   thuốc…) thì thường ĐỦ.
+- Câu hỏi "thận trọng"/"an toàn khi dùng cho bệnh nhân [tình trạng cụ thể]"
+  mà BẰNG CHỨNG có đoạn CHỐNG CHỈ ĐỊNH nói đúng tình trạng đó → ĐỦ. Chống chỉ
+  định là mức cảnh báo MẠNH HƠN thận trọng (không được dùng, thay vì dùng có
+  theo dõi) — nó TRẢ LỜI được câu hỏi "có an toàn không", chỉ là câu trả lời
+  nghiêm trọng hơn người hỏi hình dung, không phải "sách không đề cập". Đừng
+  yêu cầu đúng từ "thận trọng" xuất hiện — khớp theo Ý (tình trạng bệnh nhân),
+  không khớp theo NHÃN mục sách xếp nó vào.
 
 Trả về DUY NHẤT JSON: {"sufficient": bool, "clarifying_question": string|null,
 "quick_replies": string[]}.

apps/ai-service/rag/service.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
+from .context import pack_evidence
 from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit
 from .ports import (
     ParentStore,
@@ -32,6 +33,13 @@ class EvidencePolicy:
     # 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
+    # Token budget for the overview/rerank fallback's evidence pool (2026-08-10
+    # pipeline audit, priority #1). Replaces a flat evidence_limit COUNT: 3
+    # short chunks wastes budget a real model has, 3 long ones can silently
+    # exceed it. Never applied to the deterministic section route — a whole
+    # section is the answer there, and a truncated contraindication list
+    # reads as a complete one (see `_section_hits`'s own comment).
+    max_context_tokens: int = 6000
     # symptom_to_drug: a common symptom can match far more drugs than is
     # useful to show in one answer.
     indication_candidate_limit: int = 8
@@ -94,6 +102,9 @@ class RetrievalService:
             if find_by_section is not None:
                 hits = find_by_section(drug_id, section_key)
                 if hits:
+                    hits = hits + self._pooled_neighbour_hits(
+                        query, drug_id, section_key, find_by_section
+                    )
                     return self._decide(self._hydrate(hits, limit=None))
 
         overview_hits = self._drug_overview(drug_id)
@@ -110,9 +121,12 @@ class RetrievalService:
         # Capped even when rerank is disabled/unavailable and fails open to
         # the unfiltered list — an ordering aid must never remove the size
         # bound too, or the same 29-section explosion returns through here.
-        return self._decide(
-            self._hydrate(overview_hits, limit=self._policy.evidence_limit)
-        )
+        # Token budget (not a flat count): rerank already put the best match
+        # first, so packing in that order keeps as much of it as a real
+        # model's context can hold instead of an arbitrary fixed count.
+        hydrated = self._hydrate(overview_hits, limit=None)
+        packed = pack_evidence(hydrated, max_tokens=self._policy.max_context_tokens)
+        return self._decide(packed.evidence)
 
     def retrieve(self, query: str, drug_id: str) -> RetrievalResult:
         if not query.strip() or not drug_id.strip():
@@ -235,6 +249,28 @@ class RetrievalService:
         hits = find_by_drug(drug_id)
         return hits or None
 
+    # Bounded: a query naming a specific condition can legitimately need one
+    # neighbouring section (than_trong -> chong_chi_dinh was the case found
+    # live); more than this starts pooling tangential sections rather than
+    # the one true answer, which is precision loss dressed as recall.
+    _MAX_LEXICAL_POOLED_SECTIONS = 2
+    # Measured live 2026-08-10 on the real corpus: a genuine neighbour match
+    # (chong_chi_dinh, the true positive) scored 7 matched terms; the same
+    # query's false-positive attractor scored 6 — close enough that no
+    # threshold alone separates them (see the exclusion below instead). 5
+    # keeps clearly-incidental overlap (3-4, seen on unrelated sections in
+    # the same measurement) out while still admitting real matches.
+    _LEXICAL_POOL_MIN_SCORE = 5.0
+    # `duoc_ly_va_co_che_tac_dung` is the corpus's documented false-positive
+    # attractor (see `sections.py`'s own module docstring: it's the largest,
+    # most generic section and "sits close to any question about the drug")
+    # — true for embedding similarity there, and measured true for lexical
+    # overlap here too: it scored a close second (6) right behind the real
+    # answer (7) on the exact query that motivated this pooling mechanism.
+    # Excluded from pooling outright rather than trusting a score margin
+    # that isn't reliably wide enough on its own.
+    _LEXICAL_POOL_EXCLUDED_SECTIONS = frozenset({"duoc_ly_va_co_che_tac_dung"})
+
     def _section_hits(self, query: str, drug_id: str) -> list[SearchHit] | None:
         """Hits for an explicitly named section, or None to fall back.
 
@@ -250,19 +286,43 @@ class RetrievalService:
         if match is None:
             return None
         hits = find_by_section(drug_id, match.section_key)
-        if match.section_key == "than_trong":
-            # A "thận trọng" question about a specific condition sometimes has
-            # its real answer filed under "chống chỉ định" instead — found live
-            # 2026-08-10: Aspirin's own "thận trọng" text never says "loét dạ
-            # dày", the fact only exists in its "chống chỉ định" text ("loét
-            # dạ dày hoặc tá tràng đang hoạt động"). The two are the closest
-            # pair of "is this safe for my patient" categories the book has,
-            # and chống chỉ định text is short — pooling it costs nothing on
-            # a drug where than_trong already answers, and prevents a false
-            # "not in this source" clarify/abstain on one where it doesn't.
-            hits = hits + find_by_section(drug_id, "chong_chi_dinh")
+        hits = hits + self._pooled_neighbour_hits(query, drug_id, match.section_key, find_by_section)
         return hits or None
 
+    def _pooled_neighbour_hits(
+        self, query: str, drug_id: str, resolved_section: str, find_by_section
+    ) -> list[SearchHit]:
+        """Other sections of the SAME drug whose text lexically matches the
+        query strongly enough to suggest the resolved section alone may not
+        answer it — found live 2026-08-10: a "thận trọng" question about a
+        specific condition (loét dạ dày) had its real answer filed only
+        under "chống chỉ định" instead, a category the deterministic
+        keyword route never considers once "thận trọng" itself matched.
+        `search_lexical` generalizes the one hardcoded pairing this started
+        as into a query-driven check across every section, still bounded and
+        still whole-section (never a partial, out-of-context fragment).
+        """
+        search_lexical = getattr(self._retriever, "search_lexical", None)
+        if search_lexical is None:
+            return []
+        lexical_hits = search_lexical(query, drug_id, limit=20)
+        pooled: list[SearchHit] = []
+        ineligible = {resolved_section} | self._LEXICAL_POOL_EXCLUDED_SECTIONS
+        pooled_sections: set[str] = set()
+        for hit in lexical_hits:
+            section = hit.document.section_key
+            if (
+                section in ineligible
+                or section in pooled_sections
+                or hit.score < self._LEXICAL_POOL_MIN_SCORE
+            ):
+                continue
+            pooled_sections.add(section)
+            pooled.extend(find_by_section(drug_id, section))
+            if len(pooled_sections) >= self._MAX_LEXICAL_POOLED_SECTIONS:
+                break
+        return pooled
+
     def decide(self, evidence: tuple[Evidence, ...]) -> RetrievalResult:
         """Public entry point for a caller that assembles its own evidence
         pool across several `retrieve_framed` calls — e.g. `RagAgent`'s