Wire the guarded conversational RAG answer layer end-to-end

2026-08-05 14:33:13 +07:00
parent 834d9e51b0
commit ef08b4929e
127 changed files with 37921 additions and 169 deletions
@@ -0,0 +1,305 @@
+"""The bounded reasoning loop.
+
+Understand → plan → retrieve → assess → refine → generate → verify → repair.
+Every edge is bounded, and every budget is decremented **before** the call it
+pays for, so exhaustion degrades to the best answer so far rather than to an
+error.
+
+Two rules hold across every path and are the reason this can be added to a
+formulary at all:
+
+- `grounding.verify` still gates every generated answer. Reasoning chooses what
+  to look up and how to phrase it; it is never a source of facts.
+- A clarify signal bypasses the loop entirely. Asking beats guessing, and the
+  signals are resolver states — ambiguous drug, unresolved attribute — not a
+  model's confidence score.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from typing import Protocol
+
+from . import metrics as metric_names
+from .conversation import ConversationState, ResolvedQuestion
+from .metrics import Metrics, NullMetrics
+
+MAX_RETRIEVAL_ROUNDS = 2
+MAX_REPAIRS = 1
+MAX_LLM_CALLS = 4
+MAX_WALL_CLOCK_MS = 20_000
+
+
+class BudgetExhausted(RuntimeError):
+    """Raised only inside the loop, never surfaced; the loop catches it."""
+
+
+@dataclass
+class TurnBudget:
+    """Mutable on purpose: one budget is threaded through one turn."""
+
+    llm_calls: int = MAX_LLM_CALLS
+    retrieval_rounds: int = MAX_RETRIEVAL_ROUNDS
+    repairs: int = MAX_REPAIRS
+    wall_clock_ms: int = MAX_WALL_CLOCK_MS
+    elapsed_ms: int = 0
+
+    def spend_llm(self) -> None:
+        if self.llm_calls <= 0:
+            raise BudgetExhausted("llm_calls")
+        self.llm_calls -= 1
+
+    def spend_retrieval(self) -> None:
+        if self.retrieval_rounds <= 0:
+            raise BudgetExhausted("retrieval_rounds")
+        self.retrieval_rounds -= 1
+
+    def spend_repair(self) -> None:
+        if self.repairs <= 0:
+            raise BudgetExhausted("repairs")
+        self.repairs -= 1
+
+    def out_of_time(self) -> bool:
+        return self.elapsed_ms >= self.wall_clock_ms
+
+
+class ClarifyReason:
+    AMBIGUOUS_DRUG = "ambiguous_drug"
+    NO_ATTRIBUTE = "no_attribute"
+    MULTI_ATTRIBUTE = "multi_attribute"
+    STILL_INSUFFICIENT = "still_insufficient"
+
+
+@dataclass(frozen=True)
+class Clarification:
+    reason: str
+    question: str
+    options: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True)
+class Sufficiency:
+    """The assessor's verdict on retrieved evidence.
+
+    `missing` must name something specific — a section, a population, a second
+    drug. "Feels incomplete" does not buy a retrieval round; a round is only
+    spent when there is a concrete thing to go and fetch.
+    """
+
+    sufficient: bool
+    missing: str | None = None
+    refined_query: str | None = None
+
+
+class SufficiencyAssessor(Protocol):
+    def assess(
+        self, resolved: ResolvedQuestion, evidence_texts: tuple[str, ...]
+    ) -> Sufficiency: ...
+
+
+class DeterministicAssessor:
+    """The no-LLM default, and the reference for what the port must do.
+
+    Runs offline and is what the loop uses until a provider is enabled. It only
+    reports insufficiency it can *demonstrate* — a population was asked for and
+    no retrieved text mentions it — so it can never spin the loop on a feeling.
+    """
+
+    POPULATION_TERMS = {
+        "nguoi_lon": ("người lớn",),
+        "tre_em": ("trẻ em", "trẻ nhỏ", "trẻ "),
+        "tre_so_sinh": ("sơ sinh",),
+        "phu_nu_co_thai": ("thai", "mang thai"),
+        "phu_nu_cho_con_bu": ("cho con bú", "sữa mẹ"),
+        "nguoi_cao_tuoi": ("người cao tuổi", "người già"),
+        "suy_than": ("suy thận", "clcr"),
+        "suy_gan": ("suy gan",),
+    }
+
+    def assess(
+        self, resolved: ResolvedQuestion, evidence_texts: tuple[str, ...]
+    ) -> Sufficiency:
+        if not evidence_texts:
+            return Sufficiency(False, missing="no_evidence")
+        if resolved.population is None:
+            return Sufficiency(True)
+
+        terms = self.POPULATION_TERMS.get(resolved.population, ())
+        haystack = " ".join(evidence_texts).casefold()
+        if any(term in haystack for term in terms):
+            return Sufficiency(True)
+        return Sufficiency(
+            False,
+            missing=f"population:{resolved.population}",
+            refined_query=f"{resolved.text} {terms[0] if terms else ''}".strip(),
+        )
+
+
+@dataclass(frozen=True)
+class LoopOutcome:
+    """What one turn produced, plus what it cost."""
+
+    answer: str | None
+    clarification: Clarification | None
+    evidence_texts: tuple[str, ...]
+    retrieval_rounds_used: int
+    repairs_used: int
+    stopped_because: str
+    generated: bool = False
+
+
+@dataclass
+class LoopTrace:
+    """Ordered record of stages, for the dashboard and for debugging."""
+
+    stages: list[str] = field(default_factory=list)
+
+    def enter(self, stage: str) -> None:
+        self.stages.append(stage)
+
+
+def clarify_for(
+    reason: str, options: tuple[str, ...] = ()
+) -> Clarification:
+    questions = {
+        ClarifyReason.NO_ATTRIBUTE: (
+            "Anh/chị muốn tra thuộc tính nào của thuốc này?"
+        ),
+        ClarifyReason.AMBIGUOUS_DRUG: (
+            "Câu hỏi có thể ứng với nhiều thuốc. Anh/chị muốn tra thuốc nào?"
+        ),
+        ClarifyReason.MULTI_ATTRIBUTE: (
+            "Câu hỏi nhắc tới nhiều mục. Anh/chị muốn xem mục nào trước?"
+        ),
+        ClarifyReason.STILL_INSUFFICIENT: (
+            "Chưa tìm đủ căn cứ trong Dược thư cho ý này. "
+            "Anh/chị có thể nêu rõ hơn điều cần tra không?"
+        ),
+    }
+    return Clarification(reason, questions[reason], options)
+
+
+class Retrieve(Protocol):
+    def __call__(self, resolved: ResolvedQuestion) -> tuple[str, ...]: ...
+
+
+class Generate(Protocol):
+    def __call__(
+        self, resolved: ResolvedQuestion, evidence: tuple[str, ...], state: ConversationState
+    ) -> str | None: ...
+
+
+def run_turn(
+    state: ConversationState,
+    resolved: ResolvedQuestion,
+    retrieve: Retrieve,
+    generate: Generate,
+    clarify_signals: tuple[str, ...] = (),
+    assessor: SufficiencyAssessor | None = None,
+    budget: TurnBudget | None = None,
+    metrics: Metrics | None = None,
+    trace: LoopTrace | None = None,
+) -> LoopOutcome:
+    """One conversational turn through the bounded loop.
+
+    `clarify_signals` comes from the existing resolvers — ambiguous drug,
+    unresolved section, multi-attribute. They short-circuit before any spend,
+    because a question worth asking is cheaper and safer than a guess.
+    """
+    budget = budget or TurnBudget()
+    assessor = assessor or DeterministicAssessor()
+    metrics = metrics or NullMetrics()
+    trace = trace or LoopTrace()
+
+    trace.enter("understand")
+    if clarify_signals:
+        reason = clarify_signals[0]
+        metrics.increment(metric_names.CLARIFY_ASKED, reason=reason)
+        trace.enter("clarify")
+        return LoopOutcome(
+            answer=None,
+            clarification=clarify_for(reason),
+            evidence_texts=(),
+            retrieval_rounds_used=0,
+            repairs_used=0,
+            stopped_because="clarify_signal",
+        )
+
+    evidence: tuple[str, ...] = ()
+    rounds_used = 0
+    stopped = "sufficient"
+
+    while True:
+        try:
+            budget.spend_retrieval()
+        except BudgetExhausted:
+            stopped = "retrieval_budget"
+            break
+        trace.enter("retrieve")
+        evidence = retrieve(resolved)
+        rounds_used += 1
+
+        trace.enter("assess")
+        verdict = assessor.assess(resolved, evidence)
+        if verdict.sufficient:
+            break
+        if budget.retrieval_rounds <= 0 or budget.out_of_time():
+            stopped = "retrieval_budget"
+            break
+        # A round is spent only on a named gap with a genuinely new query.
+        if not verdict.missing or not verdict.refined_query:
+            stopped = "no_actionable_gap"
+            break
+        if verdict.refined_query == resolved.text:
+            stopped = "query_unchanged"
+            break
+        trace.enter("refine")
+        metrics.increment(metric_names.LOOP_REFINED, missing=verdict.missing)
+        resolved = replace(resolved, text=verdict.refined_query)
+
+    metrics.increment(metric_names.LOOP_ROUNDS, rounds=str(rounds_used))
+
+    if not evidence:
+        trace.enter("clarify")
+        metrics.increment(
+            metric_names.CLARIFY_ASKED, reason=ClarifyReason.STILL_INSUFFICIENT
+        )
+        return LoopOutcome(
+            answer=None,
+            clarification=clarify_for(ClarifyReason.STILL_INSUFFICIENT),
+            evidence_texts=(),
+            retrieval_rounds_used=rounds_used,
+            repairs_used=0,
+            stopped_because="no_evidence",
+        )
+
+    repairs_used = 0
+    answer: str | None = None
+    while True:
+        trace.enter("generate")
+        try:
+            budget.spend_llm()
+        except BudgetExhausted:
+            stopped = "llm_budget"
+            break
+        answer = generate(resolved, evidence, state)
+        if answer is not None:
+            break
+        # `generate` returning None means verification already refused it.
+        try:
+            budget.spend_repair()
+        except BudgetExhausted:
+            stopped = "repair_budget"
+            break
+        repairs_used += 1
+        trace.enter("repair")
+        metrics.increment(metric_names.LOOP_REPAIRED)
+
+    return LoopOutcome(
+        answer=answer,
+        clarification=None,
+        evidence_texts=evidence,
+        retrieval_rounds_used=rounds_used,
+        repairs_used=repairs_used,
+        stopped_because=stopped,
+        generated=answer is not None,
+    )