diff --git a/nanobot/utils/helpers.py b/nanobot/utils/helpers.py index a1ec6076..7a5d808c 100644 --- a/nanobot/utils/helpers.py +++ b/nanobot/utils/helpers.py @@ -8,12 +8,61 @@ import time import uuid from contextlib import suppress from datetime import datetime +from functools import lru_cache from pathlib import Path from typing import Any import tiktoken from loguru import logger +_TOOLS_TOKEN_CACHE_MAX_ENTRIES = 64 +_TOOLS_TOKEN_CACHE: dict[int, tuple[tuple[int, ...], dict[bool, int]]] = {} + + +@lru_cache(maxsize=1) +def _get_token_encoding() -> Any: + return tiktoken.get_encoding("cl100k_base") + + +def _cache_tools_token_count( + tools_id: int, + fingerprint: tuple[int, ...], + counts: dict[bool, int], +) -> None: + if ( + tools_id not in _TOOLS_TOKEN_CACHE + and len(_TOOLS_TOKEN_CACHE) >= _TOOLS_TOKEN_CACHE_MAX_ENTRIES + ): + _TOOLS_TOKEN_CACHE.pop(next(iter(_TOOLS_TOKEN_CACHE))) + _TOOLS_TOKEN_CACHE[tools_id] = (fingerprint, counts) + + +def _estimate_tools_tokens( + enc: Any, + tools: list[dict[str, Any]], + *, + leading_separator: bool, +) -> int: + """Estimate stable tool definition tokens without re-encoding every loop.""" + tools_id = id(tools) + fingerprint = tuple(id(tool) for tool in tools) + cached = _TOOLS_TOKEN_CACHE.get(tools_id) + if cached and cached[0] == fingerprint: + token_count = cached[1].get(leading_separator) + if token_count is not None: + return token_count + counts = cached[1] + else: + counts = {} + + rendered = json.dumps(tools, ensure_ascii=False) + if leading_separator: + rendered = "\n" + rendered + token_count = len(enc.encode(rendered)) + counts[leading_separator] = token_count + _cache_tools_token_count(tools_id, fingerprint, counts) + return token_count + def strip_think(text: str) -> str: """Remove thinking blocks, unclosed trailing tags, and tokenizer-level @@ -115,7 +164,7 @@ class IncrementalThinkExtractor: thinking, _ = extract_think(buf) if not thinking or thinking == self._emitted: return False - new = thinking[len(self._emitted):].strip() + new = thinking[len(self._emitted) :].strip() self._emitted = thinking if not new: return False @@ -249,7 +298,7 @@ def truncate_text_to_tokens(text: str, max_tokens: int) -> str: if max_tokens <= 0: return text try: - enc = tiktoken.get_encoding("cl100k_base") + enc = _get_token_encoding() tokens = enc.encode(text) if len(tokens) <= max_tokens: return text @@ -486,7 +535,7 @@ def estimate_prompt_tokens( reasoning_content, tool_call_id, name, plus per-message framing overhead. """ try: - enc = tiktoken.get_encoding("cl100k_base") + enc = _get_token_encoding() parts: list[str] = [] for msg in messages: content = msg.get("content") @@ -512,11 +561,13 @@ def estimate_prompt_tokens( if isinstance(value, str) and value: parts.append(value) - if tools: - parts.append(json.dumps(tools, ensure_ascii=False)) + tool_tokens = ( + _estimate_tools_tokens(enc, tools, leading_separator=bool(parts)) if tools else 0 + ) per_message_overhead = len(messages) * 4 - return len(enc.encode("\n".join(parts))) + per_message_overhead + message_tokens = len(enc.encode("\n".join(parts))) if parts else 0 + return message_tokens + tool_tokens + per_message_overhead except Exception: return 0 @@ -553,7 +604,7 @@ def estimate_message_tokens(message: dict[str, Any]) -> int: if not payload: return 4 try: - enc = tiktoken.get_encoding("cl100k_base") + enc = _get_token_encoding() return max(4, len(enc.encode(payload)) + 4) except Exception: return max(4, len(payload) // 4 + 4) diff --git a/tests/utils/test_token_estimation.py b/tests/utils/test_token_estimation.py index c3257b31..254bf530 100644 --- a/tests/utils/test_token_estimation.py +++ b/tests/utils/test_token_estimation.py @@ -1,4 +1,7 @@ -from nanobot.utils.helpers import estimate_prompt_tokens_chain +import json + +from nanobot.utils import helpers +from nanobot.utils.helpers import estimate_prompt_tokens, estimate_prompt_tokens_chain class _NoCounterProvider: @@ -30,3 +33,67 @@ def test_estimate_prompt_tokens_chain_falls_back_when_provider_counter_fails() - assert tokens > 0 assert source == "tiktoken" + + +def test_estimate_prompt_tokens_caches_tools_encoding(monkeypatch) -> None: + helpers._get_token_encoding.cache_clear() + helpers._TOOLS_TOKEN_CACHE.clear() + + class FakeEncoding: + def __init__(self) -> None: + self.encoded: list[str] = [] + + def encode(self, text: str) -> list[int]: + self.encoded.append(text) + return list(range(max(1, len(text) // 4))) + + fake_encoding = FakeEncoding() + get_encoding_calls = 0 + + def fake_get_encoding(name: str) -> FakeEncoding: + nonlocal get_encoding_calls + assert name == "cl100k_base" + get_encoding_calls += 1 + return fake_encoding + + monkeypatch.setattr(helpers.tiktoken, "get_encoding", fake_get_encoding) + tools = [{"type": "function", "function": {"name": "demo", "description": "cached"}}] + messages = [{"role": "user", "content": "hello"}] + + first = estimate_prompt_tokens(messages, tools) + second = estimate_prompt_tokens(messages, tools) + + assert first == second + assert get_encoding_calls == 1 + rendered_tools = "\n" + json.dumps(tools, ensure_ascii=False) + assert fake_encoding.encoded.count(rendered_tools) == 1 + + +def test_estimate_prompt_tokens_recomputes_when_tool_items_change(monkeypatch) -> None: + helpers._get_token_encoding.cache_clear() + helpers._TOOLS_TOKEN_CACHE.clear() + + class FakeEncoding: + def __init__(self) -> None: + self.encoded: list[str] = [] + + def encode(self, text: str) -> list[int]: + self.encoded.append(text) + return list(range(max(1, len(text) // 4))) + + fake_encoding = FakeEncoding() + monkeypatch.setattr(helpers.tiktoken, "get_encoding", lambda _name: fake_encoding) + + tools = [{"type": "function", "function": {"name": "before"}}] + messages = [{"role": "user", "content": "hello"}] + estimate_prompt_tokens(messages, tools) + + tools[0] = {"type": "function", "function": {"name": "after"}} + estimate_prompt_tokens(messages, tools) + + before_tools = "\n" + json.dumps( + [{"type": "function", "function": {"name": "before"}}], ensure_ascii=False + ) + after_tools = "\n" + json.dumps(tools, ensure_ascii=False) + assert before_tools in fake_encoding.encoded + assert after_tools in fake_encoding.encoded