tests/test_ai.py

"""Constrained generated explanation tests."""

import json
from datetime import date, datetime, timezone

from coastwise.ai import (
    build_hosted_interpreter,
    build_local_interpreter,
    explain_safety_result,
    interpret_question_with_model,
    synthesize_cached_evidence_answer,
    validate_cached_evidence_answer,
    validate_generated_text,
    validate_source_backed_wording,
    word_source_backed_answer,
)
from coastwise.data import get_rule_card
from coastwise.qna import answer_question
from coastwise.schemas import (
    CachedEvidencePacket,
    CachedEvidenceSource,
    FreshnessStatus,
    RequestedFactType,
)
from coastwise.rules import compose_safety_result, default_context
from coastwise.source_cache import DEFAULT_SEED_CACHE_DIR, load_seed_cache
from coastwise.submission import default_model_inventory


def test_explain_safety_result_returns_fallback_without_mutating_result():
    result = compose_safety_result(
        get_rule_card("lingcod"),
        default_context(today=date(2026, 6, 13)),
    )
    explanation = explain_safety_result(result, default_model_inventory(), generator=None)

    assert explanation.fallback_used
    assert explanation.model_entry_key == "minicpm5_1b"
    assert "Lingcod" in explanation.text
    assert result.status.value == "rules_apply"


def test_generated_text_with_forbidden_claims_is_rejected():
    result = compose_safety_result(
        get_rule_card("mussels"),
        default_context(today=date(2026, 6, 13)),
    )

    assert not validate_generated_text(result, "This is definitely legal and safe to eat.")
    explanation = explain_safety_result(
        result,
        default_model_inventory(),
        generator=lambda _prompt: "This is definitely legal and safe to eat.",
    )
    assert explanation.fallback_used
    assert explanation.blocked_reason == "forbidden_claim"


def test_source_backed_model_wording_cannot_add_facts_links_dates_or_next_actions():
    answer = answer_question(
        "what's the min size of lingcod?",
        "",
        load_seed_cache(DEFAULT_SEED_CACHE_DIR),
        datetime(2026, 6, 14, 13, tzinfo=timezone.utc),
    )
    unsafe_generated = (
        "Lingcod minimum size is 24 inches. See https://example.com. "
        "Retrieved 2026-06-15. You may harvest today."
    )

    assert not validate_source_backed_wording(answer, unsafe_generated)

    wording = word_source_backed_answer(
        answer,
        default_model_inventory(),
        generator=lambda _prompt: unsafe_generated,
    )

    assert wording.fallback_used
    assert wording.blocked_reason == "source_invariant"
    assert "22" in wording.text
    assert "24" not in wording.text
    assert "https://example.com" not in wording.text


def test_cached_evidence_answer_accepts_wording_bound_to_packet_excerpts():
    packet = CachedEvidencePacket(
        question="what are the restrictions on cabezone?",
        species_or_category="cabezon",
        display_name="Cabezon",
        requested_fact_type=RequestedFactType.GENERAL_SUMMARY,
        matched_place="Monterey",
        region_key="central",
        sources=(
            CachedEvidenceSource(
                source_id="cdfw_central_region",
                title="CDFW Central Region",
                url="https://wildlife.ca.gov/Fishing/Ocean/Regulations/Fishing-Map/Central",
                retrieved_at="2026-06-14T12:00:00+00:00",
                freshness_status=FreshnessStatus.CURRENT,
                excerpt="Cabezon rules vary by area and may include size, bag, and seasonal restrictions.",
            ),
        ),
        deterministic_next_safe_action="Verify current official CDFW guidance before harvest.",
    )
    generated = (
        "Cabezon has saved official source text saying rules vary by area and may include "
        "size, bag, and seasonal restrictions. Verify current official CDFW guidance before harvest."
    )

    assert validate_cached_evidence_answer(packet, generated)
    explanation = synthesize_cached_evidence_answer(
        packet,
        default_model_inventory(),
        generator=lambda _prompt: generated,
    )

    assert not explanation.fallback_used
    assert explanation.text == generated


def test_cached_evidence_answer_rejects_invented_open_status():
    packet = CachedEvidencePacket(
        question="is dungeness crab in season for Fort Bragg?",
        species_or_category="dungeness_crab",
        display_name="Dungeness crab",
        requested_fact_type=RequestedFactType.SEASON,
        matched_place="Fort Bragg",
        region_key="mendocino",
        sources=(
            CachedEvidenceSource(
                source_id="cdfw_crabs",
                title="CDFW Crabs",
                url="https://wildlife.ca.gov/Fishing/Ocean/Regulations/Fishing-Map/Crabs",
                retrieved_at="2026-06-14T12:00:00+00:00",
                freshness_status=FreshnessStatus.CURRENT,
                excerpt="Dungeness crab rules include season, trap, hoop-net, and public health closure details.",
            ),
        ),
        deterministic_next_safe_action="Verify current official CDFW guidance before harvest.",
    )
    generated = "Dungeness crab season is open in Fort Bragg today."

    assert not validate_cached_evidence_answer(packet, generated)
    explanation = synthesize_cached_evidence_answer(
        packet,
        default_model_inventory(),
        generator=lambda _prompt: generated,
    )

    assert explanation.fallback_used
    assert explanation.blocked_reason == "evidence_invariant"
    assert "season is open" not in explanation.text.casefold()


def test_cached_evidence_answer_falls_back_when_model_is_unavailable():
    packet = CachedEvidencePacket(
        question="is rock crab safe to eat?",
        species_or_category="rock_crab",
        display_name="Rock crab",
        requested_fact_type=RequestedFactType.ADVISORY,
        matched_place=None,
        region_key=None,
        sources=(
            CachedEvidenceSource(
                source_id="cdph_shellfish_advisories",
                title="CDPH Shellfish Advisories",
                url="https://www.cdph.ca.gov/Programs/OPA/Pages/Shellfish-Advisories.aspx",
                retrieved_at="2026-06-14T12:00:00+00:00",
                freshness_status=FreshnessStatus.CURRENT,
                excerpt="CDPH advisories should be checked before eating sport-harvested shellfish or crab.",
            ),
        ),
        deterministic_next_safe_action="Verify current CDPH guidance before eating sport-harvested crab.",
    )
    explanation = synthesize_cached_evidence_answer(
        packet,
        default_model_inventory(),
        generator=lambda _prompt: None,
    )

    assert explanation.fallback_used
    assert explanation.blocked_reason == "model_unavailable"
    assert "Rock crab" in explanation.text


def test_model_interpreter_accepts_only_allowlisted_candidates():
    prompts = []
    payload = {
        "species_or_category": "dungeness_crab",
        "location_name": "Mendocino",
        "region_key": "mendocino",
        "intent": "permission_status",
        "requested_fact_types": ["season", "closure"],
        "confidence": "high",
        "notes": "Normalized spelling.",
    }

    result = interpret_question_with_model(
        "im in mandocino. is it legal to harvest dungenese crabs?",
        "",
        allowed_species={"dungeness_crab", "lingcod"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino", "central"},
        generator=lambda prompt: prompts.append(prompt) or json.dumps(payload),
    )

    assert result is not None
    assert result.species_or_category == "dungeness_crab"
    assert result.location_name == "Mendocino"
    assert result.region_key == "mendocino"
    assert result.intent.value == "permission_status"
    assert [fact.value for fact in result.requested_fact_types] == ["season", "closure"]
    assert 'allowed_locations={"Mendocino": "mendocino"}' in prompts[0]


def test_model_interpreter_rejects_unknown_candidates_and_facts():
    payload = {
        "species_or_category": "maine_lobster",
        "location_name": "Mendocino",
        "region_key": "mendocino",
        "intent": "permission_status",
        "requested_fact_types": ["season"],
        "confidence": "high",
    }

    result = interpret_question_with_model(
        "can I harvest lobster in mendocino?",
        "",
        allowed_species={"dungeness_crab"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino"},
        generator=lambda _prompt: json.dumps(payload),
    )
    unknown_fact_type = interpret_question_with_model(
        "can I harvest crab in mendocino?",
        "",
        allowed_species={"dungeness_crab"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino"},
        generator=lambda _prompt: json.dumps(
            {
                "species_or_category": "dungeness_crab",
                "location_name": "Mendocino",
                "region_key": "mendocino",
                "intent": "permission_status",
                "requested_fact_types": ["unknown"],
                "confidence": "high",
            }
        ),
    )

    assert result is None
    assert unknown_fact_type is None


def test_model_interpreter_falls_back_on_invalid_json_or_low_confidence():
    invalid = interpret_question_with_model(
        "can I harvest crab?",
        "",
        allowed_species={"dungeness_crab"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino"},
        generator=lambda _prompt: "not json",
    )
    low_confidence = interpret_question_with_model(
        "can I harvest crab?",
        "",
        allowed_species={"dungeness_crab"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino"},
        generator=lambda _prompt: json.dumps(
            {
                "species_or_category": "dungeness_crab",
                "location_name": "Mendocino",
                "region_key": "mendocino",
                "intent": "permission_status",
                "requested_fact_types": ["season"],
                "confidence": "low",
            }
        ),
    )
    malformed_confidence = interpret_question_with_model(
        "can I harvest crab?",
        "",
        allowed_species={"dungeness_crab"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino"},
        generator=lambda _prompt: json.dumps(
            {
                "species_or_category": "dungeness_crab",
                "location_name": "Mendocino",
                "region_key": "mendocino",
                "intent": "permission_status",
                "requested_fact_types": ["season"],
                "confidence": ["high"],
            }
        ),
    )

    assert invalid is None
    assert low_confidence is None
    assert malformed_confidence is None


def test_build_hosted_interpreter_returns_none_without_api_key():
    assert build_hosted_interpreter(api_key="") is None


def test_hosted_interpreter_posts_openai_chat_and_returns_model_json():
    calls = []

    def fake_post(base_url, api_key, model, prompt, timeout):
        calls.append((base_url, api_key, model, prompt, timeout))
        return '{"species_or_category": "dungeness_crab"}'

    generate = build_hosted_interpreter(
        api_key="secret-key",
        base_url="https://api.modelbest.cn/v1",
        model="MiniCPM5-1B",
        post=fake_post,
    )
    out = generate("interpret this question as JSON")

    assert out == '{"species_or_category": "dungeness_crab"}'
    assert len(calls) == 1
    base_url, api_key, model, prompt, _timeout = calls[0]
    assert base_url == "https://api.modelbest.cn/v1"
    assert api_key == "secret-key"
    assert model == "MiniCPM5-1B"
    assert prompt == "interpret this question as JSON"


def test_hosted_interpreter_extracts_json_from_fenced_or_prose_output():
    fenced = build_hosted_interpreter(
        api_key="k",
        post=lambda *_a, **_k: '```json\n{"species_or_category": "lingcod"}\n```',
    )
    assert fenced("p") == '{"species_or_category": "lingcod"}'

    prose = build_hosted_interpreter(
        api_key="k",
        post=lambda *_a, **_k: 'Sure: {"species_or_category": "lingcod"} hope this helps',
    )
    assert prose("p") == '{"species_or_category": "lingcod"}'


def test_hosted_interpreter_fails_closed_on_transport_error_or_garbage():
    def boom(*_a, **_k):
        raise RuntimeError("network down")

    assert build_hosted_interpreter(api_key="k", post=boom)("p") is None
    assert build_hosted_interpreter(api_key="k", post=lambda *_a, **_k: "no json here")("p") is None


def test_hosted_interpreter_output_drives_constrained_validation():
    payload = {
        "species_or_category": "dungeness_crab",
        "location_name": "Mendocino",
        "region_key": "mendocino",
        "intent": "permission_status",
        "requested_fact_types": ["season", "closure"],
        "confidence": "high",
    }
    generate = build_hosted_interpreter(
        api_key="k",
        post=lambda *_a, **_k: f"```json\n{json.dumps(payload)}\n```",
    )
    result = interpret_question_with_model(
        "im in mandocino. is it legal to harvest dungenese crabs?",
        "",
        allowed_species={"dungeness_crab", "lingcod"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino"},
        generator=generate,
    )
    assert result is not None
    assert result.species_or_category == "dungeness_crab"
    assert result.location_name == "Mendocino"
    assert result.region_key == "mendocino"
    assert result.intent.value == "permission_status"


def test_build_local_interpreter_disabled_returns_none():
    assert build_local_interpreter(enabled=False) is None


def test_local_interpreter_lazy_loads_once_and_returns_model_json():
    loads = {"n": 0}

    class FakeModel:
        def create_chat_completion(self, messages, **kwargs):
            return {"choices": [{"message": {"content": '{"species_or_category": "lingcod"}'}}]}

    def factory():
        loads["n"] += 1
        return FakeModel()

    generate = build_local_interpreter(enabled=True, model_factory=factory)
    assert generate("first prompt") == '{"species_or_category": "lingcod"}'
    assert generate("second prompt") == '{"species_or_category": "lingcod"}'
    assert loads["n"] == 1


def test_local_interpreter_extracts_json_and_fails_closed():
    class FencedModel:
        def create_chat_completion(self, messages, **kwargs):
            return {
                "choices": [
                    {"message": {"content": '```json\n{"species_or_category": "dungeness_crab"}\n```'}}
                ]
            }

    fenced = build_local_interpreter(enabled=True, model_factory=lambda: FencedModel())
    assert fenced("p") == '{"species_or_category": "dungeness_crab"}'

    def boom_factory():
        raise RuntimeError("model load failed")

    assert build_local_interpreter(enabled=True, model_factory=boom_factory)("p") is None


def test_model_interpreter_canonicalizes_case_insensitive_location_name():
    payload = {
        "species_or_category": "dungeness_crab",
        "location_name": "mendocino",
        "region_key": "mendocino",
        "intent": "permission_status",
        "requested_fact_types": ["season"],
        "confidence": "high",
    }

    result = interpret_question_with_model(
        "im in mandocino. is it legal to harvest dungenese crabs?",
        "",
        allowed_species={"dungeness_crab"},
        allowed_locations={"Mendocino": "mendocino"},
        allowed_region_keys={"mendocino"},
        generator=lambda _prompt: json.dumps(payload),
    )

    assert result is not None
    assert result.location_name == "Mendocino"
    assert result.region_key == "mendocino"


def test_local_interpreter_constrains_output_with_json_schema():
    captured = {}

    class RecordingModel:
        def create_chat_completion(self, messages, **kwargs):
            captured.update(kwargs)
            return {"choices": [{"message": {"content": '{"species_or_category": "lingcod"}'}}]}

    generate = build_local_interpreter(enabled=True, model_factory=lambda: RecordingModel())
    generate("some prompt")

    response_format = captured.get("response_format") or {}
    assert response_format.get("type") == "json_object"
    schema = response_format.get("schema") or {}
    assert "confidence" in (schema.get("required") or [])
    assert schema.get("properties", {}).get("confidence", {}).get("enum") == ["high", "medium"]


def test_local_interpreter_schema_constrains_to_prompt_allowlists():
    captured = {}

    class RecordingModel:
        def create_chat_completion(self, messages, **kwargs):
            captured.update(kwargs)
            return {"choices": [{"message": {"content": "{}"}}]}

    generate = build_local_interpreter(enabled=True, model_factory=lambda: RecordingModel())
    generate(
        "question=x\n"
        "allowed_species=['dungeness_crab', 'lingcod']\n"
        "allowed_locations=['Mendocino', 'Monterey']\n"
        "allowed_region_keys=['central', 'mendocino']"
    )

    schema = captured["response_format"]["schema"]
    props = schema["properties"]
    assert set(props["species_or_category"]["enum"]) >= {"dungeness_crab", "lingcod"}
    assert "Mendocino" in props["location_name"]["enum"]
    assert "mendocino" in props["region_key"]["enum"]
    # species is required (model must pick from the allowlist, never null/omit)
    assert "species_or_category" in schema["required"]
    # location/region are optional so location-less questions are not forced to invent a place
    assert "location_name" not in schema["required"]
    assert "region_key" not in schema["required"]
    assert None in props["location_name"]["enum"]