feat: 成本效率跟踪基础设施
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
This commit is contained in:
sinohqb 2026-08-25 16:06:38 +08:00
parent 2c79abf1de
commit 56c56a7b4a
4 changed files with 313 additions and 0 deletions

View File

@ -0,0 +1,157 @@
"""Cost tracking and calculation for evaluation runs."""
from typing import Any
from pydantic import BaseModel, Field
class ModelPricing(BaseModel):
"""Pricing for a model (per 1M tokens)."""
model_id: str
prompt_cost_per_1m: float = Field(description="Cost per 1M prompt tokens in USD")
completion_cost_per_1m: float = Field(description="Cost per 1M completion tokens in USD")
class TokenUsage(BaseModel):
"""Token usage for a single API call."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
class CostBreakdown(BaseModel):
"""Cost breakdown for a turn, case, or run."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
cost_usd: float = 0.0
# Default pricing for common models (per 1M tokens in USD)
DEFAULT_PRICING: dict[str, ModelPricing] = {
"gpt-4o": ModelPricing(model_id="gpt-4o", prompt_cost_per_1m=5.0, completion_cost_per_1m=15.0),
"gpt-4o-mini": ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60),
"gpt-4-turbo": ModelPricing(model_id="gpt-4-turbo", prompt_cost_per_1m=10.0, completion_cost_per_1m=30.0),
"gpt-3.5-turbo": ModelPricing(model_id="gpt-3.5-turbo", prompt_cost_per_1m=0.50, completion_cost_per_1m=1.50),
"claude-3-5-sonnet": ModelPricing(model_id="claude-3-5-sonnet", prompt_cost_per_1m=3.0, completion_cost_per_1m=15.0),
"claude-3-haiku": ModelPricing(model_id="claude-3-haiku", prompt_cost_per_1m=0.25, completion_cost_per_1m=1.25),
}
def calculate_cost(
prompt_tokens: int,
completion_tokens: int,
pricing: ModelPricing,
) -> float:
"""Calculate cost in USD for given token usage and pricing.
Args:
prompt_tokens: Number of prompt tokens
completion_tokens: Number of completion tokens
pricing: Model pricing configuration
Returns:
Cost in USD
"""
prompt_cost = (prompt_tokens / 1_000_000) * pricing.prompt_cost_per_1m
completion_cost = (completion_tokens / 1_000_000) * pricing.completion_cost_per_1m
return prompt_cost + completion_cost
def aggregate_token_usage(turns: list[dict[str, Any]]) -> TokenUsage:
"""Aggregate token usage from a list of turns.
Args:
turns: List of turn dicts with optional token fields
Returns:
Aggregated token usage
"""
prompt_tokens = sum(t.get("prompt_tokens") or 0 for t in turns)
completion_tokens = sum(t.get("completion_tokens") or 0 for t in turns)
return TokenUsage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
def calculate_turn_cost(
turn: dict[str, Any],
pricing: ModelPricing,
) -> CostBreakdown:
"""Calculate cost for a single turn.
Args:
turn: Turn dict with optional token fields
pricing: Model pricing configuration
Returns:
Cost breakdown for the turn
"""
prompt_tokens = turn.get("prompt_tokens") or 0
completion_tokens = turn.get("completion_tokens") or 0
total_tokens = prompt_tokens + completion_tokens
cost_usd = calculate_cost(prompt_tokens, completion_tokens, pricing)
return CostBreakdown(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
cost_usd=cost_usd,
)
def calculate_case_cost(
turns: list[dict[str, Any]],
pricing: ModelPricing,
) -> CostBreakdown:
"""Calculate cost for a case (multiple turns).
Args:
turns: List of turn dicts
pricing: Model pricing configuration
Returns:
Aggregated cost breakdown for the case
"""
usage = aggregate_token_usage(turns)
cost_usd = calculate_cost(usage.prompt_tokens, usage.completion_tokens, pricing)
return CostBreakdown(
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
cost_usd=cost_usd,
)
def calculate_run_cost(
cases: list[dict[str, Any]],
pricing: ModelPricing,
) -> CostBreakdown:
"""Calculate total cost for a run (multiple cases).
Args:
cases: List of case dicts, each with a 'turns' field
pricing: Model pricing configuration
Returns:
Aggregated cost breakdown for the run
"""
all_turns = []
for case in cases:
all_turns.extend(case.get("turns", []))
usage = aggregate_token_usage(all_turns)
cost_usd = calculate_cost(usage.prompt_tokens, usage.completion_tokens, pricing)
return CostBreakdown(
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
cost_usd=cost_usd,
)

View File

@ -32,6 +32,17 @@ class OpenAICompatibleAdapter(ModelProtocolAdapter):
raise ProtocolAdapterError("模型接口返回内容为空")
return content
def parse_usage(self, data: dict[str, Any]) -> dict[str, int] | None:
"""Extract token usage from OpenAI-compatible response."""
usage = data.get("usage")
if not usage or not isinstance(usage, dict):
return None
return {
"prompt_tokens": usage.get("prompt_tokens", 0),
"completion_tokens": usage.get("completion_tokens", 0),
"total_tokens": usage.get("total_tokens", 0),
}
def embedding_payload(self, model_name: str | None, inputs: str | list[str]) -> dict[str, Any]:
return {"model": self.require_model(model_name), "input": inputs}

View File

@ -332,6 +332,10 @@ class Turn(BaseModel):
reply: Optional[dict[str, Any]] = None
received_at: Optional[datetime] = None
latency_ms: Optional[int] = None
# Token usage for cost tracking
prompt_tokens: Optional[int] = None
completion_tokens: Optional[int] = None
total_tokens: Optional[int] = None
class EvalResult(BaseModel):

View File

@ -0,0 +1,141 @@
"""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