From e7008f64d728d655602f787eff670e950c83d6b4 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Wed, 8 Apr 2026 17:31:00 -0700 Subject: [PATCH] Count Strands reasoningContent, image, document, video tokens (#111 follow-up) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit reasoningContent: exact counting via count_text() — pure text, no estimation image: decode with Pillow for (w*h)/750 formula, fallback by byte size document: ~1500 tokens/page heuristic (3KB/page of PDF) video: ~1000 tokens/frame heuristic (30KB/frame) Text content (reasoning, text, toolResult) uses exact tokenization. Binary content (image, document, video) uses provider formula or size-based estimates — accurate counting requires content extraction that only the provider can do. 14 tests covering all Strands content block types. --- headroom/tokenizers/base.py | 80 ++++++++++++++ tests/test_strands_tokenizer.py | 189 ++++++++++++++++++++++++++++++++ 2 files changed, 269 insertions(+) diff --git a/headroom/tokenizers/base.py b/headroom/tokenizers/base.py index c8f239367..c9f015bbc 100644 --- a/headroom/tokenizers/base.py +++ b/headroom/tokenizers/base.py @@ -175,6 +175,47 @@ class BaseTokenizer(ABC): total += self._count_content_parts(tr_content) else: total += self.count_text(json.dumps(tr_content)) + elif not part_type and "reasoningContent" in part: + # Strands SDK reasoning: {"reasoningContent": {"reasoningText": {"text": "..."}}} + # This is actual text — count it precisely. + reasoning = part["reasoningContent"] + reasoning_text = reasoning.get("reasoningText", {}) + if isinstance(reasoning_text, dict): + total += self.count_text(reasoning_text.get("text", "")) + elif isinstance(reasoning_text, str): + total += self.count_text(reasoning_text) + elif not part_type and "document" in part: + # Strands SDK document: {"document": {"source": {"bytes": ...}}} + # Provider internally extracts text from PDF/DOCX then tokenizes. + # Accurate counting would require a PDF parser — instead we use + # the Anthropic documented estimate of ~1500 tokens per page, + # with ~3KB of PDF per page as a rough heuristic. + doc = part["document"] + source = doc.get("source", {}) + doc_bytes = source.get("bytes", b"") + if isinstance(doc_bytes, bytes | bytearray): + estimated_pages = max(1, len(doc_bytes) // 3000) + total += estimated_pages * 1500 + else: + total += self.count_text(str(doc_bytes)) + elif not part_type and "image" in part: + # Strands SDK image: {"image": {"source": {"bytes": ...}}} + # Anthropic formula: tokens = (width * height) / 750. + # Decode with Pillow for exact count; fall back to estimate. + total += self._estimate_image_tokens(part["image"]) + elif not part_type and "video" in part: + # Strands SDK video: provider samples ~1 fps, each frame costs + # image tokens. We can't decode frames without heavy deps, so + # estimate from byte size assuming ~30KB per frame, ~1000 tokens + # per frame (average image). + vid = part["video"] + source = vid.get("source", {}) + vid_bytes = source.get("bytes", b"") + if isinstance(vid_bytes, bytes | bytearray): + frames = max(1, len(vid_bytes) // 30000) + total += frames * 1000 + else: + total += 3200 else: # Unknown type - estimate from JSON total += self.count_text(json.dumps(part)) @@ -183,6 +224,45 @@ class BaseTokenizer(ABC): return total + @staticmethod + def _estimate_image_tokens(image_data: dict[str, Any]) -> int: + """Estimate tokens for an image using Anthropic's formula: (w*h)/750. + + Tries to decode dimensions with Pillow. Falls back to a conservative + estimate based on byte size. + """ + source = image_data.get("source", {}) + img_bytes = source.get("bytes", b"") + + if isinstance(img_bytes, bytes | bytearray) and len(img_bytes) > 0: + try: + import io + + from PIL import Image + + img = Image.open(io.BytesIO(img_bytes)) + w, h = img.size + # Anthropic resizes to fit 1568x1568 max + max_dim = 1568 + if w > max_dim or h > max_dim: + scale = max_dim / max(w, h) + w, h = int(w * scale), int(h * scale) + return max(100, (w * h) // 750) + except Exception: + pass + + # Fallback: estimate from byte size. + # Typical screenshot: ~200KB ≈ 1200x800 ≈ 1280 tokens + if isinstance(img_bytes, bytes | bytearray): + size_kb = len(img_bytes) / 1024 + if size_kb < 50: + return 400 # Small icon/thumbnail + if size_kb < 500: + return 1200 # Typical screenshot + return 1600 # Large/high-res image + + return 1200 # Default estimate + def _count_tool_calls(self, tool_calls: list[dict[str, Any]]) -> int: """Count tokens in tool calls.""" total = 0 diff --git a/tests/test_strands_tokenizer.py b/tests/test_strands_tokenizer.py index 0bdb78b33..f80f0dc69 100644 --- a/tests/test_strands_tokenizer.py +++ b/tests/test_strands_tokenizer.py @@ -190,3 +190,192 @@ class TestMixedFormats: count = t.count_messages(messages) # Should be substantial — the tool result alone is ~700 tokens assert count > 500, f"Mixed conversation count too low: {count}" + + +class TestStrandsReasoningContent: + """Strands reasoning blocks: {"reasoningContent": {"reasoningText": {"text": "..."}}}.""" + + def test_reasoning_text_counted_as_text(self): + """reasoningContent text should be counted with count_text, not estimated.""" + t = _get_counter() + reasoning = "Let me think step by step about this problem. " * 100 + + # Strands format + msg_strands = [ + { + "role": "assistant", + "content": [{"reasoningContent": {"reasoningText": {"text": reasoning}}}], + } + ] + + # Equivalent plain text for comparison + msg_plain = [{"role": "assistant", "content": reasoning}] + + s = t.count_messages(msg_strands) + p = t.count_messages(msg_plain) + assert s == p, f"Reasoning={s} should equal plain text={p}" + + def test_reasoning_plus_text_both_counted(self): + """Message with both reasoning and text blocks.""" + t = _get_counter() + msg = [ + { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": "thinking " * 200}}}, + {"text": "Here is my answer " * 50}, + ], + } + ] + count = t.count_messages(msg) + # Should be substantial — both blocks counted + assert count > 200, f"Combined reasoning+text too low: {count}" + + +class TestStrandsMediaContent: + """Strands image, document, video blocks.""" + + def test_image_not_zero(self): + """Image block should have nonzero token count.""" + t = _get_counter() + msg = [ + { + "role": "user", + "content": [{"image": {"format": "png", "source": {"bytes": b"x" * 50000}}}], + } + ] + count = t.count_messages(msg) + assert count > 100, f"Image count too low: {count}" + + def test_document_not_zero(self): + """Document block should have nonzero token count.""" + t = _get_counter() + msg = [ + { + "role": "user", + "content": [ + { + "document": { + "format": "pdf", + "name": "report.pdf", + "source": {"bytes": b"x" * 30000}, + } + } + ], + } + ] + count = t.count_messages(msg) + assert count > 1000, f"Document count too low: {count}" + + def test_video_not_zero(self): + """Video block should have nonzero token count.""" + t = _get_counter() + msg = [ + { + "role": "user", + "content": [{"video": {"format": "mp4", "source": {"bytes": b"x" * 300000}}}], + } + ] + count = t.count_messages(msg) + assert count > 1000, f"Video count too low: {count}" + + +class TestStrandsFullConversation: + """End-to-end conversation with all Strands content types.""" + + def test_agent_conversation_with_reasoning_and_tools(self): + """Realistic Strands agent conversation.""" + t = _get_counter() + messages = [ + {"role": "user", "content": [{"text": "Analyze this code and fix the bug"}]}, + { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": {"text": "Let me examine the code carefully. " * 50} + } + }, + { + "toolUse": { + "toolUseId": "t1", + "name": "read_file", + "input": {"path": "main.py"}, + } + }, + ], + }, + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [ + { + "text": "def process():\n data = fetch()\n return transform(data)\n" + * 50 + } + ], + } + }, + ], + }, + { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": {"text": "The bug is in the transform function. " * 30} + } + }, + {"text": "I found the issue. The transform function doesn't handle None."}, + ], + }, + ] + count = t.count_messages(messages) + # Reasoning + tool result + text = should be substantial + assert count > 500, f"Full conversation too low: {count}" + + # Verify reasoning contributes meaningfully + no_reasoning = [ + {"role": "user", "content": [{"text": "Analyze this code"}]}, + { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "t1", + "name": "read_file", + "input": {"path": "main.py"}, + } + }, + ], + }, + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [ + { + "text": "def process():\n data = fetch()\n return transform(data)\n" + * 50 + } + ], + } + }, + ], + }, + { + "role": "assistant", + "content": [ + {"text": "I found the issue."}, + ], + }, + ] + count_no_reasoning = t.count_messages(no_reasoning) + assert count > count_no_reasoning + 100, ( + f"Reasoning should add significant tokens: with={count}, without={count_no_reasoning}" + )