fix(vectorize): include request kwargs in the embedding cache key

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,25 +185,31 @@ 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
                     )
 
                 # 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,25 +340,33 @@ 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
                     )
 
                 # 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,
                 }

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