diff --git a/tests/conftest_pricing.py b/tests/conftest_pricing.py new file mode 100644 index 00000000..dad59072 --- /dev/null +++ b/tests/conftest_pricing.py @@ -0,0 +1,7 @@ +# Minimal conftest for pricing tests +import pytest + +# Simple configuration without problematic imports +def pytest_configure(config): + """Configure pytest for pricing tests""" + pass \ No newline at end of file diff --git a/tests/test_pricing.py b/tests/test_pricing.py index 80abbd45..9e938a38 100644 --- a/tests/test_pricing.py +++ b/tests/test_pricing.py @@ -1,5 +1,7 @@ #!/usr/bin/env python3 +from __future__ import annotations + import sys import os import tempfile @@ -7,21 +9,43 @@ import json from pathlib import Path from unittest.mock import patch, MagicMock +# Try to import pytest, but make it optional +try: + import pytest + PYTEST_AVAILABLE = True +except ImportError: + PYTEST_AVAILABLE = False + # Create a dummy fixture decorator for when pytest is not available + def pytest_fixture(*args, **kwargs): + def decorator(func): + return func + return decorator + pytest = type('pytest', (), {'fixture': pytest_fixture}) + # 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() +if PYTEST_AVAILABLE: + @pytest.fixture(autouse=True) + def setup_test(): + """Clear the pricing cache before each test""" + COST_TRACKER.model_pricing_cache.clear() + yield + COST_TRACKER.model_pricing_cache.clear() +else: + 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() - + if not PYTEST_AVAILABLE: + setup_test() + local_models = [ "qwen3:14b", "qwen3:32b", @@ -57,7 +81,9 @@ def test_local_models_return_zero_cost(): "alias00:14b" ] - failed_models = [] + print("\n" + "="*80) + print("LOCAL MODELS PRICING TEST") + print("="*80) for model in local_models: with patch('requests.get') as mock_get: @@ -70,25 +96,26 @@ def test_local_models_return_zero_cost(): 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}") + # Show pricing information + print(f"Model: {model:<25} | Pricing: {pricing} | Cost (100/50 tokens): ${cost:.6f}") + + if PYTEST_AVAILABLE: + assert pricing == (0.0, 0.0), f"Model {model} should have zero pricing, got {pricing}" + assert cost == 0.0, f"Model {model} should have zero cost, got {cost}" + else: + if pricing != (0.0, 0.0): + raise AssertionError(f"Model {model} should have zero pricing, got {pricing}") + if cost != 0.0: + raise AssertionError(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 + print("="*80) def test_paid_models_return_nonzero_cost(): """Test that known paid models return non-zero cost""" - setup_test() - + if not PYTEST_AVAILABLE: + 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}, @@ -96,7 +123,9 @@ def test_paid_models_return_nonzero_cost(): "claude-3-5-sonnet-20241022": {"input_cost_per_token": 0.000003, "output_cost_per_token": 0.000015} } - failed_models = [] + print("\n" + "="*80) + print("PAID MODELS PRICING TEST") + print("="*80) for model, expected_pricing in paid_models_with_expected_pricing.items(): with patch('requests.get') as mock_get: @@ -111,25 +140,26 @@ def test_paid_models_return_nonzero_cost(): 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}") + # Show pricing information + print(f"Model: {model:<30} | Pricing: {pricing} | Cost (100/50 tokens): ${cost:.6f}") + + if PYTEST_AVAILABLE: + assert pricing[0] > 0 or pricing[1] > 0, f"Model {model} should have non-zero pricing, got {pricing}" + assert cost > 0, f"Model {model} should have non-zero cost, got {cost}" + else: + if not (pricing[0] > 0 or pricing[1] > 0): + raise AssertionError(f"Model {model} should have non-zero pricing, got {pricing}") + if not (cost > 0): + raise AssertionError(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 + print("="*80) def test_private_model_alias0_with_pricing_json(): """Test that alias0 works correctly when defined in pricing.json""" - setup_test() - + if not PYTEST_AVAILABLE: + setup_test() + # Create a pricing.json with alias0 configuration pricing_config = { "alias0": { @@ -145,52 +175,55 @@ def test_private_model_alias0_with_pricing_json(): } } - 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 + print("\n" + "="*80) + print("PRIVATE MODEL (alias0) PRICING TEST") + print("="*80) + + # 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 - 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) - # 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 - + expected_pricing = (5e-06, 5e-05) + expected_cost = 100 * 5e-06 + 50 * 5e-05 # 0.0030 + + # Show pricing information + print(f"Model: alias0 | Pricing: {pricing} | Cost (100/50 tokens): ${cost:.6f}") + print(f"Expected pricing: {expected_pricing} | Expected cost: ${expected_cost:.6f}") + + if PYTEST_AVAILABLE: + assert pricing == expected_pricing, f"alias0 should have pricing {expected_pricing}, got {pricing}" + assert abs(cost - expected_cost) < 1e-10, f"alias0 should have cost {expected_cost}, got {cost}" + else: 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 - + raise AssertionError(f"alias0 should have pricing {expected_pricing}, got {pricing}") 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 + raise AssertionError(f"alias0 should have cost {expected_cost}, got {cost}") + + print("="*80) def test_reset_cost_for_local_model(): """Test the reset_cost_for_local_model function""" - setup_test() - - failed_tests = [] + if not PYTEST_AVAILABLE: + setup_test() + print("\n" + "="*80) + print("RESET COST FOR LOCAL MODEL TEST") + print("="*80) + # Test with a free model with patch('requests.get') as mock_get: mock_response = MagicMock() @@ -199,8 +232,15 @@ def test_reset_cost_for_local_model(): 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") + pricing = COST_TRACKER.get_model_pricing("qwen3:14b") + + print(f"Model: qwen3:14b (free) | Pricing: {pricing} | Reset result: {result}") + + if PYTEST_AVAILABLE: + assert result == True, "qwen3:14b should be identified as a free model" + else: + if result != True: + raise AssertionError("qwen3:14b should be identified as a free model") # Test with a paid model with patch('requests.get') as mock_get: @@ -215,54 +255,178 @@ def test_reset_cost_for_local_model(): 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") + pricing = COST_TRACKER.get_model_pricing("gpt-4") + + print(f"Model: gpt-4 (paid) | Pricing: {pricing} | Reset result: {result}") + + if PYTEST_AVAILABLE: + assert result == False, "gpt-4 should not be identified as a free model" + else: + if result != False: + raise AssertionError("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 + print("="*80) +def test_model_not_found_anywhere_returns_zero(): + """Test that models not found anywhere return zero cost""" + if not PYTEST_AVAILABLE: + setup_test() + + unknown_model = "unknown-model-12345" + + print("\n" + "="*80) + print("UNKNOWN MODEL PRICING TEST") + print("="*80) + + with patch('pathlib.Path') as mock_path: + mock_path_instance = MagicMock() + mock_path_instance.exists.return_value = False # No pricing.json + mock_path.return_value = mock_path_instance + + with patch('requests.get') as mock_get: + # Mock LiteLLM API to return empty response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {} # Model not found + mock_get.return_value = mock_response + + pricing = COST_TRACKER.get_model_pricing(unknown_model) + cost = calculate_model_cost(unknown_model, 100, 50) + + # Show pricing information + print(f"Model: {unknown_model:<20} | Pricing: {pricing} | Cost (100/50 tokens): ${cost:.6f}") + + if PYTEST_AVAILABLE: + assert pricing == (0.0, 0.0), f"Unknown model should have zero pricing, got {pricing}" + assert cost == 0.0, f"Unknown model should have zero cost, got {cost}" + else: + if pricing != (0.0, 0.0): + raise AssertionError(f"Unknown model should have zero pricing, got {pricing}") + if cost != 0.0: + raise AssertionError(f"Unknown model should have zero cost, got {cost}") + + print("="*80) + + +def test_model_not_in_pricing_json_falls_back_to_litellm(): + """Test that models not in pricing.json fall back to LiteLLM API""" + if not PYTEST_AVAILABLE: + setup_test() + + # Create a pricing.json with only alias0 + pricing_config = { + "alias0": { + "input_cost_per_token": 5e-06, + "output_cost_per_token": 5e-05 + } + } + + print("\n" + "="*80) + print("PRICING.JSON vs LITELLM FALLBACK TEST") + print("="*80) + + 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_open.return_value = mock_file + + with patch('json.load', return_value=pricing_config): + with patch('requests.get') as mock_get: + # Mock LiteLLM API response for gpt-4 + 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 + + # Test that gpt-4 (not in pricing.json) uses LiteLLM + pricing = COST_TRACKER.get_model_pricing("gpt-4") + expected_gpt4_pricing = (0.00003, 0.00006) + cost_gpt4 = calculate_model_cost("gpt-4", 100, 50) + + # Test that alias0 (in pricing.json) uses local pricing + pricing_alias0 = COST_TRACKER.get_model_pricing("alias0") + expected_alias0_pricing = (5e-06, 5e-05) + cost_alias0 = calculate_model_cost("alias0", 100, 50) + + # Show pricing information + print(f"Model: gpt-4 (from LiteLLM) | Pricing: {pricing} | Cost (100/50 tokens): ${cost_gpt4:.6f}") + print(f"Model: alias0 (from pricing.json) | Pricing: {pricing_alias0} | Cost (100/50 tokens): ${cost_alias0:.6f}") + + if PYTEST_AVAILABLE: + assert pricing == expected_gpt4_pricing, f"gpt-4 should use LiteLLM pricing, got {pricing}" + assert pricing_alias0 == expected_alias0_pricing, f"alias0 should use local pricing, got {pricing_alias0}" + else: + if pricing != expected_gpt4_pricing: + raise AssertionError(f"gpt-4 should use LiteLLM pricing, got {pricing}") + if pricing_alias0 != expected_alias0_pricing: + raise AssertionError(f"alias0 should use local pricing, got {pricing_alias0}") + + print("="*80) + + +# Fallback for direct execution def run_all_tests(): - """Run all tests and report results""" + """Run all tests and report results - for direct execution""" print("Running pricing tests...") print("=" * 50) + print("NOTE: This will show detailed pricing information for all models tested") + print("=" * 50) - tests = [ + test_functions = [ 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 + test_reset_cost_for_local_model, + test_model_not_found_anywhere_returns_zero, + test_model_not_in_pricing_json_falls_back_to_litellm ] passed = 0 failed = 0 - for test in tests: + # Setup + COST_TRACKER.model_pricing_cache.clear() + + for test_func in test_functions: try: - if test(): - passed += 1 - else: - failed += 1 + COST_TRACKER.model_pricing_cache.clear() # Clear before each test + print(f"\n๐Ÿงช Running: {test_func.__name__}") + test_func() + print(f"โœ… PASSED: {test_func.__name__}") + passed += 1 except Exception as e: - print(f"FAILED: {test.__name__}") + print(f"โŒ FAILED: {test_func.__name__}") print(f" - Exception: {e}") failed += 1 print() - print("=" * 50) - print(f"Test Results: {passed} passed, {failed} failed") + print("=" * 80) + print(f"FINAL TEST RESULTS: {passed} passed, {failed} failed") + print("=" * 80) if failed == 0: - print("All tests passed! โœ…") + print("๐ŸŽ‰ All tests passed! โœ…") + print("\nSUMMARY:") + print("- Local models correctly return zero cost") + print("- Paid models correctly return non-zero cost") + print("- Private models (alias0) work with pricing.json") + print("- Reset cost function works correctly") + print("- Unknown models default to zero cost") + print("- Fallback from pricing.json to LiteLLM works") return True else: - print("Some tests failed! โŒ") + print("๐Ÿ’ฅ Some tests failed! โŒ") return False