diff --git a/packages/sdk/agentscope/_pricing.py b/packages/sdk/agentscope/_pricing.py index c5e149e..7f4794c 100644 --- a/packages/sdk/agentscope/_pricing.py +++ b/packages/sdk/agentscope/_pricing.py @@ -6,10 +6,29 @@ "gpt-4": {"input": 30.00, "output": 60.00}, "gpt-3.5-turbo": {"input": 0.50, "output": 1.50}, "claude-3-5-sonnet": {"input": 3.00, "output": 15.00}, + "claude-3-haiku": {"input": 0.25, "output": 1.25}, + "claude-3-opus": {"input": 15.00, "output": 75.00}, "gemini-1.5-pro": {"input": 1.25, "output": 5.00}, + "gemini-1.5-flash": {"input": 0.075, "output": 0.30}, } +def get_model_pricing(model_name: str) -> dict[str, float] | None: + """Lookup pricing configuration for a given model name, including minor/date-pinned versions.""" + if not model_name: + return None + + model_name_lower = model_name.lower() + if model_name_lower in PRICING_TABLE: + return PRICING_TABLE[model_name_lower] + + for key in sorted(PRICING_TABLE.keys(), key=len, reverse=True): + if model_name_lower.startswith(f"{key}-"): + return PRICING_TABLE[key] + + return None + + def calculate_cost( model_name: str, prompt_tokens: int | None, completion_tokens: int | None ) -> float: @@ -23,14 +42,20 @@ def calculate_cost( Returns: The estimated cost in USD (float). """ - if not model_name or model_name not in PRICING_TABLE: + prices = get_model_pricing(model_name) + if not prices: return 0.0 - prices = PRICING_TABLE[model_name] - input_tokens = prompt_tokens or 0 - output_tokens = completion_tokens or 0 + if prompt_tokens is not None and prompt_tokens < 0: + raise ValueError("prompt_tokens cannot be negative") + if completion_tokens is not None and completion_tokens < 0: + raise ValueError("completion_tokens cannot be negative") + + input_tokens = 0 if prompt_tokens is None else prompt_tokens + output_tokens = 0 if completion_tokens is None else completion_tokens input_cost = (input_tokens / 1_000_000) * prices["input"] output_cost = (output_tokens / 1_000_000) * prices["output"] return input_cost + output_cost + diff --git a/packages/sdk/tests/test_sdk.py b/packages/sdk/tests/test_sdk.py index 285cfba..379a5e6 100644 --- a/packages/sdk/tests/test_sdk.py +++ b/packages/sdk/tests/test_sdk.py @@ -12,6 +12,32 @@ def test_calculate_cost(): assert calculate_cost("gpt-4o", 1_000_000, 1_000_000) == 12.50 assert calculate_cost("unknown-model", 100, 100) == 0.0 + # New model pricing checks + # claude-3-haiku: 0.25 input / 1.25 output per 1M tokens + assert calculate_cost("claude-3-haiku", 1_000_000, 1_000_000) == 1.50 + # claude-3-opus: 15.00 input / 75.00 output per 1M tokens + assert calculate_cost("claude-3-opus", 1_000_000, 1_000_000) == 90.00 + # gemini-1.5-flash: 0.075 input / 0.30 output per 1M tokens + assert calculate_cost("gemini-1.5-flash", 1_000_000, 1_000_000) == pytest.approx(0.375) + + # Date-pinned & minor version matching + assert calculate_cost("claude-3-haiku-20240307", 1_000_000, 1_000_000) == 1.50 + assert calculate_cost("CLAUDE-3-HAIKU-20240307", 1_000_000, 1_000_000) == 1.50 + assert calculate_cost("claude-3-opus-20240229", 1_000_000, 1_000_000) == 90.00 + assert calculate_cost("gemini-1.5-flash-001", 1_000_000, 1_000_000) == pytest.approx(0.375) + assert calculate_cost("gpt-4o-2024-05-13", 1_000_000, 1_000_000) == 12.50 + + # Ensure non-delimited model names do not false match + assert calculate_cost("gpt-4orange", 1_000_000, 1_000_000) == 0.0 + assert calculate_cost("claude-3-haikuish", 1_000_000, 1_000_000) == 0.0 + + # Negative token count validation + with pytest.raises(ValueError, match="prompt_tokens cannot be negative"): + calculate_cost("gpt-4o", -1, 100) + with pytest.raises(ValueError, match="completion_tokens cannot be negative"): + calculate_cost("gpt-4o", 100, -1) + + @trace(name="test_sync") def sync_fn(x):