from typing import Any, Dict, List, Optional
from redis import Redis
from redisvl.extensions.llmcache.base import BaseLLMCache
from redisvl.index import SearchIndex
from redisvl.query import RangeQuery
from redisvl.redis.utils import array_to_buffer
from redisvl.schema.schema import IndexSchema
from redisvl.utils.vectorize import BaseVectorizer, HFTextVectorizer
[docs]
class SemanticCache(BaseLLMCache):
"""Semantic Cache for Large Language Models."""
entry_id_field_name: str = "id"
prompt_field_name: str = "prompt"
vector_field_name: str = "prompt_vector"
response_field_name: str = "response"
metadata_field_name: str = "metadata"
def __init__(
self,
name: str = "llmcache",
prefix: Optional[str] = None,
distance_threshold: float = 0.1,
ttl: Optional[int] = None,
vectorizer: Optional[BaseVectorizer] = None,
redis_client: Optional[Redis] = None,
redis_url: str = "redis://localhost:6379",
connection_args: Dict[str, Any] = {},
**kwargs,
):
"""Semantic Cache for Large Language Models.
Args:
name (str, optional): The name of the semantic cache search index.
Defaults to "llmcache".
prefix (Optional[str], optional): The prefix for Redis keys
associated with the semantic cache search index. Defaults to
None, and the index name will be used as the key prefix.
distance_threshold (float, optional): Semantic threshold for the
cache. Defaults to 0.1.
ttl (Optional[int], optional): The time-to-live for records cached
in Redis. Defaults to None.
vectorizer (BaseVectorizer, optional): The vectorizer for the cache.
Defaults to HFTextVectorizer.
redis_client(Redis, optional): A redis client connection instance.
Defaults to None.
redis_url (str, optional): The redis url. Defaults to
"redis://localhost:6379".
connection_args (Dict[str, Any], optional): The connection arguments
for the redis client. Defaults to None.
Raises:
TypeError: If an invalid vectorizer is provided.
TypeError: If the TTL value is not an int.
ValueError: If the threshold is not between 0 and 1.
ValueError: If the index name is not provided
"""
super().__init__(ttl)
# Use the index name as the key prefix by default
if prefix is None:
prefix = name
# Set vectorizer default
if vectorizer is None:
vectorizer = HFTextVectorizer(
model="sentence-transformers/all-mpnet-base-v2"
)
# build cache index schema
schema = IndexSchema.from_dict({"index": {"name": name, "prefix": prefix}})
# add fields
schema.add_fields(
[
{"name": self.prompt_field_name, "type": "text"},
{"name": self.response_field_name, "type": "text"},
{
"name": self.vector_field_name,
"type": "vector",
"attrs": {
"dims": vectorizer.dims,
"datatype": "float32",
"distance_metric": "cosine",
"algorithm": "flat",
},
},
]
)
# build search index
self._index = SearchIndex(schema=schema)
# handle redis connection
if redis_client:
self._index.set_client(redis_client)
else:
self._index.connect(redis_url=redis_url, **connection_args)
# initialize other components
self.default_return_fields = [
self.entry_id_field_name,
self.prompt_field_name,
self.response_field_name,
self.vector_field_name,
self.metadata_field_name,
]
self.set_vectorizer(vectorizer)
self.set_threshold(distance_threshold)
self._index.create(overwrite=False)
@property
def index(self) -> SearchIndex:
"""The underlying SearchIndex for the cache.
Returns:
SearchIndex: The search index.
"""
return self._index
@property
def distance_threshold(self) -> float:
"""The semantic distance threshold for the cache.
Returns:
float: The semantic distance threshold.
"""
return self._distance_threshold
[docs]
def set_threshold(self, distance_threshold: float) -> None:
"""Sets the semantic distance threshold for the cache.
Args:
distance_threshold (float): The semantic distance threshold for
the cache.
Raises:
ValueError: If the threshold is not between 0 and 1.
"""
if not 0 <= float(distance_threshold) <= 1:
raise ValueError(
f"Distance must be between 0 and 1, got {distance_threshold}"
)
self._distance_threshold = float(distance_threshold)
[docs]
def set_vectorizer(self, vectorizer: BaseVectorizer) -> None:
"""Sets the vectorizer for the LLM cache.
Must be a valid subclass of BaseVectorizer and have equivalent
dimensions to the vector field defined in the schema.
Args:
vectorizer (BaseVectorizer): The RedisVL vectorizer to use for
vectorizing cache entries.
Raises:
TypeError: If the vectorizer is not a valid type.
ValueError: If the vector dimensions are mismatched.
"""
if not isinstance(vectorizer, BaseVectorizer):
raise TypeError("Must provide a valid redisvl.vectorizer class.")
schema_vector_dims = self._index.schema.fields[self.vector_field_name].attrs.dims # type: ignore
if schema_vector_dims != vectorizer.dims:
raise ValueError(
"Invalid vector dimensions! "
f"Vectorizer has dims defined as {vectorizer.dims}",
f"Vector field has dims defined as {schema_vector_dims}",
)
self._vectorizer = vectorizer
[docs]
def clear(self) -> None:
"""Clear the cache of all keys while preserving the index."""
with self._index.client.pipeline(transaction=False) as pipe: # type: ignore
for key in self._index.client.scan_iter(match=f"{self._index.prefix}:*"): # type: ignore
pipe.delete(key)
pipe.execute()
[docs]
def delete(self) -> None:
"""Clear the semantic cache of all keys and remove the underlying search
index."""
self._index.delete(drop=True)
def _refresh_ttl(self, key: str) -> None:
"""Refresh the time-to-live for the specified key."""
if self.ttl:
self._index.client.expire(key, self.ttl) # type: ignore
def _vectorize_prompt(self, prompt: Optional[str]) -> List[float]:
"""Converts a text prompt to its vector representation using the
configured vectorizer."""
if not isinstance(prompt, str):
raise TypeError("Prompt must be a string.")
return self._vectorizer.embed(prompt)
def _search_cache(
self, vector: List[float], num_results: int, return_fields: Optional[List[str]]
) -> List[Dict[str, Any]]:
"""Searches the semantic cache for similar prompt vectors and returns
the specified return fields for each cache hit."""
# Setup and type checks
if not isinstance(vector, list):
raise TypeError("Vector must be a list of floats")
return_fields = return_fields or self.default_return_fields
if not isinstance(return_fields, list):
raise TypeError("return_fields must be a list of field names")
# Construct vector RangeQuery for the cache check
query = RangeQuery(
vector=vector,
vector_field_name=self.vector_field_name,
return_fields=return_fields,
distance_threshold=self._distance_threshold,
num_results=num_results,
return_score=True,
)
# Gather and return the cache hits
cache_hits: List[Dict[str, Any]] = self._index.query(query)
# Process cache hits
for hit in cache_hits:
self._refresh_ttl(hit[self.entry_id_field_name])
# Check for metadata and deserialize
if self.metadata_field_name in hit:
hit[self.metadata_field_name] = self.deserialize(
hit[self.metadata_field_name]
)
return cache_hits
[docs]
def check(
self,
prompt: Optional[str] = None,
vector: Optional[List[float]] = None,
num_results: int = 1,
return_fields: Optional[List[str]] = None,
) -> List[Dict[str, Any]]:
"""Checks the semantic cache for results similar to the specified prompt
or vector.
This method searches the cache using vector similarity with
either a raw text prompt (converted to a vector) or a provided vector as
input. It checks for semantically similar prompts and fetches the cached
LLM responses.
Args:
prompt (Optional[str], optional): The text prompt to search for in
the cache.
vector (Optional[List[float]], optional): The vector representation
of the prompt to search for in the cache.
num_results (int, optional): The number of cached results to return.
Defaults to 1.
return_fields (Optional[List[str]], optional): The fields to include
in each returned result. If None, defaults to all available
fields in the cached entry.
Returns:
List[Dict[str, Any]]: A list of dicts containing the requested
return fields for each similar cached response.
Raises:
ValueError: If neither a `prompt` nor a `vector` is specified.
TypeError: If `return_fields` is not a list when provided.
.. code-block:: python
response = cache.check(
prompt="What is the captial city of France?"
)
"""
if not (prompt or vector):
raise ValueError("Either prompt or vector must be specified.")
# Use provided vector or create from prompt
vector = vector or self._vectorize_prompt(prompt)
# Check for cache hits by searching the cache
cache_hits = self._search_cache(vector, num_results, return_fields)
return cache_hits
[docs]
def store(
self,
prompt: str,
response: str,
vector: Optional[List[float]] = None,
metadata: Optional[dict] = None,
) -> str:
"""Stores the specified key-value pair in the cache along with metadata.
Args:
prompt (str): The user prompt to cache.
response (str): The LLM response to cache.
vector (Optional[List[float]], optional): The prompt vector to
cache. Defaults to None, and the prompt vector is generated on
demand.
metadata (Optional[dict], optional): The optional metadata to cache
alongside the prompt and response. Defaults to None.
Returns:
str: The Redis key for the entries added to the semantic cache.
Raises:
ValueError: If neither prompt nor vector is specified.
TypeError: If provided metadata is not a dictionary.
.. code-block:: python
key = cache.store(
prompt="What is the captial city of France?",
response="Paris",
metadata={"city": "Paris", "country": "France"}
)
"""
# Vectorize prompt if necessary and create cache payload
vector = vector or self._vectorize_prompt(prompt)
# Construct semantic cache payload
id_field = self.entry_id_field_name
payload = {
id_field: self.hash_input(prompt),
self.prompt_field_name: prompt,
self.response_field_name: response,
self.vector_field_name: array_to_buffer(vector),
}
if metadata is not None:
if not isinstance(metadata, dict):
raise TypeError("If specified, cached metadata must be a dictionary.")
# Serialize the metadata dict and add to cache payload
payload[self.metadata_field_name] = self.serialize(metadata)
# Load LLMCache entry with TTL
keys = self._index.load(data=[payload], ttl=self._ttl, id_field=id_field)
return keys[0]
