Merge pull request #26990 from stuxf/codex/semantic-cache-tenant-isolation

chore(caching): isolate semantic cache entries
This commit is contained in:
yuneng-jiang
2026-05-04 16:02:43 -07:00
committed by GitHub
4 changed files with 882 additions and 141 deletions
+123 -71
View File
@@ -11,17 +11,23 @@ Has 4 methods:
import ast
import asyncio
import json
from typing import Any, cast
import os
from typing import Any, Dict, cast
import litellm
from litellm._logging import print_verbose
from litellm.constants import QDRANT_SCALAR_QUANTILE, QDRANT_VECTOR_SIZE
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
from litellm.types.utils import EmbeddingResponse
from .base_cache import BaseCache
class QdrantSemanticCache(BaseCache):
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
def __init__( # noqa: PLR0915
self,
qdrant_api_base=None,
@@ -33,8 +39,6 @@ class QdrantSemanticCache(BaseCache):
host_type=None,
vector_size=None,
):
import os
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
@@ -115,7 +119,9 @@ class QdrantSemanticCache(BaseCache):
print_verbose(
f"Collection already exists.\nCollection details:{self.collection_info}"
)
self._ensure_cache_key_payload_index()
else:
quantization_params: Dict[str, Any]
if quantization_config is None or quantization_config == "binary":
quantization_params = {
"binary": {
@@ -156,6 +162,7 @@ class QdrantSemanticCache(BaseCache):
print_verbose(
f"New collection created.\nCollection details:{self.collection_info}"
)
self._ensure_cache_key_payload_index()
else:
raise Exception("Error while creating new collection")
@@ -170,15 +177,94 @@ class QdrantSemanticCache(BaseCache):
cached_response = ast.literal_eval(cached_response)
return cached_response
def _get_qdrant_cache_key_filter(self, key: str) -> dict:
return {
"must": [
{
"key": self.CACHE_KEY_FIELD_NAME,
"match": {"value": str(key)},
}
]
}
def _add_cache_key_filter_to_search_data(self, data: dict, key: str) -> None:
data["filter"] = self._get_qdrant_cache_key_filter(key)
def _ensure_cache_key_payload_index(self) -> None:
try:
response = self.sync_client.put(
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/index",
headers=self.headers,
json={
"field_name": self.CACHE_KEY_FIELD_NAME,
"field_schema": "keyword",
},
)
if response.status_code not in (200, 201):
print_verbose(
"Qdrant semantic-cache could not create cache-key payload index: "
f"{response.text}"
)
except Exception as exc:
print_verbose(
"Qdrant semantic-cache could not create cache-key payload index: "
f"{str(exc)}"
)
def _payload_matches_cache_key(self, payload: dict, key: str) -> bool:
# Pre-isolation points stored only prompt + response with no cache-key
# payload field. Reassigning them to a caller's key would risk
# cross-scope hits, so they're treated as misses and re-populated on
# the next set_cache.
cached_key = payload.get(self.CACHE_KEY_FIELD_NAME)
return cached_key is not None and str(cached_key) == str(key)
async def _get_async_embedding(self, prompt: str, **kwargs) -> Any:
llm_model_list = None
llm_router = None
try:
from litellm.proxy.proxy_server import (
llm_model_list as proxy_llm_model_list,
llm_router as proxy_llm_router,
)
llm_model_list = proxy_llm_model_list
llm_router = proxy_llm_router
except ImportError:
pass
router_model_names = (
[m["model_name"] for m in llm_model_list]
if llm_model_list is not None
else []
)
if llm_router is not None and self.embedding_model in router_model_names:
user_api_key = kwargs.get("metadata", {}).get("user_api_key", "")
return await llm_router.aembedding(
model=self.embedding_model,
input=prompt,
cache={"no-store": True, "no-cache": True},
metadata={
"user_api_key": user_api_key,
"semantic-cache-embedding": True,
"trace_id": kwargs.get("metadata", {}).get("trace_id", None),
},
)
return await litellm.aembedding(
model=self.embedding_model,
input=prompt,
cache={"no-store": True, "no-cache": True},
)
def set_cache(self, key, value, **kwargs):
print_verbose(f"qdrant semantic-cache set_cache, kwargs: {kwargs}")
from litellm._uuid import uuid
# get the prompt
messages = kwargs["messages"]
prompt = ""
for message in messages:
prompt += message["content"]
prompt = get_str_from_messages(messages)
# create an embedding for prompt
embedding_response = cast(
@@ -202,6 +288,7 @@ class QdrantSemanticCache(BaseCache):
"id": str(uuid.uuid4()),
"vector": embedding,
"payload": {
self.CACHE_KEY_FIELD_NAME: str(key),
"text": prompt,
"response": value,
},
@@ -220,9 +307,7 @@ class QdrantSemanticCache(BaseCache):
# get the messages
messages = kwargs["messages"]
prompt = ""
for message in messages:
prompt += message["content"]
prompt = get_str_from_messages(messages)
# convert to embedding
embedding_response = cast(
@@ -249,6 +334,7 @@ class QdrantSemanticCache(BaseCache):
"limit": 1,
"with_payload": True,
}
self._add_cache_key_filter_to_search_data(data=data, key=key)
search_response = self.sync_client.post(
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points/search",
@@ -258,21 +344,33 @@ class QdrantSemanticCache(BaseCache):
results = search_response.json()["result"]
if results is None:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
if isinstance(results, list):
if len(results) == 0:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
similarity = results[0]["score"]
cached_prompt = results[0]["payload"]["text"]
payload = results[0]["payload"]
if not self._payload_matches_cache_key(payload=payload, key=key):
print_verbose("Qdrant semantic-cache hit did not match cache key scope")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
cached_prompt = payload["text"]
# check similarity, if more than self.similarity_threshold, return results
print_verbose(
f"semantic cache: similarity threshold: {self.similarity_threshold}, similarity: {similarity}, prompt: {prompt}, closest_cached_prompt: {cached_prompt}"
)
# update kwargs["metadata"] with similarity, don't rewrite the original metadata
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
if similarity >= self.similarity_threshold:
# cache hit !
cached_value = results[0]["payload"]["response"]
cached_value = payload["response"]
print_verbose(
f"got a cache hit, similarity: {similarity}, Current prompt: {prompt}, cached_prompt: {cached_prompt}"
)
@@ -285,40 +383,12 @@ class QdrantSemanticCache(BaseCache):
async def async_set_cache(self, key, value, **kwargs):
from litellm._uuid import uuid
from litellm.proxy.proxy_server import llm_model_list, llm_router
print_verbose(f"async qdrant semantic-cache set_cache, kwargs: {kwargs}")
# get the prompt
messages = kwargs["messages"]
prompt = ""
for message in messages:
prompt += message["content"]
# create an embedding for prompt
router_model_names = (
[m["model_name"] for m in llm_model_list]
if llm_model_list is not None
else []
)
if llm_router is not None and self.embedding_model in router_model_names:
user_api_key = kwargs.get("metadata", {}).get("user_api_key", "")
embedding_response = await llm_router.aembedding(
model=self.embedding_model,
input=prompt,
cache={"no-store": True, "no-cache": True},
metadata={
"user_api_key": user_api_key,
"semantic-cache-embedding": True,
"trace_id": kwargs.get("metadata", {}).get("trace_id", None),
},
)
else:
# convert to embedding
embedding_response = await litellm.aembedding(
model=self.embedding_model,
input=prompt,
cache={"no-store": True, "no-cache": True},
)
prompt = get_str_from_messages(messages)
embedding_response = await self._get_async_embedding(prompt, **kwargs)
# get the embedding
embedding = embedding_response["data"][0]["embedding"]
@@ -332,6 +402,7 @@ class QdrantSemanticCache(BaseCache):
"id": str(uuid.uuid4()),
"vector": embedding,
"payload": {
self.CACHE_KEY_FIELD_NAME: str(key),
"text": prompt,
"response": value,
},
@@ -348,38 +419,12 @@ class QdrantSemanticCache(BaseCache):
async def async_get_cache(self, key, **kwargs):
print_verbose(f"async qdrant semantic-cache get_cache, kwargs: {kwargs}")
from litellm.proxy.proxy_server import llm_model_list, llm_router
# get the messages
messages = kwargs["messages"]
prompt = ""
for message in messages:
prompt += message["content"]
prompt = get_str_from_messages(messages)
router_model_names = (
[m["model_name"] for m in llm_model_list]
if llm_model_list is not None
else []
)
if llm_router is not None and self.embedding_model in router_model_names:
user_api_key = kwargs.get("metadata", {}).get("user_api_key", "")
embedding_response = await llm_router.aembedding(
model=self.embedding_model,
input=prompt,
cache={"no-store": True, "no-cache": True},
metadata={
"user_api_key": user_api_key,
"semantic-cache-embedding": True,
"trace_id": kwargs.get("metadata", {}).get("trace_id", None),
},
)
else:
# convert to embedding
embedding_response = await litellm.aembedding(
model=self.embedding_model,
input=prompt,
cache={"no-store": True, "no-cache": True},
)
embedding_response = await self._get_async_embedding(prompt, **kwargs)
# get the embedding
embedding = embedding_response["data"][0]["embedding"]
@@ -396,6 +441,7 @@ class QdrantSemanticCache(BaseCache):
"limit": 1,
"with_payload": True,
}
self._add_cache_key_filter_to_search_data(data=data, key=key)
search_response = await self.async_client.post(
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points/search",
@@ -414,7 +460,13 @@ class QdrantSemanticCache(BaseCache):
return None
similarity = results[0]["score"]
cached_prompt = results[0]["payload"]["text"]
payload = results[0]["payload"]
if not self._payload_matches_cache_key(payload=payload, key=key):
print_verbose("Qdrant semantic-cache hit did not match cache key scope")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
cached_prompt = payload["text"]
# check similarity, if more than self.similarity_threshold, return results
print_verbose(
@@ -426,7 +478,7 @@ class QdrantSemanticCache(BaseCache):
if similarity >= self.similarity_threshold:
# cache hit !
cached_value = results[0]["payload"]["response"]
cached_value = payload["response"]
print_verbose(
f"got a cache hit, similarity: {similarity}, Current prompt: {prompt}, cached_prompt: {cached_prompt}"
)
+139 -33
View File
@@ -35,6 +35,7 @@ class RedisSemanticCache(BaseCache):
"""
DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index"
CACHE_KEY_FIELD_NAME: str = "litellm_cache_key"
def __init__(
self,
@@ -66,8 +67,8 @@ class RedisSemanticCache(BaseCache):
Exception: If similarity_threshold is not provided or required Redis
connection information is missing
"""
from redisvl.extensions.llmcache import SemanticCache
from redisvl.utils.vectorize import CustomTextVectorizer
from redisvl.extensions.llmcache import SemanticCache # type: ignore[import-not-found, import-untyped]
from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped]
if index_name is None:
index_name = self.DEFAULT_REDIS_INDEX_NAME
@@ -109,14 +110,94 @@ class RedisSemanticCache(BaseCache):
# Initialize the Redis vectorizer and cache
cache_vectorizer = CustomTextVectorizer(self._get_embedding)
self.llmcache = SemanticCache(
name=index_name,
self.llmcache = self._init_semantic_cache(
semantic_cache_cls=SemanticCache,
index_name=index_name,
redis_url=redis_url,
vectorizer=cache_vectorizer,
distance_threshold=self.distance_threshold,
overwrite=False,
cache_vectorizer=cache_vectorizer,
)
@classmethod
def _cache_key_filterable_field(cls) -> Dict[str, str]:
return {
"name": cls.CACHE_KEY_FIELD_NAME,
"type": "tag",
}
def _init_semantic_cache(
self,
semantic_cache_cls: Any,
index_name: str,
redis_url: str,
cache_vectorizer: Any,
) -> Any:
def _is_schema_mismatch(exc: ValueError) -> bool:
error_message = str(exc).lower()
return any(
phrase in error_message
for phrase in ("schema does not match", "index schema")
)
try:
return semantic_cache_cls(
name=index_name,
redis_url=redis_url,
vectorizer=cache_vectorizer,
distance_threshold=self.distance_threshold,
filterable_fields=[self._cache_key_filterable_field()],
overwrite=False,
)
except ValueError as exc:
if not _is_schema_mismatch(exc):
raise
isolated_index_name = f"{index_name}_isolated"
print_verbose(
"Redis semantic-cache existing index schema is not isolated; "
f"using isolated index - {isolated_index_name}"
)
try:
return semantic_cache_cls(
name=isolated_index_name,
redis_url=redis_url,
vectorizer=cache_vectorizer,
distance_threshold=self.distance_threshold,
filterable_fields=[self._cache_key_filterable_field()],
overwrite=False,
)
except ValueError as isolated_exc:
if not _is_schema_mismatch(isolated_exc):
raise
print_verbose(
"Redis semantic-cache isolated index schema is stale; "
f"recreating isolated index - {isolated_index_name}"
)
return semantic_cache_cls(
name=isolated_index_name,
redis_url=redis_url,
vectorizer=cache_vectorizer,
distance_threshold=self.distance_threshold,
filterable_fields=[self._cache_key_filterable_field()],
overwrite=True,
)
def _get_cache_filters(self, key: str) -> Dict[str, str]:
return {self.CACHE_KEY_FIELD_NAME: str(key)}
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)
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 _get_ttl(self, **kwargs) -> Optional[int]:
"""
Get the TTL (time-to-live) value for cache entries.
@@ -188,7 +269,7 @@ class RedisSemanticCache(BaseCache):
Store a value in the semantic cache.
Args:
key: The cache key (not directly used in semantic caching)
key: The cache key used to isolate semantic cache entries
value: The response value to cache
**kwargs: Additional arguments including 'messages' for the prompt
and optional 'ttl' for time-to-live
@@ -206,12 +287,15 @@ class RedisSemanticCache(BaseCache):
prompt = get_str_from_messages(messages)
value_str = str(value)
store_kwargs: Dict[str, Any] = {
"filters": self._get_cache_filters(key),
}
# Get TTL and store in Redis semantic cache
ttl = self._get_ttl(**kwargs)
if ttl is not None:
self.llmcache.store(prompt, value_str, ttl=int(ttl))
else:
self.llmcache.store(prompt, value_str)
store_kwargs["ttl"] = int(ttl)
self.llmcache.store(prompt, value_str, **store_kwargs)
except Exception as e:
print_verbose(
f"Error setting {value_str or value} in the Redis semantic cache: {str(e)}"
@@ -222,7 +306,7 @@ class RedisSemanticCache(BaseCache):
Retrieve a semantically similar cached response.
Args:
key: The cache key (not directly used in semantic caching)
key: The cache key used to isolate semantic cache entries
**kwargs: Additional arguments including 'messages' for the prompt
Returns:
@@ -235,18 +319,29 @@ class RedisSemanticCache(BaseCache):
messages = kwargs.get("messages", [])
if not messages:
print_verbose("No messages provided for semantic cache lookup")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
prompt = get_str_from_messages(messages)
# Check the cache for semantically similar prompts
results = self.llmcache.check(prompt=prompt)
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
check_kwargs: Dict[str, Any] = {
"prompt": prompt,
"filter_expression": self._get_cache_key_filter_expression(key),
}
results = self.llmcache.check(**check_kwargs)
# Return None if no similar prompts found
if not results:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
# 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")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
vector_distance = float(cache_hit["vector_distance"])
# Convert vector distance back to similarity score
@@ -257,6 +352,9 @@ class RedisSemanticCache(BaseCache):
cached_prompt = cache_hit["prompt"]
cached_response = cache_hit["response"]
# update kwargs["metadata"] with similarity, don't rewrite the original metadata
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
print_verbose(
f"Cache hit: similarity threshold: {self.similarity_threshold}, "
f"actual similarity: {similarity}, "
@@ -267,6 +365,7 @@ class RedisSemanticCache(BaseCache):
return self._get_cache_logic(cached_response=cached_response)
except Exception as e:
print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def _get_async_embedding(self, prompt: str, **kwargs) -> List[float]:
"""
@@ -321,7 +420,7 @@ class RedisSemanticCache(BaseCache):
Asynchronously store a value in the semantic cache.
Args:
key: The cache key (not directly used in semantic caching)
key: The cache key used to isolate semantic cache entries
value: The response value to cache
**kwargs: Additional arguments including 'messages' for the prompt
and optional 'ttl' for time-to-live
@@ -341,21 +440,20 @@ class RedisSemanticCache(BaseCache):
# Generate embedding for the value (response) to cache
prompt_embedding = await self._get_async_embedding(prompt, **kwargs)
store_kwargs: Dict[str, Any] = {
"vector": prompt_embedding,
"filters": self._get_cache_filters(key),
}
# Get TTL and store in Redis semantic cache
ttl = self._get_ttl(**kwargs)
if ttl is not None:
await self.llmcache.astore(
prompt,
value_str,
vector=prompt_embedding, # Pass through custom embedding
ttl=ttl,
)
else:
await self.llmcache.astore(
prompt,
value_str,
vector=prompt_embedding, # Pass through custom embedding
)
store_kwargs["ttl"] = ttl
await self.llmcache.astore(
prompt,
value_str,
**store_kwargs,
)
except Exception as e:
print_verbose(f"Error in async_set_cache: {str(e)}")
@@ -364,7 +462,7 @@ class RedisSemanticCache(BaseCache):
Asynchronously retrieve a semantically similar cached response.
Args:
key: The cache key (not directly used in semantic caching)
key: The cache key used to isolate semantic cache entries
**kwargs: Additional arguments including 'messages' for the prompt
Returns:
@@ -385,17 +483,25 @@ class RedisSemanticCache(BaseCache):
# Generate embedding for the prompt
prompt_embedding = await self._get_async_embedding(prompt, **kwargs)
# Check the cache for semantically similar prompts
results = await self.llmcache.acheck(prompt=prompt, vector=prompt_embedding)
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
check_kwargs: Dict[str, Any] = {
"prompt": prompt,
"vector": prompt_embedding,
"filter_expression": self._get_cache_key_filter_expression(key),
}
results = await self.llmcache.acheck(**check_kwargs)
# handle results / cache hit
if not results:
kwargs.setdefault("metadata", {})[
"semantic-similarity"
] = 0.0 # TODO why here but not above??
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
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")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
vector_distance = float(cache_hit["vector_distance"])
# Convert vector distance back to similarity
@@ -19,9 +19,7 @@ def test_qdrant_semantic_cache_initialization(monkeypatch):
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
@@ -31,6 +29,9 @@ def test_qdrant_semantic_cache_initialization(monkeypatch):
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_index_response = MagicMock()
mock_index_response.status_code = 200
mock_sync_client_instance.put.return_value = mock_index_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
@@ -48,6 +49,17 @@ def test_qdrant_semantic_cache_initialization(monkeypatch):
assert qdrant_cache.qdrant_api_base == "http://test.qdrant.local"
assert qdrant_cache.qdrant_api_key == "test_key"
assert qdrant_cache.similarity_threshold == 0.8
mock_sync_client_instance.put.assert_called_once_with(
url="http://test.qdrant.local/collections/test_collection/index",
headers={
"Content-Type": "application/json",
"api-key": "test_key",
},
json={
"field_name": QdrantSemanticCache.CACHE_KEY_FIELD_NAME,
"field_schema": "keyword",
},
)
# Test initialization with missing similarity_threshold
with pytest.raises(Exception, match="similarity_threshold must be provided"):
@@ -67,9 +79,7 @@ def test_qdrant_semantic_cache_get_cache_hit():
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
@@ -98,6 +108,7 @@ def test_qdrant_semantic_cache_get_cache_hit():
"result": [
{
"payload": {
QdrantSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
"text": "What is the capital of France?", # Original prompt
"response": '{"id": "test-123", "choices": [{"message": {"content": "Paris is the capital of France."}}]}',
},
@@ -127,6 +138,177 @@ def test_qdrant_semantic_cache_get_cache_hit():
# Verify search was called
qdrant_cache.sync_client.post.assert_called()
assert qdrant_cache.sync_client.post.call_args.kwargs["json"]["filter"] == {
"must": [
{
"key": QdrantSemanticCache.CACHE_KEY_FIELD_NAME,
"match": {"value": "test_key"},
}
]
}
def test_qdrant_semantic_cache_rejects_unscoped_cache_hit():
"""
Test QDRANT semantic cache rejects old or unscoped cache hits.
Legacy points have only text and response payloads, so they cannot be
safely migrated to a generated LiteLLM cache key.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
mock_search_response = MagicMock()
mock_search_response.status_code = 200
mock_search_response.json.return_value = {
"result": [
{
"payload": {
"text": "What is the capital of France?",
"response": '{"id": "test-123"}',
},
"score": 0.9,
}
]
}
qdrant_cache.sync_client.post = MagicMock(return_value=mock_search_response)
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
metadata = {}
result = qdrant_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_qdrant_semantic_cache_payload_index_failure_is_non_blocking():
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
qdrant_cache.qdrant_api_base = "http://test.qdrant.local"
qdrant_cache.collection_name = "test_collection"
qdrant_cache.headers = {"Content-Type": "application/json"}
qdrant_cache.sync_client = MagicMock()
response = MagicMock()
response.status_code = 400
response.text = "bad index"
qdrant_cache.sync_client.put.return_value = response
qdrant_cache._ensure_cache_key_payload_index()
qdrant_cache.sync_client.put.assert_called_once()
def test_qdrant_semantic_cache_payload_index_exception_is_non_blocking():
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
qdrant_cache.qdrant_api_base = "http://test.qdrant.local"
qdrant_cache.collection_name = "test_collection"
qdrant_cache.headers = {"Content-Type": "application/json"}
qdrant_cache.sync_client = MagicMock()
qdrant_cache.sync_client.put.side_effect = Exception("boom")
qdrant_cache._ensure_cache_key_payload_index()
qdrant_cache.sync_client.put.assert_called_once()
def _mock_qdrant_get_cache_result(qdrant_result):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
qdrant_cache.embedding_model = "text-embedding-ada-002"
qdrant_cache.qdrant_api_base = "http://test.qdrant.local"
qdrant_cache.collection_name = "test_collection"
qdrant_cache.headers = {
"Content-Type": "application/json",
"api-key": "test_key",
}
qdrant_cache.similarity_threshold = 0.8
qdrant_cache.sync_client = MagicMock()
mock_search_response = MagicMock()
mock_search_response.status_code = 200
mock_search_response.json.return_value = {"result": qdrant_result}
qdrant_cache.sync_client.post.return_value = mock_search_response
return qdrant_cache, QdrantSemanticCache
@pytest.mark.parametrize("qdrant_result", [None, []])
def test_qdrant_semantic_cache_get_cache_sets_metadata_on_empty_miss(qdrant_result):
qdrant_cache, _ = _mock_qdrant_get_cache_result(qdrant_result)
metadata = {}
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
result = qdrant_cache.get_cache(
key="test_key",
messages=[{"content": "What is the capital of Spain?"}],
metadata=metadata,
)
assert result is None
assert metadata["semantic-similarity"] == 0.0
def test_qdrant_semantic_cache_get_cache_sets_metadata_on_below_threshold_miss():
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache, _ = _mock_qdrant_get_cache_result(
[
{
"payload": {
QdrantSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
"text": "What is the capital of Spain?",
"response": '{"id": "test-456"}',
},
"score": 0.7,
}
]
)
metadata = {}
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
result = qdrant_cache.get_cache(
key="test_key",
messages=[{"content": "What is the capital of Spain?"}],
metadata=metadata,
)
assert result is None
assert metadata["semantic-similarity"] == 0.7
def test_qdrant_semantic_cache_get_cache_miss():
@@ -138,9 +320,7 @@ def test_qdrant_semantic_cache_get_cache_miss():
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
@@ -230,6 +410,7 @@ async def test_qdrant_semantic_cache_async_get_cache_hit():
"result": [
{
"payload": {
QdrantSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
"text": "What is the capital of Spain?", # Original prompt
"response": '{"id": "test-456", "choices": [{"message": {"content": "Madrid is the capital of Spain."}}]}',
},
@@ -262,6 +443,16 @@ async def test_qdrant_semantic_cache_async_get_cache_hit():
# Verify async search was called
qdrant_cache.async_client.post.assert_called()
assert qdrant_cache.async_client.post.call_args.kwargs["json"][
"filter"
] == {
"must": [
{
"key": QdrantSemanticCache.CACHE_KEY_FIELD_NAME,
"match": {"value": "test_key"},
}
]
}
@pytest.mark.asyncio
@@ -336,9 +527,7 @@ def test_qdrant_semantic_cache_set_cache():
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
@@ -384,6 +573,12 @@ def test_qdrant_semantic_cache_set_cache():
# Verify upsert was called
qdrant_cache.sync_client.put.assert_called()
upsert_payload = qdrant_cache.sync_client.put.call_args.kwargs["json"][
"points"
][0]["payload"]
assert (
upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key"
)
@pytest.mark.asyncio
@@ -450,6 +645,12 @@ async def test_qdrant_semantic_cache_async_set_cache():
# Verify async upsert was called
qdrant_cache.async_client.put.assert_called()
upsert_payload = qdrant_cache.async_client.put.call_args.kwargs["json"][
"points"
][0]["payload"]
assert (
upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key"
)
def test_qdrant_semantic_cache_custom_vector_size():
@@ -462,9 +663,7 @@ def test_qdrant_semantic_cache_custom_vector_size():
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection does NOT exist (so it will be created)
@@ -505,9 +704,13 @@ def test_qdrant_semantic_cache_custom_vector_size():
assert qdrant_cache.vector_size == 768
# Verify the PUT call to create the collection used vector_size=768
put_call = mock_sync_client_instance.put.call_args
assert put_call is not None
create_payload = put_call.kwargs.get("json") or put_call[1].get("json")
put_call = next(
call
for call in mock_sync_client_instance.put.call_args_list
if call.kwargs["url"]
== "http://test.qdrant.local/collections/test_collection_768"
)
create_payload = put_call.kwargs["json"]
assert create_payload["vectors"]["size"] == 768
assert create_payload["vectors"]["distance"] == "Cosine"
@@ -521,9 +724,7 @@ def test_qdrant_semantic_cache_default_vector_size():
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
@@ -559,9 +760,7 @@ def test_qdrant_semantic_cache_large_vector_size():
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection does NOT exist (so it will be created)
@@ -599,6 +798,11 @@ def test_qdrant_semantic_cache_large_vector_size():
assert qdrant_cache.vector_size == 4096
# Verify the collection was created with 4096
put_call = mock_sync_client_instance.put.call_args
create_payload = put_call.kwargs.get("json") or put_call[1].get("json")
put_call = next(
call
for call in mock_sync_client_instance.put.call_args_list
if call.kwargs["url"]
== "http://test.qdrant.local/collections/test_collection_4096"
)
create_payload = put_call.kwargs["json"]
assert create_payload["vectors"]["size"] == 4096
@@ -72,24 +72,301 @@ 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",
}
]
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]}]}
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?"}]
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()
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):
semantic_cache_mock = MagicMock()
custom_vectorizer_mock = MagicMock()
with patch.dict(
"sys.modules",
{
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
"redisvl.utils.vectorize": MagicMock(
CustomTextVectorizer=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.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 patch.dict(
"sys.modules",
{
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
"redisvl.utils.vectorize": MagicMock(
CustomTextVectorizer=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()
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'}",
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 patch.dict(
"sys.modules",
{
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
"redisvl.utils.vectorize": MagicMock(
CustomTextVectorizer=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 patch.dict(
"sys.modules",
{
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
"redisvl.utils.vectorize": MagicMock(
CustomTextVectorizer=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 patch.dict(
"sys.modules",
{
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
"redisvl.utils.vectorize": MagicMock(
CustomTextVectorizer=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")
with pytest.raises(ValueError, match="connection failed"):
RedisSemanticCache(
similarity_threshold=0.8,
index_name="existing_index",
)
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)
with patch.dict("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
@@ -123,6 +400,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",
}
]
@@ -131,16 +409,117 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch):
return_value=[0.1, 0.2, 0.3]
)
# 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={},
)
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()
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 patch.dict(
"sys.modules",
{
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
"redisvl.utils.vectorize": MagicMock(
CustomTextVectorizer=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 patch.dict(
"sys.modules",
{
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
"redisvl.utils.vectorize": MagicMock(
CustomTextVectorizer=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,
)