diff --git a/headroom/cache/compression_cache.py b/headroom/cache/compression_cache.py index d45082144..3659ef8c4 100644 --- a/headroom/cache/compression_cache.py +++ b/headroom/cache/compression_cache.py @@ -6,11 +6,15 @@ 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: @@ -20,6 +24,54 @@ class _CacheEntry: 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. @@ -84,3 +136,71 @@ class CompressionCache: else: raw = content return hashlib.sha256(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 self.get_compressed(h) is None: + 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) - len(comp_content) + self.store_compressed(h, comp_content, tokens_saved=max(tokens_saved, 0)) diff --git a/tests/test_compression_cache.py b/tests/test_compression_cache.py index aa9e710ca..1e7f5c45d 100644 --- a/tests/test_compression_cache.py +++ b/tests/test_compression_cache.py @@ -107,3 +107,177 @@ class TestCompressionCache: def test_content_hash_string_length(self) -> None: h = CompressionCache.content_hash("test") assert len(h) == 16 + + +class TestCompressionCacheFrozenCount: + def test_empty_cache_returns_zero(self, cache: CompressionCache) -> None: + assert cache.compute_frozen_count([]) == 0 + + def test_user_assistant_always_stable(self, cache: CompressionCache) -> None: + messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi there"}, + {"role": "user", "content": "how are you"}, + ] + assert cache.compute_frozen_count(messages) == 3 + + def test_tool_result_with_cache_hit_is_stable(self, cache: CompressionCache) -> None: + tool_content = "tool output data" + h = CompressionCache.content_hash(tool_content) + cache.store_compressed(h, "compressed tool output", tokens_saved=5) + + messages = [ + {"role": "user", "content": "do something"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "my_tool", "input": {}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}], + }, + ] + assert cache.compute_frozen_count(messages) == 3 + + def test_tool_result_cache_miss_stops_frozen(self, cache: CompressionCache) -> None: + messages = [ + {"role": "user", "content": "hello"}, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "uncached stuff"} + ], + }, + {"role": "user", "content": "follow up"}, + ] + assert cache.compute_frozen_count(messages) == 1 + + def test_frozen_count_with_dropped_messages(self, cache: CompressionCache) -> None: + cached_content = "cached tool output" + h = CompressionCache.content_hash(cached_content) + cache.store_compressed(h, "compressed", tokens_saved=3) + + messages = [ + {"role": "user", "content": "start"}, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": cached_content} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t2", "content": "not cached"}], + }, + ] + assert cache.compute_frozen_count(messages) == 2 + + +class TestCompressionCacheApplyAndUpdate: + def test_apply_cached_swaps_tool_results(self, cache: CompressionCache) -> None: + original_content = "big tool output" + h = CompressionCache.content_hash(original_content) + cache.store_compressed(h, "small output", tokens_saved=5) + + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": original_content} + ], + }, + ] + result = cache.apply_cached(messages) + assert result[1]["content"][0]["content"] == "small output" + + def test_apply_cached_preserves_uncached_messages(self, cache: CompressionCache) -> None: + messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + result = cache.apply_cached(messages) + assert result[0] is messages[0] + assert result[1] is messages[1] + + def test_apply_cached_never_adds_messages(self, cache: CompressionCache) -> None: + # Store something in cache that doesn't correspond to any message + cache.store_compressed("orphan_hash", "orphan_value", tokens_saved=1) + + messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + result = cache.apply_cached(messages) + assert len(result) == len(messages) + + def test_update_from_result_caches_changes(self, cache: CompressionCache) -> None: + originals = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "original output"} + ], + }, + ] + compressed = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "compressed output"} + ], + }, + ] + cache.update_from_result(originals, compressed) + + h = CompressionCache.content_hash("original output") + assert cache.get_compressed(h) == "compressed output" + + def test_update_from_result_ignores_unchanged(self, cache: CompressionCache) -> None: + originals = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "same content"} + ], + }, + ] + compressed = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "same content"} + ], + }, + ] + cache.update_from_result(originals, compressed) + h = CompressionCache.content_hash("same content") + assert cache.get_compressed(h) is None + + def test_apply_does_not_modify_original_messages(self, cache: CompressionCache) -> None: + original_content = "big tool output" + h = CompressionCache.content_hash(original_content) + cache.store_compressed(h, "small output", tokens_saved=5) + + msg = { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": original_content}], + } + messages = [msg] + cache.apply_cached(messages) + + # Original must be untouched + assert msg["content"][0]["content"] == original_content + + def test_openai_format_tool_result(self, cache: CompressionCache) -> None: + original_content = "openai tool output" + h = CompressionCache.content_hash(original_content) + cache.store_compressed(h, "compressed openai", tokens_saved=4) + + messages = [ + {"role": "tool", "tool_call_id": "tc1", "content": original_content}, + ] + result = cache.apply_cached(messages) + assert result[0]["content"] == "compressed openai" + # Original untouched + assert messages[0]["content"] == original_content