headroom-ai 0.2.13__py3-none-any.whl

headroom/tokenizer.py

@@ -0,0 +1,80 @@
+"""Token counting wrapper for Headroom SDK.
+
+This module provides a unified interface for token counting that
+delegates to provider-specific implementations.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .providers.base import TokenCounter
+
+
+class Tokenizer:
+    """
+    Token counting wrapper with model awareness.
+
+    This class wraps a provider-specific TokenCounter to provide
+    a consistent interface throughout the Headroom SDK.
+    """
+
+    def __init__(self, token_counter: TokenCounter, model: str = ""):
+        """
+        Initialize tokenizer with a provider's token counter.
+
+        Args:
+            token_counter: Provider-specific token counter.
+            model: Model name (for reference only).
+        """
+        self._counter = token_counter
+        self.model = model
+
+    def count_text(self, text: str) -> int:
+        """Count tokens in text."""
+        return self._counter.count_text(text)
+
+    def count_message(self, message: dict[str, Any]) -> int:
+        """Count tokens in a message."""
+        return self._counter.count_message(message)
+
+    def count_messages(self, messages: list[dict[str, Any]]) -> int:
+        """Count tokens in a list of messages."""
+        return self._counter.count_messages(messages)
+
+    @property
+    def available(self) -> bool:
+        """Whether token counting is available."""
+        return self._counter is not None
+
+
+# Convenience functions that require a token counter
+def count_tokens_text(text: str, token_counter: TokenCounter) -> int:
+    """
+    Count tokens in a text string.
+
+    Args:
+        text: The text to count tokens for.
+        token_counter: Provider-specific token counter.
+
+    Returns:
+        Token count.
+    """
+    return token_counter.count_text(text)
+
+
+def count_tokens_messages(
+    messages: list[dict[str, Any]],
+    token_counter: TokenCounter,
+) -> int:
+    """
+    Count total tokens for a list of messages.
+
+    Args:
+        messages: List of message dicts.
+        token_counter: Provider-specific token counter.
+
+    Returns:
+        Total token count.
+    """
+    return token_counter.count_messages(messages)

headroom/tokenizers/__init__.py

@@ -0,0 +1,75 @@
+"""Pluggable tokenizer system for universal LLM support.
+
+This module provides a registry-based tokenizer system that supports
+multiple backends:
+
+1. tiktoken - OpenAI models (GPT-3.5, GPT-4, GPT-4o)
+2. HuggingFace - Open models (Llama, Mistral, Falcon, etc.)
+3. Anthropic - Claude models (via SDK or estimation)
+4. Estimation - Fallback for unknown models
+
+Usage:
+    from headroom.tokenizers import TokenizerRegistry, get_tokenizer
+
+    # Auto-detect tokenizer from model name
+    tokenizer = get_tokenizer("gpt-4o")
+    tokens = tokenizer.count_text("Hello, world!")
+
+    # Get tokenizer for specific backend
+    tokenizer = get_tokenizer("llama-3-8b", backend="huggingface")
+
+    # Register custom tokenizer
+    TokenizerRegistry.register("my-model", my_tokenizer)
+"""
+
+from .base import BaseTokenizer, TokenCounter
+from .estimator import CharacterCounter, EstimatingTokenCounter
+from .registry import (
+    TokenizerRegistry,
+    get_tokenizer,
+    list_supported_models,
+    register_tokenizer,
+)
+from .tiktoken_counter import TiktokenCounter
+
+
+# Lazy imports for optional dependencies
+def get_huggingface_tokenizer():
+    """Get HuggingFaceTokenizer class (requires transformers)."""
+    from .huggingface import HuggingFaceTokenizer
+
+    return HuggingFaceTokenizer
+
+
+def get_mistral_tokenizer():
+    """Get MistralTokenizer class (requires mistral-common)."""
+    from .mistral import MistralTokenizer
+
+    return MistralTokenizer
+
+
+def is_mistral_tokenizer_available() -> bool:
+    """Check if Mistral tokenizer is available."""
+    from .mistral import is_mistral_available
+
+    return is_mistral_available()
+
+
+__all__ = [
+    # Registry
+    "TokenizerRegistry",
+    "get_tokenizer",
+    "register_tokenizer",
+    "list_supported_models",
+    # Base classes
+    "TokenCounter",
+    "BaseTokenizer",
+    # Implementations
+    "TiktokenCounter",
+    "EstimatingTokenCounter",
+    "CharacterCounter",
+    # Lazy loaders
+    "get_huggingface_tokenizer",
+    "get_mistral_tokenizer",
+    "is_mistral_tokenizer_available",
+]

headroom/tokenizers/base.py

@@ -0,0 +1,210 @@
+"""Base classes for tokenizer implementations.
+
+Defines the TokenCounter protocol and BaseTokenizer class that all
+tokenizer backends must implement.
+"""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class TokenCounter(Protocol):
+    """Protocol for token counting implementations.
+
+    Any class implementing this protocol can be used with Headroom
+    for token counting. This allows integration with various
+    tokenizer backends (tiktoken, HuggingFace, custom, etc.).
+    """
+
+    def count_text(self, text: str) -> int:
+        """Count tokens in a text string.
+
+        Args:
+            text: The text to count tokens for.
+
+        Returns:
+            Number of tokens in the text.
+        """
+        ...
+
+    def count_messages(self, messages: list[dict[str, Any]]) -> int:
+        """Count tokens in a list of chat messages.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'.
+
+        Returns:
+            Total token count including message overhead.
+        """
+        ...
+
+
+class BaseTokenizer(ABC):
+    """Abstract base class for tokenizer implementations.
+
+    Provides common functionality for counting messages while
+    requiring subclasses to implement text tokenization.
+    """
+
+    # Token overhead per message (role, formatting, etc.)
+    # Override in subclasses for model-specific overhead
+    MESSAGE_OVERHEAD = 4
+    REPLY_OVERHEAD = 3  # Assistant reply start tokens
+
+    @abstractmethod
+    def count_text(self, text: str) -> int:
+        """Count tokens in a text string. Must be implemented by subclasses."""
+        pass
+
+    def count_message(self, message: dict[str, Any]) -> int:
+        """Count tokens in a single message.
+
+        Args:
+            message: A message dict with 'role' and 'content'.
+
+        Returns:
+            Token count for this message.
+        """
+        return self.count_messages([message]) - self.REPLY_OVERHEAD
+
+    def count_messages(self, messages: list[dict[str, Any]]) -> int:
+        """Count tokens in a list of chat messages.
+
+        Uses OpenAI-style message counting as the baseline, which
+        works well for most models.
+
+        Args:
+            messages: List of message dicts.
+
+        Returns:
+            Total token count.
+        """
+        total = 0
+
+        for message in messages:
+            # Base message overhead
+            total += self.MESSAGE_OVERHEAD
+
+            # Count role
+            role = message.get("role", "")
+            total += self.count_text(role)
+
+            # Count content
+            content = message.get("content")
+            if content is not None:
+                if isinstance(content, str):
+                    total += self.count_text(content)
+                elif isinstance(content, list):
+                    # Multi-part content (images, tool results, etc.)
+                    total += self._count_content_parts(content)
+
+            # Count tool calls
+            tool_calls = message.get("tool_calls")
+            if tool_calls:
+                total += self._count_tool_calls(tool_calls)
+
+            # Count function call (legacy)
+            function_call = message.get("function_call")
+            if function_call:
+                total += self._count_function_call(function_call)
+
+            # Count name field
+            name = message.get("name")
+            if name:
+                total += self.count_text(name)
+                total += 1  # Name field overhead
+
+        # Reply start overhead
+        total += self.REPLY_OVERHEAD
+
+        return total
+
+    def _count_content_parts(self, parts: list[Any]) -> int:
+        """Count tokens in multi-part content."""
+        total = 0
+        for part in parts:
+            if isinstance(part, dict):
+                part_type = part.get("type", "")
+
+                if part_type == "text":
+                    total += self.count_text(part.get("text", ""))
+                elif part_type == "image_url":
+                    # Images have fixed token cost (varies by model)
+                    total += 85  # Base image token count
+                elif part_type == "tool_result":
+                    content = part.get("content", "")
+                    if isinstance(content, str):
+                        total += self.count_text(content)
+                    else:
+                        total += self.count_text(json.dumps(content))
+                elif part_type == "tool_use":
+                    total += self.count_text(part.get("name", ""))
+                    total += self.count_text(json.dumps(part.get("input", {})))
+                else:
+                    # Unknown type - estimate from JSON
+                    total += self.count_text(json.dumps(part))
+            elif isinstance(part, str):
+                total += self.count_text(part)
+
+        return total
+
+    def _count_tool_calls(self, tool_calls: list[dict[str, Any]]) -> int:
+        """Count tokens in tool calls."""
+        total = 0
+        for call in tool_calls:
+            total += 4  # Tool call overhead
+
+            if "function" in call:
+                func = call["function"]
+                total += self.count_text(func.get("name", ""))
+                total += self.count_text(func.get("arguments", ""))
+
+            if "id" in call:
+                total += self.count_text(call["id"])
+
+        return total
+
+    def _count_function_call(self, function_call: dict[str, Any]) -> int:
+        """Count tokens in legacy function call."""
+        total = 4  # Function call overhead
+        total += self.count_text(function_call.get("name", ""))
+        total += self.count_text(function_call.get("arguments", ""))
+        return total
+
+    def encode(self, text: str) -> list[int]:
+        """Encode text to token IDs.
+
+        Optional method - not all backends support encoding.
+        Default implementation raises NotImplementedError.
+
+        Args:
+            text: Text to encode.
+
+        Returns:
+            List of token IDs.
+
+        Raises:
+            NotImplementedError: If encoding is not supported.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support encoding")
+
+    def decode(self, tokens: list[int]) -> str:
+        """Decode token IDs to text.
+
+        Optional method - not all backends support decoding.
+        Default implementation raises NotImplementedError.
+
+        Args:
+            tokens: List of token IDs.
+
+        Returns:
+            Decoded text.
+
+        Raises:
+            NotImplementedError: If decoding is not supported.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support decoding")

headroom/tokenizers/estimator.py

@@ -0,0 +1,198 @@
+"""Estimation-based token counter for fallback scenarios.
+
+When no exact tokenizer is available (e.g., unknown models, missing
+dependencies), this provides a reasonable approximation based on
+character/word heuristics calibrated against real tokenizers.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any
+
+from .base import BaseTokenizer
+
+
+class EstimatingTokenCounter(BaseTokenizer):
+    """Token counter using estimation heuristics.
+
+    This is the fallback tokenizer used when:
+    - Model is unknown/unsupported
+    - Required tokenizer library not installed
+    - Speed is prioritized over accuracy
+
+    The estimation is calibrated against tiktoken cl100k_base and
+    provides ~90% accuracy for typical text. It tends to slightly
+    overestimate, which is safer for context window management.
+
+    Estimation Strategy:
+    - Base: ~4 characters per token (calibrated against GPT-4)
+    - Adjustments for code, URLs, numbers, whitespace
+    - Special handling for JSON structure
+
+    Example:
+        counter = EstimatingTokenCounter()
+        tokens = counter.count_text("Hello, world!")
+        print(f"Estimated tokens: {tokens}")
+    """
+
+    # Calibration constants (derived from tiktoken analysis)
+    CHARS_PER_TOKEN = 4.0  # Average for English text
+    CHARS_PER_TOKEN_CODE = 3.5  # Code is denser
+    CHARS_PER_TOKEN_JSON = 3.2  # JSON has more structure
+
+    # Patterns for content type detection
+    CODE_PATTERN = re.compile(
+        r"(?:def |class |function |const |let |var |import |from |"
+        r"if \(|for \(|while \(|switch \(|try \{|catch \(|"
+        r"=>|->|\{\{|\}\}|;$)",
+        re.MULTILINE,
+    )
+    JSON_PATTERN = re.compile(r"^\s*[\[\{]")
+    URL_PATTERN = re.compile(r"https?://\S+")
+    UUID_PATTERN = re.compile(
+        r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", re.IGNORECASE
+    )
+
+    def __init__(self, chars_per_token: float | None = None):
+        """Initialize estimating counter.
+
+        Args:
+            chars_per_token: Override default chars per token ratio.
+                            If None, auto-detects based on content type.
+        """
+        self._fixed_ratio = chars_per_token
+
+    def count_text(self, text: str) -> int:
+        """Estimate token count for text.
+
+        Args:
+            text: Text to count tokens for.
+
+        Returns:
+            Estimated number of tokens.
+        """
+        if not text:
+            return 0
+
+        # Use fixed ratio if provided
+        if self._fixed_ratio is not None:
+            return max(1, int(len(text) / self._fixed_ratio + 0.5))
+
+        # Auto-detect content type and adjust ratio
+        ratio = self._detect_ratio(text)
+
+        # Apply ratio with minimum of 1 token
+        base_count = int(len(text) / ratio + 0.5)
+
+        # Add overhead for special patterns
+        overhead = self._count_special_overhead(text)
+
+        return max(1, base_count + overhead)
+
+    def _detect_ratio(self, text: str) -> float:
+        """Detect optimal chars-per-token ratio based on content.
+
+        Args:
+            text: Text to analyze.
+
+        Returns:
+            Chars per token ratio.
+        """
+        # Check for JSON
+        if self.JSON_PATTERN.match(text):
+            try:
+                json.loads(text)
+                return self.CHARS_PER_TOKEN_JSON
+            except (json.JSONDecodeError, ValueError):
+                pass
+
+        # Check for code
+        code_matches = len(self.CODE_PATTERN.findall(text))
+        if code_matches > len(text) / 500:  # ~2 matches per KB
+            return self.CHARS_PER_TOKEN_CODE
+
+        return self.CHARS_PER_TOKEN
+
+    def _count_special_overhead(self, text: str) -> int:
+        """Count additional tokens for special patterns.
+
+        URLs and UUIDs often tokenize into more tokens than
+        character count would suggest.
+
+        Args:
+            text: Text to analyze.
+
+        Returns:
+            Additional token overhead.
+        """
+        overhead = 0
+
+        # URLs typically tokenize to more tokens
+        urls = self.URL_PATTERN.findall(text)
+        for url in urls:
+            # Each URL component adds overhead
+            overhead += url.count("/") + url.count("?") + url.count("&")
+
+        # UUIDs are typically 8-10 tokens despite being 36 chars
+        uuids = self.UUID_PATTERN.findall(text)
+        overhead += len(uuids) * 2  # Each UUID adds ~2 extra tokens
+
+        return overhead
+
+    def count_messages(self, messages: list[dict[str, Any]]) -> int:
+        """Estimate tokens in chat messages.
+
+        Uses the base class implementation with estimation-based
+        text counting.
+
+        Args:
+            messages: List of chat messages.
+
+        Returns:
+            Estimated total token count.
+        """
+        # Use base class implementation
+        return super().count_messages(messages)
+
+    def __repr__(self) -> str:
+        if self._fixed_ratio:
+            return f"EstimatingTokenCounter(chars_per_token={self._fixed_ratio})"
+        return "EstimatingTokenCounter(auto)"
+
+
+class CharacterCounter(BaseTokenizer):
+    """Simple character-based counter.
+
+    Uses a fixed character-to-token ratio. Useful for:
+    - Quick approximations
+    - Testing
+    - Models with unknown tokenization
+
+    This is less accurate than EstimatingTokenCounter but faster.
+    """
+
+    def __init__(self, chars_per_token: float = 4.0):
+        """Initialize character counter.
+
+        Args:
+            chars_per_token: Characters per token ratio.
+        """
+        self.chars_per_token = chars_per_token
+
+    def count_text(self, text: str) -> int:
+        """Count tokens based on character count.
+
+        Args:
+            text: Text to count.
+
+        Returns:
+            Estimated token count.
+        """
+        if not text:
+            return 0
+        return max(1, int(len(text) / self.chars_per_token + 0.5))
+
+    def __repr__(self) -> str:
+        return f"CharacterCounter(chars_per_token={self.chars_per_token})"