diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 997ad10bc33..cb122e90102 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -100,6 +100,8 @@ class Cache: gcs_path: Optional[str] = None, redis_semantic_cache_embedding_model: str = "text-embedding-ada-002", redis_semantic_cache_index_name: Optional[str] = None, + valkey_semantic_cache_embedding_model: str = "text-embedding-ada-002", + valkey_semantic_cache_index_name: str | None = None, redis_flush_size: Optional[int] = None, redis_startup_nodes: Optional[List] = None, disk_cache_dir: Optional[str] = None, @@ -208,6 +210,21 @@ class Cache: index_name=redis_semantic_cache_index_name, **kwargs, ) + elif type == LiteLLMCacheType.VALKEY_SEMANTIC: + # Imported here, not at module top, so the optional redis dependency + # is only required when this backend is actually selected. + from .valkey_semantic_cache import ValkeySemanticCache + + self.cache = ValkeySemanticCache( + host=host, + port=port, + password=password, + similarity_threshold=similarity_threshold, + embedding_model=valkey_semantic_cache_embedding_model, + index_name=valkey_semantic_cache_index_name, + startup_nodes=redis_startup_nodes, + **kwargs, + ) elif type == LiteLLMCacheType.QDRANT_SEMANTIC: self.cache = QdrantSemanticCache( qdrant_api_base=qdrant_api_base, @@ -267,12 +284,50 @@ class Cache: if ( self.type == LiteLLMCacheType.REDIS or self.type == LiteLLMCacheType.REDIS_SEMANTIC + or self.type == LiteLLMCacheType.VALKEY_SEMANTIC ) and default_in_redis_ttl is not None: self.ttl = default_in_redis_ttl if self.namespace is not None and isinstance(self.cache, RedisCache): self.cache.namespace = self.namespace + # Params whose values carry prompt content. Excluded from semantic-cache + # scope keys so differently worded prompts share a bucket and match via + # vector similarity rather than being split into per-wording buckets. + _SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS: frozenset = frozenset( + {"messages", "prompt", "input"} + ) + + # Server-set identity (from proxy auth) used to isolate semantic-cache + # buckets per tenant. Required once the prompt is out of the scope key, so a + # similar prompt from another key/team/org stays in a separate bucket. + _SEMANTIC_CACHE_TENANT_SCOPE_FIELDS: tuple[str, ...] = ( + "user_api_key", + "user_api_key_team_id", + "user_api_key_org_id", + ) + + def _is_semantic_cache(self) -> bool: + return self.type in ( + LiteLLMCacheType.REDIS_SEMANTIC, + LiteLLMCacheType.QDRANT_SEMANTIC, + LiteLLMCacheType.VALKEY_SEMANTIC, + ) + + def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str: + metadata: dict = kwargs.get("metadata") or {} + litellm_params: dict = kwargs.get("litellm_params") or {} + metadata_in_litellm_params: dict = litellm_params.get("metadata") or {} + + scope = "" + for field in self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS: + value = metadata.get(field) + if value is None: + value = metadata_in_litellm_params.get(field) + if value is not None: + scope += f"{field}: {value}" + return scope + def get_cache_key(self, **kwargs) -> str: """ Get the cache key for the given arguments. @@ -293,7 +348,15 @@ class Cache: combined_kwargs = ModelParamHelper._get_all_llm_api_params() litellm_param_kwargs = all_litellm_params + is_semantic_cache = self._is_semantic_cache() + scope_excluded_params = ( + self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS + if is_semantic_cache + else frozenset() + ) for param in kwargs: + if param in scope_excluded_params: + continue if param in combined_kwargs: param_value: Optional[str] = self._get_param_value(param, kwargs) if param_value is not None: @@ -309,6 +372,9 @@ class Cache: param_value = kwargs[param] cache_key += f"{str(param)}: {str(param_value)}" + if is_semantic_cache: + cache_key += self._get_semantic_cache_tenant_scope(kwargs) + hashed_cache_key = Cache._get_hashed_cache_key(cache_key) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) verbose_logger.debug( diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py new file mode 100644 index 00000000000..bf368b74d07 --- /dev/null +++ b/litellm/caching/valkey_semantic_cache.py @@ -0,0 +1,353 @@ +""" +Valkey Semantic Cache implementation for LiteLLM + +Backs semantic caching with Valkey (for example AWS ElastiCache for Valkey) +running the valkey-search module. + +RedisVL cannot drive valkey-search: it gates on a RediSearch module version +that valkey-search does not report, and its SemanticCache index uses a TEXT +field that valkey-search does not implement. This backend therefore talks to +valkey-search directly over redis-py, building a vector index from the field +types valkey-search does support (TAG for cache-key isolation and VECTOR for +the prompt embedding) and running KNN queries for retrieval. Prompt extraction, +embedding generation, and cached-response parsing are reused from +RedisSemanticCache since those are backend agnostic. +""" + +import asyncio +import hashlib +import os +import struct +from dataclasses import dataclass +from typing import Any + +from redis import Redis +from redis.asyncio import Redis as AsyncRedis +from redis.commands.search.field import TagField, VectorField +from redis.commands.search.indexDefinition import IndexDefinition, IndexType +from redis.commands.search.query import Query + +from litellm._logging import print_verbose +from litellm._uuid import uuid + +from .redis_semantic_cache import RedisSemanticCache + + +@dataclass(frozen=True, slots=True) +class _ValkeyCacheHit: + response: str + distance: float + + +class ValkeySemanticCache(RedisSemanticCache): + """Valkey-backed semantic cache for LLM responses.""" + + DEFAULT_VALKEY_INDEX_NAME: str = "litellm_semantic_cache_index" + EMBEDDING_FIELD_NAME: str = "embedding" + PROMPT_FIELD_NAME: str = "prompt" + RESPONSE_FIELD_NAME: str = "response" + DISTANCE_FIELD_NAME: str = "vector_distance" + + def __init__( + self, + host: str | None = None, + port: str | None = None, + password: str | None = None, + redis_url: str | None = None, + similarity_threshold: float | None = None, + embedding_model: str = "text-embedding-ada-002", + index_name: str | None = None, + ssl: bool = False, + startup_nodes: list | None = None, + sync_client: Redis | None = None, + async_client: AsyncRedis | None = None, + **kwargs: Any, + ): + if similarity_threshold is None: + raise ValueError("similarity_threshold must be provided, passed None") + + if startup_nodes: + raise ValueError( + "valkey-semantic does not support cluster-mode-enabled (multi-shard) " + "endpoints. The async cluster client cannot route the FT.* search " + "commands reliably. Point it at a cluster-mode-disabled endpoint " + "instead (a primary with replicas is fine; only horizontal sharding " + "is unsupported), or pass a single redis_url. On AWS, vector search " + "needs ElastiCache for Valkey 8.2+ on a node-based cluster." + ) + + self.similarity_threshold = similarity_threshold + self.embedding_model = embedding_model + self.index_name = index_name or self.DEFAULT_VALKEY_INDEX_NAME + self.key_prefix = f"{self.index_name}:" + self._index_dim: int | None = None + + resolved_url = None + if sync_client is None or async_client is None: + resolved_url = redis_url or self._build_valkey_url( + host, port, password, ssl + ) + self.sync_client = ( + sync_client if sync_client is not None else Redis.from_url(resolved_url) # type: ignore[arg-type] + ) + self.async_client = ( + async_client + if async_client is not None + else AsyncRedis.from_url(resolved_url) # type: ignore[arg-type] + ) + + print_verbose(f"Valkey semantic-cache initializing index - {self.index_name}") + + @staticmethod + def _build_valkey_url( + host: str | None, port: str | None, password: str | None, ssl: bool = False + ) -> str: + host = host or os.environ.get("VALKEY_HOST") or os.environ.get("REDIS_HOST") + port = port or os.environ.get("VALKEY_PORT") or os.environ.get("REDIS_PORT") + password = ( + password + or os.environ.get("VALKEY_PASSWORD") + or os.environ.get("REDIS_PASSWORD") + ) + + if not host or not port: + raise ValueError( + "Missing required Valkey configuration. Provide host and port " + "(or VALKEY_HOST/VALKEY_PORT), or pass redis_url." + ) + + credentials = f":{password}@" if password else "" + scheme = "rediss" if ssl else "redis" + return f"{scheme}://{credentials}{host}:{port}" + + @classmethod + def _scope_tag(cls, key: str) -> str: + # valkey-search TAG fields tokenize on punctuation and do not honour + # backslash escaping, so an arbitrary cache key cannot be matched + # verbatim. Hashing to hex yields a token that is always exact-match + # safe and still uniquely isolates a caller's scope. + return hashlib.sha256(str(key).encode("utf-8")).hexdigest() + + @staticmethod + def _embedding_to_bytes(embedding: list[float]) -> bytes: + return struct.pack(f"<{len(embedding)}f", *embedding) + + def _index_schema(self, dim: int) -> tuple[TagField, VectorField]: + return ( + TagField(self.CACHE_KEY_FIELD_NAME), + VectorField( + self.EMBEDDING_FIELD_NAME, + "HNSW", + {"TYPE": "FLOAT32", "DIM": dim, "DISTANCE_METRIC": "COSINE"}, + ), + ) + + def _index_definition(self) -> IndexDefinition: + return IndexDefinition(prefix=[self.key_prefix], index_type=IndexType.HASH) + + @staticmethod + def _is_index_exists_error(exc: Exception) -> bool: + return "already exists" in str(exc).lower() + + @staticmethod + def _extract_index_dim(info: dict) -> int | None: + # FT.INFO nests the vector field's "dimensions" one level inside its + # "index" block, so flatten each field descriptor a single level and + # scan for the dimensions marker. + for field in info.get("attributes") or []: + if not isinstance(field, (list, tuple)): + continue + flat = [ + sub + for item in field + for sub in (item if isinstance(item, (list, tuple)) else [item]) + ] + for i, marker in enumerate(flat): + if marker in (b"dimensions", "dimensions") and i + 1 < len(flat): + return int(flat[i + 1]) + return None + + def _assert_dim_matches(self, info: dict, dim: int) -> None: + existing_dim = self._extract_index_dim(info) + if existing_dim is not None and existing_dim != dim: + raise ValueError( + f"Valkey semantic-cache index '{self.index_name}' already exists with " + f"embedding dimension {existing_dim}, but the configured embedding " + f"model produced dimension {dim}. Use a different " + f"valkey_semantic_cache_index_name or drop the existing index." + ) + + def _ensure_index_sync(self, dim: int) -> None: + if self._index_dim == dim: + return + try: + self.sync_client.ft(self.index_name).create_index( + self._index_schema(dim), definition=self._index_definition() + ) + except Exception as exc: + if not self._is_index_exists_error(exc): + raise + self._assert_dim_matches(self.sync_client.ft(self.index_name).info(), dim) + self._index_dim = dim + + async def _ensure_index_async(self, dim: int) -> None: + if self._index_dim == dim: + return + try: + await self.async_client.ft(self.index_name).create_index( + self._index_schema(dim), definition=self._index_definition() + ) + except Exception as exc: + if not self._is_index_exists_error(exc): + raise + info = await self.async_client.ft(self.index_name).info() + self._assert_dim_matches(info, dim) + self._index_dim = dim + + def _doc_key(self, key: str) -> str: + return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}" + + def _doc_mapping( + self, key: str, prompt: str, value_str: str, embedding: list[float] + ) -> dict: + return { + self.CACHE_KEY_FIELD_NAME: self._scope_tag(key), + self.PROMPT_FIELD_NAME: prompt, + self.RESPONSE_FIELD_NAME: value_str, + self.EMBEDDING_FIELD_NAME: self._embedding_to_bytes(embedding), + } + + def _knn_query(self, key: str) -> Query: + scope = self._scope_tag(key) + query_string = ( + f"(@{self.CACHE_KEY_FIELD_NAME}:{{{scope}}})" + f"=>[KNN 1 @{self.EMBEDDING_FIELD_NAME} $vec AS {self.DISTANCE_FIELD_NAME}]" + ) + return ( + Query(query_string) + .return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME) + .dialect(2) + ) + + @classmethod + def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None: + docs = getattr(search_result, "docs", []) + if not docs: + return None + doc = docs[0] + return _ValkeyCacheHit( + response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)), + distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)), + ) + + def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any: + if hit is None: + kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + return None + + similarity = 1 - hit.distance + kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity + + if similarity < self.similarity_threshold: + return None + return self._get_cache_logic(cached_response=hit.response) + + def set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}") + try: + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic caching") + return + + embedding = self._get_embedding(prompt) + self._ensure_index_sync(len(embedding)) + + doc_key = self._doc_key(key) + self.sync_client.hset( + doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding) + ) + ttl = self._get_ttl(**kwargs) + if ttl is not None: + self.sync_client.expire(doc_key, ttl) + except Exception as e: + print_verbose(f"Error in Valkey semantic-cache set_cache: {str(e)}") + + def get_cache(self, key: str, **kwargs: Any) -> Any: + print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}") + try: + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + return None + + embedding = self._get_embedding(prompt) + self._ensure_index_sync(len(embedding)) + + search_result = self.sync_client.ft(self.index_name).search( + self._knn_query(key), + query_params={"vec": self._embedding_to_bytes(embedding)}, + ) + return self._resolve_hit(self._first_hit(search_result), key, **kwargs) + except Exception as e: + print_verbose(f"Error in Valkey semantic-cache get_cache: {str(e)}") + kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + + async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None: + print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}") + try: + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic caching") + return + + embedding = await self._get_async_embedding(prompt, **kwargs) + await self._ensure_index_async(len(embedding)) + + doc_key = self._doc_key(key) + await self.async_client.hset( + doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding) + ) + ttl = self._get_ttl(**kwargs) + if ttl is not None: + await self.async_client.expire(doc_key, ttl) + except Exception as e: + print_verbose(f"Error in async Valkey semantic-cache set_cache: {str(e)}") + + async def async_get_cache(self, key: str, **kwargs: Any) -> Any: + print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}") + try: + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + return None + + embedding = await self._get_async_embedding(prompt, **kwargs) + await self._ensure_index_async(len(embedding)) + + search_result = await self.async_client.ft(self.index_name).search( + self._knn_query(key), + query_params={"vec": self._embedding_to_bytes(embedding)}, + ) + return self._resolve_hit(self._first_hit(search_result), key, **kwargs) + except Exception as e: + print_verbose(f"Error in async Valkey semantic-cache get_cache: {str(e)}") + kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 + + async def async_set_cache_pipeline( + self, cache_list: list[tuple[str, Any]], **kwargs: Any + ) -> None: + try: + await asyncio.gather( + *[ + self.async_set_cache(key, value, **kwargs) + for key, value in cache_list + ] + ) + except Exception as e: + print_verbose( + f"Error in Valkey semantic-cache async_set_cache_pipeline: {str(e)}" + ) + + async def _index_info(self) -> dict: + return await self.async_client.ft(self.index_name).info() diff --git a/litellm/types/caching.py b/litellm/types/caching.py index 10453c74a15..eaa80c2f525 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -9,6 +9,7 @@ class LiteLLMCacheType(str, Enum): LOCAL = "local" REDIS = "redis" REDIS_SEMANTIC = "redis-semantic" + VALKEY_SEMANTIC = "valkey-semantic" S3 = "s3" DISK = "disk" QDRANT_SEMANTIC = "qdrant-semantic" diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py index 20614103ed2..eaee54bac5a 100644 --- a/tests/test_litellm/caching/test_caching.py +++ b/tests/test_litellm/caching/test_caching.py @@ -76,3 +76,73 @@ def test_get_per_item_prompt_tokens_distributes_with_remainder(): per_item = [cache._get_per_item_prompt_tokens(result, i) for i in range(3)] assert sum(per_item) == 10 # 4 + 3 + 3 assert per_item == [4, 3, 3] + + +def _semantic_cache(): + return Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + host="localhost", + port="6379", + similarity_threshold=0.8, + ) + + +def test_semantic_cache_key_excludes_prompt_so_paraphrases_share_a_bucket(): + cache = _semantic_cache() + tenant = {"user_api_key": "hash-abc"} + key_a = cache.get_cache_key( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "What color is the sky?"}], + metadata=dict(tenant), + ) + key_b = cache.get_cache_key( + model="gpt-4o-mini", + messages=[ + {"role": "user", "content": "Tell me the colour of the daytime sky."} + ], + metadata=dict(tenant), + ) + assert key_a == key_b + + +def test_semantic_cache_key_isolates_tenants(): + messages = [{"role": "user", "content": "What color is the sky?"}] + cache = _semantic_cache() + key_a = cache.get_cache_key( + model="gpt-4o-mini", messages=messages, metadata={"user_api_key": "hash-A"} + ) + key_b = cache.get_cache_key( + model="gpt-4o-mini", messages=messages, metadata={"user_api_key": "hash-B"} + ) + key_team = cache.get_cache_key( + model="gpt-4o-mini", + messages=messages, + metadata={"user_api_key": "hash-A", "user_api_key_team_id": "team-1"}, + ) + assert key_a != key_b + assert key_a != key_team + + +def test_semantic_cache_key_still_separates_models_and_params(): + cache = _semantic_cache() + messages = [{"role": "user", "content": "hi"}] + tenant = {"user_api_key": "hash-A"} + assert cache.get_cache_key( + model="gpt-4o-mini", messages=messages, metadata=dict(tenant) + ) != cache.get_cache_key(model="gpt-4o", messages=messages, metadata=dict(tenant)) + assert cache.get_cache_key( + model="gpt-4o-mini", messages=messages, temperature=0, metadata=dict(tenant) + ) != cache.get_cache_key( + model="gpt-4o-mini", messages=messages, temperature=1, metadata=dict(tenant) + ) + + +def test_exact_cache_key_still_includes_prompt(): + cache = Cache(type=LiteLLMCacheType.LOCAL) + key_a = cache.get_cache_key( + model="gpt-4o-mini", messages=[{"role": "user", "content": "a"}] + ) + key_b = cache.get_cache_key( + model="gpt-4o-mini", messages=[{"role": "user", "content": "b"}] + ) + assert key_a != key_b diff --git a/tests/test_litellm/caching/test_valkey_semantic_cache.py b/tests/test_litellm/caching/test_valkey_semantic_cache.py new file mode 100644 index 00000000000..44b9f061998 --- /dev/null +++ b/tests/test_litellm/caching/test_valkey_semantic_cache.py @@ -0,0 +1,473 @@ +import hashlib +import os +import struct +import subprocess +import sys +import textwrap +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.caching.valkey_semantic_cache import ValkeySemanticCache + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")) + + +def _make_cache(sync_client=None, async_client=None, similarity_threshold=0.8): + return ValkeySemanticCache( + similarity_threshold=similarity_threshold, + index_name="test_index", + sync_client=sync_client or MagicMock(), + async_client=async_client or AsyncMock(), + ) + + +def _search_result(distance, response='{"content": "Paris"}'): + return SimpleNamespace( + docs=[SimpleNamespace(response=response, vector_distance=str(distance))] + ) + + +def test_build_valkey_url_prefers_valkey_env(monkeypatch): + monkeypatch.setenv("REDIS_HOST", "redis-host") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "rpass") + monkeypatch.setenv("VALKEY_HOST", "valkey-host") + monkeypatch.setenv("VALKEY_PORT", "6380") + monkeypatch.setenv("VALKEY_PASSWORD", "vpass") + + assert ( + ValkeySemanticCache._build_valkey_url(None, None, None) + == "redis://:vpass@valkey-host:6380" + ) + + +def test_build_valkey_url_supports_passwordless(monkeypatch): + monkeypatch.delenv("REDIS_PASSWORD", raising=False) + monkeypatch.delenv("VALKEY_PASSWORD", raising=False) + monkeypatch.setenv("VALKEY_HOST", "valkey-host") + monkeypatch.setenv("VALKEY_PORT", "6380") + + assert ( + ValkeySemanticCache._build_valkey_url(None, None, None) + == "redis://valkey-host:6380" + ) + + +def test_build_valkey_url_falls_back_to_redis_env(monkeypatch): + monkeypatch.delenv("VALKEY_HOST", raising=False) + monkeypatch.delenv("VALKEY_PORT", raising=False) + monkeypatch.delenv("VALKEY_PASSWORD", raising=False) + monkeypatch.setenv("REDIS_HOST", "redis-host") + monkeypatch.setenv("REDIS_PORT", "6379") + monkeypatch.setenv("REDIS_PASSWORD", "rpass") + + assert ( + ValkeySemanticCache._build_valkey_url(None, None, None) + == "redis://:rpass@redis-host:6379" + ) + + +def test_build_valkey_url_requires_host_and_port(monkeypatch): + for var in ( + "VALKEY_HOST", + "VALKEY_PORT", + "VALKEY_PASSWORD", + "REDIS_HOST", + "REDIS_PORT", + "REDIS_PASSWORD", + ): + monkeypatch.delenv(var, raising=False) + + with pytest.raises(ValueError, match="Missing required Valkey configuration"): + ValkeySemanticCache._build_valkey_url(None, None, None) + + +def test_build_valkey_url_uses_rediss_scheme_when_ssl(monkeypatch): + monkeypatch.setenv("VALKEY_HOST", "valkey-host") + monkeypatch.setenv("VALKEY_PORT", "6379") + monkeypatch.setenv("VALKEY_PASSWORD", "vpass") + + assert ( + ValkeySemanticCache._build_valkey_url(None, None, None, ssl=True) + == "rediss://:vpass@valkey-host:6379" + ) + assert ValkeySemanticCache._build_valkey_url( + "h", "6379", None, ssl=False + ).startswith("redis://") + + +def test_init_requires_similarity_threshold(): + with pytest.raises(ValueError, match="similarity_threshold must be provided"): + ValkeySemanticCache(sync_client=MagicMock(), async_client=AsyncMock()) + + +def test_init_rejects_cluster_startup_nodes(): + with pytest.raises(ValueError, match="cluster-mode-enabled"): + ValkeySemanticCache( + similarity_threshold=0.8, + startup_nodes=[{"host": "shard1", "port": 6379}], + ) + + +def test_cache_dispatch_rejects_cluster_for_valkey_semantic(): + from litellm.caching.caching import Cache + from litellm.types.caching import LiteLLMCacheType + + with pytest.raises(ValueError, match="cluster-mode-enabled"): + Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + host="valkey-host", + port="6379", + similarity_threshold=0.8, + redis_startup_nodes=[{"host": "shard1", "port": 6379}], + ) + + +def test_scope_tag_is_deterministic_hex(): + tag = ValkeySemanticCache._scope_tag("model:gpt-4o::abc-123") + assert tag == hashlib.sha256(b"model:gpt-4o::abc-123").hexdigest() + assert len(tag) == 64 + assert ValkeySemanticCache._scope_tag("a") != ValkeySemanticCache._scope_tag("b") + + +def test_embedding_to_bytes_is_little_endian_float32(): + assert ValkeySemanticCache._embedding_to_bytes([1.0, 0.0]) == struct.pack( + "<2f", 1.0, 0.0 + ) + + +def test_set_cache_stores_scoped_doc_with_embedding(monkeypatch): + sync_client = MagicMock() + cache = _make_cache(sync_client=sync_client) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + cache.set_cache( + key="cache-key", + value={"content": "Paris"}, + messages=[{"role": "user", "content": "What is the capital of France?"}], + ) + + sync_client.ft.return_value.create_index.assert_called_once() + assert sync_client.hset.call_count == 1 + doc_key, kwargs = ( + sync_client.hset.call_args.args[0], + sync_client.hset.call_args.kwargs, + ) + mapping = kwargs["mapping"] + scope = ValkeySemanticCache._scope_tag("cache-key") + assert mapping[ValkeySemanticCache.CACHE_KEY_FIELD_NAME] == scope + assert mapping["prompt"] == "What is the capital of France?" + assert mapping["response"] == "{'content': 'Paris'}" + assert mapping["embedding"] == struct.pack("<3f", 0.1, 0.2, 0.3) + assert doc_key.startswith(f"test_index:{scope}:") + + +def test_set_cache_applies_ttl(): + sync_client = MagicMock() + cache = _make_cache(sync_client=sync_client) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + cache.set_cache( + key="cache-key", + value={"content": "Paris"}, + messages=[{"role": "user", "content": "What is the capital of France?"}], + ttl=60, + ) + + sync_client.expire.assert_called_once() + assert sync_client.expire.call_args.args[1] == 60 + + +def test_set_cache_skips_ttl_when_absent(): + sync_client = MagicMock() + cache = _make_cache(sync_client=sync_client) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + cache.set_cache( + key="cache-key", + value={"content": "Paris"}, + messages=[{"role": "user", "content": "What is the capital of France?"}], + ) + + sync_client.expire.assert_not_called() + + +def test_get_cache_returns_hit_above_threshold(): + sync_client = MagicMock() + sync_client.ft.return_value.search.return_value = _search_result(0.1) + cache = _make_cache(sync_client=sync_client, similarity_threshold=0.8) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + metadata = {} + result = cache.get_cache( + key="cache-key", + messages=[{"role": "user", "content": "capital of France?"}], + metadata=metadata, + ) + + assert result == {"content": "Paris"} + assert metadata["semantic-similarity"] == pytest.approx(0.9) + + +def test_get_cache_misses_below_threshold(): + sync_client = MagicMock() + sync_client.ft.return_value.search.return_value = _search_result(0.5) + cache = _make_cache(sync_client=sync_client, similarity_threshold=0.8) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + metadata = {} + result = cache.get_cache( + key="cache-key", + messages=[{"role": "user", "content": "capital of Germany?"}], + metadata=metadata, + ) + + assert result is None + assert metadata["semantic-similarity"] == pytest.approx(0.5) + + +def test_get_cache_misses_when_no_docs(): + sync_client = MagicMock() + sync_client.ft.return_value.search.return_value = SimpleNamespace(docs=[]) + cache = _make_cache(sync_client=sync_client) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + metadata = {} + result = cache.get_cache( + key="cache-key", + messages=[{"role": "user", "content": "capital of France?"}], + metadata=metadata, + ) + + assert result is None + assert metadata["semantic-similarity"] == 0.0 + + +def test_get_cache_query_filters_by_scope_tag(): + sync_client = MagicMock() + sync_client.ft.return_value.search.return_value = _search_result(0.1) + cache = _make_cache(sync_client=sync_client) + cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3]) + + cache.get_cache( + key="cache-key", + messages=[{"role": "user", "content": "capital of France?"}], + metadata={}, + ) + + query = sync_client.ft.return_value.search.call_args.args[0] + scope = ValkeySemanticCache._scope_tag("cache-key") + assert scope in query.query_string() + assert "KNN 1 @embedding" in query.query_string() + + +def _async_ft(search_distance): + search_obj = SimpleNamespace( + search=AsyncMock(return_value=_search_result(search_distance)), + create_index=AsyncMock(), + ) + return MagicMock(return_value=search_obj) + + +@pytest.mark.asyncio +async def test_async_set_and_get_roundtrip(): + async_client = AsyncMock() + async_client.ft = _async_ft(0.05) + cache = _make_cache(async_client=async_client, similarity_threshold=0.8) + cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3]) + + await cache.async_set_cache( + key="cache-key", + value={"content": "Paris"}, + messages=[{"role": "user", "content": "What is the capital of France?"}], + ttl=30, + ) + async_client.hset.assert_awaited_once() + async_client.expire.assert_awaited_once() + assert async_client.expire.call_args.args[1] == 30 + + metadata = {} + result = await cache.async_get_cache( + key="cache-key", + messages=[{"role": "user", "content": "capital city of France"}], + metadata=metadata, + ) + assert result == {"content": "Paris"} + assert metadata["semantic-similarity"] == pytest.approx(0.95) + + +@pytest.mark.asyncio +async def test_async_get_cache_misses_below_threshold(): + async_client = AsyncMock() + async_client.ft = _async_ft(0.4) + cache = _make_cache(async_client=async_client, similarity_threshold=0.8) + cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3]) + + metadata = {} + result = await cache.async_get_cache( + key="cache-key", + messages=[{"role": "user", "content": "capital of Germany?"}], + metadata=metadata, + ) + assert result is None + assert metadata["semantic-similarity"] == pytest.approx(0.6) + + +def test_ensure_index_swallows_already_exists(): + sync_client = MagicMock() + sync_client.ft.return_value.create_index.side_effect = Exception( + "Index test_index already exists." + ) + cache = _make_cache(sync_client=sync_client) + + cache._ensure_index_sync(3) + assert cache._index_dim == 3 + + +def test_ensure_index_reraises_unexpected_error(): + sync_client = MagicMock() + sync_client.ft.return_value.create_index.side_effect = Exception( + "connection refused" + ) + cache = _make_cache(sync_client=sync_client) + + with pytest.raises(Exception, match="connection refused"): + cache._ensure_index_sync(3) + + +_FT_INFO_ATTRS_DIM_1536 = [ + [b"identifier", b"litellm_cache_key", b"type", b"TAG"], + [ + b"identifier", + b"embedding", + b"type", + b"VECTOR", + b"index", + [b"capacity", 10240, b"dimensions", 1536, b"distance_metric", b"COSINE"], + ], +] + + +def test_extract_index_dim_parses_nested_ft_info(): + info = {"attributes": _FT_INFO_ATTRS_DIM_1536} + assert ValkeySemanticCache._extract_index_dim(info) == 1536 + assert ValkeySemanticCache._extract_index_dim({"attributes": []}) is None + + +def test_ensure_index_raises_on_dimension_mismatch(): + sync_client = MagicMock() + sync_client.ft.return_value.create_index.side_effect = Exception( + "Index test_index already exists." + ) + sync_client.ft.return_value.info.return_value = { + "attributes": _FT_INFO_ATTRS_DIM_1536 + } + cache = _make_cache(sync_client=sync_client) + + with pytest.raises( + ValueError, match="already exists with embedding dimension 1536" + ): + cache._ensure_index_sync(768) + assert cache._index_dim is None + + +def test_ensure_index_accepts_matching_existing_dimension(): + sync_client = MagicMock() + sync_client.ft.return_value.create_index.side_effect = Exception( + "Index test_index already exists." + ) + sync_client.ft.return_value.info.return_value = { + "attributes": _FT_INFO_ATTRS_DIM_1536 + } + cache = _make_cache(sync_client=sync_client) + + cache._ensure_index_sync(1536) + assert cache._index_dim == 1536 + + +def test_init_builds_only_missing_client_from_url(): + sync_client = MagicMock() + cache = ValkeySemanticCache( + similarity_threshold=0.8, + redis_url="redis://valkey-host:6380", + sync_client=sync_client, + ) + assert cache.sync_client is sync_client + assert cache.async_client is not None and cache.async_client is not sync_client + + +def test_init_uses_both_injected_clients_without_connection_info(monkeypatch): + for var in ("VALKEY_HOST", "VALKEY_PORT", "REDIS_HOST", "REDIS_PORT"): + monkeypatch.delenv(var, raising=False) + sync_client = MagicMock() + async_client = AsyncMock() + + cache = ValkeySemanticCache( + similarity_threshold=0.8, + sync_client=sync_client, + async_client=async_client, + ) + + assert cache.sync_client is sync_client + assert cache.async_client is async_client + + +def test_cache_dispatches_valkey_semantic_type(): + from litellm.caching.caching import Cache + from litellm.types.caching import LiteLLMCacheType + + cache = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + host="valkey-host", + port="6380", + similarity_threshold=0.8, + ) + + assert isinstance(cache.cache, ValkeySemanticCache) + + +@pytest.mark.asyncio +async def test_index_info_uses_valkey_ft_info(): + # The /health/readiness endpoint calls _index_info() on any + # RedisSemanticCache instance; since ValkeySemanticCache subclasses it, + # the inherited RedisVL implementation (which reads self.llmcache) would + # break. This override must query valkey-search FT.INFO instead. + async_client = AsyncMock() + info_namespace = SimpleNamespace(info=AsyncMock(return_value={"num_docs": 3})) + async_client.ft = MagicMock(return_value=info_namespace) + cache = _make_cache(async_client=async_client) + + result = await cache._index_info() + + assert result == {"num_docs": 3} + async_client.ft.assert_called_once_with("test_index") + + +def test_importing_caching_does_not_require_redis(): + # redis is an optional dependency (extra_proxy), so the base SDK can be + # installed without it. Selecting valkey-semantic needs redis, but merely + # importing litellm.caching.caching must not, or `import litellm` breaks for + # every base-SDK user. This runs in a subprocess with redis blocked so the + # check is not polluted by redis already being imported in this session. + code = textwrap.dedent(""" + import sys + for name in ("redis", "redis.asyncio", "redis.commands", + "redis.commands.search"): + sys.modules[name] = None + import litellm.caching.caching # must not import redis at module top + from litellm.types.caching import LiteLLMCacheType + assert LiteLLMCacheType.VALKEY_SEMANTIC == "valkey-semantic" + print("ok") + """) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env={**os.environ, "PYTHONPATH": _REPO_ROOT}, + ) + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout