| import time | |
| from threading import Lock | |
| class TokenManager: | |
| def __init__(self, token_limit: int = 50000, reset_interval: int = 30): | |
| self.token_limit = token_limit | |
| self.reset_interval = reset_interval | |
| self.token_count = 0 | |
| self.last_reset = time.time() | |
| self.lock = Lock() | |
| def track_tokens(self, tokens: int) -> None: | |
| """ | |
| Track token usage and reset if needed | |
| """ | |
| with self.lock: | |
| current_time = time.time() | |
| if current_time - self.last_reset > self.reset_interval: | |
| self.token_count = 0 | |
| self.last_reset = current_time | |
| self.token_count += tokens | |
| if self.token_count > self.token_limit: | |
| print(f"Warning: Token limit of {self.token_limit} exceeded!") | |
| def get_token_usage(self) -> int: | |
| """ | |
| Get current token usage | |
| """ | |
| with self.lock: | |
| current_time = time.time() | |
| if current_time - self.last_reset > self.reset_interval: | |
| self.token_count = 0 | |
| self.last_reset = current_time | |
| return self.token_count |