tests/test_rewards.py

import numpy as np

from fugu_lite.config import RewardConfig
from fugu_lite.rewards import (
    compute_utility,
    deterministic_quality,
    parse_judge_score,
    soft_targets,
)
from fugu_lite.schemas import GraderSpec, TaskRecord, WorkerResult


def test_deterministic_graders():
    numeric = TaskRecord(
        task_id="n",
        prompt="question",
        reference_answer="42",
        grader=GraderSpec(type="numeric", tolerance=0.01),
    )
    contains = TaskRecord(
        task_id="c",
        prompt="question",
        reference_answer="Paris",
        grader=GraderSpec(type="contains"),
    )
    assert deterministic_quality(numeric, "The answer is 42.0") == 1.0
    assert deterministic_quality(numeric, "41") == 0.0
    assert deterministic_quality(contains, "It is PARIS, France.") == 1.0


def test_utility_and_soft_targets():
    result = WorkerResult(
        worker_id="a",
        requested_model="m",
        response="ok",
        quality=1.0,
        cost_usd=0.01,
        latency_ms=30_000,
    )
    config = RewardConfig(cost_weight=0.1, latency_weight=0.2)
    assert abs(compute_utility(result, config) - 0.7) < 1e-8
    targets = soft_targets([[1.0, 0.0]], temperature=0.1)
    assert targets.shape == (1, 2)
    assert targets[0, 0] > 0.999
    assert np.isclose(targets.sum(), 1.0)


def test_parse_judge_score():
    assert parse_judge_score("0.75") == 0.75
    assert parse_judge_score("score: 1.0") == 1.0