Cost Tracking & Budgets

from pathlib import Path
import sys

ROOT = (
    Path.cwd().resolve().parent
    if Path.cwd().name == "notebooks"
    else Path.cwd().resolve()
)
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from shipit_agent.costs import CostTracker, Budget, BudgetExceededError, MODEL_PRICING
from shipit_agent.costs.pricing import MODEL_ALIASES
from IPython.display import display, Markdown

print(f"Total models: {len(MODEL_PRICING)}\n")

providers = {
    "Anthropic Claude": [k for k in MODEL_PRICING if k.startswith("claude")],
    "OpenAI": [k for k in MODEL_PRICING if k.startswith(("gpt", "o3", "o4"))],
    "Google Gemini": [k for k in MODEL_PRICING if k.startswith("gemini")],
    "Meta Llama": [k for k in MODEL_PRICING if k.startswith("llama")],
    "AWS Bedrock": [k for k in MODEL_PRICING if k.startswith("anthropic.")],
}

for provider, models in providers.items():
    print(f"\n{'='*55}")
    print(f"  {provider}")
    print(f"{'='*55}")
    print(f"  {'Model':<40s} {'Input':>8s} {'Output':>8s}")
    print(f"  {'─'*40} {'─'*8} {'─'*8}")
    for m in sorted(models):
        p = MODEL_PRICING[m]
        print(f"  {m:<40s} ${p['input']:>6.2f} ${p['output']:>6.2f}")

print(f"\nAliases: {MODEL_ALIASES}")

tracker = CostTracker()

cost = tracker.calculate_cost(
    "claude-sonnet-4", input_tokens=10_000, output_tokens=2_000
)
print(f"Claude Sonnet — 10K in / 2K out: ${cost:.4f}")

cost_cached = tracker.calculate_cost(
    "claude-sonnet-4", input_tokens=5_000, output_tokens=2_000, cache_read_tokens=50_000
)
print(f"With 50K cache read: ${cost_cached:.4f}  (saved ${cost - cost_cached:.4f})")

print("\nModel comparison (10K in / 2K out):")
for m in [
    "claude-haiku-4",
    "claude-sonnet-4",
    "claude-opus-4",
    "gpt-4o",
    "gpt-4o-mini",
    "gemini-2.5-pro",
]:
    c = tracker.calculate_cost(m, 10_000, 2_000)
    print(f"  {m:<25s} ${c:.4f}")

# Simulate a typical agent run: 3 LLM calls
workload = [
    (8_000, 1_500),  # planning
    (15_000, 3_000),  # tool-using
    (20_000, 2_500),  # synthesis
]

total_in = sum(i for i, o in workload)
total_out = sum(o for i, o in workload)

print(
    f"Workload: {total_in:,} input + {total_out:,} output tokens ({len(workload)} calls)"
)
print(f"\n{'Model':<30s} {'Per-call avg':>12s} {'Total':>10s} {'vs Haiku':>10s}")
print(f"{'─'*30} {'─'*12} {'─'*10} {'─'*10}")

haiku_total = None
for model in [
    "claude-haiku-4",
    "claude-sonnet-4",
    "claude-opus-4",
    "gpt-4o-mini",
    "gpt-4o",
    "gemini-2.5-flash",
    "gemini-2.5-pro",
    "o3-mini",
]:
    t = CostTracker()
    for inp, out in workload:
        t.record_call(model, inp, out)

    if haiku_total is None:
        haiku_total = t.total_cost

    ratio = t.total_cost / haiku_total if haiku_total > 0 else 0
    print(
        f"  {model:<28s} ${t.total_cost/len(workload):>10.4f} ${t.total_cost:>8.4f} {ratio:>8.1f}x"
    )

# Compare cached vs uncached for Claude Sonnet
print("=== Prompt Caching Savings (Claude Sonnet) ===\n")
print(f"{'Scenario':<40s} {'Cost':>10s} {'Savings':>10s}")
print(f"{'─'*40} {'─'*10} {'─'*10}")

scenarios = [
    ("No caching", 50_000, 0, 0, 5_000),
    ("50% cache hit", 25_000, 25_000, 0, 5_000),
    ("90% cache hit", 5_000, 45_000, 0, 5_000),
    ("First call (cache write)", 10_000, 0, 40_000, 5_000),
]

t = CostTracker()
base_cost = t.calculate_cost("claude-sonnet-4", 50_000, 5_000)

for name, inp, cache_r, cache_w, out in scenarios:
    cost = t.calculate_cost(
        "claude-sonnet-4",
        inp,
        out,
        cache_read_tokens=cache_r,
        cache_write_tokens=cache_w,
    )
    saved = base_cost - cost
    print(f"  {name:<38s} ${cost:>8.4f} ${saved:>+8.4f}")

print(f"\n  Baseline (no cache): ${base_cost:.4f}")

tracker = CostTracker()

calls = [
    ("claude-sonnet-4", 8000, 1500),
    ("claude-sonnet-4", 12000, 3000),
    ("claude-sonnet-4", 15000, 2000),
]

for model, inp, out in calls:
    rec = tracker.record_call(model, inp, out)
    print(f"Call #{rec.call_number}: {inp:,} in / {out:,} out = ${rec.cost_usd:.4f}")

print(f"\nTotal cost:   ${tracker.total_cost:.4f}")
print(f"Total tokens: {tracker.total_tokens}")

tracker_budget = CostTracker(budget=Budget(max_dollars=0.50, warn_at=0.80))

try:
    for i in range(20):
        tracker_budget.record_call("claude-sonnet-4", 10_000, 3_000)
        print(f"  Call #{i+1}: total = ${tracker_budget.total_cost:.4f}")
except BudgetExceededError as e:
    print("\n🛑 Budget exceeded!")
    print(f"  Spent:  ${e.spent:.4f}")
    print(f"  Budget: ${e.budget:.4f}")
    print(f"  Model:  {e.model}")

alerts = []


def on_alert(spent, budget_limit):
    alerts.append(spent)
    print(
        f"  ⚠️  ALERT: ${spent:.4f} of ${budget_limit:.2f} ({spent/budget_limit*100:.0f}%)"
    )


tracker_warn = CostTracker(
    budget=Budget(max_dollars=0.20, warn_at=0.70), on_cost_alert=on_alert
)

try:
    for i in range(10):
        tracker_warn.record_call("claude-sonnet-4", 5_000, 1_500)
        print(f"  Call #{i+1}: ${tracker_warn.total_cost:.4f}")
except BudgetExceededError:
    print(f"\n🛑 Budget exceeded after {len(tracker_warn.breakdown())} calls")
print(f"Alerts triggered: {len(alerts)}")

breakdown = tracker.breakdown()

print("Per-call breakdown:")
print(f"{'#':>3s}  {'Model':<25s}  {'Input':>8s}  {'Output':>8s}  {'Cost':>10s}")
print(f"{'─'*3}  {'─'*25}  {'─'*8}  {'─'*8}  {'─'*10}")
for c in breakdown:
    print(
        f"{c['call_number']:>3d}  {c['model']:<25s}  {c['input_tokens']:>8,d}  {c['output_tokens']:>8,d}  ${c['cost_usd']:>9.6f}"
    )
print(f"\n{'Total':>39s}  ${tracker.total_cost:.6f}")

import json

summary_tracker = CostTracker(budget=Budget(max_dollars=5.00))
summary_tracker.record_call("claude-sonnet-4", 20_000, 5_000)
summary_tracker.record_call("claude-opus-4", 10_000, 3_000)

print("Cost Summary:")
print(json.dumps(summary_tracker.summary(), indent=2, default=str))

tracker_custom = CostTracker()

tracker_custom.add_model("my-company/fine-tuned-v3", {"input": 5.00, "output": 20.00})
rec = tracker_custom.record_call("my-company/fine-tuned-v3", 8_000, 2_000)
print(f"Custom model: ${rec.cost_usd:.4f}")

rec_unknown = tracker_custom.record_call("some-unknown-model", 10_000, 5_000)
print(f"Unknown model: ${rec_unknown.cost_usd:.4f} (no pricing data)")

from shipit_agent import Agent
from examples.run_multi_tool_agent import build_llm_from_env

llm = build_llm_from_env("bedrock")

tracker = CostTracker(budget=Budget(max_dollars=2.00, warn_at=0.80))

agent = Agent.with_builtins(
    llm=llm,
    prompt="You are a helpful assistant.",
    hooks=tracker.as_hooks(),
)

result = agent.run("What are 3 benefits of AI agents in software development?")

print(f"Cost:   ${tracker.total_cost:.4f}")
print(f"Calls:  {len(tracker.breakdown())}")
print(f"Tokens: {tracker.total_tokens}")
for c in tracker.breakdown():
    print(
        f"  #{c['call_number']}: {c['input_tokens']:,} in / {c['output_tokens']:,} out = ${c['cost_usd']:.4f}"
    )

display(Markdown(result.output[:500]))

# Fresh tracker
stream_tracker = CostTracker(budget=Budget(max_dollars=2.00))

agent = Agent.with_builtins(
    llm=llm,
    prompt="You are a helpful analyst.",
    hooks=stream_tracker.as_hooks(),
)

print("=== Streaming with Live Cost ===\n")
for event in agent.stream(
    "What are 5 key metrics for evaluating AI agent performance?"
):
    if event.type == "run_started":
        print(f"🚀 Started | Cost so far: ${stream_tracker.total_cost:.4f}")
    elif event.type == "tool_called":
        print(
            f"🔧 Tool: {event.payload.get('tool_name', 'llm')[:30]:30s} | Cost: ${stream_tracker.total_cost:.4f}"
        )
    elif event.type == "tool_completed":
        print(
            f"✅ Done:  {event.payload.get('tool_name', 'llm')[:30]:30s} | Cost: ${stream_tracker.total_cost:.4f}"
        )
    elif event.type == "run_completed":
        print("\n🏁 Completed!")
        print(f"   Final cost:   ${stream_tracker.total_cost:.4f}")
        print(f"   Total calls:  {len(stream_tracker.breakdown())}")
        print(f"   Tokens:       {stream_tracker.total_tokens}")
        display(Markdown(event.payload.get("output", "")[:500]))

multi_tracker = CostTracker()

# Simulate using different models for different tasks
multi_tracker.record_call("claude-haiku-4", 5_000, 500)  # quick classification
multi_tracker.record_call("claude-sonnet-4", 15_000, 3_000)  # main reasoning
multi_tracker.record_call("claude-sonnet-4", 20_000, 5_000)  # tool-heavy step
multi_tracker.record_call("claude-opus-4", 10_000, 2_000)  # final synthesis

print("=== Multi-Model Cost Breakdown ===\n")
print(f"{'#':>3s}  {'Model':<30s}  {'In':>8s}  {'Out':>8s}  {'Cost':>10s}")
print(f"{'─'*3}  {'─'*30}  {'─'*8}  {'─'*8}  {'─'*10}")
for c in multi_tracker.breakdown():
    print(
        f"{c['call_number']:>3d}  {c['model']:<30s}  {c['input_tokens']:>8,d}  {c['output_tokens']:>8,d}  ${c['cost_usd']:>9.6f}"
    )
print(f"\n{'Total':>44s}  ${multi_tracker.total_cost:.6f}")

# Cost by model
from collections import defaultdict

by_model = defaultdict(float)
for c in multi_tracker.breakdown():
    by_model[c["model"]] += c["cost_usd"]

print("\nCost by model:")
for model, cost in sorted(by_model.items(), key=lambda x: -x[1]):
    pct = cost / multi_tracker.total_cost * 100
    print(f"  {model:<30s} ${cost:.6f} ({pct:.1f}%)")

tracker = CostTracker(budget=Budget(max_dollars=5.00, warn_at=0.80))

agent = Agent.with_builtins(
    llm=llm,
    prompt="You are a thorough code analyst.",
    hooks=tracker.as_hooks(),
    max_iterations=4,
)

result = agent.run(
    "List the main directories in the SHIPIT Agent project and their purpose."
)

print("=" * 55)
print("  COST REPORT")
print("=" * 55)
s = tracker.summary()
print(f"  Total Cost:    ${s['total_cost_usd']:.4f}")
print(f"  Total Calls:   {s['total_calls']}")
print(f"  Input Tokens:  {s['total_tokens']['input_tokens']:,}")
print(f"  Output Tokens: {s['total_tokens']['output_tokens']:,}")
if "budget" in s:
    b = s["budget"]
    print(f"  Budget:        ${b['max_dollars']:.2f}")
    print(f"  Remaining:     ${b['remaining']:.4f}")
    print(f"  Used:          {b['percent_used']:.1f}%")
print("\n  Per-call:")
for c in tracker.breakdown():
    print(f"    #{c['call_number']}: {c['model'][:30]:30s} ${c['cost_usd']:.4f}")
print("=" * 55)
display(Markdown(result.output[:600]))