from agent_budget import Budget, BudgetScope, ModelRates

# Define a scope with limits
scope = BudgetScope(
    name="my-session",
    max_tokens_per_session=50_000,
    max_cost_per_session_usd=1.00,
    max_cost_per_day_usd=5.00,
    alert_at_pct=0.8,   # warn at 80%
    hard_stop=True,     # raise BudgetExceededError at 100%
)

budget = Budget(scope)

# Optional: register an alert callback
budget.on_alert(lambda alert: print(f"[{alert.level.value}] {alert.message}"))

# Before calling the LLM, check prospectively
budget.check("gpt-4o", input_tokens=800, output_tokens=200)

# After the call, record actual usage
record = budget.record(
    session_id="user-123",
    model="gpt-4o",
    input_tokens=800,
    output_tokens=200,
)

print(f"Call cost: ${record.cost_usd:.4f}")
print(budget.summary())

from agent_budget import BudgetManager, BudgetScope

mgr = BudgetManager()

# Register different scopes per user / team
mgr.register(BudgetScope(name="user-alice", max_cost_per_day_usd=2.00))
mgr.register(BudgetScope(name="user-bob",   max_cost_per_day_usd=5.00))

# Record usage for alice
mgr.record("user-alice", session_id="s1", model="gpt-4o-mini",
           input_tokens=100, output_tokens=50)

# Global dashboard
print(mgr.all_summaries())

from agent_budget import Budget, BudgetScope, ModelRates

custom_rates = {
    "my-custom-model": ModelRates(input_per_1k=0.001, output_per_1k=0.002),
}
scope = BudgetScope(name="s", max_cost_per_session_usd=10.0)
budget = Budget(scope, rates=custom_rates)