RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'stream_pricing' into '0.4.0' 

Pricing streaming

See merge request aliasrobotics/alias_research/cai!188
This commit is contained in:
Luis Javier Navarrete Lozano 2025-05-21 14:58:29 +00:00
parent 83d8d387da cde27c8ea1
commit 0368d40bcb
2 changed files with 239 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

131
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -705,6 +705,15 @@ class OpenAIChatCompletionsModel(Model):
# Get token count estimate before API call for consistent counting
estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages)
# Pre-check price limit using estimated input tokens and a conservative estimate for output
# This prevents starting a stream that would immediately exceed the price limit
if hasattr(COST_TRACKER, "check_price_limit"):
# Use a conservative estimate for output tokens (roughly equal to input)
estimated_cost = calculate_model_cost(str(self.model),
estimated_input_tokens,
estimated_input_tokens) # Conservative estimate
COST_TRACKER.check_price_limit(estimated_cost)
response, stream = await self._fetch_response(
system_instructions,
input,
@ -807,12 +816,41 @@ class OpenAIChatCompletionsModel(Model):
# Update streaming display if enabled - always do this for text content
if streaming_context:
# Create token stats to pass with each update
# Check if this is a local model and should have zero cost
model_str = str(self.model).lower()
is_local_model = (
"ollama" in model_str or
"qwen" in model_str or
"llama" in model_str or
"mistral" in model_str or
":" in model_str or
self.is_ollama
)
# Calculate cost for current interaction (will be 0 for local models)
current_cost = calculate_model_cost(str(self.model), estimated_input_tokens, estimated_output_tokens)
# Check price limit only for non-local models
if not is_local_model and hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0:
COST_TRACKER.check_price_limit(current_cost)
# Update session total cost for real-time display
# This is a temporary estimate during streaming that will be properly updated at the end
estimated_session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0)
# For local models, don't add to the total cost
display_total_cost = estimated_session_total
if not is_local_model:
display_total_cost += current_cost
# Create token stats with both current interaction cost and updated total cost
token_stats = {
'input_tokens': estimated_input_tokens,
'output_tokens': estimated_output_tokens,
'cost': calculate_model_cost(str(self.model), estimated_input_tokens, estimated_output_tokens)
'cost': current_cost,
'total_cost': display_total_cost
}
update_agent_streaming_content(streaming_context, content, token_stats)
# More accurate token counting for text content
@ -820,6 +858,51 @@ class OpenAIChatCompletionsModel(Model):
token_count, _ = count_tokens_with_tiktoken(output_text)
estimated_output_tokens = token_count
# Periodically check price limit during streaming
# This allows early termination if price limit is reached mid-stream
if estimated_output_tokens > 0 and estimated_output_tokens % 50 == 0: # Check every ~50 tokens
# Check if this is a local model (Ollama, Qwen, etc.) that should have zero cost
model_str = str(self.model).lower()
is_local_model = (
"ollama" in model_str or
"qwen" in model_str or
"llama" in model_str or
"mistral" in model_str or
":" in model_str or
self.is_ollama
)
# For local models, cost should always be zero
if is_local_model:
current_estimated_cost = 0.0
else:
current_estimated_cost = calculate_model_cost(
str(self.model), estimated_input_tokens, estimated_output_tokens)
# Check price limit only for non-local models
if not is_local_model and hasattr(COST_TRACKER, "check_price_limit"):
COST_TRACKER.check_price_limit(current_estimated_cost)
# Update the COST_TRACKER with the running cost for accurate display
if hasattr(COST_TRACKER, "interaction_cost"):
COST_TRACKER.interaction_cost = current_estimated_cost
# Also update streaming context if available for live display
if streaming_context:
# For local models, don't add to the session total
if is_local_model:
session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0)
else:
session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + current_estimated_cost
updated_token_stats = {
'input_tokens': estimated_input_tokens,
'output_tokens': estimated_output_tokens,
'cost': current_estimated_cost,
'total_cost': session_total
}
update_agent_streaming_content(streaming_context, "", updated_token_stats)
if not state.text_content_index_and_output:
# Initialize a content tracker for streaming text
state.text_content_index_and_output = (
@ -1311,18 +1394,52 @@ class OpenAIChatCompletionsModel(Model):
total_input = getattr(self, 'total_input_tokens', 0)
total_output = getattr(self, 'total_output_tokens', 0)
# Calculate costs using the same token counts - ensure model is a string
# Check if this is a local model and should have zero cost
model_str = str(self.model).lower()
is_local_model = (
"ollama" in model_str or
"qwen" in model_str or
"llama" in model_str or
"mistral" in model_str or
":" in model_str or
self.is_ollama
)
# Calculate costs - use zero for local models
model_name = str(self.model)
interaction_cost = calculate_model_cost(model_name, interaction_input, interaction_output)
total_cost = calculate_model_cost(model_name, total_input, total_output)
if is_local_model:
# For local models, cost is always zero
interaction_cost = 0.0
total_cost = getattr(COST_TRACKER, 'session_total_cost', 0.0) # Keep existing total
# Ensure the cost tracking system knows this is a free model
if hasattr(COST_TRACKER, "reset_cost_for_local_model"):
COST_TRACKER.reset_cost_for_local_model(model_name)
else:
# For paid models, calculate as normal
interaction_cost = calculate_model_cost(model_name, interaction_input, interaction_output)
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 = 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)
# and check price limit for streaming mode (similar to non-streaming mode)
if not is_local_model and interaction_cost > 0.0:
# Check price limit before adding the new cost
if hasattr(COST_TRACKER, "check_price_limit"):
COST_TRACKER.check_price_limit(interaction_cost)
# Now add the cost to session total
if hasattr(COST_TRACKER, "update_session_cost"):
COST_TRACKER.update_session_cost(interaction_cost)
elif hasattr(COST_TRACKER, "add_interaction_cost"):
COST_TRACKER.add_interaction_cost(interaction_cost)
# Ensure the total cost includes the session total for display
if hasattr(COST_TRACKER, "session_total_cost"):
total_cost = COST_TRACKER.session_total_cost
# Store the total cost for future recording
self.total_cost = total_cost

127
src/cai/util.py
View File

@ -252,6 +252,50 @@ class CostTracker:
old_total = self.session_total_cost
self.session_total_cost += new_cost
def add_interaction_cost(self, new_cost: float) -> None:
"""
Add an interaction cost to the session total and check price limit.
This is a convenience method that combines check_price_limit and update_session_cost.
"""
# Skip updating costs if the cost is zero (common with local models)
if new_cost <= 0:
self.last_interaction_cost = 0.0
return
# Check price limit first
self.check_price_limit(new_cost)
# Then update the session cost
self.session_total_cost += new_cost
# Update the last interaction cost for tracking
self.last_interaction_cost = new_cost
def reset_cost_for_local_model(self, model_name: str) -> bool:
"""
Reset interaction cost tracking when switching to a local model.
Returns True if the model was identified as local and cost was reset.
"""
# Check if this is a local/free model
model_str = model_name.lower()
is_local_model = (
"ollama" in model_str or
"qwen" in model_str or
"llama" in model_str or
"mistral" in model_str or
":" in model_str or
(os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false')
)
if is_local_model:
# Reset the current interaction costs but keep total session costs
self.interaction_cost = 0.0
self.last_interaction_cost = 0.0
# Don't reset session_total_cost as that includes previous paid models
return True
return False
def log_final_cost(self) -> None:
"""Display final cost information at exit"""
# Skip displaying cost if already shown in the session summary
@ -263,7 +307,25 @@ class CostTracker:
# Use the centralized function to standardize model names
model_name = get_model_name(model_name)
# Check cache first
# Check if using Ollama or local model
model_str = model_name.lower()
is_local_model = (
"ollama" in model_str or
"qwen" in model_str or
"llama" in model_str or
"mistral" in model_str or
":" in model_str or # Ollama uses formats like qwen2.5:7b
(os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false')
)
# For local models, always return zero cost
if is_local_model:
# Set and cache zero cost for local models
free_pricing = (0.0, 0.0)
self.model_pricing_cache[model_name] = free_pricing
return free_pricing
# Check cache for non-local models
if model_name in self.model_pricing_cache:
return self.model_pricing_cache[model_name]
@ -317,6 +379,21 @@ class CostTracker:
# Standardize model name using the central function
model_name = get_model_name(model)
# Check if this is a local model (always free) first
model_str = model_name.lower()
is_local_model = (
"ollama" in model_str or
"qwen" in model_str or
"llama" in model_str or
"mistral" in model_str or
":" in model_str or
(os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false')
)
# For local models, always return zero cost
if is_local_model:
return 0.0
# Generate a cache key
cache_key = f"{model_name}_{input_tokens}_{output_tokens}"
@ -1518,16 +1595,26 @@ def update_agent_streaming_content(context, text_delta, token_stats=None):
return False
try:
# Parse the text_delta to get just the content if needed
parsed_delta = parse_message_content(text_delta)
# Skip empty updates to avoid showing an empty panel
if not parsed_delta or parsed_delta.strip() == "":
# Only parse and add text if we have actual content to add
# Skip when text_delta is empty and we're just updating token stats
if text_delta:
# Parse the text_delta to get just the content if needed
parsed_delta = parse_message_content(text_delta)
# Skip empty updates to avoid showing an empty panel
if not parsed_delta or parsed_delta.strip() == "":
# Update token stats if provided
if token_stats:
# Just update the footer, not the content
pass
else:
# Add the parsed text to the content
context["content"].append(parsed_delta)
# If no text_delta but we have token_stats, just update stats
elif not token_stats:
# No text and no stats - nothing to update
return True
# Add the parsed text to the content
context["content"].append(parsed_delta)
# Update the footer with token stats if provided
if token_stats:
# Create token stats display
@ -1543,13 +1630,24 @@ def update_agent_streaming_content(context, text_delta, token_stats=None):
# Add token stats
input_tokens = token_stats.get('input_tokens', 0)
output_tokens = token_stats.get('output_tokens', 0)
total_cost = token_stats.get('cost', 0.0)
interaction_cost = token_stats.get('cost', 0.0)
# Get session total cost - either from token_stats or directly from COST_TRACKER
session_total_cost = token_stats.get('total_cost', 0.0)
if session_total_cost == 0.0 and hasattr(COST_TRACKER, 'session_total_cost'):
session_total_cost = COST_TRACKER.session_total_cost
if input_tokens > 0:
footer_stats.append(" | ", style="dim")
footer_stats.append(f"I:{input_tokens} O:{output_tokens}", style="green")
if total_cost > 0:
footer_stats.append(f" (${total_cost:.4f})", style="bold cyan")
# Show both interaction cost and total session cost
if interaction_cost > 0:
footer_stats.append(f" (${interaction_cost:.4f})", style="bold cyan")
# Add the total cost information on the same line
footer_stats.append(" | Session: ", style="dim")
footer_stats.append(f"${session_total_cost:.4f}", style="bold magenta")
# Add context usage indicator
model_name = context.get("model", os.environ.get('CAI_MODEL', 'qwen2.5:14b'))
@ -1680,6 +1778,11 @@ def finish_agent_streaming(context, final_stats=None):
compact_tokens.append(f"I:{interaction_input_tokens} O:{interaction_output_tokens} ", style="green")
compact_tokens.append(f"(${interaction_cost:.4f}) ", style="bold cyan")
# Include the total session cost
session_total_cost = COST_TRACKER.session_total_cost if hasattr(COST_TRACKER, 'session_total_cost') else total_cost
compact_tokens.append(" | Session: ", style="dim")
compact_tokens.append(f"${session_total_cost:.4f}", style="bold magenta")
# Añadir un indicador de uso de contexto
context_pct = interaction_input_tokens / get_model_input_tokens(model_name) * 100
if context_pct < 50: