simonw--llm
Simon Willison
0a5a4d84a7
Test / test (macos-latest, 3.10) (push) Has been cancelled
Test / test (macos-latest, 3.11) (push) Has been cancelled
Test / test (macos-latest, 3.9) (push) Has been cancelled
Test / test (ubuntu-latest, 3.10) (push) Has been cancelled
Test / test (windows-latest, 3.11) (push) Has been cancelled
Test / test (windows-latest, 3.12) (push) Has been cancelled
Test / test (windows-latest, 3.13) (push) Has been cancelled
Test / test (windows-latest, 3.9) (push) Has been cancelled
Test / test (macos-latest, 3.12) (push) Has been cancelled
Test / test (macos-latest, 3.13) (push) Has been cancelled
Test / test (ubuntu-latest, 3.11) (push) Has been cancelled
Test / test (ubuntu-latest, 3.12) (push) Has been cancelled
Test / test (ubuntu-latest, 3.13) (push) Has been cancelled
Test / test (ubuntu-latest, 3.9) (push) Has been cancelled
Test / test (windows-latest, 3.10) (push) Has been cancelled
328 行
10 KiB
Python
328 行
10 KiB
Python
"""
|
|
Example test cases for cost estimation feature
|
|
These serve as a reference for the actual implementation
|
|
"""
|
|
|
|
import pytest
|
|
from datetime import datetime
|
|
from llm.costs import CostEstimator, Cost, PriceInfo
|
|
|
|
|
|
# Fixture: Sample pricing data for testing
|
|
@pytest.fixture
|
|
def sample_pricing_data(tmp_path):
|
|
"""Create a temporary pricing data file with sample data"""
|
|
data = {
|
|
"prices": [
|
|
{
|
|
"id": "gpt-4",
|
|
"vendor": "openai",
|
|
"name": "GPT-4",
|
|
"input": 30.0,
|
|
"output": 60.0,
|
|
"input_cached": None,
|
|
"from_date": None,
|
|
"to_date": None,
|
|
},
|
|
{
|
|
"id": "gpt-3.5-turbo",
|
|
"vendor": "openai",
|
|
"name": "GPT-3.5 Turbo",
|
|
"input": 0.5,
|
|
"output": 1.5,
|
|
"input_cached": None,
|
|
"from_date": None,
|
|
"to_date": None,
|
|
},
|
|
{
|
|
"id": "claude-3-opus",
|
|
"vendor": "anthropic",
|
|
"name": "Claude 3 Opus",
|
|
"input": 15.0,
|
|
"output": 75.0,
|
|
"input_cached": 1.5,
|
|
"from_date": None,
|
|
"to_date": None,
|
|
},
|
|
{
|
|
"id": "deepseek-chat",
|
|
"vendor": "deepseek",
|
|
"name": "DeepSeek Chat",
|
|
"input": 0.27,
|
|
"output": 1.1,
|
|
"input_cached": None,
|
|
"from_date": "2025-02-08",
|
|
"to_date": None,
|
|
},
|
|
{
|
|
"id": "deepseek-chat",
|
|
"vendor": "deepseek",
|
|
"name": "DeepSeek Chat",
|
|
"input": 0.14,
|
|
"output": 0.28,
|
|
"input_cached": None,
|
|
"from_date": None,
|
|
"to_date": "2025-02-08",
|
|
},
|
|
]
|
|
}
|
|
|
|
import json
|
|
|
|
pricing_file = tmp_path / "pricing_data.json"
|
|
pricing_file.write_text(json.dumps(data))
|
|
return pricing_file
|
|
|
|
|
|
class TestCostEstimatorBasics:
|
|
"""Test basic CostEstimator functionality"""
|
|
|
|
def test_load_pricing_data_success(self, sample_pricing_data):
|
|
"""Test successful loading of pricing data"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
assert estimator is not None
|
|
# Should have loaded 5 pricing entries
|
|
models = estimator.list_models()
|
|
assert len(models) >= 4 # At least 4 unique model IDs
|
|
|
|
def test_exact_model_match(self, sample_pricing_data):
|
|
"""Test finding price for exact model ID"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
price = estimator.get_price("gpt-4")
|
|
|
|
assert price is not None
|
|
assert price.id == "gpt-4"
|
|
assert price.vendor == "openai"
|
|
assert price.input_price == 30.0
|
|
assert price.output_price == 60.0
|
|
|
|
def test_model_not_found(self, sample_pricing_data):
|
|
"""Test handling of unknown model"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
price = estimator.get_price("unknown-model-xyz")
|
|
assert price is None
|
|
|
|
def test_fuzzy_model_match(self, sample_pricing_data):
|
|
"""Test fuzzy matching for model variations"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
|
|
# These should all match to gpt-4
|
|
for model_id in ["gpt-4-0613", "gpt-4-turbo", "gpt-4-1106-preview"]:
|
|
price = estimator.get_price(model_id)
|
|
assert price is not None
|
|
assert price.id == "gpt-4"
|
|
|
|
|
|
class TestCostCalculation:
|
|
"""Test cost calculation logic"""
|
|
|
|
def test_calculate_cost_basic(self, sample_pricing_data):
|
|
"""Test basic cost calculation"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
|
|
# 1000 input tokens, 500 output tokens for GPT-4
|
|
# input: 1000 * 30 / 1,000,000 = 0.03
|
|
# output: 500 * 60 / 1,000,000 = 0.03
|
|
# total: 0.06
|
|
cost = estimator.calculate_cost("gpt-4", 1000, 500)
|
|
|
|
assert cost is not None
|
|
assert cost.input_cost == 0.03
|
|
assert cost.output_cost == 0.03
|
|
assert cost.total_cost == 0.06
|
|
assert cost.model_id == "gpt-4"
|
|
|
|
def test_calculate_cost_with_cached(self, sample_pricing_data):
|
|
"""Test cost calculation with cached tokens"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
|
|
# Claude 3 Opus with cached tokens
|
|
# 1000 input at $15/M = 0.015
|
|
# 500 output at $75/M = 0.0375
|
|
# 2000 cached at $1.5/M = 0.003
|
|
# total: 0.0555
|
|
cost = estimator.calculate_cost(
|
|
"claude-3-opus", input_tokens=1000, output_tokens=500, cached_tokens=2000
|
|
)
|
|
|
|
assert cost is not None
|
|
assert cost.input_cost == 0.015
|
|
assert cost.output_cost == 0.0375
|
|
assert cost.cached_cost == 0.003
|
|
assert cost.total_cost == 0.0555
|
|
|
|
def test_calculate_cost_no_pricing(self, sample_pricing_data):
|
|
"""Test cost calculation when no pricing available"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
cost = estimator.calculate_cost("unknown-model", 1000, 500)
|
|
assert cost is None
|
|
|
|
def test_calculate_cost_zero_tokens(self, sample_pricing_data):
|
|
"""Test cost calculation with zero tokens"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
cost = estimator.calculate_cost("gpt-4", 0, 0)
|
|
|
|
assert cost is not None
|
|
assert cost.total_cost == 0.0
|
|
|
|
|
|
class TestHistoricalPricing:
|
|
"""Test historical pricing logic"""
|
|
|
|
def test_historical_pricing_current(self, sample_pricing_data):
|
|
"""Test getting current pricing (after price change)"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
|
|
# After Feb 8, 2025 - should get new price
|
|
date = datetime.fromisoformat("2025-03-01T12:00:00")
|
|
price = estimator.get_price("deepseek-chat", date)
|
|
|
|
assert price is not None
|
|
assert price.input_price == 0.27 # New price
|
|
assert price.output_price == 1.1
|
|
|
|
def test_historical_pricing_old(self, sample_pricing_data):
|
|
"""Test getting old pricing (before price change)"""
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
|
|
# Before Feb 8, 2025 - should get old price
|
|
date = datetime.fromisoformat("2025-01-01T12:00:00")
|
|
price = estimator.get_price("deepseek-chat", date)
|
|
|
|
assert price is not None
|
|
assert price.input_price == 0.14 # Old price
|
|
assert price.output_price == 0.28
|
|
|
|
|
|
class TestResponseIntegration:
|
|
"""Test integration with Response class"""
|
|
|
|
def test_response_cost_method(self, sample_pricing_data):
|
|
"""Test Response.cost() method"""
|
|
# This would need a mock Response object
|
|
# Pseudo-code:
|
|
"""
|
|
response = create_mock_response(
|
|
model_id="gpt-4",
|
|
input_tokens=1000,
|
|
output_tokens=500
|
|
)
|
|
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
cost = response.cost(estimator)
|
|
|
|
assert cost is not None
|
|
assert cost.total_cost == 0.06
|
|
"""
|
|
pass
|
|
|
|
def test_response_cached_tokens(self, sample_pricing_data):
|
|
"""Test extracting cached tokens from token_details"""
|
|
# Pseudo-code:
|
|
"""
|
|
response = create_mock_response(
|
|
model_id="claude-3-opus",
|
|
input_tokens=1000,
|
|
output_tokens=500,
|
|
token_details={
|
|
"cache_read_input_tokens": 2000
|
|
}
|
|
)
|
|
|
|
estimator = CostEstimator(str(sample_pricing_data))
|
|
cost = response.cost(estimator)
|
|
|
|
assert cost.cached_cost > 0
|
|
"""
|
|
pass
|
|
|
|
|
|
class TestCLICommands:
|
|
"""Test CLI command integration"""
|
|
|
|
def test_logs_cost_command(self):
|
|
"""Test 'llm logs cost' command"""
|
|
# Pseudo-code using Click testing:
|
|
"""
|
|
from click.testing import CliRunner
|
|
from llm.cli import cli
|
|
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ['logs', 'cost', '-1'])
|
|
|
|
assert result.exit_code == 0
|
|
assert 'Cost:' in result.output
|
|
assert '$' in result.output
|
|
"""
|
|
pass
|
|
|
|
def test_cost_models_command(self):
|
|
"""Test 'llm cost-models' command"""
|
|
# Pseudo-code:
|
|
"""
|
|
from click.testing import CliRunner
|
|
from llm.cli import cli
|
|
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ['cost-models'])
|
|
|
|
assert result.exit_code == 0
|
|
assert 'gpt-4' in result.output
|
|
assert 'claude' in result.output
|
|
"""
|
|
pass
|
|
|
|
|
|
# Example usage patterns
|
|
def example_usage_patterns():
|
|
"""Example code showing how the feature should be used"""
|
|
|
|
# Pattern 1: Get cost for a response
|
|
"""
|
|
import llm
|
|
|
|
model = llm.get_model("gpt-4")
|
|
response = model.prompt("Explain quantum computing")
|
|
|
|
cost = response.cost()
|
|
if cost:
|
|
print(f"Cost: ${cost.total_cost:.6f}")
|
|
"""
|
|
|
|
# Pattern 2: Custom estimator
|
|
"""
|
|
from llm.costs import CostEstimator
|
|
|
|
estimator = CostEstimator("/path/to/custom/pricing.json")
|
|
cost = response.cost(estimator)
|
|
"""
|
|
|
|
# Pattern 3: List all models with pricing
|
|
"""
|
|
from llm.costs import get_default_estimator
|
|
|
|
estimator = get_default_estimator()
|
|
models = estimator.list_models()
|
|
|
|
for model in models:
|
|
print(f"{model.name}: ${model.input_price}/M input")
|
|
"""
|
|
|
|
# Pattern 4: Calculate hypothetical cost
|
|
"""
|
|
from llm.costs import get_default_estimator
|
|
|
|
estimator = get_default_estimator()
|
|
cost = estimator.calculate_cost(
|
|
model_id="gpt-4",
|
|
input_tokens=5000,
|
|
output_tokens=1000
|
|
)
|
|
print(f"Estimated cost: ${cost.total_cost:.4f}")
|
|
"""
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("This file contains example test cases for the cost estimation feature")
|
|
print("Run with: pytest EXAMPLE_TESTS.py -v")
|