Slot-based token allocator using tiktoken for accurate counts. Production fix over v1: char/4 heuristic replaced with the actual tokenizer the model uses. Non-English text and code tokenize very differently — the heuristic causes silent overflow in prod.
| 295 | # ============================================================================= |
| 296 | |
| 297 | class TokenBudget: |
| 298 | """ |
| 299 | Slot-based token allocator using tiktoken for accurate counts. |
| 300 | |
| 301 | Production fix over v1: char/4 heuristic replaced with the actual |
| 302 | tokenizer the model uses. Non-English text and code tokenize |
| 303 | very differently — the heuristic causes silent overflow in prod. |
| 304 | """ |
| 305 | |
| 306 | def __init__(self, total_tokens: int, encoding_name: str = "cl100k_base"): |
| 307 | self.total_tokens = total_tokens |
| 308 | self._enc = None |
| 309 | try: |
| 310 | self._enc = tiktoken.get_encoding(encoding_name) |
| 311 | except Exception: |
| 312 | # Graceful offline fallback: char/4 heuristic |
| 313 | # Replace with tiktoken when network access is available |
| 314 | pass |
| 315 | self._slots: Dict[str, int] = {} |
| 316 | |
| 317 | def count(self, text: str) -> int: |
| 318 | """Exact token count via tiktoken, or char/4 heuristic if offline.""" |
| 319 | if self._enc is not None: |
| 320 | return len(self._enc.encode(text)) |
| 321 | return max(1, len(text) // 4) |
| 322 | |
| 323 | def reserve(self, name: str, text: str) -> bool: |
| 324 | tokens = self.count(text) |
| 325 | if self.remaining() < tokens: |
| 326 | return False |
| 327 | self._slots[name] = tokens |
| 328 | return True |
| 329 | |
| 330 | def reserve_tokens(self, name: str, tokens: int) -> bool: |
| 331 | if self.remaining() < tokens: |
| 332 | return False |
| 333 | self._slots[name] = tokens |
| 334 | return True |
| 335 | |
| 336 | def used(self) -> int: |
| 337 | return sum(self._slots.values()) |
| 338 | |
| 339 | def remaining(self) -> int: |
| 340 | return self.total_tokens - self.used() |
| 341 | |
| 342 | def remaining_chars(self) -> int: |
| 343 | # Rough inverse for truncation: 1 token ~ 4 chars (English) |
| 344 | return self.remaining() * 4 |
| 345 | |
| 346 | def report(self) -> Dict[str, int]: |
| 347 | return dict(self._slots) |
| 348 | |
| 349 | |
| 350 | # ============================================================================= |
no outgoing calls