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

local model fix

This commit is contained in:
luijait 2025-05-21 16:24:52 +02:00
parent c0cd502579
commit ec8179083d
2 changed files with 134 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -816,23 +816,39 @@ class OpenAIChatCompletionsModel(Model):
# Update streaming display if enabled - always do this for text content
if streaming_context:
# Calculate cost for current interaction
# 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 before updating (conservative check)
if hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0:
# 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': current_cost,
'total_cost': estimated_session_total + current_cost
'total_cost': display_total_cost
}
update_agent_streaming_content(streaming_context, content, token_stats)
@ -845,11 +861,26 @@ class OpenAIChatCompletionsModel(Model):
# 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
current_estimated_cost = calculate_model_cost(
str(self.model), estimated_input_tokens, estimated_output_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
)
# Check price limit
if hasattr(COST_TRACKER, "check_price_limit"):
# 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
@ -858,8 +889,12 @@ class OpenAIChatCompletionsModel(Model):
# Also update streaming context if available for live display
if streaming_context:
# Get the latest session total cost
session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + current_estimated_cost
# 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,
@ -1359,10 +1394,31 @@ 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)
@ -1370,7 +1426,7 @@ class OpenAIChatCompletionsModel(Model):
# Update the global COST_TRACKER with the cost of this specific interaction
# and check price limit for streaming mode (similar to non-streaming mode)
if interaction_cost > 0.0:
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)

65
src/cai/util.py
View File

@ -261,6 +261,11 @@ class CostTracker:
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)
@ -270,6 +275,31 @@ class CostTracker:
# 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
@ -281,7 +311,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]
@ -335,6 +383,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}"