feat: add count_token method to model with naive estimation using tiktoken (#2031)

Co-authored-by: opieter-aws <opieter@amazon.com>

Author: lizradway

src/strands/models/model.py

@@ -1,16 +1,19 @@
 """Abstract base class for Agent model providers."""
 
 import abc
+import functools
+import json
 import logging
-from collections.abc import AsyncGenerator, AsyncIterable
+import math
+from collections.abc import AsyncGenerator, AsyncIterable, Callable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeVar
 
 from pydantic import BaseModel
 
 from ..hooks.events import AfterInvocationEvent
 from ..plugins.plugin import Plugin
-from ..types.content import Messages, SystemContentBlock
+from ..types.content import ContentBlock, Messages, SystemContentBlock
 from ..types.streaming import StreamEvent
 from ..types.tools import ToolChoice, ToolSpec
 
@@ -21,6 +24,164 @@
 
 T = TypeVar("T", bound=BaseModel)
 
+_DEFAULT_ENCODING = "cl100k_base"
+
+
+def _heuristic_estimate_text(text: str) -> int:
+    """Estimate token count from text using characters / 4 heuristic."""
+    return math.ceil(len(text) / 4)
+
+
+def _heuristic_estimate_json(obj: Any) -> int:
+    """Estimate token count from a JSON-serializable object using characters / 2 heuristic."""
+    try:
+        return math.ceil(len(json.dumps(obj)) / 2)
+    except (TypeError, ValueError):
+        return 0
+
+
+@functools.lru_cache(maxsize=1)
+def _get_encoding() -> Any:
+    """Get the default tiktoken encoding, caching to avoid repeated lookups.
+
+    Returns:
+        The tiktoken encoding, or None if tiktoken is not installed.
+    """
+    try:
+        import tiktoken
+
+        return tiktoken.get_encoding(_DEFAULT_ENCODING)
+    except ImportError:
+        logger.debug("tiktoken not available, falling back to heuristic token estimation")
+        return None
+
+
+def _count_content_block_tokens(
+    block: ContentBlock, count_text: Callable[[str], int], count_json: Callable[[Any], int]
+) -> int:
+    """Count tokens for a single content block.
+
+    Args:
+        block: The content block to count tokens for.
+        count_text: Function that returns token count for a text string.
+        count_json: Function that returns token count for a JSON-serializable object.
+    """
+    total = 0
+
+    if "text" in block:
+        total += count_text(block["text"])
+
+    if "toolUse" in block:
+        tool_use = block["toolUse"]
+        total += count_text(tool_use.get("name", ""))
+        total += count_json(tool_use.get("input", {}))
+
+    if "toolResult" in block:
+        tool_result = block["toolResult"]
+        for item in tool_result.get("content", []):
+            if "text" in item:
+                total += count_text(item["text"])
+
+    if "reasoningContent" in block:
+        reasoning = block["reasoningContent"]
+        if "reasoningText" in reasoning:
+            reasoning_text = reasoning["reasoningText"]
+            if "text" in reasoning_text:
+                total += count_text(reasoning_text["text"])
+
+    if "guardContent" in block:
+        guard = block["guardContent"]
+        if "text" in guard and "text" in guard["text"]:
+            total += count_text(guard["text"]["text"])
+
+    if "citationsContent" in block:
+        citations = block["citationsContent"]
+        if "content" in citations:
+            for citation_item in citations["content"]:
+                if "text" in citation_item:
+                    total += count_text(citation_item["text"])
+
+    return total
+
+
+def _estimate_tokens_with_tiktoken(
+    messages: Messages,
+    tool_specs: list[ToolSpec] | None = None,
+    system_prompt: str | None = None,
+    system_prompt_content: list[SystemContentBlock] | None = None,
+) -> int:
+    """Estimate tokens by serializing messages/tools to text and counting with tiktoken.
+
+    This is a best-effort fallback for providers that don't expose native counting.
+    Accuracy varies by model but is sufficient for threshold-based decisions.
+
+    Raises:
+        ImportError: If tiktoken is not installed.
+    """
+    encoding = _get_encoding()
+    if encoding is None:
+        raise ImportError("tiktoken is not available")
+
+    def count_text(text: str) -> int:
+        return len(encoding.encode(text))
+
+    def count_json(obj: Any) -> int:
+        try:
+            return len(encoding.encode(json.dumps(obj)))
+        except (TypeError, ValueError):
+            return 0
+
+    total = 0
+
+    # Prefer system_prompt_content (structured) over system_prompt (plain string) to avoid double-counting,
+    # since providers wrap system_prompt into system_prompt_content when both are provided.
+    if system_prompt_content:
+        for block in system_prompt_content:
+            if "text" in block:
+                total += count_text(block["text"])
+    elif system_prompt:
+        total += count_text(system_prompt)
+
+    for message in messages:
+        for block in message["content"]:
+            total += _count_content_block_tokens(block, count_text, count_json)
+
+    if tool_specs:
+        for spec in tool_specs:
+            total += count_json(spec)
+
+    return total
+
+
+def _estimate_tokens_with_heuristic(
+    messages: Messages,
+    tool_specs: list[ToolSpec] | None = None,
+    system_prompt: str | None = None,
+    system_prompt_content: list[SystemContentBlock] | None = None,
+) -> int:
+    """Estimate tokens using character-based heuristics (text: chars/4, JSON: chars/2).
+
+    Dependency-free fallback when tiktoken is not installed.
+    """
+    total = 0
+
+    if system_prompt_content:
+        for block in system_prompt_content:
+            if "text" in block:
+                total += _heuristic_estimate_text(block["text"])
+    elif system_prompt:
+        total += _heuristic_estimate_text(system_prompt)
+
+    for message in messages:
+        for block in message["content"]:
+            total += _count_content_block_tokens(block, _heuristic_estimate_text, _heuristic_estimate_json)
+
+    if tool_specs:
+        for spec in tool_specs:
+            total += _heuristic_estimate_json(spec)
+
+    return total
+
 
 class BaseModelConfig(TypedDict, total=False):
     """Base configuration shared by all model providers.
@@ -151,6 +312,37 @@ def stream(
         """
         pass
 
+    async def count_tokens(
+        self,
+        messages: Messages,
+        tool_specs: list[ToolSpec] | None = None,
+        system_prompt: str | None = None,
+        system_prompt_content: list[SystemContentBlock] | None = None,
+    ) -> int:
+        """Estimate token count for the given input before sending to the model.
+
+        Used for proactive context management (e.g., triggering compression at a threshold).
+        Uses tiktoken's cl100k_base encoding when available, otherwise falls back to a
+        heuristic (characters / 4 for text, characters / 2 for JSON). Accuracy varies by
+        model provider. Not intended for billing or precise quota calculations.
+
+        Subclasses may override this method to provide model-specific token counting
+        using native APIs for improved accuracy.
+
+        Args:
+            messages: List of message objects to estimate tokens for.
+            tool_specs: List of tool specifications to include in the estimate.
+            system_prompt: Plain string system prompt. Ignored if system_prompt_content is provided.
+            system_prompt_content: Structured system prompt content blocks. Takes priority over system_prompt.
+
+        Returns:
+            Estimated total input tokens.
+        """
+        try:
+            return _estimate_tokens_with_tiktoken(messages, tool_specs, system_prompt, system_prompt_content)
+        except ImportError:
+            return _estimate_tokens_with_heuristic(messages, tool_specs, system_prompt, system_prompt_content)
+
 
 class _ModelPlugin(Plugin):
     """Plugin that manages model-related lifecycle hooks."""