Pythontoken_budget.py
Copy
1
from datetime import date
2
from dataclasses import dataclass, field
3
from openai import OpenAI
4
 
5
client = OpenAI()
6
 
7
class BudgetExceededError(Exception):
8
    pass
9
 
10
@dataclass
11
class TokenBudget:
12
    daily_limit: int = 50_000
13
    _usage: dict[tuple[str, date], int] = field(default_factory=dict)
14
 
15
    def get_usage(self, user_id: str) -> int:
16
        return self._usage.get((user_id, date.today()), 0)
17
 
18
    def check_budget(self, user_id: str, estimated_tokens: int = 0):
19
        current = self.get_usage(user_id)
20
        if current + estimated_tokens > self.daily_limit:
21
            remaining = self.daily_limit - current
22
            raise BudgetExceededError(
23
                f"Budget exceeded for {user_id}. Used: {current}/{self.daily_limit}. Remaining: {remaining}"
24
            )
25
 
26
    def record(self, user_id: str, tokens: int):
27
        key = (user_id, date.today())
28
        self._usage[key] = self._usage.get(key, 0) + tokens
29
 
30
budget = TokenBudget(daily_limit=1000)  # Low limit for testing
31
 
32
def budgeted_chat(user_id: str, message: str) -> str:
33
    budget.check_budget(user_id, estimated_tokens=100)
34
 
35
    resp = client.chat.completions.create(
36
        model="gpt-4o-mini",
37
        messages=[{"role": "user", "content": message}],
38
    )
39
    actual_tokens = resp.usage.total_tokens
40
    budget.record(user_id, actual_tokens)
41
 
42
    remaining = budget.daily_limit - budget.get_usage(user_id)
43
    print(f"[{user_id}] Used {actual_tokens} tokens. Remaining today: {remaining}")
44
    return resp.choices[0].message.content or ""
45
 
46
# Test
47
try:
48
    for i in range(20):
49
        budgeted_chat("user-123", "What is machine learning?")
50
except BudgetExceededError as e:
51
    print(f"\n✗ {e}")
52