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: