chopratejas--headroom/tests/test_strands_tokenizer.py

"""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}"
        )