diff --git a/backend/agenteval/evaluation/cost_tracking.py b/backend/agenteval/evaluation/cost_tracking.py new file mode 100644 index 0000000..ea41a44 --- /dev/null +++ b/backend/agenteval/evaluation/cost_tracking.py @@ -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, + ) diff --git a/backend/agenteval/model_protocols/openai.py b/backend/agenteval/model_protocols/openai.py index 02f8fd0..7765e9a 100644 --- a/backend/agenteval/model_protocols/openai.py +++ b/backend/agenteval/model_protocols/openai.py @@ -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} diff --git a/backend/agenteval/models.py b/backend/agenteval/models.py index 6871760..682ea31 100644 --- a/backend/agenteval/models.py +++ b/backend/agenteval/models.py @@ -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): diff --git a/tests/unit/test_cost_tracking.py b/tests/unit/test_cost_tracking.py new file mode 100644 index 0000000..10676c0 --- /dev/null +++ b/tests/unit/test_cost_tracking.py @@ -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