This commit is contained in:
luijait 2025-05-23 14:37:45 +02:00
parent 4aea8eda74
commit 699017af02
6 changed files with 438 additions and 122 deletions

View File

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

46
debug_litellm.py Normal file
View File

@ -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()

76
debug_pricing_flow.py Normal file
View File

@ -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')}")

View File

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

View File

@ -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}"

271
tests/test_pricing.py Normal file
View File

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