tests/agents/test_eligibility.py

"""Tests for agents/eligibility.py.

Wave 3: _evaluate_deterministic, _resolve_patient_value, _compute_overall (pure).
Wave 4: _parse_criteria, _assess_llm, run_eligibility_check, bulk_parse_and_strip (LLM-mocked).
"""
from __future__ import annotations

import pytest

from agents.eligibility import (
    _assess_llm,
    _compute_overall,
    _evaluate_deterministic,
    _parse_criteria,
    _resolve_patient_value,
    bulk_parse_and_strip,
    run_eligibility_check,
)
from tests.conftest import (
    FakeStream,
    make_message,
    make_text_block,
    make_tool_use_block,
)
from models import (
    CriterionAssessment,
    CriterionVerdict,
    EligibilityCriterion,
    ParsedConstraint,
    PatientProfile,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def make_criterion(
    key: str = "age_years",
    operator: str = ">=",
    value=18,
    unit: str | None = None,
    ctype: str = "inclusion",
    description: str = "Test criterion",
    raw_criteria: str = "Raw criterion text",
) -> EligibilityCriterion:
    return EligibilityCriterion(
        key=key,
        type=ctype,
        description=description,
        raw_criteria=raw_criteria,
        constraint=ParsedConstraint(key=key, operator=operator, value=value, unit=unit),
    )


def make_assessment(verdict: CriterionVerdict) -> CriterionAssessment:
    c = make_criterion()
    return CriterionAssessment(
        criterion=c,
        verdict=verdict,
        reason="test",
        patient_value="42",
        confidence="high",
    )


# ---------------------------------------------------------------------------
# _evaluate_deterministic — every operator × pass/fail
# ---------------------------------------------------------------------------

class TestEvaluateDeterministic:
    # --- >= ---
    def test_gte_passes_when_equal(self):
        a = _evaluate_deterministic(make_criterion(operator=">=", value=18), 18)
        assert a.verdict == CriterionVerdict.PASS

    def test_gte_passes_when_above(self):
        a = _evaluate_deterministic(make_criterion(operator=">=", value=18), 52)
        assert a.verdict == CriterionVerdict.PASS

    def test_gte_fails_when_below(self):
        a = _evaluate_deterministic(make_criterion(operator=">=", value=18), 16)
        assert a.verdict == CriterionVerdict.FAIL

    # --- <= ---
    def test_lte_passes_when_equal(self):
        a = _evaluate_deterministic(make_criterion(operator="<=", value=75), 75)
        assert a.verdict == CriterionVerdict.PASS

    def test_lte_passes_when_below(self):
        a = _evaluate_deterministic(make_criterion(operator="<=", value=75), 52)
        assert a.verdict == CriterionVerdict.PASS

    def test_lte_fails_when_above(self):
        a = _evaluate_deterministic(make_criterion(operator="<=", value=75), 80)
        assert a.verdict == CriterionVerdict.FAIL

    # --- == ---
    def test_eq_passes(self):
        a = _evaluate_deterministic(make_criterion(operator="==", value=0), 0)
        assert a.verdict == CriterionVerdict.PASS

    def test_eq_fails(self):
        a = _evaluate_deterministic(make_criterion(operator="==", value=0), 1)
        assert a.verdict == CriterionVerdict.FAIL

    # --- != ---
    def test_ne_passes(self):
        a = _evaluate_deterministic(make_criterion(operator="!=", value=0), 1)
        assert a.verdict == CriterionVerdict.PASS

    def test_ne_fails(self):
        a = _evaluate_deterministic(make_criterion(operator="!=", value=0), 0)
        assert a.verdict == CriterionVerdict.FAIL

    # --- in ---
    def test_in_passes(self):
        a = _evaluate_deterministic(make_criterion(operator="in", value=[0, 1]), 1)
        assert a.verdict == CriterionVerdict.PASS

    def test_in_fails(self):
        a = _evaluate_deterministic(make_criterion(operator="in", value=[0, 1]), 2)
        assert a.verdict == CriterionVerdict.FAIL

    # --- not_in ---
    def test_not_in_passes(self):
        a = _evaluate_deterministic(make_criterion(operator="not_in", value=[2, 3]), 1)
        assert a.verdict == CriterionVerdict.PASS

    def test_not_in_fails(self):
        a = _evaluate_deterministic(make_criterion(operator="not_in", value=[2, 3]), 2)
        assert a.verdict == CriterionVerdict.FAIL

    # --- between ---
    def test_between_passes_at_lower_bound(self):
        a = _evaluate_deterministic(make_criterion(operator="between", value=[18, 75]), 18)
        assert a.verdict == CriterionVerdict.PASS

    def test_between_passes_at_upper_bound(self):
        a = _evaluate_deterministic(make_criterion(operator="between", value=[18, 75]), 75)
        assert a.verdict == CriterionVerdict.PASS

    def test_between_passes_mid(self):
        a = _evaluate_deterministic(make_criterion(operator="between", value=[18, 75]), 52)
        assert a.verdict == CriterionVerdict.PASS

    def test_between_fails_below(self):
        a = _evaluate_deterministic(make_criterion(operator="between", value=[18, 75]), 17)
        assert a.verdict == CriterionVerdict.FAIL

    def test_between_fails_above(self):
        a = _evaluate_deterministic(make_criterion(operator="between", value=[18, 75]), 76)
        assert a.verdict == CriterionVerdict.FAIL

    # --- output shape ---
    def test_confidence_is_high(self):
        a = _evaluate_deterministic(make_criterion(operator=">=", value=18), 52)
        assert a.confidence == "high"

    def test_patient_value_stringified(self):
        a = _evaluate_deterministic(make_criterion(operator=">=", value=18), 52)
        assert a.patient_value == "52"

    def test_reason_contains_operator_and_value(self):
        a = _evaluate_deterministic(make_criterion(operator=">=", value=18), 52)
        assert ">=" in a.reason
        assert "18" in a.reason

    def test_reason_includes_unit_when_present(self):
        c = make_criterion(operator=">=", value=60, unit="mL/min")
        a = _evaluate_deterministic(c, 65)
        assert "mL/min" in a.reason

    def test_type_error_in_comparison_yields_fail(self):
        # Comparing int to a list value that can't be compared numerically
        c = make_criterion(operator=">=", value="not-a-number")
        a = _evaluate_deterministic(c, 52)
        assert a.verdict == CriterionVerdict.FAIL

    # --- unit normalization: years → months ---

    def test_onset_years_converted_to_months_fails_correctly(self):
        # Trial says "onset >= 18 years"; patient has 20 months — should FAIL (20 < 216)
        c = make_criterion(key="symptom_onset_months", operator=">=", value=18, unit="years")
        a = _evaluate_deterministic(c, 20)
        assert a.verdict == CriterionVerdict.FAIL
        assert "216" in a.reason
        assert "months" in a.reason

    def test_onset_months_unit_unchanged(self):
        # Trial says "onset >= 18 months"; patient has 20 months — should PASS, value stays 18
        c = make_criterion(key="symptom_onset_months", operator=">=", value=18, unit="months")
        a = _evaluate_deterministic(c, 20)
        assert a.verdict == CriterionVerdict.PASS
        assert "18 months" in a.reason

    def test_diagnosis_years_converted_to_months(self):
        # Trial says "diagnosis >= 2 years"; patient has 18 months — should FAIL (18 < 24)
        c = make_criterion(key="diagnosis_months", operator=">=", value=2, unit="years")
        a = _evaluate_deterministic(c, 18)
        assert a.verdict == CriterionVerdict.FAIL
        assert "24" in a.reason

    def test_between_years_normalized_to_months(self):
        # Trial says "onset between 1 and 2 years" → [12, 24] months; patient 20 months → PASS
        c = make_criterion(key="symptom_onset_months", operator="between", value=[1, 2], unit="years")
        a = _evaluate_deterministic(c, 20)
        assert a.verdict == CriterionVerdict.PASS
        assert "12" in a.reason and "24" in a.reason


# ---------------------------------------------------------------------------
# _resolve_patient_value
# ---------------------------------------------------------------------------

class TestResolvePatientValue:
    def test_patient_age_resolves(self, als_patient):
        val, found = _resolve_patient_value("age_years", als_patient, None)
        assert found is True
        assert val == 52

    def test_patient_onset_months_resolves(self, als_patient):
        val, found = _resolve_patient_value("symptom_onset_months", als_patient, None)
        assert found is True
        assert val == 18

    def test_non_canonical_onset_key_not_found(self, als_patient):
        # "weakness_onset_months" is not in _KEY_MAP — confirms the parser must use
        # the canonical "symptom_onset_months" key for deterministic evaluation to work
        _, found = _resolve_patient_value("weakness_onset_months", als_patient, None)
        assert found is False

    def test_patient_diagnosis_months_resolves(self, als_patient):
        val, found = _resolve_patient_value("diagnosis_months", als_patient, None)
        assert found is True
        assert val == 12

    def test_unknown_key_returns_not_found(self, als_patient):
        val, found = _resolve_patient_value("nonexistent_key_xyz", als_patient, None)
        assert found is False
        assert val is None

    def test_platform_key_returns_not_found_without_data(self, als_patient):
        val, found = _resolve_patient_value("ecog_status", als_patient, None)
        assert found is False

    def test_platform_key_resolves_with_data(self, als_patient):
        platform = {"ecog_status": 1}
        val, found = _resolve_patient_value("ecog_status", als_patient, platform)
        assert found is True
        assert val == 1

    def test_platform_key_missing_from_dict_returns_not_found(self, als_patient):
        platform = {"other_field": 99}
        val, found = _resolve_patient_value("ecog_status", als_patient, platform)
        assert found is False

    def test_sex_key_not_found_because_patient_profile_has_no_sex_field(self, als_patient):
        # criterion_keys.json maps "sex" → "patient.sex" but PatientProfile has no
        # sex attribute. This documents the current behavior so a future schema
        # addition is caught immediately.
        val, found = _resolve_patient_value("sex", als_patient, None)
        assert found is False
        assert val is None


# ---------------------------------------------------------------------------
# _compute_overall
# ---------------------------------------------------------------------------

class TestComputeOverall:
    def test_all_pass_returns_pass(self):
        assessments = [make_assessment(CriterionVerdict.PASS)] * 3
        assert _compute_overall(assessments) == CriterionVerdict.PASS

    def test_any_fail_returns_fail(self):
        assessments = [
            make_assessment(CriterionVerdict.PASS),
            make_assessment(CriterionVerdict.FAIL),
            make_assessment(CriterionVerdict.UNKNOWN),
        ]
        assert _compute_overall(assessments) == CriterionVerdict.FAIL

    def test_fail_takes_precedence_over_unknown(self):
        assessments = [
            make_assessment(CriterionVerdict.UNKNOWN),
            make_assessment(CriterionVerdict.FAIL),
        ]
        assert _compute_overall(assessments) == CriterionVerdict.FAIL

    def test_unknown_without_fail_returns_unknown(self):
        assessments = [
            make_assessment(CriterionVerdict.PASS),
            make_assessment(CriterionVerdict.UNKNOWN),
        ]
        assert _compute_overall(assessments) == CriterionVerdict.UNKNOWN

    def test_empty_list_returns_pass(self):
        # No criteria → no objections → eligible
        assert _compute_overall([]) == CriterionVerdict.PASS

    def test_single_fail_returns_fail(self):
        assert _compute_overall([make_assessment(CriterionVerdict.FAIL)]) == CriterionVerdict.FAIL

    def test_single_unknown_returns_unknown(self):
        assert _compute_overall([make_assessment(CriterionVerdict.UNKNOWN)]) == CriterionVerdict.UNKNOWN


# ---------------------------------------------------------------------------
# _parse_criteria — LLM tool-forced call
# ---------------------------------------------------------------------------

def _parse_tool_response(criteria: list[dict]):
    """Build a mock response as if parse_criteria tool was called."""
    return make_message(
        content=[make_tool_use_block("parse_criteria", {"criteria": criteria})],
        stop_reason="tool_use",
    )


class TestParseCriteria:
    def test_empty_text_returns_empty(self, mock_client):
        result = _parse_criteria(mock_client, "   ")
        assert result == []
        mock_client.messages.create.assert_not_called()

    def test_single_criterion_parsed(self, mock_client):
        raw = [{
            "key": "age_years",
            "type": "inclusion",
            "description": "Age 18-75",
            "raw_criteria": "Age between 18 and 75 years",
            "constraint": {"key": "age_years", "operator": "between", "value": [18, 75], "unit": None},
        }]
        mock_client.messages.create.return_value = _parse_tool_response(raw)
        result = _parse_criteria(mock_client, "Age between 18 and 75 years")
        assert len(result) == 1
        assert result[0].key == "age_years"
        assert result[0].constraint.operator == "between"
        assert result[0].constraint.value == [18, 75]

    def test_criterion_without_constraint_has_none(self, mock_client):
        raw = [{
            "key": "adequate_hepatic_function",
            "type": "inclusion",
            "description": "Adequate hepatic function",
            "raw_criteria": "Adequate hepatic function per investigator",
            "constraint": None,
        }]
        mock_client.messages.create.return_value = _parse_tool_response(raw)
        result = _parse_criteria(mock_client, "Adequate hepatic function per investigator")
        assert result[0].constraint is None

    def test_missing_tool_call_returns_empty(self, mock_client):
        mock_client.messages.create.return_value = make_message(
            content=[make_text_block("No tool called")],
            stop_reason="end_turn",
        )
        result = _parse_criteria(mock_client, "some eligibility text")
        assert result == []

    def test_exactly_one_api_call(self, mock_client):
        mock_client.messages.create.return_value = _parse_tool_response([])
        _parse_criteria(mock_client, "some text")
        assert mock_client.messages.create.call_count == 1


# ---------------------------------------------------------------------------
# _assess_llm — LLM tool-forced call
# ---------------------------------------------------------------------------

def _criterion(key: str = "ecog_status") -> "EligibilityCriterion":
    from models import EligibilityCriterion
    return EligibilityCriterion(
        key=key, type="inclusion",
        description="ECOG 0 or 1",
        raw_criteria="ECOG performance status 0 or 1",
        constraint=None,
    )


def _assess_tool_response(assessments: list[dict]):
    return make_message(
        content=[make_tool_use_block("assess_eligibility", {"assessments": assessments})],
        stop_reason="tool_use",
    )


class TestAssessLlm:
    def test_empty_criteria_returns_empty(self, mock_client, als_patient):
        result = _assess_llm(mock_client, [], als_patient, None)
        assert result == []
        mock_client.messages.create.assert_not_called()

    def test_pass_verdict_mapped(self, mock_client, als_patient):
        criteria = [_criterion("ecog_status")]
        raw = [{"criterion_key": "ecog_status", "verdict": "pass", "reason": "ECOG is 0", "patient_value": "0", "confidence": "medium"}]
        mock_client.messages.create.return_value = _assess_tool_response(raw)
        result = _assess_llm(mock_client, criteria, als_patient, None)
        assert result[0].verdict == CriterionVerdict.PASS
        assert result[0].confidence == "medium"

    def test_fail_verdict_mapped(self, mock_client, als_patient):
        criteria = [_criterion("ecog_status")]
        raw = [{"criterion_key": "ecog_status", "verdict": "fail", "reason": "ECOG is 3", "patient_value": "3", "confidence": "medium"}]
        mock_client.messages.create.return_value = _assess_tool_response(raw)
        result = _assess_llm(mock_client, criteria, als_patient, None)
        assert result[0].verdict == CriterionVerdict.FAIL

    def test_unknown_criterion_key_skipped(self, mock_client, als_patient):
        criteria = [_criterion("ecog_status")]
        raw = [{"criterion_key": "nonexistent_key", "verdict": "pass", "reason": "ok", "patient_value": None, "confidence": "low"}]
        mock_client.messages.create.return_value = _assess_tool_response(raw)
        result = _assess_llm(mock_client, criteria, als_patient, None)
        assert result == []

    def test_missing_tool_call_returns_unknown_for_all(self, mock_client, als_patient):
        criteria = [_criterion("ecog_status"), _criterion("prior_systemic_therapy_lines")]
        mock_client.messages.create.return_value = make_message(
            content=[make_text_block("No tool")], stop_reason="end_turn"
        )
        result = _assess_llm(mock_client, criteria, als_patient, None)
        assert len(result) == 2
        assert all(a.verdict == CriterionVerdict.UNKNOWN for a in result)


# ---------------------------------------------------------------------------
# run_eligibility_check — orchestration
# ---------------------------------------------------------------------------

class TestRunEligibilityCheck:
    def _age_criterion_raw(self):
        return [{
            "key": "age_years",
            "type": "inclusion",
            "description": "Age 18-75",
            "raw_criteria": "Age between 18 and 75",
            "constraint": {"key": "age_years", "operator": "between", "value": [18, 75], "unit": None},
        }]

    def test_deterministic_criterion_resolved_without_llm_assess(self, mock_client, als_patient):
        # age_years=52 satisfies between [18,75] deterministically
        parse_resp = make_message(
            content=[make_tool_use_block("parse_criteria", {"criteria": self._age_criterion_raw()})],
            stop_reason="tool_use",
        )
        assess_resp = make_message(
            content=[make_tool_use_block("assess_eligibility", {"assessments": []})],
            stop_reason="tool_use",
        )
        mock_client.messages.create.side_effect = [parse_resp, assess_resp]

        trial = {"nct_id": "NCT99999999", "eligibility": "Inclusion:\n- Age between 18 and 75"}
        report = run_eligibility_check(mock_client, trial, als_patient)

        assert report.nct_id == "NCT99999999"
        det = [a for a in report.assessments if a.confidence == "high"]
        assert len(det) == 1
        assert det[0].verdict == CriterionVerdict.PASS

    def test_overall_verdict_fail_when_criterion_fails(self, mock_client, als_patient):
        # Age >75 means patient (age 52) passes, but let's use a criterion patient fails
        raw = [{
            "key": "age_years",
            "type": "inclusion",
            "description": "Age <= 40",
            "raw_criteria": "Age 40 or younger",
            "constraint": {"key": "age_years", "operator": "<=", "value": 40, "unit": None},
        }]
        parse_resp = make_message(
            content=[make_tool_use_block("parse_criteria", {"criteria": raw})],
            stop_reason="tool_use",
        )
        assess_resp = make_message(
            content=[make_tool_use_block("assess_eligibility", {"assessments": []})],
            stop_reason="tool_use",
        )
        mock_client.messages.create.side_effect = [parse_resp, assess_resp]
        trial = {"nct_id": "NCT00000001", "eligibility": "Age <= 40"}
        report = run_eligibility_check(mock_client, trial, als_patient)
        assert report.overall_verdict == CriterionVerdict.FAIL

    def test_empty_eligibility_text_returns_pass(self, mock_client, als_patient):
        trial = {"nct_id": "NCT00000002", "eligibility": ""}
        report = run_eligibility_check(mock_client, trial, als_patient)
        assert report.overall_verdict == CriterionVerdict.PASS
        # _parse_criteria short-circuits on empty text; _assess_llm short-circuits
        # on empty criteria list — neither should make an API call
        mock_client.messages.create.assert_not_called()

    def test_parse_called_once_assess_called_once(self, mock_client, als_patient):
        # ecog_status is a platform key; without platform_data it goes to needs_llm,
        # so _assess_llm will make a second API call
        platform_criterion = [{
            "key": "ecog_status",
            "type": "inclusion",
            "description": "ECOG 0 or 1",
            "raw_criteria": "ECOG performance status 0 or 1",
            "constraint": {"key": "ecog_status", "operator": "in", "value": [0, 1], "unit": None},
        }]
        parse_resp = make_message(
            content=[make_tool_use_block("parse_criteria", {"criteria": platform_criterion})],
            stop_reason="tool_use",
        )
        assess_resp = make_message(
            content=[make_tool_use_block("assess_eligibility", {"assessments": []})],
            stop_reason="tool_use",
        )
        mock_client.messages.create.side_effect = [parse_resp, assess_resp]
        trial = {"nct_id": "NCT00000003", "eligibility": "some text"}
        run_eligibility_check(mock_client, trial, als_patient)
        assert mock_client.messages.create.call_count == 2


# ---------------------------------------------------------------------------
# bulk_parse_and_strip — batch parse + deterministic filter
# ---------------------------------------------------------------------------

def _bulk_parse_response(trials_payload: list[dict]):
    return make_message(
        content=[make_tool_use_block("parse_criteria_bulk", {"trials": trials_payload})],
        stop_reason="tool_use",
    )


class TestBulkParseAndStrip:
    def _trial(self, nct_id: str, eligibility: str = "Age >= 18") -> dict:
        return {"nct_id": nct_id, "eligibility": eligibility, "title": f"Trial {nct_id}"}

    def test_strips_eligibility_field(self, mock_client, als_patient):
        raw_criteria = [{"nct_id": "NCT00000001", "criteria": []}]
        mock_client.messages.create.return_value = _bulk_parse_response(raw_criteria)
        result = bulk_parse_and_strip(mock_client, [self._trial("NCT00000001")], als_patient)
        assert "eligibility" not in result[0]

    def test_strips_std_ages_and_healthy_volunteers(self, mock_client, als_patient):
        trial = {**self._trial("NCT00000001"), "std_ages": ["ADULT"], "healthy_volunteers": "No"}
        mock_client.messages.create.return_value = _bulk_parse_response(
            [{"nct_id": "NCT00000001", "criteria": []}]
        )
        result = bulk_parse_and_strip(mock_client, [trial], als_patient)
        assert "std_ages" not in result[0]
        assert "healthy_volunteers" not in result[0]

    def test_adds_parsed_criteria_key(self, mock_client, als_patient):
        mock_client.messages.create.return_value = _bulk_parse_response(
            [{"nct_id": "NCT00000001", "criteria": []}]
        )
        result = bulk_parse_and_strip(mock_client, [self._trial("NCT00000001")], als_patient)
        assert "parsed_criteria" in result[0]

    def test_adds_deterministic_verdicts_key(self, mock_client, als_patient):
        mock_client.messages.create.return_value = _bulk_parse_response(
            [{"nct_id": "NCT00000001", "criteria": []}]
        )
        result = bulk_parse_and_strip(mock_client, [self._trial("NCT00000001")], als_patient)
        assert "deterministic_verdicts" in result[0]

    def test_deterministic_verdict_computed_for_known_key(self, mock_client, als_patient):
        # age_years=52, criterion >= 18 → PASS
        age_criterion = {
            "key": "age_years", "type": "inclusion",
            "description": "Age >= 18", "raw_criteria": "Age >= 18",
            "constraint": {"key": "age_years", "operator": ">=", "value": 18, "unit": None},
        }
        mock_client.messages.create.return_value = _bulk_parse_response(
            [{"nct_id": "NCT00000001", "criteria": [age_criterion]}]
        )
        result = bulk_parse_and_strip(mock_client, [self._trial("NCT00000001")], als_patient)
        verdicts = result[0]["deterministic_verdicts"]
        assert len(verdicts) == 1
        assert verdicts[0]["verdict"] == "pass"

    def test_trials_beyond_top_n_passed_through_unmodified(self, mock_client, als_patient):
        trials = [self._trial(f"NCT{i:08d}") for i in range(7)]
        mock_client.messages.create.return_value = _bulk_parse_response(
            [{"nct_id": f"NCT{i:08d}", "criteria": []} for i in range(5)]
        )
        result = bulk_parse_and_strip(mock_client, trials, als_patient, top_n=5)
        assert len(result) == 7
        # Trials beyond top_n still have eligibility field
        assert "eligibility" in result[5]
        assert "eligibility" in result[6]

    def test_empty_trials_skipped(self, mock_client, als_patient):
        trial = self._trial("NCT00000001", eligibility="")
        result = bulk_parse_and_strip(mock_client, [trial], als_patient)
        # No API call — trial has no eligibility text
        mock_client.messages.create.assert_not_called()
        assert result == [trial]

    def test_api_failure_returns_trials_without_parsed_criteria(self, mock_client, als_patient):
        mock_client.messages.create.side_effect = Exception("API down")
        trial = self._trial("NCT00000001")
        result = bulk_parse_and_strip(mock_client, [trial], als_patient)
        # Should not raise; deterministic_verdicts should be empty
        assert result[0]["deterministic_verdicts"] == []