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

Pricing streaming

This commit is contained in:
luijait 2025-05-21 15:59:12 +02:00
parent 83d8d387da
commit 01e5adb2a8
2 changed files with 116 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

69
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,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

62
src/cai/util.py
View File

@ -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: