Pythoncost_optimizer.py
Copy
1
import os
2
import hashlib
3
import time
4
from dataclasses import dataclass, field
5
 
6
import openai
7
 
8
client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
9
 
10
# Pricing per 1M tokens (update regularly)
11
PRICING = {
12
    "gpt-4o": {"input": 2.50, "output": 10.00},
13
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
14
}
15
 
16
@dataclass
17
class CostTracker:
18
    total_cost: float = 0.0
19
    total_input_tokens: int = 0
20
    total_output_tokens: int = 0
21
    request_count: int = 0
22
    cache_hits: int = 0
23
 
24
tracker = CostTracker()
25
 
26
# Simple exact-match cache (replace with vector similarity for production)
27
response_cache: dict[str, str] = {}
28
 
29
def estimate_complexity(prompt: str) -> str:
30
    """Route to appropriate model based on prompt complexity."""
31
    SIMPLE_KEYWORDS = ["what is", "define", "list", "summarize in one sentence"]
32
    prompt_lower = prompt.lower()
33
    
34
    # Simple heuristics — replace with a classifier in production
35
    if len(prompt) < 200 and any(kw in prompt_lower for kw in SIMPLE_KEYWORDS):
36
        return "gpt-4o-mini"
37
    if "code" in prompt_lower or "debug" in prompt_lower or "explain" in prompt_lower:
38
        return "gpt-4o-mini"  # Mini handles most coding well
39
    return "gpt-4o"
40
 
41
def cached_completion(prompt: str, use_routing: bool = True) -> str:
42
    """LLM call with cost tracking, routing, and caching."""
43
    # Cache lookup
44
    cache_key = hashlib.sha256(prompt.encode()).hexdigest()
45
    if cache_key in response_cache:
46
        tracker.cache_hits += 1
47
        return f"[CACHED] {response_cache[cache_key]}"
48
    
49
    model = estimate_complexity(prompt) if use_routing else "gpt-4o"
50
    
51
    response = client.chat.completions.create(
52
        model=model,
53
        messages=[{"role": "user", "content": prompt}],
54
        max_tokens=500,
55
    )
56
    
57
    # Track cost
58
    usage = response.usage
59
    p = PRICING[model]
60
    cost = (usage.prompt_tokens * p["input"] + usage.completion_tokens * p["output"]) / 1_000_000
61
    
62
    tracker.total_cost += cost
63
    tracker.total_input_tokens += usage.prompt_tokens
64
    tracker.total_output_tokens += usage.completion_tokens
65
    tracker.request_count += 1
66
    
67
    result = response.choices[0].message.content
68
    response_cache[cache_key] = result
69
    
70
    print(f"  [{model}] ${cost:.6f} | {usage.prompt_tokens}in + {usage.completion_tokens}out tokens")
71
    return result
72
 
73
# Test routing
74
queries = [
75
    "What is a transformer model?",  # Simple → mini
76
    "Write a Python function to implement merge sort with detailed comments",  # Complex
77
    "What is a transformer model?",  # Cached
78
]
79
 
80
for q in queries:
81
    print(f"
82
Q: {q[:50]}...")
83
    cached_completion(q)
84
 
85
print(f"
86
=== Summary ===")
87
print(f"Requests: {tracker.request_count}, Cache hits: {tracker.cache_hits}")
88
print(f"Total cost: ${tracker.total_cost:.6f}")
89
print(f"Projected monthly (1000 req/day): ${tracker.total_cost / len(queries) * 1000 * 30:.2f}")
90