mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
e7008f64d7
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.
381 lines
13 KiB
Python
381 lines
13 KiB
Python
"""Tests for Strands SDK content block tokenization (#111).
|
|
|
|
Strands SDK sends content blocks without a "type" field:
|
|
{"text": "..."} instead of {"type": "text", "text": "..."}
|
|
{"toolUse": {...}} instead of {"type": "tool_use", ...}
|
|
{"toolResult": {...}} instead of {"type": "tool_result", ...}
|
|
|
|
The tokenizer must count these correctly.
|
|
"""
|
|
|
|
from headroom.tokenizers import get_tokenizer
|
|
|
|
|
|
def _get_counter():
|
|
return get_tokenizer("claude-sonnet-4-6")
|
|
|
|
|
|
class TestStrandsTextBlocks:
|
|
"""Strands text blocks: {"text": "..."} without "type" field."""
|
|
|
|
def test_strands_text_matches_anthropic_text(self):
|
|
"""Strands {"text": ...} should count same as Anthropic {"type": "text", "text": ...}."""
|
|
t = _get_counter()
|
|
text = "Hello world this is a test message " * 50
|
|
|
|
anthropic = [{"role": "user", "content": [{"type": "text", "text": text}]}]
|
|
strands = [{"role": "user", "content": [{"text": text}]}]
|
|
|
|
a = t.count_messages(anthropic)
|
|
s = t.count_messages(strands)
|
|
assert a == s, f"Anthropic={a}, Strands={s}"
|
|
|
|
def test_strands_text_matches_plain_string(self):
|
|
"""Strands text block should count same as plain string content."""
|
|
t = _get_counter()
|
|
text = "Some question " * 1000
|
|
|
|
plain = [{"role": "user", "content": text}]
|
|
strands = [{"role": "user", "content": [{"text": text}]}]
|
|
|
|
p = t.count_messages(plain)
|
|
s = t.count_messages(strands)
|
|
assert p == s, f"Plain={p}, Strands={s}"
|
|
|
|
def test_strands_multiple_text_blocks(self):
|
|
"""Multiple Strands text blocks should all be counted."""
|
|
t = _get_counter()
|
|
|
|
msg = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"text": "First block " * 100},
|
|
{"text": "Second block " * 100},
|
|
],
|
|
}
|
|
]
|
|
count = t.count_messages(msg)
|
|
|
|
# Should be roughly 2x a single block
|
|
single = [{"role": "user", "content": [{"text": "First block " * 100}]}]
|
|
single_count = t.count_messages(single)
|
|
|
|
assert count > single_count * 1.5, f"Multiple blocks={count}, single={single_count}"
|
|
|
|
def test_strands_system_message(self):
|
|
"""System message with Strands text blocks."""
|
|
t = _get_counter()
|
|
text = "You are a helpful assistant. " * 200
|
|
|
|
strands = [{"role": "system", "content": [{"text": text}]}]
|
|
plain = [{"role": "system", "content": text}]
|
|
|
|
s = t.count_messages(strands)
|
|
p = t.count_messages(plain)
|
|
assert s == p, f"Strands={s}, Plain={p}"
|
|
|
|
|
|
class TestStrandsToolBlocks:
|
|
"""Strands tool blocks: {"toolUse": {...}} and {"toolResult": {...}}."""
|
|
|
|
def test_strands_tool_use_counted(self):
|
|
"""Strands toolUse block should be counted, not zero."""
|
|
t = _get_counter()
|
|
|
|
msg = [
|
|
{
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"toolUse": {
|
|
"toolUseId": "t1",
|
|
"name": "read_file",
|
|
"input": {"path": "/src/main.py"},
|
|
}
|
|
}
|
|
],
|
|
}
|
|
]
|
|
count = t.count_messages(msg)
|
|
# Should include the tool name and input, not just message overhead
|
|
assert count > 15, f"toolUse count too low: {count}"
|
|
|
|
def test_strands_tool_result_counted(self):
|
|
"""Strands toolResult with nested text content should be counted."""
|
|
t = _get_counter()
|
|
big_content = "File contents here. " * 500
|
|
|
|
msg = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"toolResult": {
|
|
"toolUseId": "t1",
|
|
"content": [{"text": big_content}],
|
|
}
|
|
}
|
|
],
|
|
}
|
|
]
|
|
count = t.count_messages(msg)
|
|
|
|
# Should reflect the size of the content, not just overhead
|
|
plain_count = t.count_messages([{"role": "user", "content": big_content}])
|
|
assert count > plain_count * 0.5, (
|
|
f"toolResult count={count} should be close to plain={plain_count}"
|
|
)
|
|
|
|
def test_strands_tool_result_string_content(self):
|
|
"""Strands toolResult with string content."""
|
|
t = _get_counter()
|
|
|
|
msg = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"toolResult": {
|
|
"toolUseId": "t1",
|
|
"content": "Simple string result " * 100,
|
|
}
|
|
}
|
|
],
|
|
}
|
|
]
|
|
count = t.count_messages(msg)
|
|
assert count > 50, f"toolResult string count too low: {count}"
|
|
|
|
|
|
class TestMixedFormats:
|
|
"""Messages mixing Anthropic and Strands formats."""
|
|
|
|
def test_mixed_conversation(self):
|
|
"""Full conversation with mixed Strands and Anthropic blocks."""
|
|
t = _get_counter()
|
|
messages = [
|
|
# Strands system
|
|
{"role": "system", "content": [{"text": "You are helpful. " * 50}]},
|
|
# Strands user
|
|
{"role": "user", "content": [{"text": "Fix the bug in auth.py"}]},
|
|
# Strands assistant with toolUse
|
|
{
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"toolUse": {
|
|
"toolUseId": "t1",
|
|
"name": "read_file",
|
|
"input": {"path": "auth.py"},
|
|
}
|
|
}
|
|
],
|
|
},
|
|
# Strands tool result
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"toolResult": {
|
|
"toolUseId": "t1",
|
|
"content": [{"text": "def authenticate():\n pass\n" * 100}],
|
|
}
|
|
}
|
|
],
|
|
},
|
|
# Anthropic-style text (for comparison)
|
|
{"role": "assistant", "content": [{"type": "text", "text": "I found the issue."}]},
|
|
]
|
|
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}"
|
|
)
|