mirror of https://github.com/aliasrobotics/cai.git
271 lines
8.9 KiB
Python
271 lines
8.9 KiB
Python
#!/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) |