diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 11733ce4cee..6b3cebb5a45 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -511,7 +511,6 @@ class Cache: try: # never block execution if self.should_use_cache(**kwargs) is not True: return - messages = kwargs.get("messages", []) if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: @@ -523,12 +522,13 @@ class Cache: or cache_control_args.get("s-max-age") or float("inf") ) + cache_kwargs = {k: v for k, v in kwargs.items() if k != "cache_key"} if dynamic_cache_object is not None: cached_result = dynamic_cache_object.get_cache( - cache_key, messages=messages + cache_key, **cache_kwargs ) else: - cached_result = self.cache.get_cache(cache_key, messages=messages) + cached_result = self.cache.get_cache(cache_key, **cache_kwargs) return self._get_cache_logic( cached_result=cached_result, max_age=max_age ) diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index da9e7b1e587..7f651581c94 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -36,6 +36,15 @@ class RedisSemanticCache(BaseCache): DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index" CACHE_KEY_FIELD_NAME: str = "litellm_cache_key" + API_KEY_HASH_FIELD_NAME: str = "litellm_user_api_key_hash" + TEAM_ID_FIELD_NAME: str = "litellm_user_api_key_team_id" + USER_ID_FIELD_NAME: str = "litellm_user_id" + _SCOPE_FIELD_TO_METADATA_KEY: Tuple[Tuple[str, str], ...] = ( + (API_KEY_HASH_FIELD_NAME, "user_api_key_hash"), + (TEAM_ID_FIELD_NAME, "user_api_key_team_id"), + (TEAM_ID_FIELD_NAME, "team_id"), + (USER_ID_FIELD_NAME, "user_id"), + ) def __init__( self, @@ -124,6 +133,15 @@ class RedisSemanticCache(BaseCache): "type": "tag", } + @classmethod + def _filterable_fields(cls) -> List[Dict[str, str]]: + return [ + cls._cache_key_filterable_field(), + {"name": cls.API_KEY_HASH_FIELD_NAME, "type": "tag"}, + {"name": cls.TEAM_ID_FIELD_NAME, "type": "tag"}, + {"name": cls.USER_ID_FIELD_NAME, "type": "tag"}, + ] + def _init_semantic_cache( self, semantic_cache_cls: Any, @@ -144,7 +162,7 @@ class RedisSemanticCache(BaseCache): redis_url=redis_url, vectorizer=cache_vectorizer, distance_threshold=self.distance_threshold, - filterable_fields=[self._cache_key_filterable_field()], + filterable_fields=self._filterable_fields(), overwrite=False, ) except ValueError as exc: @@ -162,7 +180,7 @@ class RedisSemanticCache(BaseCache): redis_url=redis_url, vectorizer=cache_vectorizer, distance_threshold=self.distance_threshold, - filterable_fields=[self._cache_key_filterable_field()], + filterable_fields=self._filterable_fields(), overwrite=False, ) except ValueError as isolated_exc: @@ -178,25 +196,59 @@ class RedisSemanticCache(BaseCache): redis_url=redis_url, vectorizer=cache_vectorizer, distance_threshold=self.distance_threshold, - filterable_fields=[self._cache_key_filterable_field()], + filterable_fields=self._filterable_fields(), overwrite=True, ) - def _get_cache_filters(self, key: str) -> Dict[str, str]: - return {self.CACHE_KEY_FIELD_NAME: str(key)} + def _get_metadata_from_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: + metadata: Dict[str, Any] = {} + for key in ("metadata", "litellm_metadata"): + value = kwargs.get(key) + if isinstance(value, dict): + metadata.update(value) + + litellm_params = kwargs.get("litellm_params") + if isinstance(litellm_params, dict): + value = litellm_params.get("metadata") + if isinstance(value, dict): + metadata.update(value) + + return metadata + + def _get_scope_filter(self, **kwargs) -> Optional[Tuple[str, str]]: + metadata = self._get_metadata_from_kwargs(kwargs) + for field_name, metadata_key in self._SCOPE_FIELD_TO_METADATA_KEY: + value = metadata.get(metadata_key) + if value: + return field_name, str(value) + return None + + def _get_cache_filters(self, key: str, **kwargs) -> Dict[str, str]: + filters = {self.CACHE_KEY_FIELD_NAME: str(key)} + scope_filter = self._get_scope_filter(**kwargs) + if scope_filter is not None: + filters[scope_filter[0]] = scope_filter[1] + return filters + + def _get_scope_filter_expression(self, **kwargs) -> Any: + scope_filter = self._get_scope_filter(**kwargs) + if scope_filter is None: + return None - def _get_cache_key_filter_expression(self, key: str) -> Any: from redisvl.query.filter import Tag # type: ignore[import-not-found, import-untyped] - return Tag(self.CACHE_KEY_FIELD_NAME) == str(key) + return Tag(scope_filter[0]) == scope_filter[1] - def _cache_hit_matches_key(self, cache_hit: Dict[str, Any], key: str) -> bool: - # Pre-isolation entries with no ``litellm_cache_key`` field cannot be - # safely reassigned to a caller's scope and are treated as misses. - cached_key = cache_hit.get(self.CACHE_KEY_FIELD_NAME) - if isinstance(cached_key, bytes): - cached_key = cached_key.decode("utf-8") - return cached_key is not None and str(cached_key) == str(key) + def _cache_hit_matches_scope(self, cache_hit: Dict[str, Any], **kwargs) -> bool: + scope_filter = self._get_scope_filter(**kwargs) + if scope_filter is None: + return True + + field_name, expected_value = scope_filter + cached_value = cache_hit.get(field_name) + if isinstance(cached_value, bytes): + cached_value = cached_value.decode("utf-8") + return cached_value is not None and str(cached_value) == expected_value def _get_ttl(self, **kwargs) -> Optional[int]: """ @@ -288,7 +340,7 @@ class RedisSemanticCache(BaseCache): value_str = str(value) store_kwargs: Dict[str, Any] = { - "filters": self._get_cache_filters(key), + "filters": self._get_cache_filters(key, **kwargs), } # Get TTL and store in Redis semantic cache @@ -323,12 +375,15 @@ class RedisSemanticCache(BaseCache): return None prompt = get_str_from_messages(messages) - # Check the cache for semantically similar prompts in this exact - # LiteLLM cache-key scope. + # Keep caller isolation separate from prompt similarity. Filtering + # on the full request hash would make KNN unreachable for + # non-identical prompts. check_kwargs: Dict[str, Any] = { "prompt": prompt, - "filter_expression": self._get_cache_key_filter_expression(key), } + filter_expression = self._get_scope_filter_expression(**kwargs) + if filter_expression is not None: + check_kwargs["filter_expression"] = filter_expression results = self.llmcache.check(**check_kwargs) # Return None if no similar prompts found @@ -338,8 +393,8 @@ class RedisSemanticCache(BaseCache): # Process the best matching result cache_hit = results[0] - if not self._cache_hit_matches_key(cache_hit=cache_hit, key=key): - print_verbose("Redis semantic-cache hit did not match cache key scope") + if not self._cache_hit_matches_scope(cache_hit=cache_hit, **kwargs): + print_verbose("Redis semantic-cache hit did not match caller scope") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None vector_distance = float(cache_hit["vector_distance"]) @@ -442,7 +497,7 @@ class RedisSemanticCache(BaseCache): store_kwargs: Dict[str, Any] = { "vector": prompt_embedding, - "filters": self._get_cache_filters(key), + "filters": self._get_cache_filters(key, **kwargs), } # Get TTL and store in Redis semantic cache @@ -483,13 +538,16 @@ class RedisSemanticCache(BaseCache): # Generate embedding for the prompt prompt_embedding = await self._get_async_embedding(prompt, **kwargs) - # Check the cache for semantically similar prompts in this exact - # LiteLLM cache-key scope. + # Keep caller isolation separate from prompt similarity. Filtering + # on the full request hash would make KNN unreachable for + # non-identical prompts. check_kwargs: Dict[str, Any] = { "prompt": prompt, "vector": prompt_embedding, - "filter_expression": self._get_cache_key_filter_expression(key), } + filter_expression = self._get_scope_filter_expression(**kwargs) + if filter_expression is not None: + check_kwargs["filter_expression"] = filter_expression results = await self.llmcache.acheck(**check_kwargs) # handle results / cache hit @@ -498,8 +556,8 @@ class RedisSemanticCache(BaseCache): return None cache_hit = results[0] - if not self._cache_hit_matches_key(cache_hit=cache_hit, key=key): - print_verbose("Redis semantic-cache hit did not match cache key scope") + if not self._cache_hit_matches_scope(cache_hit=cache_hit, **kwargs): + print_verbose("Redis semantic-cache hit did not match caller scope") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None vector_distance = float(cache_hit["vector_distance"]) diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index b50a35ef50e..df82895d851 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -72,22 +72,15 @@ def test_redis_semantic_cache_get_cache(monkeypatch): "prompt": "What is the capital of France?", "response": '{"content": "Paris is the capital of France."}', "vector_distance": 0.1, # Distance of 0.1 means similarity of 0.9 - RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "different_request_hash", } ] redis_semantic_cache.llmcache.check = MagicMock(return_value=mock_result) # Mock the embedding function - with ( - patch( - "litellm.embedding", - return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, - ), - patch.object( - redis_semantic_cache, - "_get_cache_key_filter_expression", - return_value="cache-key-filter", - ), + with patch( + "litellm.embedding", + return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}, ): # Test get_cache with a message metadata = {} @@ -104,11 +97,53 @@ def test_redis_semantic_cache_get_cache(monkeypatch): # Verify llmcache.check was called redis_semantic_cache.llmcache.check.assert_called_once_with( prompt="What is the capital of France?", - filter_expression="cache-key-filter", ) -def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): +def test_cache_get_cache_forwards_metadata_to_semantic_cache(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="cache-key") + cache._get_cache_logic = MagicMock(return_value={"content": "Paris"}) + cache.cache = MagicMock() + cache.cache.get_cache.return_value = {"response": '{"content":"Paris"}'} + + metadata = {"user_api_key_hash": "hashed-key"} + result = cache.get_cache( + model="gpt-4o-mini", + messages=[{"content": "Which city is the capital of France?"}], + metadata=metadata, + cache={}, + ) + + assert result == {"content": "Paris"} + cache.cache.get_cache.assert_called_once_with( + "cache-key", + model="gpt-4o-mini", + messages=[{"content": "Which city is the capital of France?"}], + metadata=metadata, + cache={}, + ) + + cache.cache.get_cache.reset_mock() + cache.get_cache( + cache_key="preset-key", + messages=[{"content": "cached prompt"}], + metadata=metadata, + cache={}, + ) + + cache.cache.get_cache.assert_called_once_with( + "preset-key", + messages=[{"content": "cached prompt"}], + metadata=metadata, + cache={}, + ) + + +def test_redis_semantic_cache_rejects_cross_scope_cache_hit(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() @@ -134,16 +169,17 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): "prompt": "What is the capital of France?", "response": '{"content": "Paris"}', "vector_distance": 0.1, + RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "other-key", } ] ) with patch.object( redis_semantic_cache, - "_get_cache_key_filter_expression", - return_value="cache-key-filter", + "_get_scope_filter_expression", + return_value="caller-scope-filter", ): - metadata = {} + metadata = {"user_api_key_hash": "this-key"} result = redis_semantic_cache.get_cache( key="test_key", messages=[{"content": "What is the capital of France?"}], @@ -152,6 +188,10 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): assert result is None assert metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + filter_expression="caller-scope-filter", + ) def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch): @@ -180,13 +220,17 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch): key="test_key", value={"content": "Paris"}, messages=[{"content": "What is the capital of France?"}], + metadata={"user_api_key_hash": "hashed-key"}, ttl=60, ) redis_semantic_cache.llmcache.store.assert_called_once_with( "What is the capital of France?", "{'content': 'Paris'}", - filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + filters={ + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "hashed-key", + }, ttl=60, ) @@ -227,9 +271,10 @@ def test_redis_semantic_cache_uses_isolated_index_for_old_schema(monkeypatch): semantic_cache_mock.call_args_list[1].kwargs["name"] == "existing_index_isolated" ) - assert semantic_cache_mock.call_args_list[1].kwargs["filterable_fields"] == [ - RedisSemanticCache._cache_key_filterable_field() - ] + assert ( + semantic_cache_mock.call_args_list[1].kwargs["filterable_fields"] + == RedisSemanticCache._filterable_fields() + ) def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch): @@ -269,9 +314,10 @@ def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch): == "existing_index_isolated" ) assert semantic_cache_mock.call_args_list[2].kwargs["overwrite"] is True - assert semantic_cache_mock.call_args_list[2].kwargs["filterable_fields"] == [ - RedisSemanticCache._cache_key_filterable_field() - ] + assert ( + semantic_cache_mock.call_args_list[2].kwargs["filterable_fields"] + == RedisSemanticCache._filterable_fields() + ) def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypatch): @@ -321,20 +367,19 @@ def test_redis_semantic_cache_reraises_unexpected_index_error(): ) -def test_redis_semantic_cache_matches_bytes_cache_key(): +def test_redis_semantic_cache_matches_bytes_scope(): from litellm.caching.redis_semantic_cache import RedisSemanticCache redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) - assert redis_semantic_cache._cache_hit_matches_key( - cache_hit={RedisSemanticCache.CACHE_KEY_FIELD_NAME: b"test_key"}, - key="test_key", + assert redis_semantic_cache._cache_hit_matches_scope( + cache_hit={RedisSemanticCache.API_KEY_HASH_FIELD_NAME: b"hashed-key"}, + metadata={"user_api_key_hash": "hashed-key"}, ) -def test_redis_semantic_cache_rejects_pre_isolation_unscoped_hit(): - """Pre-isolation entries with no cache-key field cannot be safely - reassigned to a caller's scope and are treated as misses.""" +def test_redis_semantic_cache_rejects_pre_scope_hit_for_scoped_request(): + """Entries with no caller scope cannot be safely reused for a scoped request.""" from litellm.caching.redis_semantic_cache import RedisSemanticCache redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) @@ -344,13 +389,13 @@ def test_redis_semantic_cache_rejects_pre_isolation_unscoped_hit(): "response": '{"content": "Paris"}', "vector_distance": 0.1, } - assert not redis_semantic_cache._cache_hit_matches_key( + assert not redis_semantic_cache._cache_hit_matches_scope( cache_hit=cache_hit, - key="test_key", + metadata={"user_api_key_hash": "hashed-key"}, ) -def test_redis_semantic_cache_builds_filter_expression(monkeypatch): +def test_redis_semantic_cache_builds_scope_filter_expression(monkeypatch): class FakeTag: def __init__(self, field_name): self.field_name = field_name @@ -363,10 +408,13 @@ def test_redis_semantic_cache_builds_filter_expression(monkeypatch): redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) - assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == ( - RedisSemanticCache.CACHE_KEY_FIELD_NAME, - "test_key", + assert redis_semantic_cache._get_scope_filter_expression( + metadata={"user_api_key_hash": "hashed-key"} + ) == ( + RedisSemanticCache.API_KEY_HASH_FIELD_NAME, + "hashed-key", ) + assert redis_semantic_cache._get_scope_filter_expression(metadata={}) is None @pytest.mark.asyncio @@ -400,7 +448,7 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch): "prompt": "What is the capital of France?", "response": '{"content": "Paris is the capital of France."}', "vector_distance": 0.1, # Distance of 0.1 means similarity of 0.9 - RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "different_request_hash", } ] @@ -409,17 +457,12 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch): return_value=[0.1, 0.2, 0.3] ) - with patch.object( - redis_semantic_cache, - "_get_cache_key_filter_expression", - return_value="cache-key-filter", - ): - # Test async_get_cache with a message - result = await redis_semantic_cache.async_get_cache( - key="test_key", - messages=[{"content": "What is the capital of France?"}], - metadata={}, - ) + # Test async_get_cache with a message + result = await redis_semantic_cache.async_get_cache( + key="test_key", + messages=[{"content": "What is the capital of France?"}], + metadata={}, + ) # Verify result is properly parsed assert result == {"content": "Paris is the capital of France."} @@ -429,12 +472,13 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch): redis_semantic_cache.llmcache.acheck.assert_called_once_with( prompt="What is the capital of France?", vector=[0.1, 0.2, 0.3], - filter_expression="cache-key-filter", ) @pytest.mark.asyncio -async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeypatch): +async def test_redis_semantic_cache_async_get_cache_rejects_cross_scope_hit( + monkeypatch, +): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() @@ -460,6 +504,7 @@ async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeyp "prompt": "What is the capital of France?", "response": '{"content": "Paris"}', "vector_distance": 0.1, + RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "other-key", } ] ) @@ -469,13 +514,13 @@ async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeyp with patch.object( redis_semantic_cache, - "_get_cache_key_filter_expression", - return_value="cache-key-filter", + "_get_scope_filter_expression", + return_value="caller-scope-filter", ): result = await redis_semantic_cache.async_get_cache( key="test_key", messages=[{"content": "What is the capital of France?"}], - metadata={}, + metadata={"user_api_key_hash": "this-key"}, ) assert result is None @@ -513,6 +558,7 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter( key="test_key", value={"content": "Paris"}, messages=[{"content": "What is the capital of France?"}], + metadata={"user_api_key_hash": "hashed-key"}, ttl=60, ) @@ -520,6 +566,9 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter( "What is the capital of France?", "{'content': 'Paris'}", vector=[0.1, 0.2, 0.3], - filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + filters={ + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + RedisSemanticCache.API_KEY_HASH_FIELD_NAME: "hashed-key", + }, ttl=60, )