feat(sdk): Add Claude 3 and Gemini 1.5 Flash to pricing table

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
+

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):