jpmschweitzerandClaude Opus 4.5
2f97041aa9
- litellm had dependency conflicts with pydantic-ai - tiktoken is lighter and already required by pydantic-ai - Updated documentation (README.md, architecture.md) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
"""
|
|
Tests for token counting utilities.
|
|
"""
|
|
import pytest
|
|
|
|
from src.shared.tokens import count_tokens, count_message_tokens, estimate_tokens
|
|
|
|
|
|
class TestTokenCounting:
|
|
"""Tests for token counting functions."""
|
|
|
|
def test_estimate_tokens_basic(self):
|
|
"""Test basic token estimation."""
|
|
text = "Hello world"
|
|
tokens = estimate_tokens(text)
|
|
# ~4 chars per token
|
|
assert tokens == len(text) // 4
|
|
|
|
def test_estimate_tokens_empty(self):
|
|
"""Test estimation with empty string."""
|
|
assert estimate_tokens("") == 0
|
|
|
|
def test_estimate_tokens_long_text(self):
|
|
"""Test estimation with longer text."""
|
|
text = "a" * 400
|
|
tokens = estimate_tokens(text)
|
|
assert tokens == 100
|
|
|
|
def test_count_tokens_basic(self):
|
|
"""Test actual token counting."""
|
|
text = "Hello, how are you today?"
|
|
tokens = count_tokens(text)
|
|
# Should return reasonable token count
|
|
assert tokens > 0
|
|
assert tokens < len(text) # Should be fewer tokens than characters
|
|
|
|
def test_count_tokens_empty(self):
|
|
"""Test counting empty string."""
|
|
tokens = count_tokens("")
|
|
assert tokens == 0
|
|
|
|
def test_count_message_tokens_single(self):
|
|
"""Test counting tokens in single message."""
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
tokens = count_message_tokens(messages)
|
|
assert tokens > 0
|
|
|
|
def test_count_message_tokens_multiple(self):
|
|
"""Test counting tokens in multiple messages."""
|
|
messages = [
|
|
{"role": "user", "content": "Hello, how are you?"},
|
|
{"role": "assistant", "content": "I'm doing well, thank you!"},
|
|
]
|
|
tokens = count_message_tokens(messages)
|
|
# Should be more than single message
|
|
single_tokens = count_message_tokens([messages[0]])
|
|
assert tokens > single_tokens
|
|
|
|
def test_count_message_tokens_empty_list(self):
|
|
"""Test counting empty message list."""
|
|
tokens = count_message_tokens([])
|
|
# tiktoken returns small overhead for empty list (assistant priming)
|
|
assert tokens < 10
|
|
|
|
|
|
class TestTokenCountingAccuracy:
|
|
"""Tests for token counting accuracy."""
|
|
|
|
def test_code_tokens_reasonable(self):
|
|
"""Test that code is tokenized reasonably."""
|
|
code = """
|
|
def hello_world():
|
|
print("Hello, World!")
|
|
return True
|
|
"""
|
|
tokens = count_tokens(code)
|
|
# Code should have reasonable token count
|
|
assert 10 < tokens < 100
|
|
|
|
def test_special_characters(self):
|
|
"""Test tokenization of special characters."""
|
|
text = "Hello! @#$%^&*() World?"
|
|
tokens = count_tokens(text)
|
|
assert tokens > 0
|
|
|
|
def test_unicode_text(self):
|
|
"""Test tokenization of unicode text."""
|
|
text = "Hello 世界 🌍"
|
|
tokens = count_tokens(text)
|
|
assert tokens > 0
|