mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
ead221995e
Code outside the router now reaches it only through the entry points (registry, settings, health, select, apply, observe, register_plugin, retire) or the endpoint methods; the outside-member allowlist shrinks from 139 to 45. litellm/router.py no longer imports litellm.proxy: it asks the installed RouterHost (litellm/router_host.py) for vertex pass-through credentials, member auto-router authorization and auto-router prompt compression, and the proxy installs its host on import. Call sites where a model name can be missing keep the old answer instead of reaching an entry point with None.
1454 lines
51 KiB
Python
1454 lines
51 KiB
Python
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from importlib import import_module
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
|
|
|
|
|
|
@contextmanager
|
|
def _fake_redisvl_modules(semantic_cache_mock: MagicMock, custom_vectorizer_mock: MagicMock) -> Iterator[None]:
|
|
with pytest.MonkeyPatch.context() as mp:
|
|
mp.setitem(sys.modules, "redisvl.extensions.llmcache", MagicMock(SemanticCache=semantic_cache_mock))
|
|
mp.setitem(sys.modules, "redisvl.utils.vectorize", MagicMock(CustomTextVectorizer=custom_vectorizer_mock))
|
|
yield
|
|
|
|
|
|
# Tests for RedisSemanticCache
|
|
def test_redis_semantic_cache_initialization(monkeypatch):
|
|
# Mock the redisvl import
|
|
semantic_cache_mock = MagicMock()
|
|
with _fake_redisvl_modules(semantic_cache_mock, MagicMock()):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
# Set environment variables
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
# Initialize the cache with a similarity threshold
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
# Verify the semantic cache was initialized with correct parameters
|
|
assert redis_semantic_cache.similarity_threshold == 0.8
|
|
|
|
# Use pytest.approx for floating point comparison to handle precision issues
|
|
assert redis_semantic_cache.distance_threshold == pytest.approx(0.2, abs=1e-10)
|
|
assert redis_semantic_cache.embedding_model == "text-embedding-ada-002"
|
|
|
|
# Test initialization with missing similarity_threshold
|
|
with pytest.raises(ValueError, match="similarity_threshold must be provided"):
|
|
RedisSemanticCache()
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache(monkeypatch):
|
|
# Mock the redisvl import and embedding function
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
# Set environment variables
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
# Initialize cache
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
# Mock the llmcache.check method to return a result
|
|
mock_result = [
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris is the capital of France."}',
|
|
"vector_distance": 0.1, # Distance of 0.1 means similarity of 0.9
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
redis_semantic_cache.llmcache.check = MagicMock(return_value=mock_result)
|
|
|
|
# Mock the embedding function
|
|
with (
|
|
patch(
|
|
"litellm.embedding",
|
|
return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
|
|
),
|
|
patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
),
|
|
):
|
|
# Test get_cache with a message
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata=metadata,
|
|
)
|
|
|
|
# Verify result is properly parsed
|
|
assert result == {"content": "Paris is the capital of France."}
|
|
assert metadata["semantic-similarity"] == pytest.approx(0.9)
|
|
|
|
# Verify llmcache.check was called
|
|
redis_semantic_cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.check = MagicMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
}
|
|
]
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.embedding",
|
|
return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
|
|
),
|
|
patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
),
|
|
):
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result is None
|
|
assert metadata["semantic-similarity"] == 0.0
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.store = MagicMock()
|
|
|
|
with patch(
|
|
"litellm.embedding",
|
|
return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
|
|
):
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
ttl=60,
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
ttl=60,
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_uses_isolated_index_for_old_schema(monkeypatch):
|
|
fallback_cache_mock = MagicMock()
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[
|
|
ValueError("stored index schema differs from requested fields"),
|
|
fallback_cache_mock,
|
|
]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(
|
|
similarity_threshold=0.8,
|
|
index_name="existing_index",
|
|
)
|
|
|
|
assert redis_semantic_cache.llmcache is fallback_cache_mock
|
|
assert semantic_cache_mock.call_args_list[0].kwargs["name"] == "existing_index"
|
|
assert (
|
|
semantic_cache_mock.call_args_list[1].kwargs["name"]
|
|
== "existing_index_isolated"
|
|
)
|
|
assert semantic_cache_mock.call_args_list[1].kwargs["filterable_fields"] == [
|
|
RedisSemanticCache._cache_key_filterable_field()
|
|
]
|
|
|
|
|
|
def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch):
|
|
fallback_cache_mock = MagicMock()
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[
|
|
ValueError("Existing index schema does not match"),
|
|
ValueError("Existing index schema does not match"),
|
|
fallback_cache_mock,
|
|
]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(
|
|
similarity_threshold=0.8,
|
|
index_name="existing_index",
|
|
)
|
|
|
|
assert redis_semantic_cache.llmcache is fallback_cache_mock
|
|
assert (
|
|
semantic_cache_mock.call_args_list[2].kwargs["name"]
|
|
== "existing_index_isolated"
|
|
)
|
|
assert semantic_cache_mock.call_args_list[2].kwargs["overwrite"] is True
|
|
assert semantic_cache_mock.call_args_list[2].kwargs["filterable_fields"] == [
|
|
RedisSemanticCache._cache_key_filterable_field()
|
|
]
|
|
|
|
|
|
def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypatch):
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[
|
|
ValueError("Existing index schema does not match"),
|
|
ValueError("connection failed"),
|
|
]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
cache = RedisSemanticCache(
|
|
similarity_threshold=0.8,
|
|
index_name="existing_index",
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="connection failed"):
|
|
_ = cache.llmcache
|
|
|
|
|
|
def test_redis_semantic_cache_reraises_unexpected_index_error():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.distance_threshold = 0.2
|
|
semantic_cache_mock = MagicMock(side_effect=ValueError("connection failed"))
|
|
|
|
with pytest.raises(ValueError, match="connection failed"):
|
|
redis_semantic_cache._init_semantic_cache(
|
|
semantic_cache_cls=semantic_cache_mock,
|
|
index_name="existing_index",
|
|
redis_url="redis://localhost:6379",
|
|
cache_vectorizer=MagicMock(),
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_matches_bytes_cache_key():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
|
|
assert redis_semantic_cache._cache_hit_matches_key(
|
|
cache_hit={RedisSemanticCache.CACHE_KEY_FIELD_NAME: b"test_key"},
|
|
key="test_key",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_rejects_pre_isolation_unscoped_hit():
|
|
"""Pre-isolation entries with no cache-key field cannot be safely
|
|
reassigned to a caller's scope and are treated as misses."""
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
|
|
cache_hit = {
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
}
|
|
assert not redis_semantic_cache._cache_hit_matches_key(
|
|
cache_hit=cache_hit,
|
|
key="test_key",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_builds_filter_expression(monkeypatch):
|
|
class FakeTag:
|
|
def __init__(self, field_name):
|
|
self.field_name = field_name
|
|
|
|
def __eq__(self, value):
|
|
return (self.field_name, value)
|
|
|
|
monkeypatch.setitem(sys.modules, "redisvl.query.filter", MagicMock(Tag=FakeTag))
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
|
|
assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == (
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME,
|
|
"test_key",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_get_cache(monkeypatch):
|
|
# Mock the redisvl import
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
# Set environment variables
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
# Initialize cache
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
# Mock the async methods
|
|
mock_result = [
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris is the capital of France."}',
|
|
"vector_distance": 0.1, # Distance of 0.1 means similarity of 0.9
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(return_value=mock_result)
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(
|
|
return_value=[0.1, 0.2, 0.3]
|
|
)
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
# Test async_get_cache with a message
|
|
result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata={},
|
|
)
|
|
|
|
# Verify result is properly parsed
|
|
assert result == {"content": "Paris is the capital of France."}
|
|
|
|
# Verify methods were called
|
|
redis_semantic_cache._get_async_embedding.assert_called_once()
|
|
redis_semantic_cache.llmcache.acheck.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
}
|
|
]
|
|
)
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(
|
|
return_value=[0.1, 0.2, 0.3]
|
|
)
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata={},
|
|
)
|
|
|
|
assert result is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter(
|
|
monkeypatch,
|
|
):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.astore = AsyncMock()
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(
|
|
return_value=[0.1, 0.2, 0.3]
|
|
)
|
|
|
|
await redis_semantic_cache.async_set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
ttl=60,
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.astore.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
ttl=60,
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_uses_responses_string_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
redis_semantic_cache._get_ttl = MagicMock(return_value=None)
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input="What is the capital of France?",
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache_uses_responses_string_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.similarity_threshold = 0.8
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.check = MagicMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
)
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
assert metadata["semantic-similarity"] == pytest.approx(0.9)
|
|
redis_semantic_cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_flattens_structured_responses_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
redis_semantic_cache._get_ttl = MagicMock(return_value=None)
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "input_text", "text": "What is the capital of France?"},
|
|
{"type": "input_text", "text": "Answer briefly."},
|
|
{
|
|
"type": "input_image",
|
|
"image_url": "https://example.com/paris.png",
|
|
},
|
|
],
|
|
}
|
|
],
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?\nAnswer briefly.",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_reads_function_call_output_blocks():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
input=[
|
|
{"role": "user", "content": "update the config"},
|
|
{"type": "function_call", "call_id": "c1", "name": "write_file", "arguments": '{"path": "a"}'},
|
|
{"type": "function_call_output", "call_id": "c1", "output": [{"type": "input_text", "text": "wrote a"}]},
|
|
]
|
|
)
|
|
|
|
assert prompt == "update the config\nwrote a"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_prefers_messages():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
messages=[{"content": "message prompt"}],
|
|
input="responses prompt",
|
|
)
|
|
|
|
assert prompt == "message prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_handles_model_objects():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
class ModelDumpInput:
|
|
def model_dump(self):
|
|
return {"content": [{"text": "model dump prompt"}]}
|
|
|
|
class DictInput:
|
|
def dict(self):
|
|
return {"content": [{"output_text": "dict prompt"}]}
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
input=[
|
|
ModelDumpInput(),
|
|
DictInput(),
|
|
{"content": [{"input_text": "inline prompt"}]},
|
|
{"content": [{"type": "input_image", "image_url": "https://example.com"}]},
|
|
]
|
|
)
|
|
|
|
assert prompt == "model dump prompt\ndict prompt\ninline prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_returns_none_without_text():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
assert RedisSemanticCache._get_prompt_from_kwargs() is None
|
|
assert RedisSemanticCache._get_prompt_from_kwargs(input=None) is None
|
|
assert RedisSemanticCache._get_prompt_from_kwargs(input=" ") is None
|
|
assert (
|
|
RedisSemanticCache._get_prompt_from_kwargs(
|
|
input=[{"type": "input_image", "image_url": "https://example.com"}]
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_skips_blank_dict_text_keys():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
input={"text": " ", "input_text": "fallback prompt"}
|
|
)
|
|
|
|
assert prompt == "fallback prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_skips_blank_object_text_keys():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
class ResponseInput:
|
|
text = " "
|
|
input_text = "fallback prompt"
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput())
|
|
|
|
assert prompt == "fallback prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_handles_object_content():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
class ResponseInput:
|
|
content = [{"text": "object content prompt"}]
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput())
|
|
|
|
assert prompt == "object content prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_skips_blank_responses_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input=" ",
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_not_called()
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache_sets_similarity_on_blank_responses_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
metadata = {}
|
|
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
input=" ",
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result is None
|
|
assert metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.check.assert_not_called()
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.check = MagicMock(return_value=[])
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result is None
|
|
assert metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_paths_use_responses_string_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.similarity_threshold = 0.8
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.astore = AsyncMock()
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
)
|
|
redis_semantic_cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
redis_semantic_cache._get_ttl = MagicMock(return_value=None)
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
await redis_semantic_cache.async_set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input="What is the capital of France?",
|
|
)
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
metadata = {}
|
|
result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.astore.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
assert result == {"content": "Paris"}
|
|
assert metadata["semantic-similarity"] == pytest.approx(0.9)
|
|
redis_semantic_cache.llmcache.acheck.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_paths_set_similarity_on_misses():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.astore = AsyncMock()
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(return_value=[])
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
await redis_semantic_cache.async_set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input=" ",
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.astore.assert_not_called()
|
|
redis_semantic_cache._get_async_embedding.assert_not_called()
|
|
|
|
blank_metadata = {}
|
|
blank_result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
input=" ",
|
|
metadata=blank_metadata,
|
|
)
|
|
|
|
assert blank_result is None
|
|
assert blank_metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.acheck.assert_not_called()
|
|
redis_semantic_cache._get_async_embedding.assert_not_called()
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
miss_metadata = {}
|
|
miss_result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=miss_metadata,
|
|
)
|
|
|
|
assert miss_result is None
|
|
assert miss_metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.acheck.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
def test_redis_get_embedding_routes_through_router(monkeypatch):
|
|
import types
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (None, None)
|
|
router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = router
|
|
fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
with patch("litellm.embedding") as direct_embed:
|
|
vec = cache._get_embedding("hello", metadata={"user_api_key": "sk-x"})
|
|
|
|
assert vec == [0.5, 0.6]
|
|
router.embedding.assert_called_once()
|
|
assert router.embedding.call_args.kwargs["model"] == "sem-embed"
|
|
assert router.embedding.call_args.kwargs["input"] == "hello"
|
|
assert router.embedding.call_args.kwargs["cache"] == {
|
|
"no-store": True,
|
|
"no-cache": True,
|
|
}
|
|
assert router.embedding.call_args.kwargs["metadata"] == {
|
|
"user_api_key": "sk-x",
|
|
"semantic-cache-embedding": True,
|
|
}
|
|
direct_embed.assert_not_called()
|
|
|
|
|
|
def test_redis_get_embedding_falls_back_to_direct(monkeypatch):
|
|
import types
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "text-embedding-ada-002"
|
|
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = None
|
|
fake_proxy.llm_model_list = None
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
with patch(
|
|
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2]}]}
|
|
) as direct_embed:
|
|
vec = cache._get_embedding("hello")
|
|
|
|
assert vec == [0.1, 0.2]
|
|
direct_embed.assert_called_once()
|
|
|
|
|
|
def test_cache_get_cache_passes_responses_input_to_backend_cache():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.cache = MagicMock()
|
|
cache.cache.get_cache = MagicMock(return_value=None)
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
|
|
metadata = {}
|
|
cache.get_cache(
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
cache={},
|
|
)
|
|
|
|
cache.cache.get_cache.assert_called_once_with(
|
|
"test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
|
|
def test_cache_get_cache_filters_non_lookup_kwargs_from_backend_cache():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.cache = MagicMock()
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
cache.get_cache_logic = MagicMock(return_value={"content": "Paris"})
|
|
|
|
def _cache_hit(_cache_key, **cache_kwargs):
|
|
cache_kwargs["metadata"]["semantic-similarity"] = 0.7
|
|
return {"content": "Paris"}
|
|
|
|
cache.cache.get_cache = MagicMock(side_effect=_cache_hit)
|
|
|
|
metadata = {"user_api_key": "sk-secret", "trace_id": "trace-id"}
|
|
result = cache.get_cache(
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
cache={"s-maxage": 10},
|
|
api_key="sk-secret",
|
|
headers={"authorization": "Bearer sk-secret"},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
assert metadata == {
|
|
"user_api_key": "sk-secret",
|
|
"trace_id": "trace-id",
|
|
"semantic-similarity": 0.7,
|
|
}
|
|
|
|
forwarded_kwargs = cache.cache.get_cache.call_args.kwargs
|
|
assert forwarded_kwargs == {
|
|
"input": "What is the capital of France?",
|
|
"metadata": {
|
|
"user_api_key": "sk-secret",
|
|
"trace_id": "trace-id",
|
|
"semantic-similarity": 0.7,
|
|
},
|
|
}
|
|
assert forwarded_kwargs["metadata"] is not metadata
|
|
cache.get_cache_logic.assert_called_once_with(
|
|
cached_result={"content": "Paris"},
|
|
max_age=10,
|
|
)
|
|
|
|
|
|
def test_cache_get_cache_filters_sensitive_kwargs_without_metadata():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.cache = MagicMock()
|
|
cache.cache.get_cache = MagicMock(return_value={"content": "Paris"})
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
cache.get_cache_logic = MagicMock(return_value={"content": "Paris"})
|
|
|
|
result = cache.get_cache(
|
|
input="What is the capital of France?",
|
|
cache={"s-maxage": 10},
|
|
api_key="sk-secret",
|
|
headers={"authorization": "Bearer sk-secret"},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
cache.cache.get_cache.assert_called_once_with(
|
|
"test_key",
|
|
input="What is the capital of France?",
|
|
)
|
|
|
|
|
|
def test_cache_get_cache_passes_responses_input_to_dynamic_cache():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
cache.get_cache_logic = MagicMock(return_value={"content": "Paris"})
|
|
dynamic_cache_object = MagicMock()
|
|
dynamic_cache_object.get_cache = MagicMock(return_value={"content": "Paris"})
|
|
|
|
metadata = {}
|
|
result = cache.get_cache(
|
|
dynamic_cache_object=dynamic_cache_object,
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
cache={},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
dynamic_cache_object.get_cache.assert_called_once_with(
|
|
"test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
cache.get_cache_logic.assert_called_once_with(
|
|
cached_result={"content": "Paris"},
|
|
max_age=float("inf"),
|
|
)
|
|
|
|
|
|
def test_redis_sync_set_cache_passes_precomputed_vector():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.llmcache = MagicMock()
|
|
cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
cache._get_ttl = MagicMock(return_value=None)
|
|
cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
)
|
|
|
|
cache._get_embedding.assert_called_once()
|
|
cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
|
|
|
|
def test_redis_sync_get_cache_passes_precomputed_vector():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.similarity_threshold = 0.8
|
|
cache.llmcache = MagicMock()
|
|
cache.llmcache.check = MagicMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
)
|
|
cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
with patch.object(
|
|
cache, "_get_cache_key_filter_expression", return_value="cache-key-filter"
|
|
):
|
|
result = cache.get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata={},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
cache._get_embedding.assert_called_once()
|
|
cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_async_embedding_forwards_full_metadata(monkeypatch):
|
|
import types
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (None, None)
|
|
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = router
|
|
fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
await cache._get_async_embedding(
|
|
"hello",
|
|
metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"},
|
|
)
|
|
|
|
md = router.aembedding.call_args.kwargs["metadata"]
|
|
assert md["user_api_key"] == "sk-x"
|
|
assert md["user_api_key_team_id"] == "team-1" # FAILS today: team_id is dropped
|
|
assert md["semantic-cache-embedding"] is True
|
|
|
|
|
|
LONG_PROMPT = " ".join(f"token{i}" for i in range(300))
|
|
|
|
|
|
def _proxy_with_router(monkeypatch: pytest.MonkeyPatch, router: MagicMock, model_name: str) -> None:
|
|
import types
|
|
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = router
|
|
fake_proxy.llm_model_list = [{"model_name": model_name}]
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
|
|
def _token_count(model: str, text: str) -> int:
|
|
import litellm
|
|
|
|
return len(litellm.encode(model=model, text=text))
|
|
|
|
|
|
def test_redis_get_embedding_truncates_to_deployment_max_input_tokens(monkeypatch):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (5, None)
|
|
router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
|
|
_proxy_with_router(monkeypatch, router, "sem-embed")
|
|
|
|
assert cache._get_embedding(LONG_PROMPT) == [0.5, 0.6]
|
|
|
|
sent_input = router.embedding.call_args.kwargs["input"]
|
|
assert LONG_PROMPT.startswith(sent_input)
|
|
assert _token_count("sem-embed", sent_input) == 5
|
|
assert _token_count("sem-embed", LONG_PROMPT) > 5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_async_embedding_explicit_limit_beats_deployment_limit(monkeypatch):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
cache.embedding_max_input_tokens = 3
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (8191, None)
|
|
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
|
|
_proxy_with_router(monkeypatch, router, "sem-embed")
|
|
|
|
assert await cache._get_async_embedding(LONG_PROMPT) == [0.1, 0.2]
|
|
|
|
sent_input = router.aembedding.call_args.kwargs["input"]
|
|
assert _token_count("sem-embed", sent_input) == 3
|
|
|
|
|
|
def test_redis_get_embedding_truncates_direct_path_with_explicit_limit(monkeypatch):
|
|
import types
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "text-embedding-3-small"
|
|
cache.embedding_max_input_tokens = 4
|
|
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = None
|
|
fake_proxy.llm_model_list = None
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
with patch(
|
|
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2]}]}
|
|
) as direct_embed:
|
|
cache._get_embedding(LONG_PROMPT)
|
|
|
|
sent_input = direct_embed.call_args.kwargs["input"]
|
|
assert _token_count("text-embedding-3-small", sent_input) == 4
|
|
|
|
|
|
def test_redis_semantic_cache_init_stores_embedding_max_input_tokens(monkeypatch):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache(
|
|
redis_url="redis://localhost:6379",
|
|
similarity_threshold=0.8,
|
|
embedding_max_input_tokens=512,
|
|
)
|
|
assert cache.embedding_max_input_tokens == 512
|
|
default_cache = RedisSemanticCache(redis_url="redis://localhost:6379", similarity_threshold=0.8)
|
|
assert default_cache.embedding_max_input_tokens is None
|
|
|
|
|
|
def test_redis_init_defers_redisvl_construction(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
semantic_cache_mock.assert_not_called()
|
|
custom_vectorizer_mock.assert_not_called()
|
|
|
|
first = cache.llmcache
|
|
semantic_cache_mock.assert_called_once()
|
|
custom_vectorizer_mock.assert_called_once()
|
|
|
|
second = cache.llmcache
|
|
assert first is second
|
|
semantic_cache_mock.assert_called_once()
|
|
|
|
|
|
def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch):
|
|
built_cache = MagicMock()
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[ConnectionError("redis down"), built_cache]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
with pytest.raises(ConnectionError, match="redis down"):
|
|
_ = cache.llmcache
|
|
|
|
assert cache.llmcache is built_cache
|
|
assert semantic_cache_mock.call_count == 2
|
|
|
|
|
|
def test_redis_llmcache_setter_supported():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
sentinel = MagicMock()
|
|
cache.llmcache = sentinel
|
|
assert cache.llmcache is sentinel
|
|
|
|
|
|
def _router_proxy_module(router, model_name):
|
|
import types
|
|
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = router
|
|
fake_proxy.llm_model_list = [{"model_name": model_name}]
|
|
return fake_proxy
|
|
|
|
|
|
def test_redis_sync_embedding_call_is_bounded(monkeypatch):
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
cache.embedding_timeout = 1.5
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (None, None)
|
|
router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"litellm.proxy.proxy_server",
|
|
_router_proxy_module(router, "sem-embed"),
|
|
)
|
|
|
|
assert cache._get_embedding("hello") == [0.5, 0.6]
|
|
assert router.embedding.call_args.kwargs["timeout"] == 1.5
|
|
assert router.embedding.call_args.kwargs["num_retries"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_async_embedding_call_is_bounded(monkeypatch):
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
cache.embedding_timeout = 1.5
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (None, None)
|
|
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"litellm.proxy.proxy_server",
|
|
_router_proxy_module(router, "sem-embed"),
|
|
)
|
|
|
|
assert await cache._get_async_embedding("hello") == [0.5, 0.6]
|
|
assert router.aembedding.call_args.kwargs["timeout"] == 1.5
|
|
assert router.aembedding.call_args.kwargs["num_retries"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_async_embedding_gives_up_on_unresponsive_endpoint(monkeypatch):
|
|
import asyncio
|
|
import time
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
cache.embedding_timeout = 0.05
|
|
|
|
async def never_responds(**kwargs):
|
|
await asyncio.sleep(3)
|
|
return {"data": [{"embedding": [0.1, 0.2]}]}
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (None, None)
|
|
router.aembedding = never_responds
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"litellm.proxy.proxy_server",
|
|
_router_proxy_module(router, "sem-embed"),
|
|
)
|
|
|
|
started = time.monotonic()
|
|
with pytest.raises(ValueError, match="Failed to generate embedding"):
|
|
await cache._get_async_embedding("hello")
|
|
assert time.monotonic() - started < 1.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_async_get_cache_fails_open_when_embedding_hangs(monkeypatch):
|
|
import asyncio
|
|
import time
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
cache.embedding_timeout = 0.05
|
|
cache.similarity_threshold = 0.8
|
|
cache.distance_threshold = 0.2
|
|
cache.llmcache = MagicMock()
|
|
|
|
async def never_responds(**kwargs):
|
|
await asyncio.sleep(3)
|
|
return {"data": [{"embedding": [0.1, 0.2]}]}
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (None, None)
|
|
router.aembedding = never_responds
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"litellm.proxy.proxy_server",
|
|
_router_proxy_module(router, "sem-embed"),
|
|
)
|
|
|
|
metadata = {}
|
|
started = time.monotonic()
|
|
result = await cache.async_get_cache(
|
|
key="test_key",
|
|
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
|
metadata=metadata,
|
|
)
|
|
elapsed = time.monotonic() - started
|
|
|
|
assert result is None
|
|
assert metadata["semantic-similarity"] == 0.0
|
|
assert elapsed < 1.0
|
|
cache.llmcache.acheck.assert_not_called()
|
|
|
|
|
|
def test_cache_forwards_semantic_cache_embedding_timeout():
|
|
from litellm.caching.caching import Cache
|
|
from litellm.types.caching import LiteLLMCacheType
|
|
|
|
with patch.object(import_module("litellm.caching.caching"), "RedisSemanticCache") as backend:
|
|
Cache(
|
|
type=LiteLLMCacheType.REDIS_SEMANTIC,
|
|
similarity_threshold=0.8,
|
|
redis_url="redis://localhost:6379",
|
|
semantic_cache_embedding_timeout=2.5,
|
|
)
|
|
|
|
assert backend.call_args.kwargs["embedding_timeout"] == 2.5
|
|
|
|
|
|
def test_redis_semantic_cache_defaults_embedding_timeout():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
|
assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_async_embedding_truncates_off_the_event_loop(monkeypatch):
|
|
from tests.large_text import text
|
|
from tests.unit.litellm_core_utils.event_loop_lag import (
|
|
assert_loop_stayed_free,
|
|
timed_with_loop_lags,
|
|
warm_tokenizer,
|
|
)
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
warm_tokenizer("sem-embed")
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
cache.embedding_max_input_tokens = 5
|
|
cache.embedding_timeout = 5
|
|
|
|
router = route_entry_points_to_stubs(MagicMock())
|
|
router.get_configured_token_limits.return_value = (8191, None)
|
|
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
|
|
_proxy_with_router(monkeypatch, router, "sem-embed")
|
|
|
|
embedding, took, lags = await timed_with_loop_lags(lambda: cache._get_async_embedding(text * 100))
|
|
|
|
assert embedding == [0.1, 0.2]
|
|
assert _token_count("sem-embed", router.aembedding.call_args.kwargs["input"]) == 5
|
|
assert_loop_stayed_free(took, lags)
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_keeps_tool_result_text():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
messages=[
|
|
{"role": "user", "content": "list the files"},
|
|
{
|
|
"role": "assistant",
|
|
"content": [{"type": "tool_use", "id": "toolu_1", "name": "Bash", "input": {"command": "ls"}}],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "calc.py test_calc.py"}],
|
|
},
|
|
]
|
|
)
|
|
|
|
assert prompt == "list the filescalc.py test_calc.py"
|