mirror of https://github.com/aliasrobotics/cai.git
Pricing
This commit is contained in:
parent
4aea8eda74
commit
699017af02
|
|
@ -245,6 +245,10 @@
|
|||
variables:
|
||||
TEST_PATH: tests/others/test_computer_action.py
|
||||
|
||||
▪️ others test_pricing:
|
||||
<<: *run_test
|
||||
variables:
|
||||
TEST_PATH: tests/test_pricing.py
|
||||
▪️ others test_pretty_print.py:
|
||||
<<: *run_test
|
||||
variables:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,46 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
import requests
|
||||
import json
|
||||
|
||||
def debug_litellm_pricing():
|
||||
"""Debug why LiteLLM returns pricing for local models"""
|
||||
|
||||
LITELLM_URL = 'https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json'
|
||||
|
||||
try:
|
||||
response = requests.get(LITELLM_URL, timeout=5)
|
||||
data = response.json()
|
||||
|
||||
print('Checking if qwen3:14b exists in LiteLLM pricing...')
|
||||
if 'qwen3:14b' in data:
|
||||
print('Found qwen3:14b in LiteLLM:', data['qwen3:14b'])
|
||||
else:
|
||||
print('qwen3:14b NOT found in LiteLLM pricing')
|
||||
|
||||
print('\nChecking similar qwen models...')
|
||||
qwen_models = [k for k in data.keys() if 'qwen' in k.lower()]
|
||||
print(f'Total Qwen models found in LiteLLM: {len(qwen_models)}')
|
||||
print('First 10 Qwen models:')
|
||||
for model in sorted(qwen_models)[:10]:
|
||||
print(f' {model}: {data[model]}')
|
||||
|
||||
# Check if there's a pattern match or fallback
|
||||
print('\nChecking for pattern matches...')
|
||||
potential_matches = [k for k in data.keys() if 'qwen' in k.lower() and '14b' in k.lower()]
|
||||
if potential_matches:
|
||||
print('Models with qwen and 14b:')
|
||||
for model in potential_matches:
|
||||
print(f' {model}: {data[model]}')
|
||||
|
||||
# Test our current logic
|
||||
print('\nTesting our get_model_pricing logic...')
|
||||
from src.cai.util import COST_TRACKER
|
||||
pricing = COST_TRACKER.get_model_pricing('qwen3:14b')
|
||||
print(f'Our function returns for qwen3:14b: {pricing}')
|
||||
|
||||
except Exception as e:
|
||||
print(f'Error: {e}')
|
||||
|
||||
if __name__ == '__main__':
|
||||
debug_litellm_pricing()
|
||||
|
|
@ -0,0 +1,76 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
import pathlib
|
||||
import json
|
||||
import requests
|
||||
from src.cai.util import COST_TRACKER, get_model_name
|
||||
|
||||
def debug_pricing_flow():
|
||||
"""Debug step by step the get_model_pricing flow"""
|
||||
|
||||
model_name = 'qwen3:14b'
|
||||
print(f"Debugging pricing flow for: {model_name}")
|
||||
|
||||
# Step 1: Standardize model name
|
||||
standardized = get_model_name(model_name)
|
||||
print(f"1. Standardized model name: {standardized}")
|
||||
|
||||
# Step 2: Check cache
|
||||
if standardized in COST_TRACKER.model_pricing_cache:
|
||||
cached = COST_TRACKER.model_pricing_cache[standardized]
|
||||
print(f"2. Found in cache: {cached}")
|
||||
return cached
|
||||
else:
|
||||
print("2. Not found in cache")
|
||||
|
||||
# Step 3: Check local pricing.json
|
||||
print("3. Checking local pricing.json...")
|
||||
try:
|
||||
pricing_path = pathlib.Path("pricing.json")
|
||||
if pricing_path.exists():
|
||||
print(" pricing.json exists")
|
||||
with open(pricing_path, "r", encoding="utf-8") as f:
|
||||
local_pricing = json.load(f)
|
||||
print(f" Content: {local_pricing}")
|
||||
pricing_info = local_pricing.get("alias0", {})
|
||||
input_cost = pricing_info.get("input_cost_per_token", 0)
|
||||
output_cost = pricing_info.get("output_cost_per_token", 0)
|
||||
print(f" Extracted pricing: input={input_cost}, output={output_cost}")
|
||||
if input_cost or output_cost:
|
||||
print(f" Would return: ({input_cost}, {output_cost})")
|
||||
return (input_cost, output_cost)
|
||||
else:
|
||||
print(" pricing.json does not exist")
|
||||
except Exception as e:
|
||||
print(f" Error reading pricing.json: {e}")
|
||||
|
||||
# Step 4: Check LiteLLM API
|
||||
print("4. Checking LiteLLM API...")
|
||||
LITELLM_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
|
||||
try:
|
||||
response = requests.get(LITELLM_URL, timeout=2)
|
||||
if response.status_code == 200:
|
||||
model_pricing_data = response.json()
|
||||
pricing_info = model_pricing_data.get(standardized, {})
|
||||
input_cost_per_token = pricing_info.get("input_cost_per_token", 0)
|
||||
output_cost_per_token = pricing_info.get("output_cost_per_token", 0)
|
||||
print(f" LiteLLM response for {standardized}: {pricing_info}")
|
||||
print(f" Extracted: input={input_cost_per_token}, output={output_cost_per_token}")
|
||||
if input_cost_per_token or output_cost_per_token:
|
||||
print(f" Would return: ({input_cost_per_token}, {output_cost_per_token})")
|
||||
return (input_cost_per_token, output_cost_per_token)
|
||||
else:
|
||||
print(f" LiteLLM API returned status: {response.status_code}")
|
||||
except Exception as e:
|
||||
print(f" Error fetching from LiteLLM: {e}")
|
||||
|
||||
# Step 5: Default fallback
|
||||
print("5. Using default fallback: (0.0, 0.0)")
|
||||
return (0.0, 0.0)
|
||||
|
||||
if __name__ == '__main__':
|
||||
result = debug_pricing_flow()
|
||||
print(f"\nFinal result: {result}")
|
||||
|
||||
# Now test the actual function
|
||||
print(f"\nActual function result: {COST_TRACKER.get_model_pricing('qwen3:14b')}")
|
||||
|
|
@ -1066,23 +1066,11 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# Update streaming display if enabled - ALWAYS respect CAI_STREAM setting
|
||||
# Both thinking and regular content should stream if streaming is enabled
|
||||
if streaming_context:
|
||||
# Check if this is a local model and should have zero cost
|
||||
model_str = str(self.model).lower()
|
||||
is_local_model = (
|
||||
"alias" not in model_str and
|
||||
"ollama" in model_str or
|
||||
"qwen" in model_str or
|
||||
"llama" in model_str or
|
||||
"mistral" in model_str or
|
||||
":" in model_str or
|
||||
self.is_ollama
|
||||
)
|
||||
|
||||
# Calculate cost for current interaction (will be 0 for local models)
|
||||
# Calculate cost for current interaction
|
||||
current_cost = calculate_model_cost(str(self.model), estimated_input_tokens, estimated_output_tokens)
|
||||
|
||||
# Check price limit only for non-local models
|
||||
if not is_local_model and hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0:
|
||||
# Check price limit only for paid models
|
||||
if current_cost > 0 and hasattr(COST_TRACKER, "check_price_limit") and estimated_output_tokens % 50 == 0:
|
||||
try:
|
||||
COST_TRACKER.check_price_limit(current_cost)
|
||||
except Exception as e:
|
||||
|
|
@ -1101,9 +1089,9 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# 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)
|
||||
|
||||
# For local models, don't add to the total cost
|
||||
# For free models, don't add to the total cost
|
||||
display_total_cost = estimated_session_total
|
||||
if not is_local_model:
|
||||
if current_cost > 0:
|
||||
display_total_cost += current_cost
|
||||
|
||||
# Create token stats with both current interaction cost and updated total cost
|
||||
|
|
@ -1124,27 +1112,12 @@ class OpenAIChatCompletionsModel(Model):
|
|||
# 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
|
||||
# Check if this is a local model (Ollama, Qwen, etc.) that should have zero cost
|
||||
model_str = str(self.model).lower()
|
||||
is_local_model = (
|
||||
"alias" not in model_str and
|
||||
"ollama" in model_str or
|
||||
"qwen" in model_str or
|
||||
"llama" in model_str or
|
||||
"mistral" in model_str or
|
||||
":" in model_str or
|
||||
self.is_ollama
|
||||
)
|
||||
# Calculate current estimated cost
|
||||
current_estimated_cost = calculate_model_cost(
|
||||
str(self.model), estimated_input_tokens, estimated_output_tokens)
|
||||
|
||||
# For local models, cost should always be zero
|
||||
if is_local_model:
|
||||
current_estimated_cost = 0.0
|
||||
else:
|
||||
current_estimated_cost = calculate_model_cost(
|
||||
str(self.model), estimated_input_tokens, estimated_output_tokens)
|
||||
|
||||
# Check price limit only for non-local models
|
||||
if not is_local_model and hasattr(COST_TRACKER, "check_price_limit"):
|
||||
# Check price limit only for paid models
|
||||
if current_estimated_cost > 0 and hasattr(COST_TRACKER, "check_price_limit"):
|
||||
try:
|
||||
COST_TRACKER.check_price_limit(current_estimated_cost)
|
||||
except Exception as e:
|
||||
|
|
@ -1165,8 +1138,8 @@ class OpenAIChatCompletionsModel(Model):
|
|||
|
||||
# Also update streaming context if available for live display
|
||||
if streaming_context:
|
||||
# For local models, don't add to the session total
|
||||
if is_local_model:
|
||||
# For free models, don't add to the session total
|
||||
if current_estimated_cost == 0:
|
||||
session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0)
|
||||
else:
|
||||
session_total = getattr(COST_TRACKER, 'session_total_cost', 0.0) + current_estimated_cost
|
||||
|
|
@ -1678,32 +1651,17 @@ class OpenAIChatCompletionsModel(Model):
|
|||
total_input = getattr(self, 'total_input_tokens', 0)
|
||||
total_output = getattr(self, 'total_output_tokens', 0)
|
||||
|
||||
# Check if this is a local model and should have zero cost
|
||||
model_str = str(self.model).lower()
|
||||
is_local_model = (
|
||||
"alias" not in model_str and
|
||||
"ollama" in model_str or
|
||||
"qwen" in model_str or
|
||||
"llama" in model_str or
|
||||
"mistral" in model_str or
|
||||
":" in model_str or
|
||||
self.is_ollama
|
||||
)
|
||||
|
||||
# Calculate costs - use zero for local models
|
||||
# Calculate costs for this model
|
||||
model_name = str(self.model)
|
||||
if is_local_model:
|
||||
# For local models, cost is always zero
|
||||
interaction_cost = 0.0
|
||||
total_cost = getattr(COST_TRACKER, 'session_total_cost', 0.0) # Keep existing total
|
||||
|
||||
# Ensure the cost tracking system knows this is a free model
|
||||
interaction_cost = calculate_model_cost(model_name, interaction_input, interaction_output)
|
||||
total_cost = calculate_model_cost(model_name, total_input, total_output)
|
||||
|
||||
# If interaction cost is zero, this is a free model
|
||||
if interaction_cost == 0:
|
||||
# For free models, keep existing total and ensure cost tracking system knows it's free
|
||||
total_cost = getattr(COST_TRACKER, 'session_total_cost', 0.0)
|
||||
if hasattr(COST_TRACKER, "reset_cost_for_local_model"):
|
||||
COST_TRACKER.reset_cost_for_local_model(model_name)
|
||||
else:
|
||||
# For paid models, calculate as normal
|
||||
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 = float(interaction_cost if interaction_cost is not None else 0.0)
|
||||
|
|
@ -1711,7 +1669,7 @@ class OpenAIChatCompletionsModel(Model):
|
|||
|
||||
# Update the global COST_TRACKER with the cost of this specific interaction
|
||||
# and check price limit for streaming mode (similar to non-streaming mode)
|
||||
if not is_local_model and interaction_cost > 0.0:
|
||||
if interaction_cost > 0.0:
|
||||
# Check price limit before adding the new cost
|
||||
if hasattr(COST_TRACKER, "check_price_limit"):
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -84,7 +84,6 @@ def stop_active_timer():
|
|||
# Start idle timer if not already running
|
||||
if _idle_timer_start is None:
|
||||
_idle_timer_start = time.time()
|
||||
|
||||
def start_idle_timer():
|
||||
"""
|
||||
Start measuring idle time (when waiting for user input).
|
||||
|
|
@ -285,19 +284,11 @@ class CostTracker:
|
|||
Reset interaction cost tracking when switching to a local model.
|
||||
Returns True if the model was identified as local and cost was reset.
|
||||
"""
|
||||
# Check if this is a local/free model
|
||||
model_str = model_name.lower()
|
||||
is_local_model = (
|
||||
"alias" not in model_str and
|
||||
"ollama" in model_str or
|
||||
"qwen" in model_str or
|
||||
"llama" in model_str or
|
||||
"mistral" in model_str or
|
||||
":" in model_str or
|
||||
(os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false')
|
||||
)
|
||||
# Check if this is a local/free model by getting its pricing
|
||||
input_cost, output_cost = self.get_model_pricing(model_name)
|
||||
|
||||
if is_local_model:
|
||||
# If both costs are zero, it's a free/local model
|
||||
if input_cost == 0.0 and output_cost == 0.0:
|
||||
# Reset the current interaction costs but keep total session costs
|
||||
self.interaction_cost = 0.0
|
||||
self.last_interaction_cost = 0.0
|
||||
|
|
@ -318,44 +309,30 @@ class CostTracker:
|
|||
# Use the centralized function to standardize model names
|
||||
model_name = get_model_name(model_name)
|
||||
|
||||
# Check if using Ollama or local model
|
||||
model_str = model_name.lower()
|
||||
is_local_model = (
|
||||
"alias" not in model_str and
|
||||
"ollama" in model_str or
|
||||
"qwen" in model_str or
|
||||
"llama" in model_str or
|
||||
"mistral" in model_str or
|
||||
":" in model_str or # Ollama uses formats like qwen2.5:7b
|
||||
(os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false')
|
||||
)
|
||||
|
||||
# For local models, always return zero cost
|
||||
if is_local_model:
|
||||
# Set and cache zero cost for local models
|
||||
free_pricing = (0.0, 0.0)
|
||||
self.model_pricing_cache[model_name] = free_pricing
|
||||
return free_pricing
|
||||
|
||||
# Check cache for non-local models
|
||||
# Check cache first
|
||||
if model_name in self.model_pricing_cache:
|
||||
return self.model_pricing_cache[model_name]
|
||||
|
||||
# Try to load pricing from local pricing.json first
|
||||
# Only use if the specific model name exists in the file
|
||||
try:
|
||||
pricing_path = pathlib.Path("pricing.json")
|
||||
if pricing_path.exists():
|
||||
with open(pricing_path, "r", encoding="utf-8") as f:
|
||||
local_pricing = json.load(f)
|
||||
pricing_info = local_pricing.get("alias0", {})
|
||||
input_cost = pricing_info.get("input_cost_per_token", 0)
|
||||
output_cost = pricing_info.get("output_cost_per_token", 0)
|
||||
|
||||
# Cache and return local pricing
|
||||
self.model_pricing_cache[model_name] = (input_cost, output_cost)
|
||||
return input_cost, output_cost
|
||||
# Only use pricing if the exact model name exists in the file
|
||||
if model_name in local_pricing:
|
||||
pricing_info = local_pricing[model_name]
|
||||
input_cost = pricing_info.get("input_cost_per_token", 0)
|
||||
output_cost = pricing_info.get("output_cost_per_token", 0)
|
||||
|
||||
# Cache and return local pricing
|
||||
self.model_pricing_cache[model_name] = (input_cost, output_cost)
|
||||
return input_cost, output_cost
|
||||
except Exception as e:
|
||||
print(f" WARNING: Error loading local pricing.json: {str(e)}")
|
||||
# Fallback to LiteLLM API if both remote and local pricing not found
|
||||
|
||||
# Fallback to LiteLLM API if local pricing not found
|
||||
LITELLM_URL = (
|
||||
"https://raw.githubusercontent.com/BerriAI/litellm/main/"
|
||||
"model_prices_and_context_window.json"
|
||||
|
|
@ -378,8 +355,8 @@ class CostTracker:
|
|||
except Exception as e:
|
||||
print(f" WARNING: Error fetching model pricing: {str(e)}")
|
||||
|
||||
# Default values if no pricing found
|
||||
default_pricing = (0, 0)
|
||||
# Default to zero cost if no pricing found (local/free models)
|
||||
default_pricing = (0.0, 0.0)
|
||||
self.model_pricing_cache[model_name] = default_pricing
|
||||
return default_pricing
|
||||
|
||||
|
|
@ -389,22 +366,6 @@ class CostTracker:
|
|||
# Standardize model name using the central function
|
||||
model_name = get_model_name(model)
|
||||
|
||||
# Check if this is a local model (always free) first
|
||||
model_str = model_name.lower()
|
||||
is_local_model = (
|
||||
"alias" not in model_str and
|
||||
"ollama" in model_str or
|
||||
"qwen" in model_str or
|
||||
"llama" in model_str or
|
||||
"mistral" in model_str or
|
||||
":" in model_str or
|
||||
(os.getenv('OLLAMA') is not None and os.getenv('OLLAMA').lower() != 'false')
|
||||
)
|
||||
|
||||
# For local models, always return zero cost
|
||||
if is_local_model:
|
||||
return 0.0
|
||||
|
||||
# Generate a cache key
|
||||
cache_key = f"{model_name}_{input_tokens}_{output_tokens}"
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,271 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
import sys
|
||||
import os
|
||||
import tempfile
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Add the src directory to the Python path
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
|
||||
|
||||
from cai.util import COST_TRACKER, calculate_model_cost
|
||||
|
||||
|
||||
def setup_test():
|
||||
"""Clear the pricing cache before each test"""
|
||||
COST_TRACKER.model_pricing_cache.clear()
|
||||
|
||||
|
||||
def test_local_models_return_zero_cost():
|
||||
"""Test that local models return zero cost"""
|
||||
setup_test()
|
||||
|
||||
local_models = [
|
||||
"qwen3:14b",
|
||||
"qwen3:32b",
|
||||
"qwen2.5:14b",
|
||||
"qwen2.5:7b",
|
||||
"qwen2.5:72b",
|
||||
"llama3.1:8b",
|
||||
"llama3.1:70b",
|
||||
"mistral:7b",
|
||||
"mistral:latest",
|
||||
"codellama:13b",
|
||||
"ollama/llama3.1",
|
||||
"ollama/qwen2.5",
|
||||
"deepseek-coder:6.7b",
|
||||
"phi3:mini",
|
||||
"gemma:7b",
|
||||
"vicuna:13b",
|
||||
"alpaca:7b",
|
||||
"orca-mini:3b",
|
||||
"neural-chat:7b",
|
||||
"starling-lm:7b",
|
||||
"zephyr:7b",
|
||||
"openchat:7b",
|
||||
"wizard-coder:15b",
|
||||
"sqlcoder:7b",
|
||||
"magicoder:7b",
|
||||
"dolphin-mixtral:8x7b",
|
||||
"nous-hermes2:10.7b",
|
||||
"yi:34b",
|
||||
"qwq:32b",
|
||||
"alias01:14b",
|
||||
"alias01:14b-ctx-32000",
|
||||
"alias00:14b"
|
||||
]
|
||||
|
||||
failed_models = []
|
||||
|
||||
for model in local_models:
|
||||
with patch('requests.get') as mock_get:
|
||||
# Mock LiteLLM API to return empty response (model not found)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {} # Empty response, model not found
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
pricing = COST_TRACKER.get_model_pricing(model)
|
||||
cost = calculate_model_cost(model, 100, 50)
|
||||
|
||||
if pricing != (0.0, 0.0):
|
||||
failed_models.append(f"Model {model} should have zero pricing, got {pricing}")
|
||||
if cost != 0.0:
|
||||
failed_models.append(f"Model {model} should have zero cost, got {cost}")
|
||||
|
||||
if failed_models:
|
||||
print("FAILED: test_local_models_return_zero_cost")
|
||||
for failure in failed_models:
|
||||
print(f" - {failure}")
|
||||
return False
|
||||
else:
|
||||
print("PASSED: test_local_models_return_zero_cost")
|
||||
return True
|
||||
|
||||
|
||||
def test_paid_models_return_nonzero_cost():
|
||||
"""Test that known paid models return non-zero cost"""
|
||||
setup_test()
|
||||
|
||||
paid_models_with_expected_pricing = {
|
||||
"gpt-4": {"input_cost_per_token": 0.00003, "output_cost_per_token": 0.00006},
|
||||
"gpt-4o": {"input_cost_per_token": 0.0000025, "output_cost_per_token": 0.00001},
|
||||
"claude-3-sonnet-20240229": {"input_cost_per_token": 0.000003, "output_cost_per_token": 0.000015},
|
||||
"claude-3-5-sonnet-20241022": {"input_cost_per_token": 0.000003, "output_cost_per_token": 0.000015}
|
||||
}
|
||||
|
||||
failed_models = []
|
||||
|
||||
for model, expected_pricing in paid_models_with_expected_pricing.items():
|
||||
with patch('requests.get') as mock_get:
|
||||
# Mock LiteLLM API to return pricing for paid models
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
model: expected_pricing
|
||||
}
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
pricing = COST_TRACKER.get_model_pricing(model)
|
||||
cost = calculate_model_cost(model, 100, 50)
|
||||
|
||||
if not (pricing[0] > 0 or pricing[1] > 0):
|
||||
failed_models.append(f"Model {model} should have non-zero pricing, got {pricing}")
|
||||
if not (cost > 0):
|
||||
failed_models.append(f"Model {model} should have non-zero cost, got {cost}")
|
||||
|
||||
if failed_models:
|
||||
print("FAILED: test_paid_models_return_nonzero_cost")
|
||||
for failure in failed_models:
|
||||
print(f" - {failure}")
|
||||
return False
|
||||
else:
|
||||
print("PASSED: test_paid_models_return_nonzero_cost")
|
||||
return True
|
||||
|
||||
|
||||
def test_private_model_alias0_with_pricing_json():
|
||||
"""Test that alias0 works correctly when defined in pricing.json"""
|
||||
setup_test()
|
||||
|
||||
# Create a pricing.json with alias0 configuration
|
||||
pricing_config = {
|
||||
"alias0": {
|
||||
"max_tokens": 128000,
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 128000,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"output_cost_per_token": 5e-05,
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": True,
|
||||
"supports_vision": True
|
||||
}
|
||||
}
|
||||
|
||||
try:
|
||||
# Mock the file reading to simulate pricing.json with alias0
|
||||
with patch('pathlib.Path') as mock_path:
|
||||
mock_path_instance = MagicMock()
|
||||
mock_path_instance.exists.return_value = True
|
||||
mock_path.return_value = mock_path_instance
|
||||
|
||||
with patch('builtins.open', create=True) as mock_open:
|
||||
mock_file = MagicMock()
|
||||
mock_file.__enter__.return_value = mock_file
|
||||
mock_file.read.return_value = json.dumps(pricing_config)
|
||||
mock_open.return_value = mock_file
|
||||
|
||||
# Mock json.load to return our config
|
||||
with patch('json.load', return_value=pricing_config):
|
||||
pricing = COST_TRACKER.get_model_pricing("alias0")
|
||||
cost = calculate_model_cost("alias0", 100, 50)
|
||||
|
||||
expected_pricing = (5e-06, 5e-05)
|
||||
expected_cost = 100 * 5e-06 + 50 * 5e-05 # 0.0030
|
||||
|
||||
if pricing != expected_pricing:
|
||||
print(f"FAILED: test_private_model_alias0_with_pricing_json")
|
||||
print(f" - alias0 should have pricing {expected_pricing}, got {pricing}")
|
||||
return False
|
||||
|
||||
if abs(cost - expected_cost) >= 1e-10:
|
||||
print(f"FAILED: test_private_model_alias0_with_pricing_json")
|
||||
print(f" - alias0 should have cost {expected_cost}, got {cost}")
|
||||
return False
|
||||
|
||||
print("PASSED: test_private_model_alias0_with_pricing_json")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"FAILED: test_private_model_alias0_with_pricing_json")
|
||||
print(f" - Exception: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_reset_cost_for_local_model():
|
||||
"""Test the reset_cost_for_local_model function"""
|
||||
setup_test()
|
||||
|
||||
failed_tests = []
|
||||
|
||||
# Test with a free model
|
||||
with patch('requests.get') as mock_get:
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {} # Model not found, will return (0.0, 0.0)
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
result = COST_TRACKER.reset_cost_for_local_model("qwen3:14b")
|
||||
if result != True:
|
||||
failed_tests.append("qwen3:14b should be identified as a free model")
|
||||
|
||||
# Test with a paid model
|
||||
with patch('requests.get') as mock_get:
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"gpt-4": {
|
||||
"input_cost_per_token": 0.00003,
|
||||
"output_cost_per_token": 0.00006
|
||||
}
|
||||
}
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
result = COST_TRACKER.reset_cost_for_local_model("gpt-4")
|
||||
if result != False:
|
||||
failed_tests.append("gpt-4 should not be identified as a free model")
|
||||
|
||||
if failed_tests:
|
||||
print("FAILED: test_reset_cost_for_local_model")
|
||||
for failure in failed_tests:
|
||||
print(f" - {failure}")
|
||||
return False
|
||||
else:
|
||||
print("PASSED: test_reset_cost_for_local_model")
|
||||
return True
|
||||
|
||||
|
||||
def run_all_tests():
|
||||
"""Run all tests and report results"""
|
||||
print("Running pricing tests...")
|
||||
print("=" * 50)
|
||||
|
||||
tests = [
|
||||
test_local_models_return_zero_cost,
|
||||
test_paid_models_return_nonzero_cost,
|
||||
test_private_model_alias0_with_pricing_json,
|
||||
test_reset_cost_for_local_model
|
||||
]
|
||||
|
||||
passed = 0
|
||||
failed = 0
|
||||
|
||||
for test in tests:
|
||||
try:
|
||||
if test():
|
||||
passed += 1
|
||||
else:
|
||||
failed += 1
|
||||
except Exception as e:
|
||||
print(f"FAILED: {test.__name__}")
|
||||
print(f" - Exception: {e}")
|
||||
failed += 1
|
||||
print()
|
||||
|
||||
print("=" * 50)
|
||||
print(f"Test Results: {passed} passed, {failed} failed")
|
||||
|
||||
if failed == 0:
|
||||
print("All tests passed! ✅")
|
||||
return True
|
||||
else:
|
||||
print("Some tests failed! ❌")
|
||||
return False
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
success = run_all_tests()
|
||||
sys.exit(0 if success else 1)
|
||||
Loading…
Reference in New Issue