wehub-resource-sync
644 lines
25 KiB
Python
644 lines
25 KiB
Python
"""Unit tests for token-aware truncation functionality.
|
|
|
|
This test suite defines the contract for token truncation functions that prevent
|
|
500 errors from Ollama when text exceeds model token limits. These tests verify:
|
|
|
|
1. Model token limit retrieval (known and unknown models)
|
|
2. Text truncation behavior for single and multiple texts
|
|
3. Token counting and truncation accuracy using tiktoken
|
|
|
|
All tests are written in Red Phase - they should FAIL initially because the
|
|
implementation does not exist yet.
|
|
"""
|
|
|
|
import pytest
|
|
import tiktoken
|
|
from leann.embedding_compute import (
|
|
EMBEDDING_MODEL_LIMITS,
|
|
get_model_token_limit,
|
|
truncate_to_token_limit,
|
|
)
|
|
|
|
|
|
class TestModelTokenLimits:
|
|
"""Tests for retrieving model-specific token limits."""
|
|
|
|
def test_get_model_token_limit_known_model(self):
|
|
"""Verify correct token limit is returned for known models.
|
|
|
|
Known models should return their specific token limits from
|
|
EMBEDDING_MODEL_LIMITS dictionary.
|
|
"""
|
|
# Test nomic-embed-text (2048 tokens)
|
|
limit = get_model_token_limit("nomic-embed-text")
|
|
assert limit == 2048, "nomic-embed-text should have 2048 token limit"
|
|
|
|
# Test nomic-embed-text-v1.5 (2048 tokens)
|
|
limit = get_model_token_limit("nomic-embed-text-v1.5")
|
|
assert limit == 2048, "nomic-embed-text-v1.5 should have 2048 token limit"
|
|
|
|
# Test nomic-embed-text-v2 (512 tokens)
|
|
limit = get_model_token_limit("nomic-embed-text-v2")
|
|
assert limit == 512, "nomic-embed-text-v2 should have 512 token limit"
|
|
|
|
# Test OpenAI models (8192 tokens)
|
|
limit = get_model_token_limit("text-embedding-3-small")
|
|
assert limit == 8192, "text-embedding-3-small should have 8192 token limit"
|
|
|
|
def test_get_model_token_limit_unknown_model(self):
|
|
"""Verify default token limit is returned for unknown models.
|
|
|
|
Unknown models should return the default limit (2048) to allow
|
|
operation with reasonable safety margin.
|
|
"""
|
|
# Test with completely unknown model
|
|
limit = get_model_token_limit("unknown-model-xyz")
|
|
assert limit == 2048, "Unknown models should return default 2048"
|
|
|
|
# Test with empty string
|
|
limit = get_model_token_limit("")
|
|
assert limit == 2048, "Empty model name should return default 2048"
|
|
|
|
def test_get_model_token_limit_custom_default(self):
|
|
"""Verify custom default can be specified for unknown models.
|
|
|
|
Allow callers to specify their own default token limit when
|
|
model is not in the known models dictionary.
|
|
"""
|
|
limit = get_model_token_limit("unknown-model", default=4096)
|
|
assert limit == 4096, "Should return custom default for unknown models"
|
|
|
|
# Known model should ignore custom default
|
|
limit = get_model_token_limit("nomic-embed-text", default=4096)
|
|
assert limit == 2048, "Known model should ignore custom default"
|
|
|
|
def test_embedding_model_limits_dictionary_exists(self):
|
|
"""Verify EMBEDDING_MODEL_LIMITS dictionary contains expected models.
|
|
|
|
The dictionary should be importable and contain at least the
|
|
known nomic models with correct token limits.
|
|
"""
|
|
assert isinstance(EMBEDDING_MODEL_LIMITS, dict), "Should be a dictionary"
|
|
assert "nomic-embed-text" in EMBEDDING_MODEL_LIMITS, "Should contain nomic-embed-text"
|
|
assert "nomic-embed-text-v1.5" in EMBEDDING_MODEL_LIMITS, (
|
|
"Should contain nomic-embed-text-v1.5"
|
|
)
|
|
assert EMBEDDING_MODEL_LIMITS["nomic-embed-text"] == 2048
|
|
assert EMBEDDING_MODEL_LIMITS["nomic-embed-text-v1.5"] == 2048
|
|
assert EMBEDDING_MODEL_LIMITS["nomic-embed-text-v2"] == 512
|
|
# OpenAI models
|
|
assert EMBEDDING_MODEL_LIMITS["text-embedding-3-small"] == 8192
|
|
|
|
|
|
class TestTokenTruncation:
|
|
"""Tests for truncating texts to token limits."""
|
|
|
|
@pytest.fixture
|
|
def tokenizer(self):
|
|
"""Provide tiktoken tokenizer for token counting verification."""
|
|
return tiktoken.get_encoding("cl100k_base")
|
|
|
|
def test_truncate_single_text_under_limit(self, tokenizer):
|
|
"""Verify text under token limit remains unchanged.
|
|
|
|
When text is already within the token limit, it should be
|
|
returned unchanged with no truncation.
|
|
"""
|
|
text = "This is a short text that is well under the token limit."
|
|
token_count = len(tokenizer.encode(text))
|
|
assert token_count < 100, f"Test setup: text should be short (has {token_count} tokens)"
|
|
|
|
# Truncate with generous limit
|
|
result = truncate_to_token_limit([text], token_limit=512)
|
|
|
|
assert len(result) == 1, "Should return same number of texts"
|
|
assert result[0] == text, "Text under limit should be unchanged"
|
|
|
|
def test_truncate_single_text_over_limit(self, tokenizer):
|
|
"""Verify text over token limit is truncated correctly.
|
|
|
|
When text exceeds the token limit, it should be truncated to
|
|
fit within the limit while maintaining valid token boundaries.
|
|
"""
|
|
# Create a text that definitely exceeds limit
|
|
text = "word " * 200 # ~200 tokens (each "word " is typically 1-2 tokens)
|
|
original_token_count = len(tokenizer.encode(text))
|
|
assert original_token_count > 50, (
|
|
f"Test setup: text should be long (has {original_token_count} tokens)"
|
|
)
|
|
|
|
# Truncate to 50 tokens
|
|
result = truncate_to_token_limit([text], token_limit=50)
|
|
|
|
assert len(result) == 1, "Should return same number of texts"
|
|
assert result[0] != text, "Text over limit should be truncated"
|
|
assert len(result[0]) < len(text), "Truncated text should be shorter"
|
|
|
|
# Verify truncated text is within token limit
|
|
truncated_token_count = len(tokenizer.encode(result[0]))
|
|
assert truncated_token_count <= 50, (
|
|
f"Truncated text should be ≤50 tokens, got {truncated_token_count}"
|
|
)
|
|
|
|
def test_truncate_multiple_texts_mixed_lengths(self, tokenizer):
|
|
"""Verify multiple texts with mixed lengths are handled correctly.
|
|
|
|
When processing multiple texts:
|
|
- Texts under limit should remain unchanged
|
|
- Texts over limit should be truncated independently
|
|
- Output list should maintain same order and length
|
|
"""
|
|
texts = [
|
|
"Short text.", # Under limit
|
|
"word " * 200, # Over limit
|
|
"Another short one.", # Under limit
|
|
"token " * 150, # Over limit
|
|
]
|
|
|
|
# Verify test setup
|
|
for i, text in enumerate(texts):
|
|
token_count = len(tokenizer.encode(text))
|
|
if i in [1, 3]:
|
|
assert token_count > 50, f"Text {i} should be over limit (has {token_count} tokens)"
|
|
else:
|
|
assert token_count < 50, (
|
|
f"Text {i} should be under limit (has {token_count} tokens)"
|
|
)
|
|
|
|
# Truncate with 50 token limit
|
|
result = truncate_to_token_limit(texts, token_limit=50)
|
|
|
|
assert len(result) == len(texts), "Should return same number of texts"
|
|
|
|
# Verify each text individually
|
|
for i, (original, truncated) in enumerate(zip(texts, result)):
|
|
token_count = len(tokenizer.encode(truncated))
|
|
assert token_count <= 50, f"Text {i} should be ≤50 tokens, got {token_count}"
|
|
|
|
# Short texts should be unchanged
|
|
if i in [0, 2]:
|
|
assert truncated == original, f"Short text {i} should be unchanged"
|
|
# Long texts should be truncated
|
|
else:
|
|
assert len(truncated) < len(original), f"Long text {i} should be truncated"
|
|
|
|
def test_truncate_empty_list(self):
|
|
"""Verify empty input list returns empty output list.
|
|
|
|
Edge case: empty list should return empty list without errors.
|
|
"""
|
|
result = truncate_to_token_limit([], token_limit=512)
|
|
assert result == [], "Empty input should return empty output"
|
|
|
|
def test_truncate_preserves_order(self, tokenizer):
|
|
"""Verify truncation preserves original text order.
|
|
|
|
Output list should maintain the same order as input list,
|
|
regardless of which texts were truncated.
|
|
"""
|
|
texts = [
|
|
"First text " * 50, # Will be truncated
|
|
"Second text.", # Won't be truncated
|
|
"Third text " * 50, # Will be truncated
|
|
]
|
|
|
|
result = truncate_to_token_limit(texts, token_limit=20)
|
|
|
|
assert len(result) == 3, "Should preserve list length"
|
|
# Check that order is maintained by looking for distinctive words
|
|
assert "First" in result[0], "First text should remain in first position"
|
|
assert "Second" in result[1], "Second text should remain in second position"
|
|
assert "Third" in result[2], "Third text should remain in third position"
|
|
|
|
def test_truncate_extremely_long_text(self, tokenizer):
|
|
"""Verify extremely long texts are truncated efficiently.
|
|
|
|
Test with text that far exceeds token limit to ensure
|
|
truncation handles extreme cases without performance issues.
|
|
"""
|
|
# Create very long text (simulate real-world scenario)
|
|
text = "token " * 5000 # ~5000+ tokens
|
|
original_token_count = len(tokenizer.encode(text))
|
|
assert original_token_count > 1000, "Test setup: text should be very long"
|
|
|
|
# Truncate to small limit
|
|
result = truncate_to_token_limit([text], token_limit=100)
|
|
|
|
assert len(result) == 1
|
|
truncated_token_count = len(tokenizer.encode(result[0]))
|
|
assert truncated_token_count <= 100, (
|
|
f"Should truncate to ≤100 tokens, got {truncated_token_count}"
|
|
)
|
|
assert len(result[0]) < len(text) // 10, "Should significantly reduce text length"
|
|
|
|
def test_truncate_exact_token_limit(self, tokenizer):
|
|
"""Verify text at exactly token limit is handled correctly.
|
|
|
|
Edge case: text with exactly the token limit should either
|
|
remain unchanged or be safely truncated by 1 token.
|
|
"""
|
|
# Create text with approximately 50 tokens
|
|
# We'll adjust to get exactly 50
|
|
target_tokens = 50
|
|
text = "word " * 50
|
|
tokens = tokenizer.encode(text)
|
|
|
|
# Adjust to get exactly target_tokens
|
|
if len(tokens) > target_tokens:
|
|
tokens = tokens[:target_tokens]
|
|
text = tokenizer.decode(tokens)
|
|
elif len(tokens) < target_tokens:
|
|
# Add more words
|
|
while len(tokenizer.encode(text)) < target_tokens:
|
|
text += "word "
|
|
tokens = tokenizer.encode(text)[:target_tokens]
|
|
text = tokenizer.decode(tokens)
|
|
|
|
# Verify we have exactly target_tokens
|
|
assert len(tokenizer.encode(text)) == target_tokens, (
|
|
"Test setup: should have exactly 50 tokens"
|
|
)
|
|
|
|
result = truncate_to_token_limit([text], token_limit=target_tokens)
|
|
|
|
assert len(result) == 1
|
|
result_tokens = len(tokenizer.encode(result[0]))
|
|
assert result_tokens <= target_tokens, (
|
|
f"Should be ≤{target_tokens} tokens, got {result_tokens}"
|
|
)
|
|
|
|
|
|
class TestLMStudioHybridDiscovery:
|
|
"""Tests for LM Studio integration in get_model_token_limit() hybrid discovery.
|
|
|
|
These tests verify that get_model_token_limit() properly integrates with
|
|
the LM Studio SDK bridge for dynamic token limit discovery. The integration
|
|
should:
|
|
|
|
1. Detect LM Studio URLs (port 1234 or 'lmstudio'/'lm.studio' in URL)
|
|
2. Convert HTTP URLs to WebSocket format for SDK queries
|
|
3. Query LM Studio SDK and use discovered limit
|
|
4. Fall back to registry when SDK returns None
|
|
5. Execute AFTER Ollama detection but BEFORE registry fallback
|
|
|
|
All tests are written in Red Phase - they should FAIL initially because the
|
|
LM Studio detection and integration logic does not exist yet in get_model_token_limit().
|
|
"""
|
|
|
|
def test_get_model_token_limit_lmstudio_success(self, monkeypatch):
|
|
"""Verify LM Studio SDK query succeeds and returns detected limit.
|
|
|
|
When a LM Studio base_url is detected and the SDK query succeeds,
|
|
get_model_token_limit() should return the dynamically discovered
|
|
context length without falling back to the registry.
|
|
"""
|
|
|
|
# Mock _query_lmstudio_context_limit to return successful SDK query
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
# Verify WebSocket URL was passed (not HTTP)
|
|
assert base_url.startswith("ws://"), (
|
|
f"Should convert HTTP to WebSocket format, got: {base_url}"
|
|
)
|
|
return 8192 # Successful SDK query
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
# Test with HTTP URL that should be converted to WebSocket
|
|
limit = get_model_token_limit(
|
|
model_name="custom-model", base_url="http://localhost:1234/v1"
|
|
)
|
|
|
|
assert limit == 8192, "Should return limit from LM Studio SDK query"
|
|
|
|
def test_get_model_token_limit_lmstudio_fallback_to_registry(self, monkeypatch):
|
|
"""Verify fallback to registry when LM Studio SDK returns None.
|
|
|
|
When LM Studio SDK query fails (returns None), get_model_token_limit()
|
|
should fall back to the EMBEDDING_MODEL_LIMITS registry.
|
|
"""
|
|
|
|
# Mock _query_lmstudio_context_limit to return None (SDK failure)
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
return None # SDK query failed
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
# Test with known model that exists in registry
|
|
limit = get_model_token_limit(
|
|
model_name="nomic-embed-text", base_url="http://localhost:1234/v1"
|
|
)
|
|
|
|
# Should fall back to registry value
|
|
assert limit == 2048, "Should fall back to registry when SDK returns None"
|
|
|
|
def test_get_model_token_limit_lmstudio_port_detection(self, monkeypatch):
|
|
"""Verify detection of LM Studio via port 1234.
|
|
|
|
get_model_token_limit() should recognize port 1234 as a LM Studio
|
|
server and attempt SDK query, regardless of hostname.
|
|
"""
|
|
query_called = False
|
|
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
nonlocal query_called
|
|
query_called = True
|
|
return 4096
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
# Test with port 1234 (default LM Studio port)
|
|
limit = get_model_token_limit(model_name="test-model", base_url="http://127.0.0.1:1234/v1")
|
|
|
|
assert query_called, "Should detect port 1234 and call LM Studio SDK query"
|
|
assert limit == 4096, "Should return SDK query result"
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_url,expected_limit,keyword",
|
|
[
|
|
("http://lmstudio.local:8080/v1", 16384, "lmstudio"),
|
|
("http://api.lm.studio:5000/v1", 32768, "lm.studio"),
|
|
],
|
|
)
|
|
def test_get_model_token_limit_lmstudio_url_keyword_detection(
|
|
self, monkeypatch, test_url, expected_limit, keyword
|
|
):
|
|
"""Verify detection of LM Studio via keywords in URL.
|
|
|
|
get_model_token_limit() should recognize 'lmstudio' or 'lm.studio'
|
|
in the URL as indicating a LM Studio server.
|
|
"""
|
|
query_called = False
|
|
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
nonlocal query_called
|
|
query_called = True
|
|
return expected_limit
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
limit = get_model_token_limit(model_name="test-model", base_url=test_url)
|
|
|
|
assert query_called, f"Should detect '{keyword}' keyword and call SDK query"
|
|
assert limit == expected_limit, f"Should return SDK query result for {keyword}"
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_url,expected_protocol,expected_host",
|
|
[
|
|
("http://localhost:1234/v1", "ws://", "localhost:1234"),
|
|
("https://lmstudio.example.com:1234/v1", "wss://", "lmstudio.example.com:1234"),
|
|
],
|
|
)
|
|
def test_get_model_token_limit_protocol_conversion(
|
|
self, monkeypatch, input_url, expected_protocol, expected_host
|
|
):
|
|
"""Verify HTTP/HTTPS URL is converted to WebSocket format for SDK query.
|
|
|
|
LM Studio SDK requires WebSocket URLs. get_model_token_limit() should:
|
|
1. Convert 'http://' to 'ws://'
|
|
2. Convert 'https://' to 'wss://'
|
|
3. Remove '/v1' or other path suffixes (SDK expects base URL)
|
|
4. Preserve host and port
|
|
"""
|
|
conversions_tested = []
|
|
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
conversions_tested.append(base_url)
|
|
return 8192
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
get_model_token_limit(model_name="test-model", base_url=input_url)
|
|
|
|
# Verify conversion happened
|
|
assert len(conversions_tested) == 1, "Should have called SDK query once"
|
|
assert conversions_tested[0].startswith(expected_protocol), (
|
|
f"Should convert to {expected_protocol}"
|
|
)
|
|
assert expected_host in conversions_tested[0], (
|
|
f"Should preserve host and port: {expected_host}"
|
|
)
|
|
|
|
def test_get_model_token_limit_lmstudio_executes_after_ollama(self, monkeypatch):
|
|
"""Verify LM Studio detection happens AFTER Ollama detection.
|
|
|
|
The hybrid discovery order should be:
|
|
1. Ollama dynamic discovery (port 11434 or 'ollama' in URL)
|
|
2. LM Studio dynamic discovery (port 1234 or 'lmstudio' in URL)
|
|
3. Registry fallback
|
|
|
|
If both Ollama and LM Studio patterns match, Ollama should take precedence.
|
|
This test verifies that LM Studio is checked but doesn't interfere with Ollama.
|
|
"""
|
|
ollama_called = False
|
|
lmstudio_called = False
|
|
|
|
def mock_query_ollama(model_name, base_url):
|
|
nonlocal ollama_called
|
|
ollama_called = True
|
|
return 2048 # Ollama query succeeds
|
|
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
nonlocal lmstudio_called
|
|
lmstudio_called = True
|
|
return None # Should not be reached if Ollama succeeds
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_ollama_context_limit",
|
|
mock_query_ollama,
|
|
)
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
# Test with Ollama URL
|
|
limit = get_model_token_limit(
|
|
model_name="test-model", base_url="http://localhost:11434/api"
|
|
)
|
|
|
|
assert ollama_called, "Should attempt Ollama query first"
|
|
assert not lmstudio_called, "Should not attempt LM Studio query when Ollama succeeds"
|
|
assert limit == 2048, "Should return Ollama result"
|
|
|
|
def test_get_model_token_limit_lmstudio_not_detected_for_non_lmstudio_urls(self, monkeypatch):
|
|
"""Verify LM Studio SDK query is NOT called for non-LM Studio URLs.
|
|
|
|
Only URLs with port 1234 or 'lmstudio'/'lm.studio' keywords should
|
|
trigger LM Studio SDK queries. Other URLs should skip to registry fallback.
|
|
"""
|
|
lmstudio_called = False
|
|
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
nonlocal lmstudio_called
|
|
lmstudio_called = True
|
|
return 8192
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
# Test with non-LM Studio URLs
|
|
test_cases = [
|
|
"http://localhost:8080/v1", # Different port
|
|
"http://openai.example.com/v1", # Different service
|
|
"http://localhost:3000/v1", # Another port
|
|
]
|
|
|
|
for base_url in test_cases:
|
|
lmstudio_called = False # Reset for each test
|
|
get_model_token_limit(model_name="nomic-embed-text", base_url=base_url)
|
|
assert not lmstudio_called, f"Should NOT call LM Studio SDK for URL: {base_url}"
|
|
|
|
def test_get_model_token_limit_lmstudio_case_insensitive_detection(self, monkeypatch):
|
|
"""Verify LM Studio detection is case-insensitive for keywords.
|
|
|
|
Keywords 'lmstudio' and 'lm.studio' should be detected regardless
|
|
of case (LMStudio, LMSTUDIO, LmStudio, etc.).
|
|
"""
|
|
query_called = False
|
|
|
|
def mock_query_lmstudio(model_name, base_url):
|
|
nonlocal query_called
|
|
query_called = True
|
|
return 8192
|
|
|
|
monkeypatch.setattr(
|
|
"leann.embedding_compute._query_lmstudio_context_limit",
|
|
mock_query_lmstudio,
|
|
)
|
|
|
|
# Test various case variations
|
|
test_cases = [
|
|
"http://LMStudio.local:8080/v1",
|
|
"http://LMSTUDIO.example.com/v1",
|
|
"http://LmStudio.local/v1",
|
|
"http://api.LM.STUDIO:5000/v1",
|
|
]
|
|
|
|
for base_url in test_cases:
|
|
query_called = False # Reset for each test
|
|
limit = get_model_token_limit(model_name="test-model", base_url=base_url)
|
|
assert query_called, f"Should detect LM Studio in URL: {base_url}"
|
|
assert limit == 8192, f"Should return SDK result for URL: {base_url}"
|
|
|
|
|
|
class TestTokenLimitCaching:
|
|
"""Tests for token limit caching to prevent repeated SDK/API calls.
|
|
|
|
Caching prevents duplicate SDK/API calls within the same Python process,
|
|
which is important because:
|
|
1. LM Studio SDK load() can load duplicate model instances
|
|
2. Ollama /api/show queries add latency
|
|
3. Registry lookups are pure overhead
|
|
|
|
Cache is process-scoped and resets between leann build invocations.
|
|
"""
|
|
|
|
def setup_method(self):
|
|
"""Clear cache before each test."""
|
|
from leann.embedding_compute import _token_limit_cache
|
|
|
|
_token_limit_cache.clear()
|
|
|
|
def test_registry_lookup_is_cached(self):
|
|
"""Verify that registry lookups are cached."""
|
|
from leann.embedding_compute import _token_limit_cache
|
|
|
|
# First call
|
|
limit1 = get_model_token_limit("text-embedding-3-small")
|
|
assert limit1 == 8192
|
|
|
|
# Verify it's in cache
|
|
cache_key = ("text-embedding-3-small", "")
|
|
assert cache_key in _token_limit_cache
|
|
assert _token_limit_cache[cache_key] == 8192
|
|
|
|
# Second call should use cache
|
|
limit2 = get_model_token_limit("text-embedding-3-small")
|
|
assert limit2 == 8192
|
|
|
|
def test_default_fallback_is_cached(self):
|
|
"""Verify that default fallbacks are cached."""
|
|
from leann.embedding_compute import _token_limit_cache
|
|
|
|
# First call with unknown model
|
|
limit1 = get_model_token_limit("unknown-model-xyz", default=512)
|
|
assert limit1 == 512
|
|
|
|
# Verify it's in cache
|
|
cache_key = ("unknown-model-xyz", "")
|
|
assert cache_key in _token_limit_cache
|
|
assert _token_limit_cache[cache_key] == 512
|
|
|
|
# Second call should use cache
|
|
limit2 = get_model_token_limit("unknown-model-xyz", default=512)
|
|
assert limit2 == 512
|
|
|
|
def test_different_urls_create_separate_cache_entries(self):
|
|
"""Verify that different base_urls create separate cache entries."""
|
|
from leann.embedding_compute import _token_limit_cache
|
|
|
|
# Same model, different URLs
|
|
limit1 = get_model_token_limit("nomic-embed-text", base_url="http://localhost:11434")
|
|
limit2 = get_model_token_limit("nomic-embed-text", base_url="http://localhost:1234/v1")
|
|
|
|
# Both should find the model in registry (2048)
|
|
assert limit1 == 2048
|
|
assert limit2 == 2048
|
|
|
|
# But they should be separate cache entries
|
|
cache_key1 = ("nomic-embed-text", "http://localhost:11434")
|
|
cache_key2 = ("nomic-embed-text", "http://localhost:1234/v1")
|
|
|
|
assert cache_key1 in _token_limit_cache
|
|
assert cache_key2 in _token_limit_cache
|
|
assert len(_token_limit_cache) == 2
|
|
|
|
def test_cache_prevents_repeated_lookups(self):
|
|
"""Verify that cache prevents repeated registry/API lookups."""
|
|
from leann.embedding_compute import _token_limit_cache
|
|
|
|
model_name = "text-embedding-ada-002"
|
|
|
|
# First call - should add to cache
|
|
assert len(_token_limit_cache) == 0
|
|
limit1 = get_model_token_limit(model_name)
|
|
|
|
cache_size_after_first = len(_token_limit_cache)
|
|
assert cache_size_after_first == 1
|
|
|
|
# Multiple subsequent calls - cache size should not change
|
|
for _ in range(5):
|
|
limit = get_model_token_limit(model_name)
|
|
assert limit == limit1
|
|
assert len(_token_limit_cache) == cache_size_after_first
|
|
|
|
def test_versioned_model_names_cached_correctly(self):
|
|
"""Verify that versioned model names (e.g., model:tag) are cached."""
|
|
from leann.embedding_compute import _token_limit_cache
|
|
|
|
# Model with version tag
|
|
limit = get_model_token_limit("nomic-embed-text:latest", base_url="http://localhost:11434")
|
|
assert limit == 2048
|
|
|
|
# Should be cached with full name including version
|
|
cache_key = ("nomic-embed-text:latest", "http://localhost:11434")
|
|
assert cache_key in _token_limit_cache
|
|
assert _token_limit_cache[cache_key] == 2048
|