perf(tokens): cache tool schema estimates

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)

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