feat: add datrie dependency and implement LRUTrieCache for prefix caching

- Added datrie as a dependency in pyproject.toml.
- Introduced LRUTrieCache class for efficient caching of token sequence prefixes.
- Implemented methods for inserting and looking up cached prefixes.
- Added unit tests for LRUTrieCache to ensure functionality and eviction behavior.
- Updated server code to handle tool calls with type safety.

Author: bradhilton

pyproject.toml

@@ -61,6 +61,7 @@ tinker = [
     "torch>=2.8.0",
     "transformers>=4.55.2,<=4.57.3",
     "uvicorn>=0.35.0",
+    "datrie>=0.8.3",
 ]
 
 [project.scripts]

src/art/tinker/prefix_cache.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass
+import struct
+from typing import Sequence
+
+import datrie
+
+_TOKEN_BYTES = 4
+_TRIE_ALPHABET = "".join(chr(i) for i in range(256))
+
+
+@dataclass(frozen=True)
+class PrefixEntry:
+    rendered_len: int
+    raw_prefix: tuple[int, ...]
+
+
+class LRUTrieCache:
+    """LRU-bounded prefix trie for token sequence rewrites."""
+
+    def __init__(self, max_entries: int = 1024) -> None:
+        if max_entries <= 0:
+            raise ValueError("max_entries must be positive")
+        self._trie: datrie.Trie = datrie.Trie(_TRIE_ALPHABET)
+        self._lru: OrderedDict[str, None] = OrderedDict()
+        self._max_entries = max_entries
+
+    @staticmethod
+    def _encode_tokens(tokens: Sequence[int]) -> str:
+        packed = bytearray()
+        for token in tokens:
+            packed.extend(struct.pack(">I", token))
+        return packed.decode("latin-1")
+
+    def lookup(self, rendered_tokens: Sequence[int]) -> PrefixEntry | None:
+        key = self._encode_tokens(rendered_tokens)
+        match: tuple[str, PrefixEntry] | None = None
+        for item in self._trie.prefix_items(key):
+            match = item
+        if match is None:
+            return None
+        match_key, entry = match
+        self._lru.move_to_end(match_key)
+        return entry
+
+    def insert(self, rendered_prefix: Sequence[int], raw_prefix: Sequence[int]) -> None:
+        key = self._encode_tokens(rendered_prefix)
+        entry = PrefixEntry(
+            rendered_len=len(rendered_prefix), raw_prefix=tuple(raw_prefix)
+        )
+        self._trie[key] = entry
+        self._lru[key] = None
+        self._lru.move_to_end(key)
+        self._evict()
+
+    def _evict(self) -> None:
+        while len(self._lru) > self._max_entries:
+            old_key, _ = self._lru.popitem(last=False)
+            if old_key in self._trie:
+                del self._trie[old_key]

src/art/tinker/server.py

@@ -168,7 +168,7 @@ async def chat_completions(
                         message=ChatCompletionMessage(
                             content=openai_message.get("content") or None,
                             role="assistant",
-                            tool_calls=tool_calls,
+                            tool_calls=tool_calls,  # type: ignore
                         ),
                         logprobs=ChoiceLogprobs(
                             content=[

tests/unit/test_prefix_cache.py

@@ -0,0 +1,36 @@
+"""Tests for the LRUTrieCache prefix rewrite helper."""
+
+import pytest
+
+pytest.importorskip("datrie")
+
+from art.tinker.prefix_cache import LRUTrieCache
+
+
+class TestLRUTrieCache:
+    def test_longest_prefix_match(self) -> None:
+        cache = LRUTrieCache(max_entries=10)
+        cache.insert([1, 2], [10, 11])
+        cache.insert([1, 2, 3], [20, 21, 22])
+
+        entry = cache.lookup([1, 2, 3, 4])
+
+        assert entry is not None
+        assert entry.rendered_len == 3
+        assert entry.raw_prefix == (20, 21, 22)
+
+    def test_lru_eviction(self) -> None:
+        cache = LRUTrieCache(max_entries=2)
+        cache.insert([1], [10])
+        cache.insert([2], [20])
+
+        assert cache.lookup([1, 99]) is not None
+
+        cache.insert([3], [30])
+
+        assert cache.lookup([2, 0]) is None
+        assert cache.lookup([1, 0]) is not None
+
+    def test_invalid_size(self) -> None:
+        with pytest.raises(ValueError):
+            LRUTrieCache(max_entries=0)