diff --git a/project/paperbench/paperbench/judge/token_usage.py b/project/paperbench/paperbench/judge/token_usage.py index 1cd8cacc..91c929e8 100644 --- a/project/paperbench/paperbench/judge/token_usage.py +++ b/project/paperbench/paperbench/judge/token_usage.py @@ -45,10 +45,11 @@ def _get_leaf_node_token_usages(task: GradedTaskNode) -> list[TokenUsage]: if task.is_leaf(): # need this check because judge_metadata may be malformed in case of node errors - if task.judge_metadata is not None and "token_usage" in task.judge_metadata: - return [TokenUsage.from_dict(task.judge_metadata["token_usage"])] - else: - return [] + if task.judge_metadata is not None: + usage = task.judge_metadata.get("token_usage") + if isinstance(usage, dict): + return [TokenUsage.from_dict(usage)] + return [] token_usages = [] diff --git a/project/paperbench/tests/unit/test_token_usage.py b/project/paperbench/tests/unit/test_token_usage.py new file mode 100644 index 00000000..0d0d255f --- /dev/null +++ b/project/paperbench/tests/unit/test_token_usage.py @@ -0,0 +1,56 @@ +from paperbench.judge.graded_task_node import GradedTaskNode +from paperbench.judge.token_usage import get_total_token_usage +from paperbench.rubric.tasks import TaskNode + + +def _leaf(node_id: str, token_usage): + task = TaskNode( + id=node_id, + requirements=node_id, + weight=1, + task_category="Code Development", + ) + return GradedTaskNode.from_task( + task, + score=1.0, + valid_score=True, + explanation="graded", + judge_metadata={"token_usage": token_usage}, + ) + + +def test_total_token_usage_skips_null_leaf_usage() -> None: + root = GradedTaskNode( + id="root", + requirements="root", + weight=1, + sub_tasks=[ + _leaf("without-usage", None), + _leaf("with-usage", {"gpt-test": {"in": 12, "out": 5}}), + ], + score=1.0, + valid_score=True, + explanation="aggregate", + ) + + total = get_total_token_usage(root) + + assert total.to_dict() == {"gpt-test": {"in": 12, "out": 5}} + + +def test_total_token_usage_skips_missing_leaf_usage_key() -> None: + task = TaskNode( + id="leaf", + requirements="leaf", + weight=1, + task_category="Code Development", + ) + leaf = GradedTaskNode.from_task( + task, + score=1.0, + valid_score=True, + explanation="graded", + judge_metadata={}, + ) + + assert get_total_token_usage(leaf).to_dict() == {}