scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-30 01:52:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(redis-semantic): preserve knn hits with scoped filters

Browse source
This commit is contained in:
Ritwij Aryan Parmar 2026-05-29 13:39:44 -04:00
parent 68852ef165
commit 63a491b1d7
3 changed files with 190 additions and 83 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/caching/caching.py
View file

@ -511,7 +511,6 @@ class Cache:
try: # never block execution
if self.should_use_cache(**kwargs) is not True:
return
messages = kwargs.get("messages", [])
if "cache_key" in kwargs:
cache_key = kwargs["cache_key"]
else:
@ -523,12 +522,13 @@ class Cache:
or cache_control_args.get("s-max-age")
or float("inf")
)
cache_kwargs = {k: v for k, v in kwargs.items() if k != "cache_key"}
if dynamic_cache_object is not None:
cached_result = dynamic_cache_object.get_cache(
cache_key, messages=messages
cache_key, **cache_kwargs
)
else:
cached_result = self.cache.get_cache(cache_key, messages=messages)
cached_result = self.cache.get_cache(cache_key, **cache_kwargs)
return self._get_cache_logic(
cached_result=cached_result, max_age=max_age
)

110
litellm/caching/redis_semantic_cache.py
View file

@ -36,6 +36,15 @@ class RedisSemanticCache(BaseCache):
DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index"
CACHE_KEY_FIELD_NAME: str = "litellm_cache_key"
API_KEY_HASH_FIELD_NAME: str = "litellm_user_api_key_hash"
TEAM_ID_FIELD_NAME: str = "litellm_user_api_key_team_id"
USER_ID_FIELD_NAME: str = "litellm_user_id"
_SCOPE_FIELD_TO_METADATA_KEY: Tuple[Tuple[str, str], ...] = (
(API_KEY_HASH_FIELD_NAME, "user_api_key_hash"),
(TEAM_ID_FIELD_NAME, "user_api_key_team_id"),
(TEAM_ID_FIELD_NAME, "team_id"),
(USER_ID_FIELD_NAME, "user_id"),
)
def __init__(
self,
@ -124,6 +133,15 @@ class RedisSemanticCache(BaseCache):
"type": "tag",
}
@classmethod
def _filterable_fields(cls) -> List[Dict[str, str]]:
return [
cls._cache_key_filterable_field(),
{"name": cls.API_KEY_HASH_FIELD_NAME, "type": "tag"},
{"name": cls.TEAM_ID_FIELD_NAME, "type": "tag"},
{"name": cls.USER_ID_FIELD_NAME, "type": "tag"},
]
def _init_semantic_cache(
self,
semantic_cache_cls: Any,
@ -144,7 +162,7 @@ class RedisSemanticCache(BaseCache):
redis_url=redis_url,
vectorizer=cache_vectorizer,
distance_threshold=self.distance_threshold,
filterable_fields=[self._cache_key_filterable_field()],
filterable_fields=self._filterable_fields(),
overwrite=False,
)
except ValueError as exc:
@ -162,7 +180,7 @@ class RedisSemanticCache(BaseCache):
redis_url=redis_url,
vectorizer=cache_vectorizer,
distance_threshold=self.distance_threshold,
filterable_fields=[self._cache_key_filterable_field()],
filterable_fields=self._filterable_fields(),
overwrite=False,
)
except ValueError as isolated_exc:
@ -178,25 +196,59 @@ class RedisSemanticCache(BaseCache):
redis_url=redis_url,
vectorizer=cache_vectorizer,
distance_threshold=self.distance_threshold,
filterable_fields=[self._cache_key_filterable_field()],
filterable_fields=self._filterable_fields(),
overwrite=True,
)
def _get_cache_filters(self, key: str) -> Dict[str, str]:
return {self.CACHE_KEY_FIELD_NAME: str(key)}
def _get_metadata_from_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
metadata: Dict[str, Any] = {}
for key in ("metadata", "litellm_metadata"):
value = kwargs.get(key)
if isinstance(value, dict):
metadata.update(value)
litellm_params = kwargs.get("litellm_params")
if isinstance(litellm_params, dict):
value = litellm_params.get("metadata")
if isinstance(value, dict):
metadata.update(value)
return metadata
def _get_scope_filter(self, **kwargs) -> Optional[Tuple[str, str]]:
metadata = self._get_metadata_from_kwargs(kwargs)
for field_name, metadata_key in self._SCOPE_FIELD_TO_METADATA_KEY:
value = metadata.get(metadata_key)
if value:
return field_name, str(value)
return None
def _get_cache_filters(self, key: str, **kwargs) -> Dict[str, str]:
filters = {self.CACHE_KEY_FIELD_NAME: str(key)}
scope_filter = self._get_scope_filter(**kwargs)
if scope_filter is not None:
filters[scope_filter[0]] = scope_filter[1]
return filters
def _get_scope_filter_expression(self, **kwargs) -> Any:
scope_filter = self._get_scope_filter(**kwargs)
if scope_filter is None:
return None
def _get_cache_key_filter_expression(self, key: str) -> Any:
from redisvl.query.filter import Tag # type: ignore[import-not-found, import-untyped]
return Tag(self.CACHE_KEY_FIELD_NAME) == str(key)
return Tag(scope_filter[0]) == scope_filter[1]
def _cache_hit_matches_key(self, cache_hit: Dict[str, Any], key: str) -> bool:
# Pre-isolation entries with no ``litellm_cache_key`` field cannot be
# safely reassigned to a caller's scope and are treated as misses.
cached_key = cache_hit.get(self.CACHE_KEY_FIELD_NAME)
if isinstance(cached_key, bytes):
cached_key = cached_key.decode("utf-8")
return cached_key is not None and str(cached_key) == str(key)
def _cache_hit_matches_scope(self, cache_hit: Dict[str, Any], **kwargs) -> bool:
scope_filter = self._get_scope_filter(**kwargs)
if scope_filter is None:
return True
field_name, expected_value = scope_filter
cached_value = cache_hit.get(field_name)
if isinstance(cached_value, bytes):
cached_value = cached_value.decode("utf-8")
return cached_value is not None and str(cached_value) == expected_value
def _get_ttl(self, **kwargs) -> Optional[int]:
"""
@ -288,7 +340,7 @@ class RedisSemanticCache(BaseCache):
value_str = str(value)
store_kwargs: Dict[str, Any] = {
"filters": self._get_cache_filters(key),
"filters": self._get_cache_filters(key, **kwargs),
}
# Get TTL and store in Redis semantic cache
@ -323,12 +375,15 @@ class RedisSemanticCache(BaseCache):
return None
prompt = get_str_from_messages(messages)
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
# Keep caller isolation separate from prompt similarity. Filtering
# on the full request hash would make KNN unreachable for
# non-identical prompts.
check_kwargs: Dict[str, Any] = {
"prompt": prompt,
"filter_expression": self._get_cache_key_filter_expression(key),
}
filter_expression = self._get_scope_filter_expression(**kwargs)
if filter_expression is not None:
check_kwargs["filter_expression"] = filter_expression
results = self.llmcache.check(**check_kwargs)
# Return None if no similar prompts found
@ -338,8 +393,8 @@ class RedisSemanticCache(BaseCache):
# Process the best matching result
cache_hit = results[0]
if not self._cache_hit_matches_key(cache_hit=cache_hit, key=key):
print_verbose("Redis semantic-cache hit did not match cache key scope")
if not self._cache_hit_matches_scope(cache_hit=cache_hit, **kwargs):
print_verbose("Redis semantic-cache hit did not match caller scope")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
vector_distance = float(cache_hit["vector_distance"])
@ -442,7 +497,7 @@ class RedisSemanticCache(BaseCache):
store_kwargs: Dict[str, Any] = {
"vector": prompt_embedding,
"filters": self._get_cache_filters(key),
"filters": self._get_cache_filters(key, **kwargs),
}
# Get TTL and store in Redis semantic cache
@ -483,13 +538,16 @@ class RedisSemanticCache(BaseCache):
# Generate embedding for the prompt
prompt_embedding = await self._get_async_embedding(prompt, **kwargs)
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
# Keep caller isolation separate from prompt similarity. Filtering
# on the full request hash would make KNN unreachable for
# non-identical prompts.
check_kwargs: Dict[str, Any] = {
"prompt": prompt,
"vector": prompt_embedding,
"filter_expression": self._get_cache_key_filter_expression(key),
}
filter_expression = self._get_scope_filter_expression(**kwargs)
if filter_expression is not None:
check_kwargs["filter_expression"] = filter_expression
results = await self.llmcache.acheck(**check_kwargs)
# handle results / cache hit
@ -498,8 +556,8 @@ class RedisSemanticCache(BaseCache):
return None
cache_hit = results[0]
if not self._cache_hit_matches_key(cache_hit=cache_hit, key=key):
print_verbose("Redis semantic-cache hit did not match cache key scope")
if not self._cache_hit_matches_scope(cache_hit=cache_hit, **kwargs):
print_verbose("Redis semantic-cache hit did not match caller scope")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
vector_distance = float(cache_hit["vector_distance"])

157
tests/test_litellm/caching/test_redis_semantic_cache.py
View file

@ -72,22 +72,15 @@ def test_redis_semantic_cache_get_cache(monkeypatch):
"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",
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "different_request_hash",
}
]
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",
),
with patch(
"litellm.embedding",
return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
):
# Test get_cache with a message
metadata = {}
@ -104,11 +97,53 @@ def test_redis_semantic_cache_get_cache(monkeypatch):
# Verify llmcache.check was called
redis_semantic_cache.llmcache.check.assert_called_once_with(
prompt="What is the capital of France?",
filter_expression="cache-key-filter",
)
def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch):
def test_cache_get_cache_forwards_metadata_to_semantic_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="cache-key")
cache._get_cache_logic = MagicMock(return_value={"content": "Paris"})
cache.cache = MagicMock()
cache.cache.get_cache.return_value = {"response": '{"content":"Paris"}'}
metadata = {"user_api_key_hash": "hashed-key"}
result = cache.get_cache(
model="gpt-4o-mini",
messages=[{"content": "Which city is the capital of France?"}],
metadata=metadata,
cache={},
)
assert result == {"content": "Paris"}
cache.cache.get_cache.assert_called_once_with(
"cache-key",
model="gpt-4o-mini",
messages=[{"content": "Which city is the capital of France?"}],
metadata=metadata,
cache={},
)
cache.cache.get_cache.reset_mock()
cache.get_cache(
cache_key="preset-key",
messages=[{"content": "cached prompt"}],
metadata=metadata,
cache={},
)
cache.cache.get_cache.assert_called_once_with(
"preset-key",
messages=[{"content": "cached prompt"}],
metadata=metadata,
cache={},
)
def test_redis_semantic_cache_rejects_cross_scope_cache_hit(monkeypatch):
semantic_cache_mock = MagicMock()
custom_vectorizer_mock = MagicMock()
@ -134,16 +169,17 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch):
"prompt": "What is the capital of France?",
"response": '{"content": "Paris"}',
"vector_distance": 0.1,
RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "other-key",
}
]
)
with patch.object(
redis_semantic_cache,
"_get_cache_key_filter_expression",
return_value="cache-key-filter",
"_get_scope_filter_expression",
return_value="caller-scope-filter",
):
metadata = {}
metadata = {"user_api_key_hash": "this-key"}
result = redis_semantic_cache.get_cache(
key="test_key",
messages=[{"content": "What is the capital of France?"}],
@ -152,6 +188,10 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch):
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?",
filter_expression="caller-scope-filter",
)
def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch):
@ -180,13 +220,17 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch):
key="test_key",
value={"content": "Paris"},
messages=[{"content": "What is the capital of France?"}],
metadata={"user_api_key_hash": "hashed-key"},
ttl=60,
)
redis_semantic_cache.llmcache.store.assert_called_once_with(
"What is the capital of France?",
"{'content': 'Paris'}",
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
filters={
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "hashed-key",
},
ttl=60,
)
@ -227,9 +271,10 @@ def test_redis_semantic_cache_uses_isolated_index_for_old_schema(monkeypatch):
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()
]
assert (
semantic_cache_mock.call_args_list[1].kwargs["filterable_fields"]
== RedisSemanticCache._filterable_fields()
)
def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch):
@ -269,9 +314,10 @@ def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch):
== "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()
]
assert (
semantic_cache_mock.call_args_list[2].kwargs["filterable_fields"]
== RedisSemanticCache._filterable_fields()
)
def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypatch):
@ -321,20 +367,19 @@ def test_redis_semantic_cache_reraises_unexpected_index_error():
)
def test_redis_semantic_cache_matches_bytes_cache_key():
def test_redis_semantic_cache_matches_bytes_scope():
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",
assert redis_semantic_cache._cache_hit_matches_scope(
cache_hit={RedisSemanticCache.API_KEY_HASH_FIELD_NAME: b"hashed-key"},
metadata={"user_api_key_hash": "hashed-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."""
def test_redis_semantic_cache_rejects_pre_scope_hit_for_scoped_request():
"""Entries with no caller scope cannot be safely reused for a scoped request."""
from litellm.caching.redis_semantic_cache import RedisSemanticCache
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
@ -344,13 +389,13 @@ def test_redis_semantic_cache_rejects_pre_isolation_unscoped_hit():
"response": '{"content": "Paris"}',
"vector_distance": 0.1,
}
assert not redis_semantic_cache._cache_hit_matches_key(
assert not redis_semantic_cache._cache_hit_matches_scope(
cache_hit=cache_hit,
key="test_key",
metadata={"user_api_key_hash": "hashed-key"},
)
def test_redis_semantic_cache_builds_filter_expression(monkeypatch):
def test_redis_semantic_cache_builds_scope_filter_expression(monkeypatch):
class FakeTag:
def __init__(self, field_name):
self.field_name = field_name
@ -363,10 +408,13 @@ def test_redis_semantic_cache_builds_filter_expression(monkeypatch):
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == (
RedisSemanticCache.CACHE_KEY_FIELD_NAME,
"test_key",
assert redis_semantic_cache._get_scope_filter_expression(
metadata={"user_api_key_hash": "hashed-key"}
) == (
RedisSemanticCache.API_KEY_HASH_FIELD_NAME,
"hashed-key",
)
assert redis_semantic_cache._get_scope_filter_expression(metadata={}) is None
@pytest.mark.asyncio
@ -400,7 +448,7 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch):
"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",
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "different_request_hash",
}
]
@ -409,17 +457,12 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch):
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={},
)
# 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."}
@ -429,12 +472,13 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch):
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):
async def test_redis_semantic_cache_async_get_cache_rejects_cross_scope_hit(
monkeypatch,
):
semantic_cache_mock = MagicMock()
custom_vectorizer_mock = MagicMock()
@ -460,6 +504,7 @@ async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeyp
"prompt": "What is the capital of France?",
"response": '{"content": "Paris"}',
"vector_distance": 0.1,
RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "other-key",
}
]
)
@ -469,13 +514,13 @@ async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeyp
with patch.object(
redis_semantic_cache,
"_get_cache_key_filter_expression",
return_value="cache-key-filter",
"_get_scope_filter_expression",
return_value="caller-scope-filter",
):
result = await redis_semantic_cache.async_get_cache(
key="test_key",
messages=[{"content": "What is the capital of France?"}],
metadata={},
metadata={"user_api_key_hash": "this-key"},
)
assert result is None
@ -513,6 +558,7 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter(
key="test_key",
value={"content": "Paris"},
messages=[{"content": "What is the capital of France?"}],
metadata={"user_api_key_hash": "hashed-key"},
ttl=60,
)
@ -520,6 +566,9 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter(
"What is the capital of France?",
"{'content': 'Paris'}",
vector=[0.1, 0.2, 0.3],
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
filters={
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "hashed-key",
},
ttl=60,
)