mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-18 06:26:16 +00:00
Merge pull request #26990 from stuxf/codex/semantic-cache-tenant-isolation
chore(caching): isolate semantic cache entries
This commit is contained in:
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user