simonw--llm
12 KiB
12 KiB
Cost Estimation Feature - FINAL IMPLEMENTATION PLAN
🎯 Overview
Add cost estimation to LLM project by:
- Fetching pricing data from llm-prices.com on first use
- Caching in user directory (
~/.local/share/io.datasette.llm/historical-v1.json) - Integrating cost display into existing
-u/--usageflag - Supporting both sync and async Response classes
🔄 Key Changes from Initial Discussion
| Aspect | Initial Plan | Final Plan |
|---|---|---|
| Pricing data | Bundle with package | Fetch on demand |
| Cache location | Package directory | user_dir() / "historical-v1.json" |
| Cache refresh | 7 days | 24 hours |
| Filename | pricing_data.json |
historical-v1.json (matches source) |
| CLI | Separate commands | Integrate with -u flag |
| Async support | Not planned | AsyncResponse.cost() method |
💡 User Experience
Before (Current)
$ llm "Hello" -m gpt-4 -u
Hello! How can I help you?
Token usage: 10 input, 5 output
After (With Cost Estimation)
$ llm "Hello" -m gpt-4 -u
Hello! How can I help you?
Token usage: 10 input, 5 output, Cost: $0.000450 ($0.000300 input, $0.000150 output)
🏗️ Architecture
┌─────────────────────────────────────────────────────────────┐
│ User runs: llm "prompt" -m gpt-4 -u │
└────────────────────┬────────────────────────────────────────┘
│
▼
┌───────────────────────┐
│ token_usage_string() │
└───────────┬───────────┘
│
▼
┌───────────────────────┐
│ CostEstimator │
│ (singleton) │
└───────────┬───────────┘
│
┌───────────┴───────────┐
│ │
First use Subsequent uses
│ │
▼ ▼
┌────────────┐ ┌────────────┐
│ Fetch from │ │ Load from │
│ llm-prices │ │ cache file │
│ .com │ │ (instant) │
└─────┬──────┘ └─────┬──────┘
│ │
▼ │
┌────────────┐ │
│ Save to │ │
│ cache │ │
└─────┬──────┘ │
│ │
└───────────┬───────────┘
│
▼
┌─────────────────────┐
│ Calculate cost │
└──────────┬──────────┘
│
▼
┌─────────────────────┐
│ Display with cost │
└─────────────────────┘
📦 What Gets Built
New Files
llm/
└── costs.py # ~400 lines
├── CostEstimator (sync)
├── AsyncCostEstimator (async)
├── Cost dataclass
├── PriceInfo dataclass
└── Helper functions
tests/
└── test_costs.py # ~500 lines
├── Unit tests
├── Integration tests
└── Mock network tests
Modified Files
llm/
├── models.py # +30 lines
│ ├── Response.cost() - sync method
│ └── AsyncResponse.cost() - async method
│
├── utils.py # +40 lines
│ └── token_usage_string() - enhanced with cost
│
├── cli.py # +10 lines
│ └── Update usage display calls (2 places)
│
└── __init__.py # +5 lines
└── Export Cost, CostEstimator, etc.
User Directory (Runtime)
~/.local/share/io.datasette.llm/
├── logs.db
├── keys.json
├── aliases.json
└── historical-v1.json # NEW: Cached pricing data
🔧 Implementation Details
1. CostEstimator (Sync)
class CostEstimator:
"""Synchronous cost estimator with lazy loading."""
_instance = None
_pricing_data = None
def __init__(self, force_refresh: bool = False):
self._ensure_pricing_data(force_refresh)
def _ensure_pricing_data(self, force_refresh: bool = False):
"""Fetch/load pricing data with 24-hour cache."""
cache_path = llm.user_dir() / "historical-v1.json"
# Check if cache exists and is fresh (< 24 hours)
if not force_refresh and cache_path.exists():
age = time.time() - cache_path.stat().st_mtime
if age < 24 * 3600:
self._pricing_data = json.loads(cache_path.read_text())
return
# Fetch fresh data
try:
response = httpx.get(
"https://www.llm-prices.com/historical-v1.json",
timeout=10.0
)
data = response.json()
cache_path.write_text(json.dumps(data))
self._pricing_data = data
except Exception:
# Fall back to cache if available
if cache_path.exists():
self._pricing_data = json.loads(cache_path.read_text())
else:
raise
def calculate_cost(...) -> Optional[Cost]:
"""Calculate cost for given tokens and model."""
# Implementation...
2. AsyncCostEstimator
class AsyncCostEstimator:
"""Async version for AsyncResponse."""
async def _ensure_pricing_data(self, force_refresh: bool = False):
"""Async fetch/load with same logic."""
# Similar logic but with:
# - httpx.AsyncClient for fetching
# - await for async operations
3. Response Integration
# llm/models.py
class Response(_BaseResponse):
def cost(self, estimator=None) -> Optional[Cost]:
"""Sync cost calculation."""
if estimator is None:
from .costs import get_default_estimator
estimator = get_default_estimator() # May block on first use
return estimator.calculate_cost(
model_id=self.resolved_model or self.model.model_id,
input_tokens=self.input_tokens or 0,
output_tokens=self.output_tokens or 0,
cached_tokens=self._get_cached_tokens(),
date=self.datetime_utc()
)
class AsyncResponse(_BaseResponse):
async def cost(self, estimator=None) -> Optional[Cost]:
"""Async cost calculation."""
if estimator is None:
from .costs import get_async_estimator
estimator = await get_async_estimator()
return await estimator.calculate_cost(...)
4. CLI Integration
# llm/utils.py
def token_usage_string(
input_tokens,
output_tokens,
token_details,
model_id: Optional[str] = None, # NEW
datetime_utc: Optional[datetime] = None, # NEW
show_cost: bool = True # NEW
) -> str:
"""Enhanced to include cost when model_id provided."""
bits = []
# ... existing token formatting ...
if show_cost and model_id:
try:
from .costs import get_default_estimator
cost = get_default_estimator().calculate_cost(...)
if cost:
bits.append(f"Cost: ${cost.total_cost:.6f} (...)")
except Exception:
pass # Silently skip cost on error
return ", ".join(bits)
🚦 Error Handling
Scenario Matrix
| Scenario | Cache Exists | Network | Behavior |
|---|---|---|---|
| First use | ❌ No | ✅ Yes | Fetch → Save → Show cost |
| First use | ❌ No | ❌ No | Skip cost silently |
| Fresh cache | ✅ Yes (< 24h) | N/A | Load cache → Show cost |
| Stale cache | ✅ Yes (> 24h) | ✅ Yes | Fetch → Update → Show cost |
| Stale cache | ✅ Yes (> 24h) | ❌ No | Use stale cache → Show cost |
Key Principles
- Never break the tool - Missing pricing doesn't prevent usage display
- Fail gracefully - Network errors just skip cost
- Use stale data - Better than nothing if network unavailable
- Silent failures - Don't spam users with pricing errors
📊 Performance
First Use (Cold Start)
With network:
├── HTTP request: ~200-500ms
├── JSON parse: ~10ms
├── File write: ~5ms
└── Total delay: ~215-515ms (once)
Without network:
└── No delay (cost skipped)
Subsequent Uses (Warm Cache)
Cache fresh (< 24h):
├── File read: ~2ms
├── JSON parse: ~10ms
└── Total delay: ~12ms (negligible)
Cache Refresh (Stale Cache)
Background refresh attempt:
├── Try fetch (with timeout)
├── If success: update cache
└── If failure: use stale cache
(user sees cost either way)
✅ Implementation Phases
Phase 1: Core (2 days)
- Create
llm/costs.py - Implement
CostEstimatorwith lazy loading - Implement
AsyncCostEstimator - Add cache management (24-hour TTL)
- Implement cost calculation logic
- Add model ID matching (exact + fuzzy)
- Write unit tests
Phase 2: Integration (0.5 day)
- Add
Response.cost()method - Add
AsyncResponse.cost()method - Add
_get_cached_tokens()helper - Update exports in
__init__.py - Write integration tests
Phase 3: CLI Enhancement (0.5 day)
- Enhance
token_usage_string()in utils.py - Update
llm promptcommand (line ~901) - Update
llm logscommand (line ~2182) - Write CLI tests
Phase 4: Documentation & Polish (1 day)
- Add docstrings
- Error handling review
- Performance testing
- Write user documentation
- Update README
Total: 4 days
📝 Testing Strategy
Unit Tests
test_fetch_pricing_data()
test_cache_saves_correctly()
test_cache_age_detection()
test_load_from_stale_cache_on_network_error()
test_exact_model_match()
test_fuzzy_model_match()
test_calculate_cost_basic()
test_calculate_cost_with_cached_tokens()
test_historical_pricing()
Integration Tests
test_response_cost_sync()
test_response_cost_async()
test_first_use_with_network()
test_first_use_without_network()
test_usage_string_includes_cost()
Manual Testing
# Test first use
rm ~/.local/share/io.datasette.llm/historical-v1.json
llm "Test" -m gpt-4 -u
# Test cached use
llm "Test" -m gpt-4 -u # Should be instant
# Test stale cache
touch -t 202401010000 ~/.local/share/io.datasette.llm/historical-v1.json
llm "Test" -m gpt-4 -u # Should refresh
# Test network failure
# (disconnect network)
llm "Test" -m gpt-4 -u # Should use stale cache or skip
🎯 Success Criteria
- Planning complete
- Cost appears with
-uflag - First fetch works correctly
- Cache is created in user_dir()
- 24-hour refresh works
- Network failures handled gracefully
- Sync and async versions both work
- Tests >90% coverage
- No breaking changes
- Documentation complete
📚 Documentation Files
See INDEX.md for complete navigation. Key files:
- CACHING_UPDATE.md - Detailed caching implementation
- IMPLEMENTATION_CHECKLIST_REVISED.md - Step-by-step tasks
- COMPARISON.md - Why this approach vs alternatives
🚀 Next Steps
- Review
CACHING_UPDATE.mdfor implementation details - Start Phase 1: Create
llm/costs.py - Follow
IMPLEMENTATION_CHECKLIST_REVISED.md - Reference
EXAMPLE_TESTS.pyfor test patterns
Ready to implement! This approach is simple, fast, and robust. 🎉