mirror of https://github.com/aliasrobotics/cai.git
FIX show pricing for CAI_STREAM=true
This commit is contained in:
Lidia
parent
2cf6bc01c3
commit
d4909c5288
|
|
@ -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):
|
||||
|
|
@ -358,11 +358,16 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# Create streaming context if needed
|
||||
streaming_context = None
|
||||
if should_show_rich_stream:
|
||||
print(f"\nDEBUG: Creating streaming context with initial stats:")
|
||||
print(f"DEBUG: Model: {str(self.model)}")
|
||||
print(f"DEBUG: Agent name: {self.agent_name}")
|
||||
print(f"DEBUG: Counter: {self.interaction_counter}")
|
||||
streaming_context = create_agent_streaming_context(
|
||||
agent_name=self.agent_name,
|
||||
counter=self.interaction_counter,
|
||||
model=str(self.model)
|
||||
)
|
||||
print(f"DEBUG: Created streaming context: {streaming_context}")
|
||||
|
||||
with generation_span(
|
||||
model=str(self.model),
|
||||
|
|
@ -753,49 +758,74 @@ 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)
|
||||
|
||||
print(f"DEBUG: Final direct cost calculations - Interaction: ${interaction_cost:.6f}, Total: ${total_cost:.6f}")
|
||||
|
||||
# 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),
|
||||
}
|
||||
|
||||
print(f"DEBUG: Final stats costs (from dictionary) - Interaction: ${final_stats['interaction_cost']:.6f}, Total: ${final_stats['total_cost']:.6f}")
|
||||
print(f"DEBUG: Cost types in dictionary - Interaction: {type(final_stats['interaction_cost'])}, Total: {type(final_stats['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)
|
||||
|
||||
print(f"\nDEBUG: Final stats before finish_agent_streaming:")
|
||||
print(f"DEBUG: Direct stats costs - Interaction: ${direct_stats['interaction_cost']:.6f}, Total: ${direct_stats['total_cost']:.6f}")
|
||||
print(f"DEBUG: Direct stats types - Interaction: {type(direct_stats['interaction_cost'])}, Total: {type(direct_stats['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
|
||||
|
||||
|
|
|
|||
|
|
@ -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=None, # before 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,10 +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 - calculate if None
|
||||
# Current cost - only calculate if not provided
|
||||
if interaction_cost is None:
|
||||
interaction_cost = calculate_model_cost(model, interaction_input_tokens, interaction_output_tokens)
|
||||
current_cost = float(interaction_cost) if interaction_cost is not None else 0.0
|
||||
# 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
|
||||
print(f"DEBUG _create_token_display: Current cost after conversion: {current_cost}")
|
||||
tokens_text.append(f"(${current_cost:.4f}) ", style="bold")
|
||||
|
||||
# Separator
|
||||
|
|
@ -372,10 +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 - calculate if None
|
||||
# Total cost - only calculate if not provided
|
||||
if total_cost is None:
|
||||
total_cost = calculate_model_cost(model, total_input_tokens, total_output_tokens)
|
||||
total_cost_value = float(total_cost) if total_cost is not None else 0.0
|
||||
# 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
|
||||
print(f"DEBUG _create_token_display: Total cost after conversion: {total_cost_value}")
|
||||
tokens_text.append(f"(${total_cost_value:.4f}) ", style="bold")
|
||||
|
||||
# Separator
|
||||
|
|
@ -565,14 +577,21 @@ 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")
|
||||
|
||||
# CRITICAL FIX: 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))
|
||||
|
||||
print(f"\nDEBUG finish_agent_streaming: Received costs from final_stats - Interaction: {interaction_cost}, Total: {total_cost}")
|
||||
print(f"DEBUG finish_agent_streaming: Type of interaction_cost: {type(interaction_cost)}, Type of total_cost: {type(total_cost)}")
|
||||
|
||||
if (interaction_input_tokens is not None and
|
||||
interaction_output_tokens is not None and
|
||||
|
|
@ -581,6 +600,15 @@ 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)
|
||||
|
||||
print(f"DEBUG finish_agent_streaming: Costs before passing to _create_token_display - Interaction: {interaction_cost}, Total: {total_cost}")
|
||||
print(f"DEBUG finish_agent_streaming: Type of costs before passing - Interaction: {type(interaction_cost)}, Total: {type(total_cost)}")
|
||||
|
||||
tokens_text = _create_token_display(
|
||||
interaction_input_tokens,
|
||||
interaction_output_tokens,
|
||||
|
|
|
|||
Loading…
Reference in New Issue