51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
from rag.conversation import (
|
|
ConversationState,
|
|
DeterministicSummariser,
|
|
InMemoryConversationStore,
|
|
Turn,
|
|
)
|
|
|
|
|
|
def test_store_returns_fresh_state_for_unknown_id():
|
|
store = InMemoryConversationStore()
|
|
state = store.load("conv-new")
|
|
assert state.conversation_id == "conv-new"
|
|
assert state.turn_count == 0
|
|
assert state.recent == ()
|
|
|
|
|
|
def test_store_round_trips_saved_state():
|
|
store = InMemoryConversationStore()
|
|
state = ConversationState("conv-1", summary="s", turn_count=3)
|
|
store.save(state)
|
|
assert store.load("conv-1") is state
|
|
|
|
|
|
def test_summariser_records_topic_labels_only():
|
|
s = DeterministicSummariser()
|
|
dropped = (
|
|
Turn("user", "Chống chỉ định của metformin?", "t0",
|
|
drug_id="metformin", section_key="chong_chi_dinh"),
|
|
Turn("assistant", "Quá mẫn với metformin, suy thận Clcr < 60...", "t1",
|
|
drug_id="metformin", section_key="chong_chi_dinh"),
|
|
)
|
|
out = s.fold("", dropped)
|
|
# The label line is present...
|
|
assert "chong_chi_dinh của metformin" in out
|
|
# ...and no clinical value leaked from the assistant turn.
|
|
assert "Clcr" not in out
|
|
assert "60" not in out
|
|
|
|
|
|
def test_summariser_stays_within_budget_dropping_oldest():
|
|
s = DeterministicSummariser()
|
|
dropped = tuple(
|
|
Turn("user", f"q{i}", f"t{i}", drug_id=f"drug{i}", section_key="lieu_luong")
|
|
for i in range(400)
|
|
)
|
|
out = s.fold("", dropped)
|
|
assert len(out) <= DeterministicSummariser.MAX_CHARS
|
|
# Most-recent topic survives, oldest is dropped.
|
|
assert "drug399" in out
|
|
assert "drug0 " not in out
|