""" Unit tests for TokenCounter classes covering: - BaseTokenCounter (rule-based estimation) - OpenAITokenCounter (tiktoken-based) - HFTokenCounter (HuggingFace tokenizer-based) Usage: python test_token_counter.py --base # Test BaseTokenCounter only python test_token_counter.py --openai # Test OpenAITokenCounter only python test_token_counter.py --hf # Test HFTokenCounter only python test_token_counter.py --all # Test all token counters """ import argparse from typing import Type, List from reme.core.enumeration import Role from reme.core.schema import Message, ToolCall from reme.core.token_counter import BaseTokenCounter, OpenAITokenCounter, HFTokenCounter def get_token_counter(counter_class: Type[BaseTokenCounter], **kwargs) -> BaseTokenCounter: """Create and return a token counter instance.""" default_kwargs = { "model_name": "gpt-4o", } default_kwargs.update(kwargs) return counter_class(**default_kwargs) def get_test_messages() -> List[Message]: """Create test messages for token counting.""" return [ Message(role=Role.SYSTEM, content="You are a helpful assistant."), Message(role=Role.USER, content="Hello, how are you today?"), Message(role=Role.ASSISTANT, content="I'm doing well, thank you for asking! How can I help you?"), Message(role=Role.USER, content="Can you explain what machine learning is?"), Message( role=Role.ASSISTANT, content="Machine learning is a subset of artificial intelligence that enables computers to " "learn from data without being explicitly programmed.", ), ] def get_chinese_messages() -> List[Message]: """Create test messages with Chinese content.""" return [ Message(role=Role.SYSTEM, content="你是一个有帮助的助手。"), Message(role=Role.USER, content="你好，今天天气怎么样？"), Message(role=Role.ASSISTANT, content="今天天气很好，阳光明媚，适合外出活动。"), Message(role=Role.USER, content="能给我推荐一些好看的电影吗？"), Message(role=Role.ASSISTANT, content="当然可以！我推荐《肖申克的救赎》、《阿甘正传》和《泰坦尼克号》。"), ] def get_mixed_messages() -> List[Message]: """Create test messages with mixed English and Chinese content.""" return [ Message(role=Role.SYSTEM, content="You are a bilingual assistant. 你是一个双语助手。"), Message(role=Role.USER, content="What is AI? 什么是人工智能？"), Message( role=Role.ASSISTANT, content="AI (Artificial Intelligence) 是人工智能的英文缩写，它是计算机科学的一个分支。", ), ] def get_messages_with_reasoning() -> List[Message]: """Create test messages with reasoning content.""" return [ Message(role=Role.USER, content="What is 2 + 2?"), Message( role=Role.ASSISTANT, content="The answer is 4.", reasoning_content="Let me think about this step by step. 2 + 2 " "equals 4 because addition combines two quantities.", ), ] def get_test_tools() -> List[ToolCall]: """Create test tool calls for token counting.""" return [ ToolCall( **{ "type": "function", "function": { "name": "get_weather", "description": "Get the current weather for a specified location.", "parameters": { "type": "object", "properties": { "location": { "type": "string", "description": "The city and country, e.g., 'Beijing, China'", }, "unit": { "type": "string", "description": "Temperature unit: 'celsius' or 'fahrenheit'", "enum": ["celsius", "fahrenheit"], }, }, "required": ["location"], }, }, }, ), ToolCall( **{ "type": "function", "function": { "name": "search_web", "description": "Search the web for information.", "parameters": { "type": "object", "properties": { "query": { "type": "string", "description": "The search query", }, "num_results": { "type": "integer", "description": "Number of results to return", }, }, "required": ["query"], }, }, }, ), ] def get_tool_call_messages() -> List[Message]: """Create messages with tool call responses.""" return [ Message(role=Role.USER, content="What's the weather in Beijing?"), Message( role=Role.ASSISTANT, content="", tool_calls=[ ToolCall( id="call_123", name="get_weather", arguments='{"location": "Beijing, China", "unit": "celsius"}', ), ], ), Message( role=Role.TOOL, content='{"temperature": 25, "condition": "sunny", "humidity": 60}', tool_call_id="call_123", ), Message( role=Role.ASSISTANT, content="The weather in Beijing is sunny with a temperature of 25°C and 60% humidity.", ), ] def test_basic_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test basic token counting with simple messages.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Basic Token Count") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) messages = get_test_messages() print(f"Input: {len(messages)} messages") for i, msg in enumerate(messages, 1): content_preview = msg.content[:50] + "..." if len(msg.content) > 50 else msg.content print(f" {i}. [{msg.role.value}] {content_preview}") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" print(f"\n✓ Token count: {token_count}") print(f"✓ PASSED: {counter_name} basic token count") def test_chinese_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with Chinese content.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Chinese Token Count") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) messages = get_chinese_messages() print(f"Input: {len(messages)} Chinese messages") for i, msg in enumerate(messages, 1): content_preview = msg.content[:30] + "..." if len(msg.content) > 30 else msg.content print(f" {i}. [{msg.role.value}] {content_preview}") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" print(f"\n✓ Token count: {token_count}") print(f"✓ PASSED: {counter_name} Chinese token count") def test_mixed_language_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with mixed English and Chinese content.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Mixed Language Token Count") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) messages = get_mixed_messages() print(f"Input: {len(messages)} mixed language messages") for i, msg in enumerate(messages, 1): content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content print(f" {i}. [{msg.role.value}] {content_preview}") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" print(f"\n✓ Token count: {token_count}") print(f"✓ PASSED: {counter_name} mixed language token count") def test_reasoning_content_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with reasoning content.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Reasoning Content Token Count") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) messages = get_messages_with_reasoning() print(f"Input: {len(messages)} messages with reasoning content") for i, msg in enumerate(messages, 1): print(f" {i}. [{msg.role.value}] content: {msg.content[:30]}...") if msg.reasoning_content: print(f" reasoning: {msg.reasoning_content[:30]}...") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" print(f"\n✓ Token count: {token_count}") print(f"✓ PASSED: {counter_name} reasoning content token count") def test_token_count_with_tools(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with tool definitions.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Token Count with Tools") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) messages = get_test_messages()[:2] tools = get_test_tools() print(f"Input: {len(messages)} messages, {len(tools)} tools") for tool in tools: print(f" Tool: {tool.name} - {tool.description[:40]}...") token_count = counter.count_token(messages, tools=tools) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" # Token count with tools should be higher than without token_count_no_tools = counter.count_token(messages) assert token_count > token_count_no_tools, f"{counter_name}: Token count with tools should be higher" print(f"\n✓ Token count without tools: {token_count_no_tools}") print(f"✓ Token count with tools: {token_count}") print(f"✓ Tools added {token_count - token_count_no_tools} tokens") print(f"✓ PASSED: {counter_name} token count with tools") def test_tool_call_messages_token_count(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with messages containing tool calls.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Tool Call Messages Token Count") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) messages = get_tool_call_messages() print(f"Input: {len(messages)} messages with tool calls") for i, msg in enumerate(messages, 1): if msg.tool_calls: print(f" {i}. [{msg.role.value}] tool_calls: {[tc.name for tc in msg.tool_calls]}") else: content_preview = msg.content[:40] + "..." if len(msg.content) > 40 else msg.content print(f" {i}. [{msg.role.value}] {content_preview}") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" print(f"\n✓ Token count: {token_count}") print(f"✓ PASSED: {counter_name} tool call messages token count") def test_empty_messages(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with empty message list.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Empty Messages") print(f"{'=' * 60}") # HFTokenCounter does not support empty message list (apply_chat_template requires at least one message) if counter_class == HFTokenCounter: print("⊘ SKIPPED: HFTokenCounter does not support empty message list") return counter = get_token_counter(counter_class, **kwargs) messages: List[Message] = [] print("Input: 0 messages") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" # OpenAITokenCounter adds 3 tokens for reply priming even with empty messages if counter_class == OpenAITokenCounter: assert token_count == 3, f"{counter_name}: Empty messages should have 3 tokens (reply priming)" print(f"\n✓ Token count: {token_count} (includes 3 tokens for reply priming)") else: assert token_count == 0, f"{counter_name}: Empty messages should have 0 tokens" print(f"\n✓ Token count: {token_count}") print(f"✓ PASSED: {counter_name} empty messages") def test_single_message(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with a single message.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Single Message") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) messages = [Message(role=Role.USER, content="Hello!")] print(f"Input: 1 message - '{messages[0].content}'") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" print(f"\n✓ Token count: {token_count}") print(f"✓ PASSED: {counter_name} single message") def test_long_content(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Test token counting with long content.""" print(f"\n{'=' * 60}") print(f"Testing {counter_name}: Long Content") print(f"{'=' * 60}") counter = get_token_counter(counter_class, **kwargs) # Create a long message long_text = "This is a test sentence. " * 100 messages = [Message(role=Role.USER, content=long_text)] print(f"Input: 1 message with {len(long_text)} characters") token_count = counter.count_token(messages) assert token_count is not None, f"{counter_name}: Token count is None" assert isinstance(token_count, int), f"{counter_name}: Token count is not an integer" assert token_count > 0, f"{counter_name}: Token count should be positive" # Long content should have more tokens short_messages = [Message(role=Role.USER, content="This is a test sentence.")] short_token_count = counter.count_token(short_messages) assert token_count > short_token_count, f"{counter_name}: Long content should have more tokens" print(f"\n✓ Short content token count: {short_token_count}") print(f"✓ Long content token count: {token_count}") print(f"✓ PASSED: {counter_name} long content") def run_all_tests_for_counter(counter_class: Type[BaseTokenCounter], counter_name: str, **kwargs): """Run all tests for a specific token counter class.""" print(f"\n\n{'#' * 60}") print(f"# Running all tests for: {counter_name}") print(f"{'#' * 60}") test_basic_token_count(counter_class, counter_name, **kwargs) test_chinese_token_count(counter_class, counter_name, **kwargs) test_mixed_language_token_count(counter_class, counter_name, **kwargs) test_reasoning_content_token_count(counter_class, counter_name, **kwargs) test_token_count_with_tools(counter_class, counter_name, **kwargs) test_tool_call_messages_token_count(counter_class, counter_name, **kwargs) test_empty_messages(counter_class, counter_name, **kwargs) test_single_message(counter_class, counter_name, **kwargs) test_long_content(counter_class, counter_name, **kwargs) print(f"\n{'=' * 60}") print(f"✓ All tests passed for {counter_name}!") print(f"{'=' * 60}") def main(): """Main entry point for running tests.""" parser = argparse.ArgumentParser( description="Run token counter tests", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: python test_token_counter.py --base # Test BaseTokenCounter only python test_token_counter.py --openai # Test OpenAITokenCounter only python test_token_counter.py --hf # Test HFTokenCounter only python test_token_counter.py --all # Test all token counters """, ) parser.add_argument( "--base", action="store_true", help="Test BaseTokenCounter (rule-based)", ) parser.add_argument( "--openai", action="store_true", help="Test OpenAITokenCounter (tiktoken-based)", ) parser.add_argument( "--hf", action="store_true", help="Test HFTokenCounter (HuggingFace tokenizer-based)", ) parser.add_argument( "--hf-model", type=str, default="Qwen/Qwen2.5-0.5B-Instruct", help="HuggingFace model name for HFTokenCounter (default: Qwen/Qwen2.5-0.5B-Instruct)", ) parser.add_argument( "--all", action="store_true", help="Run tests for all available token counters", ) args = parser.parse_args() # Determine which counters to test counters_to_test = [] if args.all: counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {})) counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {})) counters_to_test.append( (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}), ) else: if args.base: counters_to_test.append((BaseTokenCounter, "BaseTokenCounter", {})) if args.openai: counters_to_test.append((OpenAITokenCounter, "OpenAITokenCounter", {})) if args.hf: counters_to_test.append( (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}), ) if not counters_to_test: # Default to all counters if no argument provided counters_to_test = [ (BaseTokenCounter, "BaseTokenCounter", {}), (OpenAITokenCounter, "OpenAITokenCounter", {}), (HFTokenCounter, "HFTokenCounter", {"model_name": args.hf_model, "trust_remote_code": True}), ] print("No counter specified, defaulting to test all counters") print("Use --base/--openai/--hf to test specific ones\n") # Run tests for each counter for counter_class, counter_name, kwargs in counters_to_test: try: run_all_tests_for_counter(counter_class, counter_name, **kwargs) except Exception as e: print(f"\n✗ FAILED: {counter_name} tests failed with error:") print(f" {type(e).__name__}: {e}") raise # Final summary print(f"\n\n{'#' * 60}") print("# TEST SUMMARY") print(f"{'#' * 60}") print(f"✓ All tests passed for {len(counters_to_test)} token counter(s):") for _, counter_name, _ in counters_to_test: print(f" - {counter_name}") print(f"{'#' * 60}\n") if __name__ == "__main__": main()