budget-ai

Token usage tracking, cost estimation, and budget enforcement for LLM applications.

Features

• Token tracking: Record per-model, per-provider, per-task token consumption
• Cost estimation: Configurable per-model rates with fallback defaults
• Budget enforcement: Set token and cost limits, raise BudgetExceeded when exhausted
• Decorator patterns: @track_budget and @enforce_budget for async LLM calls
• Zero dependencies: Python stdlib only (ADR-0008)
• Protocol-based: BudgetTracker and CostEstimator via typing.Protocol (ADR-0020)

Install

pip install "git+https://github.com/FlossWare/budget-ai.git"

Quick Start

importasynciofrombudget_aiimportInMemoryBudgetTracker, SimpleCostEstimator, TokenUsageasyncdefmain():
tracker=InMemoryBudgetTracker(
model_rates={"gpt-4o": 0.03, "claude-sonnet": 0.015}
)
awaittracker.set_budget(max_tokens=100_000, max_cost=5.00)
usage=TokenUsage(prompt_tokens=500, completion_tokens=200, total_tokens=700)
awaittracker.record_usage("gpt-4o", usage, task_id="summarize")
status=awaittracker.remaining()
print(f"Tokens remaining: {status.tokens_remaining}")
print(f"Cost remaining: ${status.cost_remaining:.4f}")
report=awaittracker.cost_report()
print(f"Total cost: ${report.total_cost:.4f}")
print(f"By model: {report.by_model}")
asyncio.run(main())

Decorators

frombudget_aiimporttrack_budget, enforce_budget, InMemoryBudgetTrackertracker=InMemoryBudgetTracker()
@track_budget(tracker=tracker)@enforce_budget(tracker=tracker)asyncdefmy_llm_call(prompt: str, *, model: str="gpt-4o"):
returnawaitbackend.chat([{"role": "user", "content": prompt}], model=model)

License

MIT