scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

feat(caching): add valkey-semantic cache backend and fix semantic cache scope keys (#30675)

Browse source 
Adds a "valkey-semantic" cache type so semantic prompt caching can run
against Valkey clusters (for example AWS ElastiCache for Valkey) using the
valkey-search module.

The existing "redis-semantic" backend cannot drive valkey-search. RedisVL
gates the connection on a RediSearch module version that valkey-search does
not report, and its SemanticCache index declares the prompt as a TEXT field,
which valkey-search does not implement. ValkeySemanticCache therefore talks to
valkey-search directly over redis-py: it builds a vector index from the field
types valkey-search supports (TAG for caller scope, VECTOR for the prompt
embedding) and runs KNN queries for retrieval. Prompt extraction, embedding
generation, and cached-response parsing are reused from RedisSemanticCache
since those are backend agnostic. The redis dependency is imported lazily in
the cache dispatch so importing litellm without redis installed still works.

It also fixes semantic-cache scope keys so similarity matching works across
reworded prompts. get_cache_key() hashed messages / prompt / input into the
litellm_cache_key that every semantic backend filters its KNN search on, so a
paraphrase landed in a different bucket and never matched, even far above the
similarity threshold. For semantic cache types the prompt-bearing params are
now excluded from the scope key and the server-set tenant identity
(user_api_key, team, org) is appended instead, restoring embedding matching
within a tenant while keeping cache entries scoped to the authenticated
key / team / org. The three semantic backends share this key, so the same
change fixes redis-semantic and qdrant-semantic.

Connections resolve from VALKEY_HOST / VALKEY_PORT / VALKEY_PASSWORD, falling
back to REDIS_* for drop-in compatibility, and passwordless clusters (IAM or
no-auth) are supported.

Resolves #29121
Fixes #29086
This commit is contained in:
Yassin Kortam 2026-06-19 17:09:17 -07:00 • committed by GitHub
parent ea17236a1e
commit 9c3ad1b094
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 963 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm/caching/caching.py
View file

@ -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(

353
litellm/caching/valkey_semantic_cache.py Normal file
View file

@ -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()

1
litellm/types/caching.py
View file

@ -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"

70
tests/test_litellm/caching/test_caching.py
View file

@ -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

473
tests/test_litellm/caching/test_valkey_semantic_cache.py Normal file
View file

@ -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