Spaces:
Sleeping
Sleeping
File size: 4,966 Bytes
d82f721 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """Budget tracking for API spending.
Abstracts over storage backends to track cumulative token usage.
Supports in-memory (default, no persistence) and HF Datasets (persistent).
"""
import json
import os
import time
from datetime import date
from threading import Lock
from dotenv import load_dotenv
load_dotenv()
# ---------------------------------------------------------------------------
# Abstract store
# ---------------------------------------------------------------------------
class BudgetStore:
"""Base class for budget storage backends."""
def get(self) -> dict:
"""Return {'tokens': int, 'date': str, 'reset_hour': int}."""
raise NotImplementedError
def save(self, data: dict) -> None:
raise NotImplementedError
# ---------------------------------------------------------------------------
# In-memory store (works everywhere, resets on restart)
# ---------------------------------------------------------------------------
class MemoryBudgetStore(BudgetStore):
def __init__(self):
self._lock = Lock()
self._data = {
"tokens": 0,
"date": str(date.today()),
"reset_hour": 0,
}
def get(self) -> dict:
with self._lock:
return dict(self._data)
def save(self, data: dict) -> None:
with self._lock:
self._data.update(data)
# ---------------------------------------------------------------------------
# HF Datasets store (persistent across restarts and deployments)
# ---------------------------------------------------------------------------
class HFBudgetStore(BudgetStore):
"""Store budget data in a Hugging Face Dataset repo as a JSON file.
Requires:
pip install huggingface_hub
huggingface-cli login (or HF_TOKEN env var)
A dataset repo created at repo_id (private recommended).
"""
def __init__(self, repo_id: str):
self.repo_id = repo_id
self.token = os.getenv("HF_TOKEN")
self._lock = Lock()
self._cache: dict | None = None # cache to avoid HF API on every call
def get(self) -> dict:
with self._lock:
if self._cache is not None:
return dict(self._cache)
try:
from huggingface_hub import hf_hub_download
path = hf_hub_download(
repo_id=self.repo_id,
filename="budget.json",
repo_type="dataset",
token=self.token,
)
self._cache = json.load(open(path))
return dict(self._cache)
except Exception:
default = {
"tokens": 0,
"date": str(date.today()),
"reset_hour": 0,
}
self._cache = default
return dict(default)
def save(self, data: dict) -> None:
with self._lock:
self._cache = dict(data)
try:
from huggingface_hub import upload_file
tmp = "/tmp/rag_budget.json"
with open(tmp, "w") as f:
json.dump(data, f)
upload_file(
path_or_fileobj=tmp,
path_in_repo="budget.json",
repo_id=self.repo_id,
repo_type="dataset",
token=self.token,
)
except Exception:
pass # silent fail — in-memory cache still tracks
# ---------------------------------------------------------------------------
# Budget manager (token counter + daily reset)
# ---------------------------------------------------------------------------
class BudgetManager:
def __init__(
self,
store: BudgetStore | None = None,
daily_cap: int = 500_000,
daily_reset_hour: int = 0, # UTC
):
self.store = store or MemoryBudgetStore()
self.daily_cap = daily_cap
self.daily_reset_hour = daily_reset_hour
def _today(self) -> str:
return str(date.today())
def _current_hour_utc(self) -> int:
return int(time.strftime("%H", time.gmtime()))
def track(self, tokens: int) -> None:
"""Record *tokens* used, respecting daily reset."""
data = self.store.get()
today = self._today()
hour = self._current_hour_utc()
# Daily reset
if data["date"] != today and hour >= self.daily_reset_hour:
data["tokens"] = 0
data["date"] = today
data["reset_hour"] = self.daily_reset_hour
data["tokens"] += tokens
self.store.save(data)
def remaining(self) -> int:
return max(0, self.daily_cap - self.store.get()["tokens"])
def used(self) -> int:
return self.store.get()["tokens"]
def exceeded(self) -> bool:
return self.remaining() <= 0
|