diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 5e3713e5a1..cb521efca0 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -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}" ) diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index c76f27377d..da9e7b1e58 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -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 diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index 13dc4b5812..949e6ccc29 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -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 diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index f9946e266f..b50a35ef50 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -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, + )