Spaces:
Sleeping
Sleeping
File size: 10,009 Bytes
a6f4696 3d5e588 a6f4696 | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 | """Tests for agents/intake.py.
_months_from_date and _resolve_months: pure date helpers (freezegun).
stream_intake_turn: FakeStream pattern for the streaming generator.
"""
from __future__ import annotations
from unittest.mock import MagicMock, patch
import pytest
from freezegun import freeze_time
from agents.intake import _months_from_date, _resolve_months, stream_intake_turn
from models import PatientProfile
from tests.conftest import FakeStream, make_message, make_text_block, make_tool_use_block
# ---------------------------------------------------------------------------
# _months_from_date
# ---------------------------------------------------------------------------
class TestMonthsFromDate:
@freeze_time("2026-05-01")
def test_same_month_is_zero(self):
assert _months_from_date("2026-05") == 0
@freeze_time("2026-05-01")
def test_one_month_ago(self):
assert _months_from_date("2026-04") == 1
@freeze_time("2026-05-01")
def test_twelve_months_ago(self):
assert _months_from_date("2025-05") == 12
@freeze_time("2026-05-01")
def test_crosses_year_boundary(self):
assert _months_from_date("2025-11") == 6
@freeze_time("2026-05-01")
def test_multi_year(self):
assert _months_from_date("2024-05") == 24
def test_invalid_string_returns_zero(self):
assert _months_from_date("not-a-date") == 0
def test_empty_string_returns_zero(self):
assert _months_from_date("") == 0
@freeze_time("2026-05-01")
def test_future_date_returns_negative(self):
# The function does not clamp — a future onset_date from the LLM produces
# a negative onset_months, which can cause eligibility criteria like
# "onset_months >= 6" to pass incorrectly. Documented here so a clamping
# fix is validated when added.
result = _months_from_date("2030-01")
assert result < 0
# ---------------------------------------------------------------------------
# _resolve_months
# ---------------------------------------------------------------------------
class TestResolveMonths:
@freeze_time("2026-05-01")
def test_prefers_date_key(self):
data = {"onset_date": "2025-05", "onset_months": 99}
assert _resolve_months(data, "onset_date", "onset_months") == 12
def test_falls_back_to_months_key(self):
data = {"onset_months": 18}
assert _resolve_months(data, "onset_date", "onset_months") == 18
def test_missing_both_returns_zero(self):
assert _resolve_months({}, "onset_date", "onset_months") == 0
def test_none_months_returns_zero(self):
data = {"onset_months": None}
assert _resolve_months(data, "onset_date", "onset_months") == 0
# ---------------------------------------------------------------------------
# stream_intake_turn — FakeStream pattern
# ---------------------------------------------------------------------------
def _initial_messages() -> list:
return [
{"role": "user", "content": "Please begin."},
{"role": "assistant", "content": [make_text_block("Hello, I'm Beacon.")]},
{"role": "user", "content": "I have ALS, age 52."},
]
class TestStreamIntakeTurnText:
def test_yields_token_events(self, mock_client):
final_msg = make_message(
content=[make_text_block("What is your ZIP code?")],
stop_reason="end_turn",
)
mock_client.messages.stream.return_value = FakeStream(
tokens=["What ", "is ", "your ", "ZIP?"],
final_message=final_msg,
)
events = list(stream_intake_turn(mock_client, _initial_messages()))
token_events = [e for e in events if e[0] == "token"]
assert token_events == [
("token", "What "), ("token", "is "), ("token", "your "), ("token", "ZIP?")
]
def test_yields_text_event_with_updated_messages(self, mock_client):
final_msg = make_message(
content=[make_text_block("What is your ZIP code?")],
stop_reason="end_turn",
)
mock_client.messages.stream.return_value = FakeStream(
tokens=["What is your ZIP code?"],
final_message=final_msg,
)
events = list(stream_intake_turn(mock_client, _initial_messages()))
text_events = [e for e in events if e[0] == "text"]
assert len(text_events) == 1
kind, text, msgs = text_events[0]
assert "ZIP" in text
assert isinstance(msgs, list)
def test_single_stream_call_for_text_response(self, mock_client):
final_msg = make_message(
content=[make_text_block("OK")], stop_reason="end_turn"
)
mock_client.messages.stream.return_value = FakeStream(tokens=["OK"], final_message=final_msg)
list(stream_intake_turn(mock_client, _initial_messages()))
assert mock_client.messages.stream.call_count == 1
class TestStreamIntakeTurnIdentifyDisease:
def test_identify_disease_yields_reset_stream_and_continues(self, mock_client):
identify_block = make_tool_use_block(
"identify_disease",
{"standardized_name": "Amyotrophic Lateral Sclerosis"},
tool_use_id="tu_id1",
)
# First stream: identify_disease tool_use
first_msg = make_message(content=[identify_block], stop_reason="tool_use")
# Second stream: text response after tool result
second_msg = make_message(
content=[make_text_block("Got it, collecting ALS benchmarks.")],
stop_reason="end_turn",
)
mock_client.messages.stream.side_effect = [
FakeStream(tokens=[], final_message=first_msg),
FakeStream(tokens=["Got it"], final_message=second_msg),
]
events = list(stream_intake_turn(mock_client, _initial_messages()))
assert ("reset_stream",) in events
assert mock_client.messages.stream.call_count == 2
def test_messages_contain_tool_result_after_identify(self, mock_client):
identify_block = make_tool_use_block(
"identify_disease",
{"standardized_name": "Amyotrophic Lateral Sclerosis"},
tool_use_id="tu_id1",
)
first_msg = make_message(content=[identify_block], stop_reason="tool_use")
second_msg = make_message(
content=[make_text_block("Collecting ALSFRS-R.")], stop_reason="end_turn"
)
mock_client.messages.stream.side_effect = [
FakeStream(tokens=[], final_message=first_msg),
FakeStream(tokens=["Collecting"], final_message=second_msg),
]
events = list(stream_intake_turn(mock_client, _initial_messages()))
text_event = next(e for e in events if e[0] == "text")
msgs = text_event[2]
roles = [m["role"] for m in msgs]
assert "user" in roles # tool_result is in a user message
class TestStreamIntakeTurnSubmitProfile:
def _submit_block(self) -> MagicMock:
return make_tool_use_block(
"submit_profile",
{
"disease": "Amyotrophic Lateral Sclerosis",
"age": 52,
"onset_date": "2024-11",
"diagnosis_date": "2025-05",
"zip_code": "02115",
"country_code": "US",
"radius_miles": 100,
"phases": ["2", "3"],
"include_eap": False,
"include_observational": False,
},
tool_use_id="tu_sub1",
)
def test_yields_profile_event(self, mock_client):
submit_block = self._submit_block()
final_msg = make_message(content=[submit_block], stop_reason="tool_use")
mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)
with patch("agents.intake.geocode_zip", return_value=(42.33, -71.10)):
events = list(stream_intake_turn(mock_client, _initial_messages()))
profile_events = [e for e in events if e[0] == "profile"]
assert len(profile_events) == 1
def test_profile_event_contains_patient_profile(self, mock_client):
submit_block = self._submit_block()
final_msg = make_message(content=[submit_block], stop_reason="tool_use")
mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)
with patch("agents.intake.geocode_zip", return_value=(42.33, -71.10)):
events = list(stream_intake_turn(mock_client, _initial_messages()))
_, profile, msgs = next(e for e in events if e[0] == "profile")
assert isinstance(profile, PatientProfile)
assert profile.disease == "Amyotrophic Lateral Sclerosis"
assert profile.age == 52
def test_geocode_failure_sets_zero_coords(self, mock_client):
submit_block = self._submit_block()
final_msg = make_message(content=[submit_block], stop_reason="tool_use")
mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)
with patch("agents.intake.geocode_zip", side_effect=Exception("timeout")):
events = list(stream_intake_turn(mock_client, _initial_messages()))
_, profile, _ = next(e for e in events if e[0] == "profile")
assert profile.lat == 0.0
assert profile.lon == 0.0
def test_generator_terminates_after_profile(self, mock_client):
submit_block = self._submit_block()
final_msg = make_message(content=[submit_block], stop_reason="tool_use")
mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)
with patch("agents.intake.geocode_zip", return_value=(42.33, -71.10)):
events = list(stream_intake_turn(mock_client, _initial_messages()))
# After profile, no more events
profile_idx = next(i for i, e in enumerate(events) if e[0] == "profile")
assert profile_idx == len(events) - 1
|