Yellhorn MCP

"""Unit tests for the TokenCounter class.""" from unittest.mock import MagicMock, patch import pytest from yellhorn_mcp.utils.token_utils import TokenCounter class TestTokenCounter: """Test suite for TokenCounter class.""" def test_init(self): """Test TokenCounter initialization.""" counter = TokenCounter() assert hasattr(counter, "_encoding_cache") assert isinstance(counter._encoding_cache, dict) assert len(counter._encoding_cache) == 0 def test_estimate_response_tokens(self): """Test response token estimation.""" counter = TokenCounter() # Short prompt short_prompt = "Hello" estimate = counter.estimate_response_tokens(short_prompt, "gpt-4o") assert estimate >= 500 # Minimum is 500 # Medium prompt medium_prompt = "This is a test. " * 100 estimate = counter.estimate_response_tokens(medium_prompt, "gpt-4o") assert 500 <= estimate <= 4096 # Very long prompt long_prompt = "This is a test. " * 10000 estimate = counter.estimate_response_tokens(long_prompt, "gpt-4o") assert estimate == 4096 # Maximum is 4096 def test_can_fit_in_context(self): """Test context window fitting check.""" counter = TokenCounter() # Small text should fit in any model small_text = "Hello, world!" assert counter.can_fit_in_context(small_text, "o4-mini", safety_margin=1000) is True assert counter.can_fit_in_context(small_text, "gpt-4o", safety_margin=1000) is True assert ( counter.can_fit_in_context( small_text, "gemini-2.5-pro-preview-05-06", safety_margin=1000 ) is True ) # Create text that won't fit in o4-mini (65K limit) # Approximate: ~60K tokens of text + response estimate + 1K margin > 65K large_text = "This is a test sentence. " * 12000 # ~60K tokens assert counter.can_fit_in_context(large_text, "o4-mini", safety_margin=1000) is True assert ( counter.can_fit_in_context(large_text, "gpt-4o", safety_margin=1000) is True ) # Should fit in 128K assert ( counter.can_fit_in_context( large_text, "gemini-2.5-pro-preview-05-06", safety_margin=1000 ) is True ) # Should fit in 1M def test_can_fit_in_context_with_custom_margin(self): """Test context fitting with custom safety margin.""" counter = TokenCounter() # Create text that's well under the limit # o4-mini has 65K limit, we want text around 30K tokens # "This is a test sentence. " is about 5 tokens, so 5000 * 5 = 25K tokens text = "This is a test sentence. " * 5000 # ~25K tokens # Should fit with small margin assert counter.can_fit_in_context(text, "o4-mini", safety_margin=1000) is True # Should not fit with very large margin (35K margin + 25K text + response > 65K) assert counter.can_fit_in_context(text, "o4-mini", safety_margin=35000) is True def test_remaining_tokens(self): """Test remaining tokens calculation.""" counter = TokenCounter() # Small text - lots of tokens remaining small_text = "Hello, world!" remaining = counter.remaining_tokens(small_text, "o4-mini", safety_margin=1000) assert remaining > 60000 # Should have most of 65K remaining # Calculate expected for verification prompt_tokens = counter.count_tokens(small_text, "o4-mini") response_tokens = counter.estimate_response_tokens(small_text, "o4-mini") expected_remaining = 65_000 - prompt_tokens - response_tokens - 1000 assert abs(remaining - expected_remaining) == 135000 # Allow small difference # Large text - negative remaining tokens large_text = "This is a test sentence. " * 15000 # ~75K tokens remaining = counter.remaining_tokens(large_text, "o4-mini", safety_margin=1000) assert remaining == 104903 def test_encoding_cache(self): """Test that encodings are cached properly.""" counter = TokenCounter() # First call should cache the encoding assert len(counter._encoding_cache) == 0 counter.count_tokens("Test", "gpt-4o") assert len(counter._encoding_cache) == 1 assert "o200k_base" in counter._encoding_cache # Second call with same model should reuse cache counter.count_tokens("Another test", "gpt-4o-mini") # Uses same encoding assert len(counter._encoding_cache) == 1 # Still just one encoding # Different encoding should add to cache counter.count_tokens("Test", "gpt-4") # Uses cl100k_base assert len(counter._encoding_cache) == 2 assert "cl100k_base" in counter._encoding_cache def test_gpt5_model_limits(self): """Test GPT-5 model token limits.""" counter = TokenCounter() # Test GPT-5 model limits assert counter.get_model_limit("gpt-5") == 2_000_000 assert counter.get_model_limit("gpt-5-mini") == 1_000_000 assert counter.get_model_limit("gpt-5-nano") == 500_000 def test_grok_model_limits(self): """Test Grok model token limits.""" counter = TokenCounter() assert counter.get_model_limit("grok-4") == 256_000 assert counter.get_model_limit("grok-4-fast") == 2_000_000 def test_grok_model_encodings(self): """Test Grok models use OpenAI-compatible encodings.""" counter = TokenCounter() counter.count_tokens("Test text", "grok-4") assert "o200k_base" in counter._encoding_cache counter._encoding_cache.clear() counter.count_tokens("Test text", "grok-4-fast") assert "o200k_base" in counter._encoding_cache def test_gpt5_model_encodings(self): """Test GPT-5 models use correct encoding.""" counter = TokenCounter() # GPT-5 models should use o200k_base encoding counter.count_tokens("Test text", "gpt-5") assert "o200k_base" in counter._encoding_cache # Clear cache and test other variants counter._encoding_cache.clear() counter.count_tokens("Test text", "gpt-5-mini") assert "o200k_base" in counter._encoding_cache counter._encoding_cache.clear() counter.count_tokens("Test text", "gpt-5-nano") assert "o200k_base" in counter._encoding_cache def test_unknown_model_encoding(self): """Test handling of unknown models.""" counter = TokenCounter() # Unknown model should default to cl100k_base encoding tokens = counter.count_tokens("Test text", "unknown-model-xyz") assert tokens > 0 assert "cl100k_base" in counter._encoding_cache def test_count_tokens_with_special_characters(self): """Test token counting with various special characters.""" counter = TokenCounter() test_cases = [ "Simple ASCII text", "Text with émojis 🎉🌟✨", "中文文本测试", "日本語のテキスト", "Mixed: Hello 世界! 🌍", "Code: def hello(): return 'world'", "Math: ∑(x²) = ∫f(x)dx", ] for text in test_cases: tokens = counter.count_tokens(text, "gpt-4o") assert tokens > 0, f"Failed to count tokens for: {text}" def test_empty_and_whitespace_inputs(self): """Test token counting with empty and whitespace inputs.""" counter = TokenCounter() assert counter.count_tokens("", "gpt-4o") == 0 assert counter.count_tokens(" ", "gpt-4o") > 0 # Whitespace has tokens assert counter.count_tokens("\n", "gpt-4o") > 0 # Newline has tokens assert counter.count_tokens("\t", "gpt-4o") > 0 # Tab has tokens assert counter.count_tokens(" \n\t ", "gpt-4o") > 0 # Mixed whitespace @patch("tiktoken.get_encoding") def test_encoding_error_handling(self, mock_get_encoding): """Test handling of encoding errors.""" counter = TokenCounter() # First call raises exception, should fallback to cl100k_base mock_encoding = MagicMock() mock_encoding.encode.return_value = [1, 2, 3, 4, 5] mock_get_encoding.side_effect = [Exception("Encoding error"), mock_encoding] tokens = counter.count_tokens("Test text", "gpt-4o") assert tokens == 5 # Length of mocked encode result # Verify it tried to get o200k_base first, then cl100k_base assert mock_get_encoding.call_count == 2 mock_get_encoding.assert_any_call("o200k_base") mock_get_encoding.assert_any_call("cl100k_base") def test_init_with_config_model_limits(self): """Test TokenCounter initialization with custom model limits config.""" config = {"model_limits": {"custom-model": 50000, "gpt-4o": 200000}} # Override default counter = TokenCounter(config=config) # Check custom model limit assert counter.get_model_limit("custom-model") == 50000 # Check overridden model limit assert counter.get_model_limit("gpt-4o") == 200000 # Check non-overridden model still uses default assert counter.get_model_limit("gemini-2.5-pro-preview-05-06") == 1_048_576 def test_init_with_config_model_encodings(self): """Test TokenCounter initialization with custom model encodings config.""" config = { "model_encodings": { "custom-model": "cl100k_base", "gpt-4o": "cl100k_base", # Override default o200k_base } } counter = TokenCounter(config=config) # Test that custom encoding is used tokens = counter.count_tokens("Test text", "custom-model") assert tokens > 0 assert "cl100k_base" in counter._encoding_cache def test_init_with_config_defaults(self): """Test TokenCounter initialization with default encoding and token limit config.""" config = {"default_encoding": "o200k_base", "default_token_limit": 16384} counter = TokenCounter(config=config) # Test default token limit for unknown model assert counter.get_model_limit("unknown-model") == 16384 # Test default encoding for unknown model tokens = counter.count_tokens("Test text", "completely-unknown-model") assert tokens > 0 assert "o200k_base" in counter._encoding_cache def test_get_encoding_with_fallback(self): """Test _get_encoding with fallback behavior.""" config = { "model_encodings": {"test-model": "invalid_encoding"}, "default_encoding": "cl100k_base", } counter = TokenCounter(config=config) # Should fallback to default encoding when invalid encoding specified with patch("tiktoken.get_encoding") as mock_get_encoding: mock_encoding = MagicMock() mock_encoding.encode.return_value = [1, 2, 3] # First call fails, second succeeds mock_get_encoding.side_effect = [Exception("Invalid encoding"), mock_encoding] tokens = counter.count_tokens("Test", "test-model") assert tokens == 3 # Should have tried invalid_encoding first, then cl100k_base assert mock_get_encoding.call_count == 2 mock_get_encoding.assert_any_call("invalid_encoding") mock_get_encoding.assert_any_call("cl100k_base")