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

Merge branch 'v0.4.0_pricing' into 'v0.4.0' 

Fix pricing for CAI_STREAM=true and CAI_STREAM=false

See merge request aliasrobotics/alias_research/cai!140
This commit is contained in:
Víctor Mayoral Vilches 2025-04-22 07:38:48 +00:00
parent fa198de051 0bced4cbf3
commit 76cdba8236
4 changed files with 118 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/cai/cli.py
View File

@ -107,7 +107,7 @@ from cai.sdk.agents import set_default_openai_client, set_tracing_disabled
from openai.types.responses import ResponseTextDeltaEvent
from rich.console import Console
import asyncio
from cai.util import fix_litellm_transcription_annotations, color
from cai.util import fix_litellm_transcription_annotations, color, calculate_model_cost
from cai.util import create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming
# Import modules from cai.repl
@ -322,8 +322,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
"total_input_tokens": getattr(model, 'total_input_tokens', interaction_input),
"total_output_tokens": getattr(model, 'total_output_tokens', output_tokens),
"total_reasoning_tokens": getattr(model, 'total_reasoning_tokens', 0),
"interaction_cost": None,
"total_cost": None
"interaction_cost": calculate_model_cost(str(model), interaction_input, output_tokens),
"total_cost": calculate_model_cost(str(model), getattr(model, 'total_input_tokens', interaction_input), getattr(model, 'total_output_tokens', output_tokens))
}
finish_agent_streaming(current_streaming_context, token_stats)

1
src/cai/repl/commands/model.py
View File

@ -201,7 +201,6 @@ class ModelCommand(Command):
console.print(
"[yellow]Warning: Could not fetch model pricing data[/yellow]"
)
# Create a flat list of all models for numeric selection
# pylint: disable=invalid-name
ALL_MODELS = []

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

@ -61,7 +61,7 @@ from openai.types.responses import (
)
from openai.types.responses.response_input_param import FunctionCallOutput, ItemReference, Message
from openai.types.responses.response_usage import OutputTokensDetails
from cai.util import calculate_model_cost
# Create custom InputTokensDetails class since it's not available in current OpenAI version
from openai._models import BaseModel
class InputTokensDetails(BaseModel):
@ -753,49 +753,65 @@ class OpenAIChatCompletionsModel(Model):
self.total_reasoning_tokens += final_response.usage.output_tokens_details.reasoning_tokens
# Prepare final statistics for display
interaction_input = final_response.usage.input_tokens if final_response.usage else 0
interaction_output = final_response.usage.output_tokens if final_response.usage else 0
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
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)
# Explicit conversion to float with fallback to ensure they're never None or 0
interaction_cost = max(float(interaction_cost if interaction_cost is not None else 0.0), 0.00001)
total_cost = max(float(total_cost if total_cost is not None else 0.0), 0.00001)
# Create final stats with explicit type conversion for all values
final_stats = {
"interaction_input_tokens": final_response.usage.input_tokens if final_response.usage else 0,
"interaction_output_tokens": final_response.usage.output_tokens if final_response.usage else 0,
"interaction_reasoning_tokens": (
"interaction_input_tokens": int(interaction_input),
"interaction_output_tokens": int(interaction_output),
"interaction_reasoning_tokens": int(
final_response.usage.output_tokens_details.reasoning_tokens
if final_response.usage and final_response.usage.output_tokens_details
and hasattr(final_response.usage.output_tokens_details, 'reasoning_tokens')
else 0
),
"total_input_tokens": getattr(self, 'total_input_tokens', 0),
"total_output_tokens": getattr(self, 'total_output_tokens', 0),
"total_reasoning_tokens": getattr(self, 'total_reasoning_tokens', 0),
"interaction_cost": None,
"total_cost": None,
"total_input_tokens": int(total_input),
"total_output_tokens": int(total_output),
"total_reasoning_tokens": int(getattr(self, 'total_reasoning_tokens', 0)),
"interaction_cost": float(interaction_cost),
"total_cost": float(total_cost),
}
# At the end of streaming, finish the streaming context if we were using it
if streaming_context:
finish_agent_streaming(streaming_context, final_stats)
# Create a direct copy of the costs to ensure they remain as floats
direct_stats = final_stats.copy()
direct_stats["interaction_cost"] = float(interaction_cost)
direct_stats["total_cost"] = float(total_cost)
# Use the direct copy with guaranteed float costs
finish_agent_streaming(streaming_context, direct_stats)
# If we're not using rich streaming and not suppressing output, use old method
elif not self.suppress_final_output and final_response.output and any(isinstance(item, ResponseOutputMessage) for item in final_response.output):
# Find the assistant message to print
for item in final_response.output:
if isinstance(item, ResponseOutputMessage) and item.role == 'assistant':
cli_print_agent_messages(
agent_name=getattr(self, 'agent_name', 'Agent'), # Default to 'Agent' if not available
agent_name=getattr(self, 'agent_name', 'Agent'),
message=item,
counter=getattr(self, 'interaction_counter', 0), # Default to 0 if not available
counter=getattr(self, 'interaction_counter', 0),
model=str(self.model),
debug=False,
interaction_input_tokens=final_response.usage.input_tokens if final_response.usage else 0,
interaction_output_tokens=final_response.usage.output_tokens if final_response.usage else 0,
interaction_reasoning_tokens=(
final_response.usage.output_tokens_details.reasoning_tokens
if final_response.usage and final_response.usage.output_tokens_details
and hasattr(final_response.usage.output_tokens_details, 'reasoning_tokens')
else 0
),
total_input_tokens=getattr(self, 'total_input_tokens', 0),
total_output_tokens=getattr(self, 'total_output_tokens', 0),
total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0),
interaction_cost=None,
total_cost=None,
interaction_input_tokens=interaction_input,
interaction_output_tokens=interaction_output,
interaction_reasoning_tokens=final_stats["interaction_reasoning_tokens"],
total_input_tokens=total_input,
total_output_tokens=total_output,
total_reasoning_tokens=final_stats["total_reasoning_tokens"],
interaction_cost=interaction_cost,
total_cost=total_cost,
)
break

82
src/cai/util.py
View File

@ -339,13 +339,15 @@ def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals
total_output_tokens,
total_reasoning_tokens,
model,
interaction_cost=0.0,
interaction_cost=None,
total_cost=None
) -> Text: # noqa: E501
"""
Create a Text object displaying token usage information
with enhanced formatting.
"""
# print(f"\nDEBUG _create_token_display: Received costs - Interaction: {interaction_cost}, Total: {total_cost}")
tokens_text = Text(justify="left")
# Create a more compact, horizontal display
@ -357,8 +359,15 @@ def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals
tokens_text.append(f"O:{interaction_output_tokens} ", style="red")
tokens_text.append(f"R:{interaction_reasoning_tokens} ", style="yellow")
# Current cost
current_cost = float(interaction_cost) if interaction_cost is not None else 0.0
# Current cost - only calculate if not provided
if interaction_cost is None:
interaction_cost = calculate_model_cost(model, interaction_input_tokens, interaction_output_tokens)
# Ensure interaction_cost is a float
try:
current_cost = float(interaction_cost) if interaction_cost is not None else 0.0
except (ValueError, TypeError):
current_cost = 0.0
tokens_text.append(f"(${current_cost:.4f}) ", style="bold")
# Separator
@ -370,8 +379,15 @@ def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals
tokens_text.append(f"O:{total_output_tokens} ", style="red")
tokens_text.append(f"R:{total_reasoning_tokens} ", style="yellow")
# Total cost
total_cost_value = float(total_cost) if total_cost is not None else 0.0
# Total cost - only calculate if not provided
if total_cost is None:
total_cost = calculate_model_cost(model, total_input_tokens, total_output_tokens)
# Ensure total_cost is a float
try:
total_cost_value = float(total_cost) if total_cost is not None else 0.0
except (ValueError, TypeError):
total_cost_value = 0.0
tokens_text.append(f"(${total_cost_value:.4f}) ", style="bold")
# Separator
@ -561,14 +577,18 @@ def finish_agent_streaming(context, final_stats=None):
# If we have token stats, add them
tokens_text = None
if final_stats:
#print(f"\nDEBUG finish_agent_streaming: Received final_stats: {final_stats}")
interaction_input_tokens = final_stats.get("interaction_input_tokens")
interaction_output_tokens = final_stats.get("interaction_output_tokens")
interaction_reasoning_tokens = final_stats.get("interaction_reasoning_tokens")
total_input_tokens = final_stats.get("total_input_tokens")
total_output_tokens = final_stats.get("total_output_tokens")
total_reasoning_tokens = final_stats.get("total_reasoning_tokens")
interaction_cost = final_stats.get("interaction_cost")
total_cost = final_stats.get("total_cost")
# Ensure costs are properly extracted and preserved as floats
interaction_cost = float(final_stats.get("interaction_cost", 0.0))
total_cost = float(final_stats.get("total_cost", 0.0))
if (interaction_input_tokens is not None and
interaction_output_tokens is not None and
@ -577,6 +597,12 @@ def finish_agent_streaming(context, final_stats=None):
total_output_tokens is not None and
total_reasoning_tokens is not None):
# Only calculate costs if they weren't provided or are zero
if interaction_cost is None or interaction_cost == 0.0:
interaction_cost = calculate_model_cost(context["model"], interaction_input_tokens, interaction_output_tokens)
if total_cost is None or total_cost == 0.0:
total_cost = calculate_model_cost(context["model"], total_input_tokens, total_output_tokens)
tokens_text = _create_token_display(
interaction_input_tokens,
interaction_output_tokens,
@ -615,4 +641,44 @@ def finish_agent_streaming(context, final_stats=None):
time.sleep(0.5)
# Stop the live display
context["live"].stop()
context["live"].stop()
def calculate_model_cost(model_name, input_tokens, output_tokens):
"""
Calculate the cost for a given model based on token usage.
Args:
model_name: The name of the model being used
input_tokens: Number of input tokens used
output_tokens: Number of output tokens used
Returns:
float: The calculated cost in dollars
"""
# Fetch model pricing data from LiteLLM GitHub repository
LITELLM_URL = (
"https://raw.githubusercontent.com/BerriAI/litellm/main/"
"model_prices_and_context_window.json"
)
try:
import requests
response = requests.get(LITELLM_URL, timeout=2)
if response.status_code == 200:
model_pricing_data = response.json()
# Get pricing info for the model
pricing_info = model_pricing_data.get(model_name, {})
input_cost_per_token = pricing_info.get("input_cost_per_token", 0)
output_cost_per_token = pricing_info.get("output_cost_per_token", 0)
# Calculate costs
input_cost = input_tokens * input_cost_per_token
output_cost = output_tokens * output_cost_per_token
return input_cost + output_cost
except Exception:
# If we can't fetch pricing data, return 0
pass
return 0.0