Spaces:
Build error
Build error
File size: 2,149 Bytes
9c7d451 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """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)
|