diff --git a/redisvl/utils/vectorize/base.py b/redisvl/utils/vectorize/base.py index 73bb3452..b17e54f3 100644 --- a/redisvl/utils/vectorize/base.py +++ b/redisvl/utils/vectorize/base.py @@ -1,4 +1,5 @@ import io +import json import logging from enum import Enum from pathlib import Path @@ -7,7 +8,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator from redisvl.extensions.cache.embeddings import EmbeddingsCache -from redisvl.redis.utils import array_to_buffer +from redisvl.redis.utils import array_to_buffer, hashify from redisvl.schema.fields import VectorDataType from redisvl.utils.utils import deprecated_argument from redisvl.utils.vectorize import _telemetry @@ -105,13 +106,16 @@ def embed( if preprocess is not None: content = preprocess(content) + cache_metadata = kwargs.pop("metadata", {}) + cache_model_name = self._cache_model_name(kwargs) + with _telemetry.embeddings(self) as call: # Check cache if available and not skipped if self.cache is not None and not skip_cache: try: cache_result = self.cache.get( content=self._serialize_for_cache(content), - model_name=self.model, + model_name=cache_model_name, ) if cache_result: logger.debug(f"Cache hit for content with model {self.model}") @@ -125,7 +129,6 @@ def embed( call.cache_result(0, 1) # Generate embedding using provider-specific implementation - cache_metadata = kwargs.pop("metadata", {}) with call.provider_call(): embedding = self._embed(content, **kwargs) @@ -134,7 +137,7 @@ def embed( try: self.cache.set( content=self._serialize_for_cache(content), - model_name=self.model, + model_name=cache_model_name, embedding=embedding, metadata=cache_metadata, ) @@ -182,17 +185,19 @@ def embed_many( else: processed_contents = contents + cache_metadata = kwargs.pop("metadata", {}) + cache_model_name = self._cache_model_name(kwargs) + with _telemetry.embeddings(self) as call: # Get cached embeddings and identify misses results, cache_misses, cache_miss_indices = self._get_from_cache_batch( - processed_contents, skip_cache + processed_contents, skip_cache, cache_model_name ) if self.cache is not None and not skip_cache: call.cache_result(len(contents) - len(cache_misses), len(contents)) # Generate embeddings for cache misses if cache_misses: - cache_metadata = kwargs.pop("metadata", {}) with call.provider_call(): new_embeddings = self._embed_many( cache_misses, batch_size=batch_size, **kwargs @@ -200,7 +205,11 @@ def embed_many( # Store new embeddings in cache self._store_in_cache_batch( - cache_misses, new_embeddings, cache_metadata, skip_cache + cache_misses, + new_embeddings, + cache_metadata, + skip_cache, + cache_model_name, ) # Insert new embeddings into results array @@ -246,13 +255,16 @@ async def aembed( if preprocess is not None: content = preprocess(content) + cache_metadata = kwargs.pop("metadata", {}) + cache_model_name = self._cache_model_name(kwargs) + with _telemetry.embeddings(self) as call: # Check cache if available and not skipped if self.cache is not None and not skip_cache: try: cache_result = await self.cache.aget( content=self._serialize_for_cache(content), - model_name=self.model, + model_name=cache_model_name, ) if cache_result: logger.debug( @@ -270,7 +282,6 @@ async def aembed( call.cache_result(0, 1) # Generate embedding using provider-specific implementation - cache_metadata = kwargs.pop("metadata", {}) with call.provider_call(): embedding = await self._aembed(content, **kwargs) @@ -279,7 +290,7 @@ async def aembed( try: await self.cache.aset( content=self._serialize_for_cache(content), - model_name=self.model, + model_name=cache_model_name, embedding=embedding, metadata=cache_metadata, ) @@ -329,17 +340,21 @@ async def aembed_many( else: processed_contents = contents + cache_metadata = kwargs.pop("metadata", {}) + cache_model_name = self._cache_model_name(kwargs) + with _telemetry.embeddings(self) as call: # Get cached embeddings and identify misses results, cache_misses, cache_miss_indices = ( - await self._aget_from_cache_batch(processed_contents, skip_cache) + await self._aget_from_cache_batch( + processed_contents, skip_cache, cache_model_name + ) ) if self.cache is not None and not skip_cache: call.cache_result(len(contents) - len(cache_misses), len(contents)) # Generate embeddings for cache misses if cache_misses: - cache_metadata = kwargs.pop("metadata", {}) with call.provider_call(): new_embeddings = await self._aembed_many( cache_misses, batch_size=batch_size, **kwargs @@ -347,7 +362,11 @@ async def aembed_many( # Store new embeddings in cache await self._astore_in_cache_batch( - cache_misses, new_embeddings, cache_metadata, skip_cache + cache_misses, + new_embeddings, + cache_metadata, + skip_cache, + cache_model_name, ) # Insert new embeddings into results array @@ -390,14 +409,27 @@ async def _aembed_many( ) return self._embed_many(contents, batch_size=batch_size, **kwargs) + def _cache_model_name(self, kwargs: dict[str, Any]) -> str: + """Return the model name used as part of the embedding cache key. + + Provider kwargs (e.g. Cohere's ``input_type``) change the embedding the + model returns, so they must be part of the key. Without kwargs the key + is just ``self.model``, so entries cached before are still found. + """ + if not kwargs: + return self.model + digest = hashify(json.dumps(kwargs, sort_keys=True, default=str)) + return f"{self.model}:{digest}" + def _get_from_cache_batch( - self, contents: list[Any], skip_cache: bool + self, contents: list[Any], skip_cache: bool, model_name: str | None = None ) -> tuple[list[list[float] | None], list[str], list[int]]: """Get vector embeddings from cache and track cache misses. Args: contents: List of content to get from cache skip_cache: Whether to skip cache lookup + model_name: Cache key model name (defaults to ``self.model``) Returns: Tuple of (results, cache_misses, cache_miss_indices) @@ -414,7 +446,7 @@ def _get_from_cache_batch( # Efficient batch cache lookup cache_results = self.cache.mget( contents=(self._serialize_for_cache(c) for c in contents), - model_name=self.model, + model_name=model_name or self.model, ) # Process cache hits and collect misses @@ -437,13 +469,14 @@ def _get_from_cache_batch( return results, cache_misses, cache_miss_indices # type: ignore async def _aget_from_cache_batch( - self, contents: list[Any], skip_cache: bool + self, contents: list[Any], skip_cache: bool, model_name: str | None = None ) -> tuple[list[list[float] | None], list[str], list[int]]: """Asynchronously get vector embeddings from cache and track cache misses. Args: contents: List of content to get from cache skip_cache: Whether to skip cache lookup + model_name: Cache key model name (defaults to ``self.model``) Returns: Tuple of (results, cache_misses, cache_miss_indices) @@ -460,7 +493,7 @@ async def _aget_from_cache_batch( # Efficient batch cache lookup cache_results = await self.cache.amget( contents=(self._serialize_for_cache(c) for c in contents), - model_name=self.model, + model_name=model_name or self.model, ) # Process cache hits and collect misses @@ -490,6 +523,7 @@ def _store_in_cache_batch( embeddings: list[list[float]], metadata: dict[str, Any], skip_cache: bool, + model_name: str | None = None, ) -> None: """Store a batch of vector embeddings in the cache. @@ -498,6 +532,7 @@ def _store_in_cache_batch( embeddings: List of vector embeddings metadata: Metadata to store with the embeddings skip_cache: Whether to skip cache storage + model_name: Cache key model name (defaults to ``self.model``) """ if skip_cache or self.cache is None: return @@ -507,7 +542,7 @@ def _store_in_cache_batch( cache_items = [ { "content": self._serialize_for_cache(content), - "model_name": self.model, + "model_name": model_name or self.model, "embedding": emb, "metadata": metadata, } @@ -523,6 +558,7 @@ async def _astore_in_cache_batch( embeddings: list[list[float]], metadata: dict[str, Any], skip_cache: bool, + model_name: str | None = None, ) -> None: """Asynchronously store a batch of vector embeddings in the cache. @@ -531,6 +567,7 @@ async def _astore_in_cache_batch( embeddings: List of vector embeddings metadata: Metadata to store with the embeddings skip_cache: Whether to skip cache storage + model_name: Cache key model name (defaults to ``self.model``) """ if skip_cache or self.cache is None: return @@ -540,7 +577,7 @@ async def _astore_in_cache_batch( cache_items = [ { "content": self._serialize_for_cache(content), - "model_name": self.model, + "model_name": model_name or self.model, "embedding": emb, "metadata": metadata, } diff --git a/tests/unit/test_base_vectorizer.py b/tests/unit/test_base_vectorizer.py index 4f8ec8e8..e7aac7bb 100644 --- a/tests/unit/test_base_vectorizer.py +++ b/tests/unit/test_base_vectorizer.py @@ -43,3 +43,100 @@ def embed_many( assert vectorizer.model == "simple" assert vectorizer.dims == 10 assert vectorizer.dtype == "float32" + + +class _InputTypeVectorizer(BaseVectorizer): + """Fake provider whose output depends on the ``input_type`` kwarg.""" + + model: str = "fake-model" + dims: int = 1 + + def _value(self, kwargs: dict) -> float: + return 2.0 if kwargs.get("input_type") == "search_document" else 1.0 + + def _embed(self, content: Any, **kwargs) -> list[float]: + return [self._value(kwargs)] + + def _embed_many( + self, contents: list[Any], batch_size: int = 10, **kwargs + ) -> list[list[float]]: + return [[self._value(kwargs)] for _ in contents] + + async def _aembed(self, content: Any, **kwargs) -> list[float]: + return [self._value(kwargs)] + + async def _aembed_many( + self, contents: list[Any], batch_size: int = 10, **kwargs + ) -> list[list[float]]: + return [[self._value(kwargs)] for _ in contents] + + +class _FakeCache: + """In-memory stand-in for EmbeddingsCache keyed by (content, model_name).""" + + def __init__(self): + self.store: dict = {} + + def get(self, content, model_name): + return self.store.get((content, model_name)) + + def set(self, content, model_name, embedding, metadata=None): + self.store[(content, model_name)] = {"embedding": embedding} + + def mget(self, contents, model_name): + return [self.get(c, model_name) for c in contents] + + def mset(self, items): + for i in items: + self.set(i["content"], i["model_name"], i["embedding"]) + + async def aget(self, content, model_name): + return self.get(content, model_name) + + async def aset(self, content, model_name, embedding, metadata=None): + self.set(content, model_name, embedding, metadata) + + async def amget(self, contents, model_name): + return self.mget(contents, model_name) + + async def amset(self, items): + self.mset(items) + + +def _cached_vectorizer() -> _InputTypeVectorizer: + vec = _InputTypeVectorizer() + vec.cache = _FakeCache() + return vec + + +def test_embed_cache_distinguishes_request_kwargs(): + vec = _cached_vectorizer() + assert vec.embed("x", input_type="search_query") == [1.0] + assert vec.embed("x", input_type="search_document") == [2.0] + # Same kwargs hit the cache + assert vec.embed("x", input_type="search_document") == [2.0] + + +def test_embed_many_cache_distinguishes_request_kwargs(): + vec = _cached_vectorizer() + assert vec.embed_many(["x"], input_type="search_query") == [[1.0]] + assert vec.embed_many(["x"], input_type="search_document") == [[2.0]] + + +async def test_aembed_cache_distinguishes_request_kwargs(): + vec = _cached_vectorizer() + assert await vec.aembed("x", input_type="search_query") == [1.0] + assert await vec.aembed("x", input_type="search_document") == [2.0] + + +async def test_aembed_many_cache_distinguishes_request_kwargs(): + vec = _cached_vectorizer() + assert await vec.aembed_many(["x"], input_type="search_query") == [[1.0]] + assert await vec.aembed_many(["x"], input_type="search_document") == [[2.0]] + + +def test_embed_cache_key_without_kwargs_is_unchanged(): + vec = _cached_vectorizer() + vec.embed("x", metadata={"a": 1}) + vec.embed_many(["y"]) + assert {k[1] for k in vec.cache.store} == {"fake-model"}