102 lines
3.6 KiB
Python
102 lines
3.6 KiB
Python
"""
|
|
Lab 6.9: Cost Optimization -- SOLUTION
|
|
|
|
Profiles an agent session, implements cascade routing, calculates savings.
|
|
"""
|
|
|
|
MODEL_COSTS = {
|
|
"gemini-flash": {"input": 0.15, "output": 0.60},
|
|
"deepseek-v3": {"input": 0.27, "output": 1.10},
|
|
"claude-sonnet": {"input": 3.00, "output": 15.00},
|
|
"claude-opus": {"input": 15.00, "output": 75.00},
|
|
}
|
|
|
|
SESSION_LOG = [
|
|
("retrieve context", 8000, 200, "retrieval"),
|
|
("analyze data", 5000, 1500, "analysis"),
|
|
("make decision", 3000, 800, "decision"),
|
|
("format output", 2000, 400, "formatting"),
|
|
("verify result", 2500, 600, "verification"),
|
|
]
|
|
|
|
CASCADE_MAP = {
|
|
"retrieval": "gemini-flash",
|
|
"formatting": "gemini-flash",
|
|
"verification": "deepseek-v3",
|
|
"analysis": "claude-sonnet",
|
|
"decision": "claude-opus",
|
|
}
|
|
|
|
|
|
def calculate_session_cost(session: list[tuple], model: str) -> float:
|
|
costs = MODEL_COSTS[model]
|
|
total = 0.0
|
|
for step_name, in_tokens, out_tokens, _ in session:
|
|
step_cost = (in_tokens * costs["input"] + out_tokens * costs["output"]) / 1_000_000
|
|
total += step_cost
|
|
return total
|
|
|
|
|
|
def calculate_cascade_cost(session: list[tuple]) -> float:
|
|
total = 0.0
|
|
for step_name, in_tokens, out_tokens, step_type in session:
|
|
model = CASCADE_MAP[step_type]
|
|
costs = MODEL_COSTS[model]
|
|
step_cost = (in_tokens * costs["input"] + out_tokens * costs["output"]) / 1_000_000
|
|
total += step_cost
|
|
print(f" {step_name:<25} {model:<16} in={in_tokens:>5} out={out_tokens:>5} -> ${step_cost:.6f}")
|
|
return total
|
|
|
|
|
|
def find_optimizations(session: list[tuple], current_cost: float) -> list[dict]:
|
|
optimizations = []
|
|
for step_name, in_tokens, out_tokens, step_type in session:
|
|
single_model = "claude-opus"
|
|
cascade_model = CASCADE_MAP[step_type]
|
|
|
|
sc = MODEL_COSTS[single_model]
|
|
cc = MODEL_COSTS[cascade_model]
|
|
|
|
single_cost = (in_tokens * sc["input"] + out_tokens * sc["output"]) / 1_000_000
|
|
cascade_cost = (in_tokens * cc["input"] + out_tokens * cc["output"]) / 1_000_000
|
|
savings = single_cost - cascade_cost
|
|
|
|
if savings > 0.001:
|
|
optimizations.append({
|
|
"description": f"'{step_name}' from {single_model} to {cascade_model}",
|
|
"savings": round(savings, 4),
|
|
"savings_pct": round((savings / single_cost) * 100, 0)
|
|
})
|
|
|
|
optimizations.sort(key=lambda x: x["savings"], reverse=True)
|
|
return optimizations
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 55)
|
|
print(" COST OPTIMIZATION ANALYSIS")
|
|
print("=" * 55)
|
|
|
|
opus_cost = calculate_session_cost(SESSION_LOG, "claude-opus")
|
|
print(f"\nSingle model (all Opus): ${opus_cost:.4f}")
|
|
|
|
print(f"\nCascade routing:")
|
|
cascade_cost = calculate_cascade_cost(SESSION_LOG)
|
|
|
|
savings_pct = ((opus_cost - cascade_cost) / opus_cost * 100)
|
|
print(f"\nCascade total: ${cascade_cost:.4f}")
|
|
print(f"All Opus total: ${opus_cost:.4f}")
|
|
print(f"Savings: {savings_pct:.0f}% (${opus_cost - cascade_cost:.4f})")
|
|
|
|
optimizations = find_optimizations(SESSION_LOG, opus_cost)
|
|
if optimizations:
|
|
print(f"\nOptimization suggestions:")
|
|
for opt in optimizations:
|
|
print(f" - {opt['description']}: save {opt['savings_pct']:.0f}% (${opt['savings']:.4f})")
|
|
|
|
print(f"\n3x Rule estimate:")
|
|
print(f" Prototype cost: ${cascade_cost:.4f}")
|
|
print(f" With retries: ${cascade_cost * 1.5:.4f}")
|
|
print(f" Production: ${cascade_cost * 3.0:.4f}")
|
|
print(f" At 1000 tasks: ${cascade_cost * 3.0 * 1000:.2f}/month")
|