Hermes-agent

tests/agent/test_compressor_image_tokens.py

@@ -0,0 +1,140 @@
+"""Tests for image-token accounting in the context compressor.
+
+Covers the native-image-routing PR's companion change: the compressor's
+multimodal message length counter now charges ~1600 tokens per attached
+image part instead of 0, so tail-cut / prune decisions are accurate for
+creative workflows that iterate on images across many turns.
+"""
+
+from __future__ import annotations
+
+
+from agent.context_compressor import (
+    _CHARS_PER_TOKEN,
+    _IMAGE_CHAR_EQUIVALENT,
+    _IMAGE_TOKEN_ESTIMATE,
+    _content_length_for_budget,
+)
+
+
+class TestContentLengthForBudget:
+    def test_plain_string(self):
+        assert _content_length_for_budget("hello world") == 11
+
+    def test_empty_string(self):
+        assert _content_length_for_budget("") == 0
+
+    def test_none_coerces_to_zero(self):
+        assert _content_length_for_budget(None) == 0
+
+    def test_text_only_list(self):
+        content = [
+            {"type": "text", "text": "first"},
+            {"type": "text", "text": "second"},
+        ]
+        assert _content_length_for_budget(content) == 5 + 6
+
+    def test_single_image_part_charges_fixed_budget(self):
+        content = [
+            {"type": "text", "text": "look"},
+            {"type": "image_url", "image_url": {"url": "data:image/png;base64,XXXX"}},
+        ]
+        # 4 chars of text + 1 image at fixed char-equivalent
+        assert _content_length_for_budget(content) == 4 + _IMAGE_CHAR_EQUIVALENT
+
+    def test_image_url_raw_base64_is_not_counted_as_chars(self):
+        """A 1MB base64 blob inside an image_url must NOT inflate token count.
+
+        The flat image estimate is what the provider actually bills; the raw
+        base64 is transport payload, not context tokens.
+        """
+        huge_url = "data:image/png;base64," + ("A" * 1_000_000)
+        content = [
+            {"type": "image_url", "image_url": {"url": huge_url}},
+        ]
+        # Exactly one image's worth, not 1M + something.
+        assert _content_length_for_budget(content) == _IMAGE_CHAR_EQUIVALENT
+
+    def test_multiple_image_parts(self):
+        content = [
+            {"type": "text", "text": "compare"},
+            {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
+            {"type": "image_url", "image_url": {"url": "data:image/png;base64,BBB"}},
+            {"type": "image_url", "image_url": {"url": "data:image/png;base64,CCC"}},
+        ]
+        assert _content_length_for_budget(content) == 7 + 3 * _IMAGE_CHAR_EQUIVALENT
+
+    def test_openai_responses_input_image_shape(self):
+        """Responses API uses type=input_image with top-level image_url string."""
+        content = [
+            {"type": "input_text", "text": "hey"},
+            {"type": "input_image", "image_url": "data:image/png;base64,XX"},
+        ]
+        # input_text has .text "hey" (3 chars) + 1 image
+        assert _content_length_for_budget(content) == 3 + _IMAGE_CHAR_EQUIVALENT
+
+    def test_anthropic_native_image_shape(self):
+        """Anthropic native shape: {type: image, source: {...}}."""
+        content = [
+            {"type": "text", "text": "hi"},
+            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "XX"}},
+        ]
+        assert _content_length_for_budget(content) == 2 + _IMAGE_CHAR_EQUIVALENT
+
+    def test_bare_string_part_in_list(self):
+        """Older code paths sometimes produce mixed list-of-strings content."""
+        content = ["hello", {"type": "text", "text": "world"}]
+        assert _content_length_for_budget(content) == 5 + 5
+
+    def test_image_estimate_constant_is_reasonable(self):
+        """Sanity-check the estimate aligns with real provider billing.
+
+        Anthropic ≈ width*height/750 → ~1600 for 1000×1200.
+        OpenAI GPT-4o high-detail 2048×2048 ≈ 1445.
+        Gemini 258/tile × 6 tiles for a 2048×2048 ≈ 1548.
+        Anything in the 800-2000 range is defensible. Enforce bounds so an
+        accidental edit doesn't drop it to e.g. 16.
+        """
+        assert 800 <= _IMAGE_TOKEN_ESTIMATE <= 2500
+        assert _IMAGE_CHAR_EQUIVALENT == _IMAGE_TOKEN_ESTIMATE * _CHARS_PER_TOKEN
+
+
+class TestTokenBudgetWithImages:
+    """Integration: the compressor's tail-cut decision now respects image cost."""
+
+    def test_image_heavy_turns_count_toward_budget(self):
+        """A tail with 5 image-bearing turns should blow past a 5K token budget."""
+        from agent.context_compressor import ContextCompressor
+
+        # Minimal compressor fixture — just enough to call _find_tail_cut_by_tokens
+        cc = object.__new__(ContextCompressor)
+        cc.tail_token_budget = 5000
+
+        # Build 10 messages: 5 with images, 5 with short text. Without the
+        # image-tokens fix, the compressor would think all 10 fit in 5K and
+        # protect them all. With the fix, images alone cost 5 × 1600 = 8K,
+        # so the tail should be trimmed.
+        messages = [{"role": "system", "content": "sys"}]
+        for i in range(5):
+            messages.append({
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": f"turn {i}"},
+                    {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
+                ],
+            })
+            messages.append({
+                "role": "assistant",
+                "content": f"response {i}",
+            })
+
+        cut = cc._find_tail_cut_by_tokens(messages, head_end=0, token_budget=5000)
+
+        # Budget is 5K, soft ceiling 7.5K. 5 images alone = 8000 image-tokens.
+        # Walking backward, the compressor should stop before including all 5.
+        # Exact cut depends on text lengths and min_tail, but it MUST be > 1
+        # (at least some head-side messages should be compressible).
+        assert cut > 1, (
+            f"Expected image-heavy tail to be trimmed; compressor placed cut at "
+            f"{cut} out of {len(messages)} (image tokens were likely ignored)."
+        )