File size: 1,179 Bytes
ba54ea9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | from datetime import datetime
from recap.models import Event
from recap.retrieval import retrieve
def _ev(eid, body, date_iso="2022-01-01"):
return Event(
id=eid,
date=datetime.fromisoformat(date_iso),
category="lab",
title=body,
source="x",
body=body,
)
def test_retrieves_relevant_events_for_question():
events = [
_ev("a", "Creatinine 1.4 mg/dL — first abnormal reading"),
_ev("b", "Influenza vaccination administered"),
_ev("c", "Hemoglobin A1c 8.2%"),
]
hits = retrieve("when did creatinine become abnormal", events, top_k=2)
assert hits[0].id == "a"
def test_retrieve_returns_at_most_top_k():
events = [_ev(str(i), f"event {i}") for i in range(20)]
hits = retrieve("event", events, top_k=5)
assert len(hits) == 5
def test_retrieve_handles_empty_event_list():
assert retrieve("anything", [], top_k=5) == []
def test_retrieve_falls_back_to_first_k_when_no_match():
events = [_ev("a", "alpha"), _ev("b", "beta"), _ev("c", "gamma")]
hits = retrieve("zzzzz", events, top_k=2)
assert len(hits) == 2 # falls back to first k rather than empty
|