From cde27c8ea1ab5e517b680b2e8720d01c2c87bbdf Mon Sep 17 00:00:00 2001 From: luijait Date: Wed, 21 May 2025 16:24:52 +0200 Subject: [PATCH] local model fix --- .../agents/models/openai_chatcompletions.py | 84 +++++++++++++++---- src/cai/util.py | 65 +++++++++++++- 2 files changed, 134 insertions(+), 15 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 73ac1c76..21ffad1b 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -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) diff --git a/src/cai/util.py b/src/cai/util.py index 06c2dfa5..13e5f482 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -257,6 +257,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) @@ -266,6 +271,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 @@ -277,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] @@ -331,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}"