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