rewards/token_budget.py

"""Token budget reward component — REW-01 + REW-04.

α-parameterized shaping with 3 variants for the Mercor Pareto sub-prize.
Pure Python, stdlib-only. No eval/exec/compile/os/subprocess.
"""


def token_budget(
    completion: str,
    prompt_token_count: int,
    alpha: float = 0.2,
    variant: str = "capped_linear",
    **_,
) -> float:
    """α-parameterized token budget reward.  REW-04.

    Variants:
        - capped_linear:    max(0.0, 1 - α * ratio)   in [0, 1]
        - capped_quadratic: max(0.0, 1 - α * ratio²)  in [0, 1]
        - uncapped:         max(-1.0, 1 - α * ratio)   in [-1, 1]

    Raises:
        ValueError: for unknown variant strings.
    """
    completion_tokens = max(1, len(completion.split()))
    ratio = completion_tokens / max(1, prompt_token_count)

    if variant == "capped_linear":
        return max(0.0, min(1.0, 1.0 - alpha * ratio))
    elif variant == "capped_quadratic":
        return max(0.0, min(1.0, 1.0 - alpha * ratio * ratio))
    elif variant == "uncapped":
        return max(-1.0, 1.0 - alpha * ratio)
    else:
        raise ValueError(
            f"REW-04: unknown variant '{variant}' "
            f"(expected capped_linear|capped_quadratic|uncapped)"
        )


__all__ = ["token_budget"]