From 01e5adb2a87e06c0fd3785fdfc638d8a472cf30a Mon Sep 17 00:00:00 2001 From: luijait Date: Wed, 21 May 2025 15:59:12 +0200 Subject: [PATCH 1/2] Pricing streaming --- .../agents/models/openai_chatcompletions.py | 69 +++++++++++++++++-- src/cai/util.py | 62 ++++++++++++++--- 2 files changed, 116 insertions(+), 15 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 95e4be9f..73ac1c76 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -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,25 @@ 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 + # Calculate cost for current interaction + 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: + 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) + + # 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': estimated_session_total + current_cost } + update_agent_streaming_content(streaming_context, content, token_stats) # More accurate token counting for text content @@ -820,6 +842,32 @@ 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 + current_estimated_cost = calculate_model_cost( + str(self.model), estimated_input_tokens, estimated_output_tokens) + + # Check price limit + if 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: + # Get the latest session total cost + 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 = ( @@ -1321,8 +1369,21 @@ class OpenAIChatCompletionsModel(Model): 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 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 diff --git a/src/cai/util.py b/src/cai/util.py index 98b95150..06c2dfa5 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -252,6 +252,20 @@ 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. + """ + # 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 log_final_cost(self) -> None: """Display final cost information at exit""" # Skip displaying cost if already shown in the session summary @@ -1518,16 +1532,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 +1567,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 +1715,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: From cde27c8ea1ab5e517b680b2e8720d01c2c87bbdf Mon Sep 17 00:00:00 2001 From: luijait Date: Wed, 21 May 2025 16:24:52 +0200 Subject: [PATCH 2/2] 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}"