tests/test_prompt_builder_contract.py

"""Prompt contract tests for load-bearing navigator model calls."""

from __future__ import annotations

import json
from typing import Any

from figment.prompt_builder import build_prompt


REQUIRED_TOP_LEVEL_KEYS = {
    "protocol_urgency",
    "red_flags",
    "intake_facts",
    "candidate_protocol_pathways",
    "missing_info_to_collect",
    "next_observations_to_collect",
    "conflicts_or_uncertainties",
    "responder_checklist",
    "do_not_do",
    "source_cards",
    "handoff_note_sbar",
    "responder_plain_language_script",
    "safety_boundary",
}


def _context_from_prompt(prompt: str) -> dict[str, Any]:
    return json.loads(prompt.split("\n\nCONTEXT:\n", maxsplit=1)[1])


def _confirmed_chest_pain_intake() -> dict[str, Any]:
    return {
        "setting": "mobile clinic",
        "patient_age": "52",
        "pregnancy_status": "not_applicable",
        "chief_concern": "Chest pain",
        "symptoms": "Crushing chest pain and shortness of breath",
        "vitals": "HR 118; blood pressure pending",
        "allergies": "unknown",
        "medications": "unknown",
        "available_supplies": "AED, radio",
        "responder_note": "Adult reports chest pain after cleanup work.",
        "confirmed": True,
    }


def _emergency_chest_pain_rules() -> list[dict[str, Any]]:
    return [
        {
            "rule_id": "red_flag_chest_pain",
            "label": "Chest pain escalation cue",
            "urgency": "emergency",
            "evidence": "chest pain",
            "card_id": "CHEST-PAIN-ESCALATION-v1",
        }
    ]


def _retrieved_cards() -> list[dict[str, Any]]:
    return [
        {
            "card_id": "CHEST-PAIN-ESCALATION-v1",
            "score": 1.0,
            "source": "test",
            "card": {
                "card_id": "CHEST-PAIN-ESCALATION-v1",
                "title": "Chest pain escalation",
                "required_observations": [
                    "chest pain description",
                    "onset and duration",
                    "available vital signs",
                ],
                "red_flags": ["chest pain or pressure"],
                "escalation_criteria": [
                    "Chest pain with shortness of breath requires emergency escalation.",
                ],
                "local_actions": ["Document onset, duration, and vital signs."],
                "forbidden_actions": ["Do not diagnose the cause of chest pain."],
            },
        },
        {
            "card_id": "REFERRAL-SBAR-v1",
            "score": 0.5,
            "source": "test",
            "card": {
                "card_id": "REFERRAL-SBAR-v1",
                "title": "Referral and SBAR format",
                "required_observations": [
                    "situation or reason for handoff",
                    "objective observations only",
                    "source protocol card IDs",
                ],
                "local_actions": ["Use Situation, Background, Assessment, and Request."],
                "forbidden_actions": ["Do not add a diagnosis as the assessment."],
            },
        },
    ]


def test_prompt_includes_literal_required_json_skeleton_with_all_schema_keys() -> None:
    prompt, _ = build_prompt(
        _confirmed_chest_pain_intake(),
        _retrieved_cards(),
        _emergency_chest_pain_rules(),
        "emergency",
    )

    assert "REQUIRED_JSON_SKELETON:" in prompt
    assert "Do not discharge" in prompt
    assert "autonomous routing" in prompt
    context = _context_from_prompt(prompt)

    skeleton = context["required_json_skeleton"]
    assert set(skeleton) == REQUIRED_TOP_LEVEL_KEYS
    assert set(skeleton["handoff_note_sbar"]) == {
        "situation",
        "background",
        "assessment_observations_only",
        "handoff_request",
    }
    assert isinstance(skeleton["source_cards"], list)
    assert isinstance(skeleton["missing_info_to_collect"], list)


def test_prompt_context_lists_allowed_facts_and_required_observations() -> None:
    prompt, _ = build_prompt(
        _confirmed_chest_pain_intake(),
        _retrieved_cards(),
        _emergency_chest_pain_rules(),
        "emergency",
        audio_draft={
            "confirmation_status": "unconfirmed",
            "transcript": "Audio-only chest pain phrase must not become an allowed fact.",
        },
    )
    context = _context_from_prompt(prompt)

    allowed_facts_text = json.dumps(context["allowed_facts_inventory"], sort_keys=True)
    assert "confirmed_intake" in allowed_facts_text
    assert "Crushing chest pain and shortness of breath" in allowed_facts_text
    assert "deterministic_rule" in allowed_facts_text
    assert "red_flag_chest_pain" in allowed_facts_text
    assert "retrieved_protocol_card" in allowed_facts_text
    assert "Chest pain escalation" in allowed_facts_text
    assert "Audio-only chest pain phrase" not in allowed_facts_text
    assert "Audio-only chest pain phrase" not in prompt

    observations = context["required_observations_inventory"]
    chest = next(item for item in observations if item["card_id"] == "CHEST-PAIN-ESCALATION-v1")
    assert chest["required_observations"] == [
        "chest pain description",
        "onset and duration",
        "available vital signs",
    ]

    targets = context["required_observation_targets"]
    chest_targets = [target for target in targets if target["card_id"] == "CHEST-PAIN-ESCALATION-v1"]
    assert chest_targets == [
        {
            "id": "CHEST-PAIN-ESCALATION-v1::required_observation::1",
            "card_id": "CHEST-PAIN-ESCALATION-v1",
            "title": "Chest pain escalation",
            "display_text": "chest pain description",
            "cue_tokens": ["chest", "pain", "description"],
        },
        {
            "id": "CHEST-PAIN-ESCALATION-v1::required_observation::2",
            "card_id": "CHEST-PAIN-ESCALATION-v1",
            "title": "Chest pain escalation",
            "display_text": "onset and duration",
            "cue_tokens": ["onset", "duration"],
        },
        {
            "id": "CHEST-PAIN-ESCALATION-v1::required_observation::3",
            "card_id": "CHEST-PAIN-ESCALATION-v1",
            "title": "Chest pain escalation",
            "display_text": "available vital signs",
            "cue_tokens": ["available", "vital", "signs"],
        },
    ]


def test_prompt_context_guides_routine_and_negated_cases() -> None:
    prompt, _ = build_prompt(
        {
            **_confirmed_chest_pain_intake(),
            "chief_concern": "routine cough check",
            "symptoms": "mild cough, no chest pain, no shortness of breath, speaking normally",
        },
        _retrieved_cards(),
        [],
        "routine",
    )
    context = _context_from_prompt(prompt)

    guidance_text = " ".join(context["routine_or_negated_case_guidance"])
    assert "Do not convert denied or absent symptoms into red_flags" in guidance_text
    assert "keep protocol_urgency routine" in guidance_text
    assert "nearby emergency card language" in guidance_text

    ledger = context["case_fact_ledger"]
    absent_text = json.dumps(ledger["absent_or_denied"], sort_keys=True).lower()
    present_text = json.dumps(ledger["present"], sort_keys=True).lower()
    assert "no chest pain" in absent_text
    assert "no shortness of breath" in absent_text
    assert "routine cough check" in present_text


def test_prompt_context_includes_sbar_template_and_internal_target_id_contract() -> None:
    prompt, _ = build_prompt(
        _confirmed_chest_pain_intake(),
        _retrieved_cards(),
        _emergency_chest_pain_rules(),
        "emergency",
    )
    context = _context_from_prompt(prompt)

    sbar_template = context["handoff_note_sbar_template"]
    assert sbar_template == {
        "situation": "Chest pain",
        "background": "Setting: mobile clinic. Age: 52. Pregnancy status: not_applicable.",
        "assessment_observations_only": "Symptoms: Crushing chest pain and shortness of breath. Vitals: HR 118; blood pressure pending. Red flags: Chest pain escalation cue.",
        "handoff_request": "Request emergency review/escalation per cited local protocol cards.",
    }

    internal_contract = context["internal_generation_contract"]
    assert internal_contract["trace_only_keys"] == ["selected_required_observation_ids"]
    assert internal_contract["required_when_required_observation_targets_selected"] == [
        "selected_required_observation_ids"
    ]
    assert internal_contract["strip_before_user_display"] is True
    assert "selected_required_observation_ids" in internal_contract["selected_required_observation_ids"]
    assert "Choose required observation IDs before writing observation text" in prompt
    assert "recognizable responder-facing observation text" in internal_contract["selected_required_observation_ids"]
    assert "selected_required_observation_ids" not in context["navigator_output_schema"]
    assert "selected_required_observation_ids" not in context["required_json_skeleton"]


def test_prompt_context_names_mandatory_source_cards_and_model_owned_observation_policy() -> None:
    prompt, _ = build_prompt(
        _confirmed_chest_pain_intake(),
        _retrieved_cards(),
        _emergency_chest_pain_rules(),
        "emergency",
    )
    context = _context_from_prompt(prompt)

    assert context["mandatory_source_card_ids"] == [
        "CHEST-PAIN-ESCALATION-v1",
        "REFERRAL-SBAR-v1",
    ]
    assert context["mandatory_required_observation_target_ids"] == [
        "CHEST-PAIN-ESCALATION-v1::required_observation::1",
        "CHEST-PAIN-ESCALATION-v1::required_observation::2",
        "CHEST-PAIN-ESCALATION-v1::required_observation::3",
    ]
    assert context["mandatory_required_observation_targets"] == [
        {
            "id": "CHEST-PAIN-ESCALATION-v1::required_observation::1",
            "card_id": "CHEST-PAIN-ESCALATION-v1",
            "title": "Chest pain escalation",
            "display_text": "chest pain description",
        },
        {
            "id": "CHEST-PAIN-ESCALATION-v1::required_observation::2",
            "card_id": "CHEST-PAIN-ESCALATION-v1",
            "title": "Chest pain escalation",
            "display_text": "onset and duration",
        },
        {
            "id": "CHEST-PAIN-ESCALATION-v1::required_observation::3",
            "card_id": "CHEST-PAIN-ESCALATION-v1",
            "title": "Chest pain escalation",
            "display_text": "available vital signs",
        },
    ]
    policy = context["required_observation_generation_policy"]
    assert policy["model_owned_not_scaffold_filled"] is True
    assert "mandatory_required_observation_targets" in policy
    assert "missing_info_to_collect" in policy["source_card_scope"]
    assert "mandatory_source_card_id" in policy["source_cards"]
    assert "mandatory_required_observation_target_id" in policy["selected_required_observation_ids"]
    assert "confirm/document" in policy["text_requirement"]
    assert "Emit selected_required_observation_ids as a trace-only key" in prompt
    assert "Include every mandatory_source_card_id in source_cards" in prompt
    assert "retrieved support cards used for safety/SBAR" in prompt
    assert "Do not add SAFETY-BOUNDARIES-v1 or REFERRAL-SBAR-v1 as candidate pathways" in prompt
    assert "make the target display_text visible in missing_info_to_collect" in prompt
    assert "Copy each display_text into missing_info_to_collect" in prompt