diff --git a/packages/sdk/agentscope/_pricing.py b/packages/sdk/agentscope/_pricing.py index c5e149e..b5c30eb 100644 --- a/packages/sdk/agentscope/_pricing.py +++ b/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 diff --git a/packages/sdk/agentscope/decorators.py b/packages/sdk/agentscope/decorators.py index 28cf17b..26917ca 100644 --- a/packages/sdk/agentscope/decorators.py +++ b/packages/sdk/agentscope/decorators.py @@ -88,7 +88,12 @@ def _make_decorator( ) -> Callable[..., Any]: @functools.wraps(func) def sync_wrapper(*args: Any, **kwargs: Any) -> Any: - client = get_global_client() + client = None + try: + client = get_global_client() + except Exception: + pass + run_id = str(uuid.uuid4()) parent_id = current_parent_run_id.get() token = current_parent_run_id.set(run_id) @@ -100,25 +105,29 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: } # Start Event - client.emit( - { - "event_id": run_id, - "session_id": "", - "parent_event_id": parent_id, - "event_type": "chain_start", - "agent_name": name, - "agent_type": agent_type, - "timestamp": datetime.now(timezone.utc).isoformat(), - "latency_ms": None, - "status": "running", - "payload": { - "chain_type": name, - "inputs": inputs, - "outputs": None, - "error": None, - }, - } - ) + if client: + try: + client.emit( + { + "event_id": run_id, + "session_id": "", + "parent_event_id": parent_id, + "event_type": "chain_start", + "agent_name": name, + "agent_type": agent_type, + "timestamp": datetime.now(timezone.utc).isoformat(), + "latency_ms": None, + "status": "running", + "payload": { + "chain_type": name, + "inputs": inputs, + "outputs": None, + "error": None, + }, + } + ) + except Exception: + pass start_time = time.perf_counter() try: @@ -126,56 +135,69 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: latency = int((time.perf_counter() - start_time) * 1000) # End Event - client.emit( - { - "event_id": run_id, - "session_id": "", - "parent_event_id": parent_id, - "event_type": "chain_end", - "agent_name": name, - "agent_type": agent_type, - "timestamp": datetime.now(timezone.utc).isoformat(), - "latency_ms": latency, - "status": "completed", - "payload": { - "chain_type": name, - "inputs": {}, - "outputs": {"result": str(result)}, - "error": None, - }, - } - ) + if client: + try: + client.emit( + { + "event_id": run_id, + "session_id": "", + "parent_event_id": parent_id, + "event_type": "chain_end", + "agent_name": name, + "agent_type": agent_type, + "timestamp": datetime.now(timezone.utc).isoformat(), + "latency_ms": latency, + "status": "completed", + "payload": { + "chain_type": name, + "inputs": {}, + "outputs": {"result": str(result)}, + "error": None, + }, + } + ) + except Exception: + pass return result except Exception as e: latency = int((time.perf_counter() - start_time) * 1000) # Error Event - client.emit( - { - "event_id": run_id, - "session_id": "", - "parent_event_id": parent_id, - "event_type": "chain_error", - "agent_name": name, - "agent_type": agent_type, - "timestamp": datetime.now(timezone.utc).isoformat(), - "latency_ms": latency, - "status": "error", - "payload": { - "chain_type": name, - "inputs": {}, - "outputs": None, - "error": str(e), - }, - } - ) + if client: + try: + client.emit( + { + "event_id": run_id, + "session_id": "", + "parent_event_id": parent_id, + "event_type": "chain_error", + "agent_name": name, + "agent_type": agent_type, + "timestamp": datetime.now(timezone.utc).isoformat(), + "latency_ms": latency, + "status": "error", + "payload": { + "chain_type": name, + "inputs": {}, + "outputs": None, + "error": str(e), + }, + } + ) + except Exception: + pass raise e finally: current_parent_run_id.reset(token) @functools.wraps(func) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: - client = get_global_client() + client = None + try: + client = get_global_client() + except Exception: + pass + run_id = str(uuid.uuid4()) parent_id = current_parent_run_id.get() token = current_parent_run_id.set(run_id) @@ -186,25 +208,29 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: } # Start Event - client.emit( - { - "event_id": run_id, - "session_id": "", - "parent_event_id": parent_id, - "event_type": "chain_start", - "agent_name": name, - "agent_type": agent_type, - "timestamp": datetime.now(timezone.utc).isoformat(), - "latency_ms": None, - "status": "running", - "payload": { - "chain_type": name, - "inputs": inputs, - "outputs": None, - "error": None, - }, - } - ) + if client: + try: + client.emit( + { + "event_id": run_id, + "session_id": "", + "parent_event_id": parent_id, + "event_type": "chain_start", + "agent_name": name, + "agent_type": agent_type, + "timestamp": datetime.now(timezone.utc).isoformat(), + "latency_ms": None, + "status": "running", + "payload": { + "chain_type": name, + "inputs": inputs, + "outputs": None, + "error": None, + }, + } + ) + except Exception: + pass start_time = time.perf_counter() try: @@ -212,49 +238,57 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: latency = int((time.perf_counter() - start_time) * 1000) # End Event - client.emit( - { - "event_id": run_id, - "session_id": "", - "parent_event_id": parent_id, - "event_type": "chain_end", - "agent_name": name, - "agent_type": agent_type, - "timestamp": datetime.now(timezone.utc).isoformat(), - "latency_ms": latency, - "status": "completed", - "payload": { - "chain_type": name, - "inputs": {}, - "outputs": {"result": str(result)}, - "error": None, - }, - } - ) + if client: + try: + client.emit( + { + "event_id": run_id, + "session_id": "", + "parent_event_id": parent_id, + "event_type": "chain_end", + "agent_name": name, + "agent_type": agent_type, + "timestamp": datetime.now(timezone.utc).isoformat(), + "latency_ms": latency, + "status": "completed", + "payload": { + "chain_type": name, + "inputs": {}, + "outputs": {"result": str(result)}, + "error": None, + }, + } + ) + except Exception: + pass return result except Exception as e: latency = int((time.perf_counter() - start_time) * 1000) # Error Event - client.emit( - { - "event_id": run_id, - "session_id": "", - "parent_event_id": parent_id, - "event_type": "chain_error", - "agent_name": name, - "agent_type": agent_type, - "timestamp": datetime.now(timezone.utc).isoformat(), - "latency_ms": latency, - "status": "error", - "payload": { - "chain_type": name, - "inputs": {}, - "outputs": None, - "error": str(e), - }, - } - ) + if client: + try: + client.emit( + { + "event_id": run_id, + "session_id": "", + "parent_event_id": parent_id, + "event_type": "chain_error", + "agent_name": name, + "agent_type": agent_type, + "timestamp": datetime.now(timezone.utc).isoformat(), + "latency_ms": latency, + "status": "error", + "payload": { + "chain_type": name, + "inputs": {}, + "outputs": None, + "error": str(e), + }, + } + ) + except Exception: + pass raise e finally: current_parent_run_id.reset(token) diff --git a/packages/sdk/tests/test_sdk.py b/packages/sdk/tests/test_sdk.py index d8c6dd3..e07f799 100644 --- a/packages/sdk/tests/test_sdk.py +++ b/packages/sdk/tests/test_sdk.py @@ -8,10 +8,30 @@ def test_calculate_cost(): - # gpt-4o price: 2.50 input / 10.00 output per 1M tokens + # Base match assert calculate_cost("gpt-4o", 1_000_000, 1_000_000) == 12.50 + # Versioned/pinned match (Issue #25) + assert calculate_cost("gpt-4o-2024-05-13", 1_000_000, 1_000_000) == 12.50 + assert calculate_cost("anthropic/claude-3-5-sonnet", 1_000_000, 1_000_000) == 18.00 + # Newly added models (Issue #29) + assert calculate_cost("claude-3-haiku-20240307", 1_000_000, 1_000_000) == 1.50 + assert calculate_cost("gemini-1.5-flash-latest", 1_000_000, 1_000_000) == 0.375 + # Unknown model fallback assert calculate_cost("unknown-model", 100, 100) == 0.0 + # Custom programmatic updates + from agentscope._pricing import update_pricing_table + update_pricing_table({"custom-model-x": {"input": 1.0, "output": 2.0}}) + assert calculate_cost("custom-model-x-v1", 1_000_000, 1_000_000) == 3.0 + + # Environment variable overrides + import os + from agentscope._pricing import _load_env_overrides + os.environ["AGENTSCOPE_CUSTOM_PRICING"] = '{"custom-env-y": {"input": 5.0, "output": 10.0}}' + _load_env_overrides() + assert calculate_cost("custom-env-y", 1_000_000, 1_000_000) == 15.0 + del os.environ["AGENTSCOPE_CUSTOM_PRICING"] + @trace(name="test_sync") def sync_fn(x): @@ -245,3 +265,35 @@ def test_client_pending_status_patch(): finally: urllib.request.urlopen = original_urlopen + +def test_decorators_resilience(): + client = get_global_client() + + # Mock emit to throw an error + def failing_emit(event): + raise RuntimeError("Websocket connection lost") + + original_emit = client.emit + client.emit = failing_emit + + try: + # Should execute successfully and not crash the user function + @trace(name="resilient_sync") + def sync_hello(x): + return x * 10 + + assert sync_hello(4) == 40 + + # Test async wrapper resilience + @trace(name="resilient_async") + async def async_hello(y): + return y + 10 + + import asyncio + + res = asyncio.run(async_hello(5)) + assert res == 15 + finally: + client.emit = original_emit + +