diff --git a/pricing.json b/pricing.json new file mode 100644 index 00000000..f0443799 --- /dev/null +++ b/pricing.json @@ -0,0 +1,27 @@ +{ + "alias0": { + "max_tokens": 128000, + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "input_cost_per_token": 0.000005, + "output_cost_per_token": 0.00005, + "cache_creation_input_token_cost": 0.000005, + "cache_read_input_token_cost": 0.0000005, + "search_context_cost_per_query": { + "search_context_size_low": 1e-2, + "search_context_size_medium": 1e-2, + "search_context_size_high": 1e-2 + }, + "litellm_provider": "openai", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 159, + "supports_assistant_prefill": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "deprecation_date": "2026-02-01", + "supports_tool_choice": true + } +} \ No newline at end of file diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py index eb44d9a0..6f3edde6 100644 --- a/src/cai/repl/ui/toolbar.py +++ b/src/cai/repl/ui/toolbar.py @@ -123,6 +123,8 @@ def update_toolbar_in_background(): os.getenv('CAI_MODEL', 'default')} | " f"Max Turns: { os.getenv('CAI_MAX_TURNS', 'inf')} | " + f"Price Limit: { + os.getenv('CAI_PRICE_LIMIT', 'inf')} | " f"{current_time_with_tz}" ) toolbar_cache['last_update'] = datetime.datetime.now() diff --git a/src/cai/sdk/agents/exceptions.py b/src/cai/sdk/agents/exceptions.py index 820390db..170bdd76 100644 --- a/src/cai/sdk/agents/exceptions.py +++ b/src/cai/sdk/agents/exceptions.py @@ -61,3 +61,9 @@ class OutputGuardrailTripwireTriggered(AgentsException): super().__init__( f"Guardrail {guardrail_result.guardrail.__class__.__name__} triggered tripwire" ) + + +class PriceLimitExceeded(AgentsException): + """Raised when the maximum price limit is exceeded.""" + def __init__(self, current_cost: float, price_limit: float): + super().__init__(f"Maximum price limit (${price_limit:.4f}) exceeded. Current cost: ${current_cost:.4f}") diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 40899e5a..f92436b3 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -14,7 +14,7 @@ import asyncio from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload -from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost +from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost, COST_TRACKER from cai.util import start_idle_timer, stop_idle_timer, start_active_timer, stop_active_timer from wasabi import color from cai.sdk.agents.run_to_jsonl import get_session_recorder @@ -1288,8 +1288,12 @@ class OpenAIChatCompletionsModel(Model): total_cost = calculate_model_cost(model_name, total_input, total_output) # Explicit conversion to float with fallback to ensure they're never None or 0 - interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0.00001) - total_cost = max(float(total_cost if total_cost is not None else 0.0), 0.00001) + interaction_cost = float(interaction_cost if interaction_cost is not None else 0.0) + total_cost = float(total_cost if total_cost is not None else 0.0) + + # Update the global COST_TRACKER with the cost of this specific interaction + if hasattr(COST_TRACKER, "add_interaction_cost") and interaction_cost > 0.0: + COST_TRACKER.add_interaction_cost(interaction_cost) # Store the total cost for future recording self.total_cost = total_cost diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 6089d331..15d6eb70 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -58,6 +58,15 @@ if max_turns_env is not None: else: DEFAULT_MAX_TURNS = float("inf") +price_limit_env = os.getenv("CAI_PRICE_LIMIT") +if price_limit_env is not None: + try: + DEFAULT_PRICE_LIMIT = float(price_limit_env) + except ValueError: + DEFAULT_PRICE_LIMIT = float("inf") +else: + DEFAULT_PRICE_LIMIT = float("inf") + @dataclass class RunConfig: diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index e31aba92..8b9e3e46 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -1252,7 +1252,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg command, stdout, timeout, - stream=stream, + stream=True, call_id=call_id, tool_name=tool_name, workspace_dir=_get_workspace_dir(), diff --git a/src/cai/util.py b/src/cai/util.py index 46f07ea1..82081f95 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -231,16 +231,27 @@ class CostTracker: # Track the last calculation to debug inconsistencies last_interaction_cost: float = 0.0 last_total_cost: float = 0.0 - - def reset_interaction_stats(self): - """Reset stats for a new interaction""" - self.interaction_input_tokens = 0 - self.interaction_output_tokens = 0 - self.interaction_reasoning_tokens = 0 - self.interaction_cost = 0.0 + + def check_price_limit(self, new_cost: float) -> None: + """Check if adding the new cost would exceed the price limit.""" + from cai.sdk.agents.exceptions import PriceLimitExceeded + import os + price_limit_env = os.getenv("CAI_PRICE_LIMIT") + try: + price_limit = float(price_limit_env) if price_limit_env is not None else float("inf") + except ValueError: + price_limit = float("inf") + + if price_limit != float("inf"): + total_cost = self.session_total_cost + new_cost + if total_cost > price_limit: + raise PriceLimitExceeded(total_cost, price_limit) def update_session_cost(self, new_cost: float) -> None: """Add cost to session total and log the update""" + # Check price limit before updating + self.check_price_limit(new_cost) + old_total = self.session_total_cost self.session_total_cost += new_cost @@ -250,7 +261,6 @@ class CostTracker: if os.environ.get("CAI_COST_DISPLAYED", "").lower() == "true": return print(f"\nTotal CAI Session Cost: ${self.session_total_cost:.6f}") - def get_model_pricing(self, model_name: str) -> tuple: """Get and cache pricing information for a model""" # Use the centralized function to standardize model names @@ -259,8 +269,24 @@ class CostTracker: # Check cache first if model_name in self.model_pricing_cache: return self.model_pricing_cache[model_name] + + # Try to load pricing from local pricing.json first + try: + pricing_path = pathlib.Path("pricing.json") + if pricing_path.exists(): + with open(pricing_path, "r", encoding="utf-8") as f: + local_pricing = json.load(f) + pricing_info = local_pricing.get("alias0", {}) + input_cost = pricing_info.get("input_cost_per_token", 0) + output_cost = pricing_info.get("output_cost_per_token", 0) + + # Cache and return local pricing + self.model_pricing_cache[model_name] = (input_cost, output_cost) + return input_cost, output_cost + except Exception as e: + print(f" WARNING: Error loading local pricing.json: {str(e)}") - # Fetch from LiteLLM API + # Fallback to LiteLLM API if local pricing not found LITELLM_URL = ( "https://raw.githubusercontent.com/BerriAI/litellm/main/" "model_prices_and_context_window.json" @@ -283,7 +309,7 @@ class CostTracker: except Exception as e: print(f" WARNING: Error fetching model pricing: {str(e)}") - # Default values if pricing not found + # Default values if no pricing found default_pricing = (0, 0) self.model_pricing_cache[model_name] = default_pricing return default_pricing