56c56a7b4a
All checks were successful
CI / test (pull_request) Successful in 4m9s
新增成本跟踪模块，为对话级和任务级成本计算提供基础。 - Turn 模型新增 prompt_tokens、completion_tokens、total_tokens 字段 - OpenAI 协议适配器新增 parse_usage() 提取 token 使用量 - 新增 evaluation/cost_tracking.py 模块： - ModelPricing：模型定价配置 - TokenUsage：token 使用量聚合 - CostBreakdown：成本明细 - calculate_cost()：根据 token 使用量和定价计算费用 - calculate_turn_cost()、calculate_case_cost()、calculate_run_cost() - 内置常见模型定价（GPT-4o、GPT-4o-mini、Claude 等） - 新增 11 项单元测试（677 tests passed） 注：引擎集成（实际捕获 API 调用的 token 使用量）留待后续实现。 Closes #25
142 lines
4.7 KiB
Python
142 lines
4.7 KiB
Python
"""Tests for cost tracking module."""
|
|
|
|
import pytest
|
|
|
|
from agenteval.evaluation.cost_tracking import (
|
|
CostBreakdown,
|
|
ModelPricing,
|
|
TokenUsage,
|
|
aggregate_token_usage,
|
|
calculate_case_cost,
|
|
calculate_cost,
|
|
calculate_run_cost,
|
|
calculate_turn_cost,
|
|
)
|
|
|
|
|
|
def test_model_pricing_model():
|
|
"""ModelPricing model should work correctly."""
|
|
pricing = ModelPricing(
|
|
model_id="gpt-4o",
|
|
prompt_cost_per_1m=5.0,
|
|
completion_cost_per_1m=15.0,
|
|
)
|
|
assert pricing.model_id == "gpt-4o"
|
|
assert pricing.prompt_cost_per_1m == 5.0
|
|
|
|
|
|
def test_token_usage_model():
|
|
"""TokenUsage model should work correctly."""
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
|
|
assert usage.prompt_tokens == 100
|
|
assert usage.completion_tokens == 50
|
|
assert usage.total_tokens == 150
|
|
|
|
|
|
def test_calculate_cost_gpt4o_mini():
|
|
"""Cost calculation for gpt-4o-mini should be correct."""
|
|
pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
|
|
# 1000 prompt tokens + 500 completion tokens
|
|
# Cost = (1000/1M * 0.15) + (500/1M * 0.60) = 0.00015 + 0.0003 = 0.00045
|
|
cost = calculate_cost(1000, 500, pricing)
|
|
assert abs(cost - 0.00045) < 0.00001
|
|
|
|
|
|
def test_calculate_cost_zero_tokens():
|
|
"""Zero tokens should result in zero cost."""
|
|
pricing = ModelPricing(model_id="gpt-4o", prompt_cost_per_1m=5.0, completion_cost_per_1m=15.0)
|
|
cost = calculate_cost(0, 0, pricing)
|
|
assert cost == 0.0
|
|
|
|
|
|
def test_aggregate_token_usage():
|
|
"""Aggregate token usage from multiple turns."""
|
|
turns = [
|
|
{"prompt_tokens": 100, "completion_tokens": 50},
|
|
{"prompt_tokens": 200, "completion_tokens": 100},
|
|
{"prompt_tokens": None, "completion_tokens": None}, # Missing data
|
|
]
|
|
usage = aggregate_token_usage(turns)
|
|
assert usage.prompt_tokens == 300
|
|
assert usage.completion_tokens == 150
|
|
assert usage.total_tokens == 450
|
|
|
|
|
|
def test_aggregate_token_usage_empty():
|
|
"""Empty turns should result in zero usage."""
|
|
usage = aggregate_token_usage([])
|
|
assert usage.prompt_tokens == 0
|
|
assert usage.completion_tokens == 0
|
|
assert usage.total_tokens == 0
|
|
|
|
|
|
def test_calculate_turn_cost():
|
|
"""Calculate cost for a single turn."""
|
|
pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
|
|
turn = {"prompt_tokens": 1000, "completion_tokens": 500}
|
|
breakdown = calculate_turn_cost(turn, pricing)
|
|
assert breakdown.prompt_tokens == 1000
|
|
assert breakdown.completion_tokens == 500
|
|
assert breakdown.total_tokens == 1500
|
|
assert abs(breakdown.cost_usd - 0.00045) < 0.00001
|
|
|
|
|
|
def test_calculate_turn_cost_missing_tokens():
|
|
"""Turn with missing token data should have zero cost."""
|
|
pricing = ModelPricing(model_id="gpt-4o", prompt_cost_per_1m=5.0, completion_cost_per_1m=15.0)
|
|
turn = {} # No token data
|
|
breakdown = calculate_turn_cost(turn, pricing)
|
|
assert breakdown.prompt_tokens == 0
|
|
assert breakdown.completion_tokens == 0
|
|
assert breakdown.cost_usd == 0.0
|
|
|
|
|
|
def test_calculate_case_cost():
|
|
"""Calculate cost for a case with multiple turns."""
|
|
pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
|
|
turns = [
|
|
{"prompt_tokens": 1000, "completion_tokens": 500},
|
|
{"prompt_tokens": 2000, "completion_tokens": 1000},
|
|
]
|
|
breakdown = calculate_case_cost(turns, pricing)
|
|
assert breakdown.prompt_tokens == 3000
|
|
assert breakdown.completion_tokens == 1500
|
|
assert breakdown.total_tokens == 4500
|
|
# Cost = (3000/1M * 0.15) + (1500/1M * 0.60) = 0.00045 + 0.0009 = 0.00135
|
|
assert abs(breakdown.cost_usd - 0.00135) < 0.00001
|
|
|
|
|
|
def test_calculate_run_cost():
|
|
"""Calculate total cost for a run with multiple cases."""
|
|
pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
|
|
cases = [
|
|
{
|
|
"case_id": "c1",
|
|
"turns": [
|
|
{"prompt_tokens": 1000, "completion_tokens": 500},
|
|
],
|
|
},
|
|
{
|
|
"case_id": "c2",
|
|
"turns": [
|
|
{"prompt_tokens": 2000, "completion_tokens": 1000},
|
|
],
|
|
},
|
|
]
|
|
breakdown = calculate_run_cost(cases, pricing)
|
|
assert breakdown.prompt_tokens == 3000
|
|
assert breakdown.completion_tokens == 1500
|
|
assert breakdown.total_tokens == 4500
|
|
|
|
|
|
def test_cost_breakdown_model():
|
|
"""CostBreakdown model should work correctly."""
|
|
breakdown = CostBreakdown(
|
|
prompt_tokens=1000,
|
|
completion_tokens=500,
|
|
total_tokens=1500,
|
|
cost_usd=0.001,
|
|
)
|
|
assert breakdown.prompt_tokens == 1000
|
|
assert breakdown.cost_usd == 0.001
|