From 23d2a0f22dc875c777f9eb50b62812e40a9f3e16 Mon Sep 17 00:00:00 2001 From: Lidia Date: Wed, 9 Apr 2025 15:48:59 +0200 Subject: [PATCH] FIX show pricing - calculate pricing --- src/cai/cli.py | 6 +-- src/cai/repl/commands/model.py | 5 +- src/cai/sdk/agents/models/openai_responses.py | 9 ++-- src/cai/util.py | 52 +++++++++++++++++-- 4 files changed, 60 insertions(+), 12 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index 3eb1fafd..677914db 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -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) diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py index 31cd8e9e..94baea75 100644 --- a/src/cai/repl/commands/model.py +++ b/src/cai/repl/commands/model.py @@ -201,7 +201,10 @@ class ModelCommand(Command): console.print( "[yellow]Warning: Could not fetch model pricing data[/yellow]" ) - + print("--------------------------------") + print(LITELLM_URL) + print(model_pricing_data) + print("--------------------------------") # Create a flat list of all models for numeric selection # pylint: disable=invalid-name ALL_MODELS = [] diff --git a/src/cai/sdk/agents/models/openai_responses.py b/src/cai/sdk/agents/models/openai_responses.py index ce587045..b25df689 100644 --- a/src/cai/sdk/agents/models/openai_responses.py +++ b/src/cai/sdk/agents/models/openai_responses.py @@ -28,6 +28,7 @@ from ..tracing import SpanError, response_span from ..usage import Usage from ..version import __version__ from .interface import Model, ModelTracing +from cai.util import calculate_model_cost if TYPE_CHECKING: from ..model_settings import ModelSettings @@ -153,8 +154,8 @@ class OpenAIResponsesModel(Model): 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_cost=calculate_model_cost(str(self.model), usage.input_tokens, usage.output_tokens), + total_cost=calculate_model_cost(str(self.model), getattr(self, 'total_input_tokens', 0), getattr(self, 'total_output_tokens', 0)), ) # Update token totals @@ -254,8 +255,8 @@ class OpenAIResponsesModel(Model): 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_cost=calculate_model_cost(str(self.model), final_response.usage.input_tokens, final_response.usage.output_tokens), + total_cost=calculate_model_cost(str(self.model), getattr(self, 'total_input_tokens', 0), getattr(self, 'total_output_tokens', 0)), ) # Update token totals diff --git a/src/cai/util.py b/src/cai/util.py index 1cdb95ed..33375956 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -339,7 +339,7 @@ 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, # before 0.0 total_cost=None ) -> Text: # noqa: E501 """ @@ -357,7 +357,9 @@ 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 - calculate if None + 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 tokens_text.append(f"(${current_cost:.4f}) ", style="bold") @@ -370,7 +372,9 @@ 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 - calculate if None + 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 tokens_text.append(f"(${total_cost_value:.4f}) ", style="bold") @@ -615,4 +619,44 @@ def finish_agent_streaming(context, final_stats=None): time.sleep(0.5) # Stop the live display - context["live"].stop() \ No newline at end of file + 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 \ No newline at end of file