""" Test cost calculation service. """ import pytest from app.services.cost_calculator import CostCalculator class TestCostCalculator: """Test cost calculation functionality.""" def test_openai_gpt35_turbo_cost(self): """Test OpenAI GPT-3.5 Turbo cost calculation.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="gpt-3.5-turbo", input_tokens=1000, output_tokens=500 ) # GPT-3.5 Turbo: $0.0005/1K input tokens, $0.0015/1K output tokens # Input: 1000 tokens * 5 units/1K = 5 units # Output: 500 tokens * 15 units/1K = 7 units expected_cost_cents = 5 + 7 assert cost_cents == expected_cost_cents def test_openai_gpt4_cost(self): """Test OpenAI GPT-4 cost calculation.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="gpt-4", input_tokens=1000, output_tokens=500 ) # GPT-4: $0.03/1K input tokens, $0.06/1K output tokens # Input: 1000 tokens * 300 units/1K = 300 units # Output: 500 tokens * 600 units/1K = 300 units expected_cost_cents = 300 + 300 assert cost_cents == expected_cost_cents def test_anthropic_claude_cost(self): """Test Anthropic Claude cost calculation.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="claude-3-haiku", input_tokens=1000, output_tokens=500 ) # Claude-3 Haiku: $0.00025/1K input tokens, $0.00125/1K output tokens # Input: 1000 tokens * 25 cents/1K = 25 cents # Output: 500 tokens * 125 cents/1K = 62.5 cents -> 62 cents (integer division) expected_cost_cents = 25 + 62 assert cost_cents == expected_cost_cents def test_google_gemini_cost(self): """Test Google Gemini cost calculation.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="gemini-pro", input_tokens=1000, output_tokens=500 ) # Gemini Pro: $0.0005/1K input tokens, $0.0015/1K output tokens # Input: 1000 tokens * 5 units/1K = 5 units # Output: 500 tokens * 15 units/1K = 7 units expected_cost_cents = 5 + 7 assert cost_cents == expected_cost_cents def test_privatemode_tee_cost(self): """Test Privatemode.ai TEE cost calculation.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="privatemode-llama-70b", input_tokens=1000, output_tokens=500 ) # Privatemode Llama 70B: $0.004/1K input tokens, $0.008/1K output tokens # Input: 1000 tokens * 40 units/1K = 40 units # Output: 500 tokens * 80 units/1K = 40 units expected_cost_cents = 40 + 40 assert cost_cents == expected_cost_cents def test_embedding_cost(self): """Test embedding cost calculation.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="text-embedding-ada-002", input_tokens=1000, output_tokens=0 ) # text-embedding-ada-002: $0.0001/1K tokens # Input: 1000 tokens * 1 unit/1K = 1 unit # Output: 0 tokens * 0 units/1K = 0 units expected_cost_cents = 1 assert cost_cents == expected_cost_cents def test_unknown_model_cost(self): """Test cost calculation for unknown model.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="unknown-model", input_tokens=1000, output_tokens=500 ) # Should use default pricing: $0.001/1K input, $0.002/1K output # Input: 1000 tokens * 10 units/1K = 10 units # Output: 500 tokens * 20 units/1K = 10 units expected_cost_cents = 10 + 10 assert cost_cents == expected_cost_cents def test_zero_tokens_cost(self): """Test cost calculation with zero tokens.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="gpt-3.5-turbo", input_tokens=0, output_tokens=0 ) assert cost_cents == 0 def test_large_token_count_cost(self): """Test cost calculation with large token counts.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="gpt-3.5-turbo", input_tokens=100000, output_tokens=50000 ) # GPT-3.5 Turbo: $0.0005/1K input tokens, $0.0015/1K output tokens # Input: 100000 tokens * 5 units/1K = 500 units # Output: 50000 tokens * 15 units/1K = 750 units expected_cost_cents = 500 + 750 assert cost_cents == expected_cost_cents def test_estimate_request_cost(self): """Test request cost estimation.""" cost_cents = CostCalculator.estimate_cost_cents( model_name="gpt-3.5-turbo", estimated_tokens=1000 ) # Should return a reasonable estimate # 1000 tokens * 70% input * 5 units/1K + 1000 tokens * 30% output * 15 units/1K # = 700 * 5/1000 + 300 * 15/1000 = 3 + 4 = 7 units expected_cost_cents = 3 + 4 assert cost_cents == expected_cost_cents def test_cost_precision(self): """Test cost calculation precision.""" cost_cents = CostCalculator.calculate_cost_cents( model_name="gpt-3.5-turbo", input_tokens=1, output_tokens=1 ) # Should maintain precision and handle small amounts # Input: 1 token * 5 units/1K = 0.005 units -> 0 units (integer division) # Output: 1 token * 15 units/1K = 0.015 units -> 0 units (integer division) expected_cost_cents = 0 assert cost_cents == expected_cost_cents def test_model_normalization(self): """Test that model names are normalized correctly.""" cost1 = CostCalculator.calculate_cost_cents( model_name="gpt-3.5-turbo", input_tokens=1000, output_tokens=500 ) cost2 = CostCalculator.calculate_cost_cents( model_name="openai/gpt-3.5-turbo", input_tokens=1000, output_tokens=500 ) assert cost1 == cost2 def test_get_model_pricing(self): """Test getting model pricing information.""" pricing = CostCalculator.get_model_pricing("gpt-3.5-turbo") assert "input" in pricing assert "output" in pricing assert pricing["input"] == 5 # $0.0005/1K in 1/10000ths assert pricing["output"] == 15 # $0.0015/1K in 1/10000ths def test_get_cost_per_1k_tokens(self): """Test getting cost per 1K tokens in dollars.""" pricing = CostCalculator.get_cost_per_1k_tokens("gpt-3.5-turbo") assert "input" in pricing assert "output" in pricing assert "currency" in pricing assert pricing["input"] == 0.0005 # $0.0005/1K assert pricing["output"] == 0.0015 # $0.0015/1K assert pricing["currency"] == "USD" def test_format_cost_display(self): """Test cost formatting for display.""" # Test zero cost assert CostCalculator.format_cost_display(0) == "$0.00" # Test small cost assert CostCalculator.format_cost_display(1) == "$0.0010" # Test larger cost assert CostCalculator.format_cost_display(1500) == "$1.50"