213 lines
8.0 KiB
Python
213 lines
8.0 KiB
Python
"""`rag/agent.py` — the new orchestrator wired live in F-03.
|
|
|
|
Codex's 2026-08-06 review (F-03/F-10) flagged that this module had zero test
|
|
coverage despite being built to replace the live front end. This is the
|
|
first coverage: routing branches with fakes, not an exhaustive golden set
|
|
(that is F-10's job — a production-path regression suite against real
|
|
Qdrant/Bedrock fixtures).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from rag.agent import RagAgent
|
|
from rag.answer import GroundedAnswerService
|
|
from rag.models import Evidence, EvidenceDecision, RetrievalResult, SourceRef
|
|
from rag.understanding import QueryFrame
|
|
|
|
SOURCE = SourceRef(physical_page=100, precision="region", printed_page=100)
|
|
|
|
|
|
def _evidence(text: str) -> Evidence:
|
|
return Evidence(
|
|
evidence_id="e0", matched_doc_id="e0", kind="prose", text=text, score=1.0,
|
|
source_refs=(SOURCE,), hydrated_from_parent=False, requires_visual_check=False,
|
|
)
|
|
|
|
|
|
class _FixedUnderstander:
|
|
def __init__(self, frame: QueryFrame) -> None:
|
|
self._frame = frame
|
|
|
|
def understand(self, turn, history=()):
|
|
return self._frame
|
|
|
|
|
|
class _FixedRetrieval:
|
|
"""Stands in for `RetrievalService.retrieve_framed` — a canned result per
|
|
drug_id regardless of section/query, so these tests assert routing, not
|
|
retrieval (that's `test_retrieval_service.py`'s job)."""
|
|
|
|
def __init__(self, results: dict[str, RetrievalResult]) -> None:
|
|
self._results = results
|
|
|
|
def retrieve_framed(self, drug_id, section_key, query, is_overview=False):
|
|
return self._results.get(
|
|
drug_id, RetrievalResult(EvidenceDecision.ABSTAIN, "not_configured")
|
|
)
|
|
|
|
|
|
def _agent(frame: QueryFrame, results: dict[str, RetrievalResult] | None = None) -> RagAgent:
|
|
# `routing=None`: `answer_from_result` (the only method this path calls)
|
|
# never touches it — see `rag/answer.py`.
|
|
answers = GroundedAnswerService(routing=None)
|
|
return RagAgent(_FixedUnderstander(frame), _FixedRetrieval(results or {}), answers)
|
|
|
|
|
|
def test_smalltalk_does_not_touch_retrieval():
|
|
agent = _agent(QueryFrame(turn_type="smalltalk"))
|
|
reply = agent.handle("chào bạn")
|
|
assert reply.decision == "answerable"
|
|
assert reply.turn_type == "smalltalk"
|
|
assert "Dược thư" in reply.answer
|
|
|
|
|
|
def test_out_of_scope_turn_type_abstains():
|
|
agent = _agent(QueryFrame(turn_type="out_of_scope"))
|
|
reply = agent.handle("cách tiêm truyền tĩnh mạch")
|
|
assert reply.decision == "abstain"
|
|
assert reply.reason == "out_of_scope"
|
|
|
|
|
|
def test_veterinary_phrase_abstains_even_if_the_model_missed_it():
|
|
# The model returned an ordinary-looking frame; the keyword backstop
|
|
# (`rag.policy.looks_non_human`, the same one F-02 wired server-side)
|
|
# still catches it before any retrieval happens.
|
|
agent = _agent(QueryFrame(turn_type="drug_attribute", drugs=("metformin",)))
|
|
reply = agent.handle("liều metformin cho chó bao nhiêu")
|
|
assert reply.decision == "abstain"
|
|
assert reply.reason == "out_of_scope"
|
|
|
|
|
|
def test_unknown_drug_name_is_reported_not_substituted():
|
|
agent = _agent(QueryFrame(turn_type="drug_attribute", unknown_drugs=("aspirinol",)))
|
|
reply = agent.handle("liều aspirinol")
|
|
assert reply.decision == "abstain"
|
|
assert reply.reason == "drug_not_in_formulary"
|
|
assert "aspirinol" in reply.answer
|
|
|
|
|
|
def test_no_drug_named_asks_which_one():
|
|
agent = _agent(QueryFrame(turn_type="drug_attribute"))
|
|
reply = agent.handle("liều dùng bao nhiêu")
|
|
assert reply.decision == "clarify"
|
|
assert reply.reason == "no_drug"
|
|
|
|
|
|
def test_needs_clarify_frame_is_surfaced_directly():
|
|
agent = _agent(QueryFrame(
|
|
turn_type="dosing_calc", drugs=("paracetamol",),
|
|
needs_clarify=True, clarify_reason="Bé mấy tuổi, cân nặng bao nhiêu kg?",
|
|
))
|
|
reply = agent.handle("liều paracetamol cho trẻ em")
|
|
assert reply.decision == "clarify"
|
|
assert reply.clarification == "Bé mấy tuổi, cân nặng bao nhiêu kg?"
|
|
|
|
|
|
def test_single_drug_attribute_retrieves_and_answers():
|
|
result = RetrievalResult(
|
|
EvidenceDecision.ANSWERABLE, "grounded_evidence_available",
|
|
(_evidence("Liều 500 mg mỗi ngày."),), resolved_drug_id="metformin",
|
|
)
|
|
agent = _agent(
|
|
QueryFrame(turn_type="drug_attribute", drugs=("metformin",),
|
|
attribute="lieu_luong_va_cach_dung"),
|
|
{"metformin": result},
|
|
)
|
|
reply = agent.handle("liều metformin")
|
|
assert reply.decision == "answerable"
|
|
assert reply.drugs == ("metformin",)
|
|
assert "500 mg" in reply.answer
|
|
|
|
|
|
def test_interaction_combines_both_drugs_evidence():
|
|
warfarin = RetrievalResult(
|
|
EvidenceDecision.ANSWERABLE, "grounded_evidence_available",
|
|
(_evidence("Tương tác với aspirin làm tăng nguy cơ chảy máu."),),
|
|
)
|
|
aspirin = RetrievalResult(
|
|
EvidenceDecision.ANSWERABLE, "grounded_evidence_available",
|
|
(_evidence("Tương tác với warfarin làm tăng nguy cơ chảy máu."),),
|
|
)
|
|
agent = _agent(
|
|
QueryFrame(turn_type="interaction", drugs=("warfarin", "aspirin")),
|
|
{"warfarin": warfarin, "aspirin": aspirin},
|
|
)
|
|
reply = agent.handle("warfarin với aspirin có dùng chung được không")
|
|
assert reply.decision == "answerable"
|
|
assert reply.turn_type == "interaction"
|
|
assert "chảy máu" in reply.answer
|
|
|
|
|
|
def test_interaction_with_no_evidence_abstains_and_does_not_imply_safety():
|
|
agent = _agent(
|
|
QueryFrame(turn_type="interaction", drugs=("drug_a", "drug_b")), {},
|
|
)
|
|
reply = agent.handle("drug_a và drug_b dùng chung được không")
|
|
assert reply.decision == "abstain"
|
|
assert reply.reason == "no_interaction_evidence"
|
|
assert "KHÔNG có nghĩa là an toàn" in reply.answer
|
|
|
|
|
|
def test_symptom_to_drug_without_a_drug_name_asks_honestly_not_wired_yet():
|
|
agent = _agent(QueryFrame(turn_type="symptom_to_drug", indication="sốt cao"))
|
|
reply = agent.handle("sốt cao uống thuốc gì")
|
|
assert reply.decision == "clarify"
|
|
assert reply.reason == "reverse_lookup_not_ready"
|
|
|
|
|
|
def test_history_is_passed_to_the_understander_on_the_next_turn():
|
|
received_history: list[tuple[str, ...]] = []
|
|
|
|
class _RecordingUnderstander:
|
|
def understand(self, turn, history=()):
|
|
received_history.append(tuple(history))
|
|
return QueryFrame(turn_type="smalltalk")
|
|
|
|
answers = GroundedAnswerService(routing=None)
|
|
agent = RagAgent(_RecordingUnderstander(), _FixedRetrieval({}), answers)
|
|
|
|
agent.handle("chào bạn", conversation_id="c1")
|
|
agent.handle("còn liều thì sao?", conversation_id="c1")
|
|
|
|
assert received_history[0] == ()
|
|
assert any("chào bạn" in line for line in received_history[1])
|
|
|
|
|
|
def test_history_is_isolated_per_conversation_id():
|
|
received_history: list[tuple[str, ...]] = []
|
|
|
|
class _RecordingUnderstander:
|
|
def understand(self, turn, history=()):
|
|
received_history.append(tuple(history))
|
|
return QueryFrame(turn_type="smalltalk")
|
|
|
|
answers = GroundedAnswerService(routing=None)
|
|
agent = RagAgent(_RecordingUnderstander(), _FixedRetrieval({}), answers)
|
|
|
|
agent.handle("xin chào", conversation_id="a")
|
|
agent.handle("liều dùng bao nhiêu", conversation_id="b")
|
|
|
|
# The second call, on a different conversation id, must not see "a"'s turn.
|
|
assert received_history[1] == ()
|
|
|
|
|
|
def test_autocomplete_delegates_to_the_configured_source():
|
|
class _Source:
|
|
def complete(self, prefix, k):
|
|
return ["metformin_id", "metoprolol_id"]
|
|
|
|
answers = GroundedAnswerService(routing=None)
|
|
agent = RagAgent(
|
|
_FixedUnderstander(QueryFrame(turn_type="smalltalk")),
|
|
_FixedRetrieval({}), answers, autocomplete=_Source(),
|
|
)
|
|
assert agent.complete("met") == ["Metformin Id", "Metoprolol Id"]
|
|
|
|
|
|
def test_autocomplete_with_no_source_configured_returns_empty():
|
|
answers = GroundedAnswerService(routing=None)
|
|
agent = RagAgent(
|
|
_FixedUnderstander(QueryFrame(turn_type="smalltalk")),
|
|
_FixedRetrieval({}), answers,
|
|
)
|
|
assert agent.complete("met") == []
|