mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Count Strands reasoningContent, image, document, video tokens (#111 follow-up)
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.
This commit is contained in:
chopratejas
parent
c92a6b9fa1
commit
e7008f64d7
2 changed files with 269 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue