diff --git a/ci/test/.test.yml b/ci/test/.test.yml index 1aac3be5..1e66a9c6 100644 --- a/ci/test/.test.yml +++ b/ci/test/.test.yml @@ -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: diff --git a/debug_litellm.py b/debug_litellm.py new file mode 100644 index 00000000..7c2d39c0 --- /dev/null +++ b/debug_litellm.py @@ -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() \ No newline at end of file diff --git a/debug_pricing_flow.py b/debug_pricing_flow.py new file mode 100644 index 00000000..8e004db9 --- /dev/null +++ b/debug_pricing_flow.py @@ -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')}") \ No newline at end of file diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 3f2d0076..62c43dcd 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -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: diff --git a/src/cai/util.py b/src/cai/util.py index c2a02c89..5b1a3ef7 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -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}" diff --git a/tests/test_pricing.py b/tests/test_pricing.py new file mode 100644 index 00000000..80abbd45 --- /dev/null +++ b/tests/test_pricing.py @@ -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) \ No newline at end of file