| # quread/cost_guard.py | |
| import time | |
| # -------- Configuration -------- | |
| DAILY_TOKEN_LIMIT = 200_000 # total tokens per day | |
| REQUEST_TOKEN_LIMIT = 1_200 # max tokens per request | |
| # -------- Internal state -------- | |
| _state = { | |
| "day": time.strftime("%Y-%m-%d"), | |
| "tokens": 0, | |
| } | |
| def allow_request(estimated_tokens: int) -> bool: | |
| """ | |
| Returns True if the request is allowed under cost limits. | |
| Resets automatically each day. | |
| """ | |
| global _state | |
| today = time.strftime("%Y-%m-%d") | |
| if _state["day"] != today: | |
| _state = {"day": today, "tokens": 0} | |
| if estimated_tokens > REQUEST_TOKEN_LIMIT: | |
| return False | |
| if _state["tokens"] + estimated_tokens > DAILY_TOKEN_LIMIT: | |
| return False | |
| _state["tokens"] += estimated_tokens | |
| return True |