library/tests/test_runtime.py

import importlib.metadata
import json
import shutil
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest
from safetensors.torch import load_file, save_file

import jev_urdu
from jev_urdu._checkpoint import build_sequence
from jev_urdu.runtime import _question


QUESTIONS = {
    "topic": jev_urdu.choice("topic", ["billing", "technical"]),
    "urgent": jev_urdu.yes_no("urgent"),
}


def test_import_has_no_heavy_imports_or_laya():
    script = "import sys, jev_urdu; assert all(x not in sys.modules for x in ('torch', 'transformers', 'laya')); assert jev_urdu.__version__ == '0.1.0'"
    library_root = Path(__file__).resolve().parents[1]
    subprocess.run([sys.executable, "-c", script], cwd=library_root, check=True)
    assert all(not requirement.lower().startswith("laya")
               for requirement in importlib.metadata.requires("jev-urdu"))


def test_checkpoint_token_contract(model):
    question = _question("topic", QUESTIONS["topic"])
    ids, markers = build_sequence(model.agent.tok, "hello world", question,
                                  max_len=96, head_max_len=48)
    # Published contract: CLS, typed question, SEP, marked options, SEP, state, SEP.
    vocabulary = model.agent.tok.get_vocab()
    expected = [vocabulary[x] for x in ["<bos>", "choice", "question:", "topic",
                "<eos>", "<mask>", "billing", "<mask>", "technical",
                "<eos>", "hello", "world", "<eos>"]]
    assert ids == expected
    assert markers == [5, 7]


def test_real_offline_inference_and_simplified_api(model):
    raw = model.predict("میرا آرڈر نہیں پہنچا", QUESTIONS)
    assert set(raw["answers"]) == set(QUESTIONS)
    assert raw["truncated"] is False
    answer = raw["answers"]["topic"]
    assert answer["choice"] in ["billing", "technical"]
    assert sum(answer["probabilities"].values()) == pytest.approx(1, abs=0.0002)
    assert 0 <= raw["answers"]["urgent"]["noul"] <= 1
    simplified = model.ask("میرا آرڈر نہیں پہنچا", QUESTIONS)
    assert simplified["topic"]["answer"] == answer["choice"]
    assert simplified["urgent"]["probability"] == raw["answers"]["urgent"]["noul"]


def test_sindhi_text_with_existing_public_api(model):
    questions = {
        "topic": jev_urdu.choice(
            "هي پيغام ڪهڙي موضوع سان لاڳاپيل آهي؟",
            ["تعليم", "ٽيڪنالاجي", "زراعت", "صحت", "ٻيو"],
        ),
        "urgent": jev_urdu.yes_no("ڇا هي معاملو فوري آهي؟"),
    }

    result = model.ask(
        "اڄ ڪڻڪ جو اگهه وڌي ويو آهي.",
        questions,
    )

    assert set(result) == {"topic", "urgent"}
    assert result["topic"]["answer"] in [
        "تعليم", "ٽيڪنالاجي", "زراعت", "صحت", "ٻيو"
    ]
    assert 0 <= result["topic"]["probability"] <= 1
    assert 0 <= result["urgent"]["probability"] <= 1


def test_batch_matches_single_and_keeps_state_order(model):
    states = ["hello", "میرا آرڈر نہیں پہنچا", "world hello"]
    batched = model.predict_batch(states, QUESTIONS, batch_size=2)
    singles = [model.predict(state, QUESTIONS) for state in states]
    for batch, single in zip(batched, singles):
        for qid in QUESTIONS:
            assert batch["answers"][qid] == single["answers"][qid]
    assert model.ask_many([], QUESTIONS) == []


def test_single_option_and_ordinal_score(model):
    answers = model.ask("hello", {
        "one": jev_urdu.choice("topic", ["billing"]),
        "rating": {"type": "score", "instructions": "topic",
                   "criteria": ["negative", "neutral", "positive"]},
    })
    assert answers["one"]["answer"] == "billing"
    assert answers["one"]["probability"] == 1.0
    assert 0 <= answers["rating"]["answer"] <= 2


def test_option_order_decodes_into_canonical_labels(model):
    question = _question("topic", {**QUESTIONS["topic"], "option_order": [1, 0]})
    item = {"markers": [5, 7], "state_stats": {"truncated": False}}
    output = model.agent._decode([np.array([5.0, -5.0])], [np.array([0.7, 0.3])],
                                [item], {"topic": question})
    assert output["answers"]["topic"]["choice"] == "technical"


def test_reports_long_input_truncation(model):
    raw = model.predict("hello " * 200, QUESTIONS)
    assert raw["truncated"] is True


@pytest.mark.parametrize("definition", [
    {"type": "other", "instructions": "topic"},
    {"type": "choice", "instructions": "topic", "criteria": []},
    {"type": "choice", "instructions": "topic", "criteria": ["billing", "billing"]},
    {"type": "choice", "instructions": "topic", "criteria": ["billing"], "option_order": [1]},
    {"type": "choice", "instructions": "", "criteria": ["billing"]},
    {"type": "score", "instructions": "topic", "criteria": ["billing"]},
])
def test_invalid_questions_fail_clearly(model, definition):
    with pytest.raises(ValueError):
        model.predict("hello", {"invalid": definition})


def test_distinct_options_that_collapse_after_tokenization_fail(model):
    with pytest.raises(ValueError, match="identical"):
        model.ask("hello", {"topic": jev_urdu.choice("topic", ["unknown-label-A", "unknown-label-B"])})


@pytest.mark.parametrize("kwargs", [{"batch_size": 0}, {"max_len": 1000}, {"head_max_len": 3}])
def test_invalid_limits_fail(model, kwargs):
    with pytest.raises(ValueError):
        model.predict_batch(["hello"], QUESTIONS, **kwargs)


def test_missing_local_checkpoint_is_not_treated_as_hub_repo(tmp_path):
    with pytest.raises(FileNotFoundError, match="Local checkpoint"):
        jev_urdu.load(tmp_path / "missing")


def test_strict_weights_refuse_an_incompatible_checkpoint(checkpoint, tmp_path):
    directory = tmp_path / "broken"
    shutil.copytree(checkpoint, directory)
    path = directory / "model.safetensors"
    weights = load_file(str(path))
    del weights["scorer.3.weight"]
    save_file(weights, str(path))
    with pytest.raises(RuntimeError, match="scorer.3.weight"):
        jev_urdu.load(directory)


def test_invalid_temperature_fails_on_load(checkpoint, tmp_path):
    directory = tmp_path / "temperature"
    shutil.copytree(checkpoint, directory)
    path = directory / "rl_agent_config.json"
    cfg = json.loads(path.read_text(encoding="utf-8"))
    cfg["temperature"] = [1.0, float("nan"), 1.0]
    path.write_text(json.dumps(cfg), encoding="utf-8")
    with pytest.raises(ValueError, match="finite positive"):
        jev_urdu.load(directory)

def test_sindhi_convenience_methods(model):
    # Test sentiment helper with lang="sd"
    sent = model.sentiment("اڄ جو ڏينهن تمام سٺو گذريو.", lang="sd")
    assert "sentiment" in sent
    assert sent["sentiment"]["answer"] in ["مثبت", "منفي", "غير جانبدار"]
    assert 0 <= sent["sentiment"]["probability"] <= 1
    assert "dissatisfied" in sent

    # Test topic helper with default lang="sd"
    top = model.topic("اسان جي زمين ۾ ڪڻڪ سٺي ٿي آهي.", lang="sd")
    assert "topic" in top
    assert top["topic"]["answer"] in ["تعليم", "ٽيڪنالاجي", "زراعت", "صحت", "ٻيو"]
    assert 0 <= top["topic"]["probability"] <= 1

    # Test triage helper with lang="sd"
    tri = model.triage("منهنجو اڪائونٽ ڪم نٿو ڪري.", lang="sd")
    assert set(tri) == {"issue", "wants_human", "priority"}
    assert tri["priority"]["answer"] in ["فوري", "جلد", "معمول"]