Spaces:
Build error
Build error
File size: 7,713 Bytes
29928b8 60741b5 29928b8 60741b5 29928b8 60741b5 29928b8 60741b5 29928b8 78d41df 60741b5 9b5e783 60741b5 9b5e783 60741b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | """Content-addressed compression cache with LRU eviction.
Used in "token headroom mode" to avoid re-compressing messages across turns.
Maps original content hashes to their compressed versions.
"""
from __future__ import annotations
import copy
import hashlib
import json
import logging
from collections import OrderedDict
from dataclasses import dataclass
logger = logging.getLogger(__name__)
@dataclass
class _CacheEntry:
"""Internal cache entry storing compressed text and metadata."""
compressed: str
tokens_saved: int
def _is_tool_result_message(msg: dict) -> bool:
"""Check if a message is a tool result in Anthropic or OpenAI format."""
# OpenAI format: role="tool"
if msg.get("role") == "tool":
return True
# Anthropic format: role="user" with content list containing tool_result blocks
content = msg.get("content")
if isinstance(content, list):
return any(
isinstance(block, dict) and block.get("type") == "tool_result" for block in content
)
return False
def _extract_tool_result_content(msg: dict) -> str | None:
"""Extract text content from a tool result message (both formats)."""
# OpenAI format
if msg.get("role") == "tool":
content = msg.get("content")
return content if isinstance(content, str) else None
# Anthropic format
content = msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
inner = block.get("content")
if isinstance(inner, str):
return inner
return None
def _swap_tool_result_content(msg: dict, new_content: str) -> dict:
"""Deep copy msg and replace tool result content with new_content."""
new_msg = copy.deepcopy(msg)
# OpenAI format
if new_msg.get("role") == "tool":
new_msg["content"] = new_content
return new_msg
# Anthropic format
content = new_msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
block["content"] = new_content
break
return new_msg
class CompressionCache:
"""Content-addressed cache mapping content hashes to compressed versions.
Uses an OrderedDict for O(1) LRU eviction. Entries are evicted oldest-first
when the cache exceeds max_entries.
"""
def __init__(self, max_entries: int = 10000) -> None:
self.max_entries = max_entries
self._cache: OrderedDict[str, _CacheEntry] = OrderedDict()
self._hits: int = 0
self._misses: int = 0
self._total_tokens_saved: int = 0
def get_compressed(self, hash: str) -> str | None:
"""Retrieve compressed content by hash, refreshing LRU position on hit."""
entry = self._cache.get(hash)
if entry is None:
self._misses += 1
return None
self._hits += 1
self._cache.move_to_end(hash)
return entry.compressed
def store_compressed(self, hash: str, compressed: str, tokens_saved: int) -> None:
"""Store a compressed version keyed by content hash.
If the hash already exists, the entry is overwritten and moved to the
end (most recently used). When the cache exceeds max_entries, the oldest
entry is evicted.
"""
if hash in self._cache:
old_entry = self._cache[hash]
self._total_tokens_saved -= old_entry.tokens_saved
del self._cache[hash]
self._cache[hash] = _CacheEntry(compressed=compressed, tokens_saved=tokens_saved)
self._total_tokens_saved += tokens_saved
while len(self._cache) > self.max_entries:
_, evicted = self._cache.popitem(last=False)
self._total_tokens_saved -= evicted.tokens_saved
def get_stats(self) -> dict:
"""Return cache statistics."""
return {
"entries": len(self._cache),
"hits": self._hits,
"misses": self._misses,
"tokens_saved": self._total_tokens_saved,
}
@staticmethod
def content_hash(content: str | list) -> str:
"""Compute a truncated SHA-256 hash for string or list content.
For list content (Anthropic-format messages with type/text/content fields),
the list is JSON-serialized with sorted keys for deterministic hashing.
"""
if isinstance(content, list):
raw = json.dumps(content, sort_keys=True, ensure_ascii=False)
else:
raw = content
return hashlib.md5(raw.encode("utf-8")).hexdigest()[:16]
def compute_frozen_count(self, messages: list[dict]) -> int:
"""Count consecutive stable messages from the start.
A message is stable if it is a plain user/assistant/system message,
an assistant message with tool_use blocks, or a tool_result whose
content hash is already in the cache. The first unstable tool_result
(cache miss) stops the count.
"""
count = 0
for msg in messages:
if _is_tool_result_message(msg):
content = _extract_tool_result_content(msg)
if content is not None:
h = self.content_hash(content)
if h not in self._cache:
break
else:
# tool_result with non-string content; treat as unstable
break
# Regular user/assistant/system messages and assistant+tool_use
# are always stable — fall through.
count += 1
return count
def apply_cached(self, messages: list[dict]) -> list[dict]:
"""Return a new list with cached compressions swapped into tool results.
Never mutates the input list or any message dict within it.
Output always has the same length as input.
"""
result: list[dict] = []
for msg in messages:
if _is_tool_result_message(msg):
content = _extract_tool_result_content(msg)
if content is not None:
h = self.content_hash(content)
compressed = self.get_compressed(h)
if compressed is not None:
result.append(_swap_tool_result_content(msg, compressed))
continue
result.append(msg)
return result
def update_from_result(self, originals: list[dict], compressed: list[dict]) -> None:
"""Cache new compressions by comparing original and compressed messages.
Index-aligned: for each position, if both are tool results and the
content differs, store the mapping original_hash -> compressed_content.
"""
if len(originals) != len(compressed):
logger.warning(
"update_from_result: length mismatch (originals=%d, compressed=%d), skipping",
len(originals),
len(compressed),
)
return
for orig, comp in zip(originals, compressed):
orig_content = _extract_tool_result_content(orig)
comp_content = _extract_tool_result_content(comp)
if orig_content is None or comp_content is None:
continue
if orig_content == comp_content:
continue
h = self.content_hash(orig_content)
tokens_saved = len(orig_content) // 4 - len(comp_content) // 4
self.store_compressed(h, comp_content, tokens_saved=max(tokens_saved, 0))
|