143 lines
4.2 KiB
Python
143 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
token_budget.py — Daily token budget tracker for pipeline orchestration.
|
|
|
|
Tracks token usage per pipeline per day, enforces daily limits,
|
|
and provides a query interface for the orchestrator.
|
|
|
|
Data: ~/.hermes/pipeline_budget.json
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
BUDGET_FILE = Path.home() / ".hermes" / "pipeline_budget.json"
|
|
DEFAULT_DAILY_LIMIT = 500_000
|
|
|
|
|
|
def _load() -> dict:
|
|
if BUDGET_FILE.exists():
|
|
try:
|
|
return json.loads(BUDGET_FILE.read_text())
|
|
except (json.JSONDecodeError, OSError):
|
|
pass
|
|
return {}
|
|
|
|
|
|
def _save(data: dict):
|
|
BUDGET_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
BUDGET_FILE.write_text(json.dumps(data, indent=2))
|
|
|
|
|
|
def today_key() -> str:
|
|
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
|
|
|
|
def get_daily_usage(pipeline: str = None) -> dict:
|
|
"""Get token usage for today. If pipeline specified, return just that pipeline."""
|
|
data = _load()
|
|
day = data.get("daily", {}).get(today_key(), {"tokens_used": 0, "pipelines": {}})
|
|
if pipeline:
|
|
return {
|
|
"pipeline": pipeline,
|
|
"tokens_used": day.get("pipelines", {}).get(pipeline, 0),
|
|
"daily_total": day.get("tokens_used", 0),
|
|
}
|
|
return day
|
|
|
|
|
|
def get_remaining(limit: int = DEFAULT_DAILY_LIMIT) -> int:
|
|
"""Get remaining token budget for today."""
|
|
usage = get_daily_usage()
|
|
return max(0, limit - usage.get("tokens_used", 0))
|
|
|
|
|
|
def can_afford(tokens: int, limit: int = DEFAULT_DAILY_LIMIT) -> bool:
|
|
"""Check if we have budget for a token spend."""
|
|
return get_remaining(limit) >= tokens
|
|
|
|
|
|
def record_usage(pipeline: str, input_tokens: int, output_tokens: int) -> dict:
|
|
"""
|
|
Record token usage for a pipeline task.
|
|
|
|
Called automatically by the orchestrator after each pipeline task completes.
|
|
Returns the updated daily state.
|
|
"""
|
|
total = input_tokens + output_tokens
|
|
data = _load()
|
|
today = today_key()
|
|
|
|
daily = data.setdefault("daily", {})
|
|
day = daily.setdefault(today, {"tokens_used": 0, "pipelines": {}})
|
|
|
|
day["tokens_used"] = day.get("tokens_used", 0) + total
|
|
pipes = day.setdefault("pipelines", {})
|
|
pipes[pipeline] = pipes.get(pipeline, 0) + total
|
|
|
|
# Track breakdown
|
|
breakdown = day.setdefault("breakdown", {})
|
|
pb = breakdown.setdefault(pipeline, {"input": 0, "output": 0, "calls": 0})
|
|
pb["input"] += input_tokens
|
|
pb["output"] += output_tokens
|
|
pb["calls"] += 1
|
|
|
|
# Track lifetime stats
|
|
lifetime = data.setdefault("lifetime", {"total_tokens": 0, "total_days": 0})
|
|
lifetime["total_tokens"] = lifetime.get("total_tokens", 0) + total
|
|
|
|
_save(data)
|
|
|
|
return {
|
|
"pipeline": pipeline,
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"total": total,
|
|
"daily_used": day["tokens_used"],
|
|
"daily_remaining": get_remaining(),
|
|
}
|
|
|
|
|
|
def get_report() -> str:
|
|
"""Generate a human-readable budget report."""
|
|
data = _load()
|
|
today = today_key()
|
|
day = data.get("daily", {}).get(today, {"tokens_used": 0, "pipelines": {}})
|
|
|
|
lines = []
|
|
lines.append(f"Token Budget — {today}")
|
|
lines.append(f" Daily usage: {day.get('tokens_used', 0):,} / {DEFAULT_DAILY_LIMIT:,}")
|
|
lines.append(f" Remaining: {get_remaining():,}")
|
|
lines.append("")
|
|
lines.append(" Pipelines:")
|
|
|
|
breakdown = day.get("breakdown", {})
|
|
for name, stats in sorted(breakdown.items(), key=lambda x: -x[1]["output"]):
|
|
total = stats["input"] + stats["output"]
|
|
lines.append(f" {name}: {total:,} tokens ({stats['calls']} calls)")
|
|
|
|
if not breakdown:
|
|
lines.append(" (no pipelines run today)")
|
|
|
|
lifetime = data.get("lifetime", {})
|
|
lines.append("")
|
|
lines.append(f" Lifetime: {lifetime.get('total_tokens', 0):,} total tokens")
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
if "--report" in sys.argv:
|
|
print(get_report())
|
|
elif "--remaining" in sys.argv:
|
|
print(get_remaining())
|
|
elif "--can-afford" in sys.argv:
|
|
idx = sys.argv.index("--can-afford")
|
|
tokens = int(sys.argv[idx + 1])
|
|
print("yes" if can_afford(tokens) else "no")
|
|
else:
|
|
print(get_report())
|