Spaces:
Running
Running
| """Author RAG Chatbot SaaS — Token Counter Utility. | |
| Counts tokens for gpt-4o using tiktoken before sending to OpenAI. | |
| RULE: Always call count_tokens() before assembling the context | |
| to ensure we stay within RAG_MAX_CONTEXT_TOKENS. | |
| """ | |
| import structlog | |
| import tiktoken | |
| logger = structlog.get_logger(__name__) | |
| # gpt-4o uses the same tokenizer as gpt-4 | |
| _ENCODING_NAME = "cl100k_base" | |
| _encoding: tiktoken.Encoding | None = None | |
| def _get_encoding() -> tiktoken.Encoding: | |
| """Lazily load and cache the tiktoken encoding.""" | |
| global _encoding | |
| if _encoding is None: | |
| _encoding = tiktoken.get_encoding(_ENCODING_NAME) | |
| return _encoding | |
| def count_tokens(text: str) -> int: | |
| """Count the number of tokens in a text string. | |
| Args: | |
| text: Input string to tokenize. | |
| Returns: | |
| Integer token count. | |
| """ | |
| return len(_get_encoding().encode(text)) | |
| def count_messages_tokens(messages: list[dict]) -> int: | |
| """Count tokens for a list of OpenAI chat messages. | |
| Accounts for per-message overhead (role + separators). | |
| Args: | |
| messages: List of dicts with 'role' and 'content' keys. | |
| Returns: | |
| Total token count including overhead. | |
| """ | |
| encoding = _get_encoding() | |
| total = 0 | |
| for msg in messages: | |
| # Each message has: 4 overhead tokens + role + content | |
| total += 4 | |
| total += len(encoding.encode(msg.get("role", ""))) | |
| total += len(encoding.encode(msg.get("content", ""))) | |
| total += 2 # Reply priming | |
| return total | |
| def trim_messages_to_budget( | |
| messages: list[dict], | |
| max_tokens: int, | |
| preserve_system: bool = True, | |
| ) -> list[dict]: | |
| """Trim conversation history to fit within a token budget. | |
| Removes oldest messages first. System message is always preserved | |
| if preserve_system=True. | |
| Args: | |
| messages: Full list of chat messages (oldest first). | |
| max_tokens: Maximum allowed token count. | |
| preserve_system: If True, never remove the system message. | |
| Returns: | |
| Trimmed list of messages that fits within max_tokens. | |
| """ | |
| if count_messages_tokens(messages) <= max_tokens: | |
| return messages | |
| system_msgs = [m for m in messages if m["role"] == "system"] if preserve_system else [] | |
| history_msgs = [m for m in messages if m["role"] != "system"] | |
| while history_msgs and count_messages_tokens(system_msgs + history_msgs) > max_tokens: | |
| history_msgs.pop(0) # Remove oldest non-system message | |
| logger.debug("Trimmed oldest message from context", remaining=len(history_msgs)) | |
| return system_msgs + history_msgs | |
| def fits_budget(text: str, budget: int) -> bool: | |
| """Check whether a text fits within a token budget. | |
| Args: | |
| text: Input string to evaluate. | |
| budget: Maximum allowed token count. | |
| Returns: | |
| True if token count is within budget, False if it exceeds it. | |
| """ | |
| return count_tokens(text) <= budget | |