feat: robust pattern-matching cost estimation and trace decorator resilience

packages/sdk/agentscope/_pricing.py

@@ -1,36 +1,83 @@
-# Hardcoded token pricing table for LLM cost estimation
-# Prices are represented in USD per 1,000,000 (1M) tokens.
+import json
+import os
+from typing import Dict, Optional
 
-PRICING_TABLE = {
+# Base pricing table (in USD per 1M tokens)
+# We support prefix-matching so pinned/versioned variants resolve correctly.
+PRICING_TABLE: Dict[str, Dict[str, float]] = {
     "gpt-4o": {"input": 2.50, "output": 10.00},
     "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 update_pricing_table(custom_pricing: Dict[str, Dict[str, float]]) -> None:
+    """Programmatically update or override the model pricing table.
+
+    Args:
+        custom_pricing: Dictionary of model names to dicts containing 'input' and 'output' rates.
+    """
+    for model_name, rates in custom_pricing.items():
+        if "input" in rates and "output" in rates:
+            PRICING_TABLE[model_name] = {
+                "input": float(rates["input"]),
+                "output": float(rates["output"]),
+            }
+
+
+def _load_env_overrides() -> None:
+    """Load custom pricing overrides from AGENTSCOPE_CUSTOM_PRICING environment variable."""
+    env_pricing = os.getenv("AGENTSCOPE_CUSTOM_PRICING")
+    if env_pricing:
+        try:
+            custom_pricing = json.loads(env_pricing)
+            if isinstance(custom_pricing, dict):
+                update_pricing_table(custom_pricing)
+        except Exception:
+            # Silently ignore parsing errors in environment configurations
+            pass
+
+
+# Automatically load any environment overrides on import
+_load_env_overrides()
+
+
 def calculate_cost(
     model_name: str, prompt_tokens: int | None, completion_tokens: int | None
 ) -> float:
     """Calculate the estimated USD cost of an LLM call.
 
     Args:
-        model_name: The name of the LLM model used.
+        model_name: The name/identifier of the LLM model.
         prompt_tokens: Number of prompt (input) tokens.
         completion_tokens: Number of completion (output) tokens.
 
     Returns:
         The estimated cost in USD (float).
     """
-    if not model_name or model_name not in PRICING_TABLE:
+    if not model_name:
+        return 0.0
+
+    # Resolve pricing mapping using prefix/pattern matching.
+    # We sort keys by length descending to match the most specific pattern first.
+    matched_prices = None
+    for base_model in sorted(PRICING_TABLE.keys(), key=len, reverse=True):
+        if base_model in model_name:
+            matched_prices = PRICING_TABLE[base_model]
+            break
+
+    if not matched_prices:
         return 0.0
 
-    prices = PRICING_TABLE[model_name]
     input_tokens = prompt_tokens or 0
     output_tokens = completion_tokens or 0
 
-    input_cost = (input_tokens / 1_000_000) * prices["input"]
-    output_cost = (output_tokens / 1_000_000) * prices["output"]
+    input_cost = (input_tokens / 1_000_000) * matched_prices["input"]
+    output_cost = (output_tokens / 1_000_000) * matched_prices["output"]
 
     return input_cost + output_cost