agentic-ai-engineering/course/labs/L6-cost-optimization/starter.py

78 lines
2.5 KiB
Python

"""
Lab 6.9: Cost Optimization
Profile a session, identify savings, and implement cascade routing.
"""
# TODO 1: Define cost profiles for different models
MODEL_COSTS = {
"gemini-flash": {"input": 0.15, "output": 0.60}, # $/M tokens
"deepseek-v3": {"input": 0.27, "output": 1.10},
"claude-sonnet": {"input": 3.00, "output": 15.00},
"claude-opus": {"input": 15.00, "output": 75.00},
}
# TODO 2: Profile a real agent session
SESSION_LOG = [
# (step, model, input_tokens, output_tokens, step_type)
("retrieve context", "claude-opus", 8000, 200, "retrieval"),
("analyze data", "claude-opus", 5000, 1500, "analysis"),
("make decision", "claude-opus", 3000, 800, "decision"),
("format output", "claude-opus", 2000, 400, "formatting"),
("verify result", "claude-opus", 2500, 600, "verification"),
]
# TODO 3: Implement cost calculation
def calculate_session_cost(session: list[tuple], model: str) -> float:
"""Calculate total cost of a session using given model's pricing."""
pass # TODO
# TODO 4: Implement cascade routing optimization
CASCADE_MAP = {
"retrieval": "gemini-flash", # Cheap: simple retrieval
"formatting": "gemini-flash", # Cheap: formatting
"verification": "deepseek-v3", # Medium: structured checks
"analysis": "claude-sonnet", # Expensive: reasoning
"decision": "claude-opus", # Most expensive: critical thinking
}
def calculate_cascade_cost(session: list[tuple]) -> float:
"""Calculate session cost with cascade routing."""
pass # TODO
# TODO 5: Find additional savings
def find_optimizations(session: list[tuple], current_cost: float) -> list[dict]:
"""
Analyze the session and suggest optimizations.
Look for:
- Steps where cascade routing saves >50%
- Steps with unusually high token counts (context waste)
- Opportunities to batch small retrievals
"""
pass # TODO
if __name__ == "__main__":
print("=" * 55)
print(" COST OPTIMIZATION ANALYSIS")
print("=" * 55)
opus_cost = calculate_session_cost(SESSION_LOG, "claude-opus")
cascade_cost = calculate_cascade_cost(SESSION_LOG)
print(f"\nAll Opus: ${opus_cost:.4f}")
print(f"Cascade: ${cascade_cost:.4f}")
print(f"Savings: {((opus_cost - cascade_cost) / opus_cost * 100):.0f}%")
optimizations = find_optimizations(SESSION_LOG, opus_cost)
if optimizations:
print(f"\nOptimization suggestions:")
for opt in optimizations:
print(f" - {opt['description']}: save ${opt['savings']:.4f}")