fix(caching): keep embedding cache hits aligned with request inputs (#42571)

* fix(caching): keep embedding cache hits aligned with request inputs

Partial hits now send only the uncached inputs to the provider and merge
fresh vectors back into their original positions. Responses whose item
count differs from the input count (one input scoring many documents)
are no longer written to the per-input cache, since a later hit would
return a single item.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(caching): drop mutable collection builds flagged by the type discipline gate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(caching): bypass embedding cache entries written before the per input cardinality check

Embedding cache entries now carry format_version and readers treat entries without it as
misses, so entries that only hold the first row of a multi row response are refetched instead
of served until their TTL expires. The provider call also receives a copy of the request kwargs
with the uncached inputs rather than mutating the caller's mapping

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(caching): assert a partial embedding cache hit becomes a full hit on repeat

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(caching): await pending embedding cache writes before asserting on cache hits

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(caching): validate cached embeddings without mutating responses or request kwargs

Validate cache rows through a frozen pydantic model so import does not depend on
TypeAdapter support for ReadOnly TypedDicts, accept string embeddings, build the
merged partial hit response instead of mutating the cached one, and hand the
provider request mapping to post call hooks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(caching): keep cache_hit and response_ms on merged partial embedding hits

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/caching/caching.py

@@ -781,35 +781,23 @@ class Cache:
         Convert any embedding response into the standardized CachedEmbedding TypedDict format.
         """
         try:
-            if isinstance(embedding_response, dict):
-                return {
-                    "embedding": embedding_response.get("embedding"),
-                    "index": embedding_response.get("index"),
-                    "object": embedding_response.get("object"),
-                    "model": model,
-                    "prompt_tokens": prompt_tokens,
-                    "prompt_tokens_details": prompt_tokens_details,
-                }
-            elif hasattr(embedding_response, "model_dump"):
-                data = embedding_response.model_dump()
-                return {
-                    "embedding": data.get("embedding"),
-                    "index": data.get("index"),
-                    "object": data.get("object"),
-                    "model": model,
-                    "prompt_tokens": prompt_tokens,
-                    "prompt_tokens_details": prompt_tokens_details,
-                }
-            else:
-                data = vars(embedding_response)
-                return {
-                    "embedding": data.get("embedding"),
-                    "index": data.get("index"),
-                    "object": data.get("object"),
-                    "model": model,
-                    "prompt_tokens": prompt_tokens,
-                    "prompt_tokens_details": prompt_tokens_details,
-                }
+            data: Final = (
+                embedding_response
+                if isinstance(embedding_response, dict)
+                else embedding_response.model_dump()
+                if hasattr(embedding_response, "model_dump")
+                else vars(embedding_response)
+            )
+            cached: Final[CachedEmbedding] = {
+                "embedding": data.get("embedding"),
+                "index": data.get("index"),
+                "object": data.get("object"),
+                "model": model,
+                "prompt_tokens": prompt_tokens,
+                "prompt_tokens_details": prompt_tokens_details,
+                "format_version": EMBEDDING_CACHE_FORMAT_VERSION,
+            }
+            return cached
         except KeyError as e:
             raise ValueError(f"Missing expected key in embedding response: {e}")
 
@@ -925,6 +913,15 @@ class Cache:
             if self.should_use_cache(**kwargs) is not True:
                 return
 
+            input_count: Final = len(kwargs["input"]) if isinstance(kwargs["input"], list) else 1
+            if len(result.data) != input_count:
+                verbose_logger.debug(
+                    "LiteLLM Cache: skipping embedding cache write, %d inputs but %d embeddings in the response",
+                    input_count,
+                    len(result.data),
+                )
+                return
+
             # set default ttl if not set
             if self.ttl is not None:
                 kwargs["ttl"] = self.ttl

litellm/caching/caching_handler.py

@@ -21,7 +21,7 @@ import time
 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, ValidationError
 
 import litellm
 from litellm._logging import print_verbose, verbose_logger
@@ -34,7 +34,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
 from litellm.litellm_core_utils.logging_utils import (
     _assemble_complete_response_from_streaming_chunks,
 )
-from litellm.types.caching import CachedEmbedding
+from litellm.types.caching import EMBEDDING_CACHE_FORMAT_VERSION, CachedEmbedding
 from litellm.types.integrations.custom_logger import converted_stream_requested
 from litellm.types.llms.openai import ResponsesAPIResponse
 from litellm.types.rerank import RerankResponse
@@ -77,6 +77,7 @@ class CachingHandlerResponse(BaseModel):
     cached_result: object | None = None
     final_embedding_cached_response: EmbeddingResponse | None = None
     embedding_all_elements_cache_hit: bool = False  # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
+    embedding_uncached_input: list[str | list[int]] | None = None
 
 
 in_memory_cache_obj: Final = InMemoryCache()
@@ -168,6 +169,37 @@ def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
     return request_kwargs.get("cache_key", None)
 
 
+class _CachedEmbeddingRecord(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    embedding: list[float] | str | None
+    index: int | None
+    object: str | None
+    model: str | None
+    prompt_tokens: int | None
+    prompt_tokens_details: dict | None
+    format_version: int
+
+
+def _current_format_embedding_entry(entry: object) -> CachedEmbedding | None:
+    try:
+        record: Final = _CachedEmbeddingRecord.model_validate(entry)
+    except ValidationError:
+        return None
+    if record.format_version != EMBEDDING_CACHE_FORMAT_VERSION:
+        return None
+    cached: Final[CachedEmbedding] = {
+        "embedding": record.embedding,
+        "index": record.index,
+        "object": record.object,
+        "model": record.model,
+        "prompt_tokens": record.prompt_tokens,
+        "prompt_tokens_details": record.prompt_tokens_details,
+        "format_version": record.format_version,
+    }
+    return cached
+
+
 class LLMCachingHandler:
     def __init__(
         self,
@@ -320,6 +352,7 @@ class LLMCachingHandler:
                     return CachingHandlerResponse(
                         final_embedding_cached_response=final_embedding_cached_response,
                         embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
+                        embedding_uncached_input=self.handle_kwargs_input_list_or_str(kwargs),
                     )
 
             verbose_logger.debug("CACHE RESULT: %s", cached_result)
@@ -657,32 +690,30 @@ class LLMCachingHandler:
         if _caching_handler_response.final_embedding_cached_response is None:
             return embedding_response
 
-        idx = 0
-        final_data_list: Final = []
-        for item in _caching_handler_response.final_embedding_cached_response.data:
-            if item is None and embedding_response.data is not None:
-                final_data_list.append(embedding_response.data[idx])
-                idx += 1
-            else:
-                final_data_list.append(item)
-
-        _caching_handler_response.final_embedding_cached_response.data = final_data_list
-        _caching_handler_response.final_embedding_cached_response._hidden_params["cache_hit"] = True
-        _caching_handler_response.final_embedding_cached_response._response_ms = (
-            end_time - start_time
-        ).total_seconds() * 1000
-
-        ## USAGE
-        if (
-            _caching_handler_response.final_embedding_cached_response.usage is not None
-            and embedding_response.usage is not None
-        ):
-            _caching_handler_response.final_embedding_cached_response.usage = self.combine_usage(
-                usage1=_caching_handler_response.final_embedding_cached_response.usage,
-                usage2=embedding_response.usage,
-            )
-
-        return _caching_handler_response.final_embedding_cached_response
+        cached: Final = _caching_handler_response.final_embedding_cached_response
+        fresh_items: Final = iter(embedding_response.data or ())
+        merged_usage: Final = (
+            self.combine_usage(usage1=cached.usage, usage2=embedding_response.usage)
+            if cached.usage is not None and embedding_response.usage is not None
+            else cached.usage
+        )
+        merged: Final = EmbeddingResponse(
+            model=cached.model,
+            data=[  # mutable-ok: EmbeddingResponse.data is a pydantic list field
+                item
+                if item is not None
+                else Embedding(embedding=next(fresh_items)["embedding"], index=position, object="embedding")
+                for position, item in enumerate(cached.data)
+            ],
+            usage=merged_usage,
+            hidden_params={  # mutable-ok: EmbeddingResponse._hidden_params is a mutable dict field
+                **cached._hidden_params,
+                "cache_hit": True,
+            },
+            _response_headers=cached._response_headers,
+        )
+        merged._response_ms = (end_time - start_time).total_seconds() * 1000
+        return merged
 
     def _async_log_cache_hit_on_callbacks(
         self,
@@ -770,7 +801,7 @@ class LLMCachingHandler:
                         dynamic_cache_object=self.dual_cache,
                     )
                 )
-            cached_result = await asyncio.gather(*tasks)
+            cached_result = [_current_format_embedding_entry(entry) for entry in await asyncio.gather(*tasks)]
             ## check if cached result is None ##
             if cached_result is not None and isinstance(cached_result, list):
                 # set cached_result to None if all elements are None

litellm/types/caching.py

@@ -3,7 +3,7 @@ from enum import Enum
 from typing import Any, Final, Literal, Optional, Union
 
 from pydantic import BaseModel
-from typing_extensions import TypedDict
+from typing_extensions import ReadOnly, TypedDict
 
 
 class LiteLLMCacheType(str, Enum):
@@ -137,12 +137,16 @@ class HealthCheckCacheParams(BaseModel):
     redis_version: str | int | float | None = None
 
 
+EMBEDDING_CACHE_FORMAT_VERSION: Final = 2
+
+
 class CachedEmbedding(TypedDict):
     """Type definition for cached embedding objects"""
 
-    embedding: list[float] | None
-    index: int | None
-    object: str | None
-    model: str | None
-    prompt_tokens: int | None
-    prompt_tokens_details: dict | None
+    embedding: ReadOnly[list[float] | str | None]
+    index: ReadOnly[int | None]
+    object: ReadOnly[str | None]
+    model: ReadOnly[str | None]
+    prompt_tokens: ReadOnly[int | None]
+    prompt_tokens_details: ReadOnly[dict | None]
+    format_version: ReadOnly[int]

litellm/utils.py

@@ -2015,13 +2015,19 @@ def client(original_function):
                     print_verbose(f"Error while checking max token limit: {e}")
 
             # MODEL CALL
+            call_kwargs: Final = (
+                {**kwargs, "input": _caching_handler_response.embedding_uncached_input}
+                if _caching_handler_response is not None
+                and _caching_handler_response.embedding_uncached_input is not None
+                else kwargs
+            )
             try:
-                result = await original_function(*args, **kwargs)
+                result = await original_function(*args, **call_kwargs)
             except Exception as deployment_error:
                 _deployment_call_end_time = datetime.datetime.now()  # noqa: DTZ005  # matches the naive datetimes this whole function already times start_time/end_time with
                 try:
                     await async_post_call_failure_deployment_hook(
-                        request_data=kwargs,
+                        request_data=call_kwargs,
                         exception=deployment_error,
                         call_type=call_type,
                     )
@@ -2062,7 +2068,7 @@ def client(original_function):
             post_call_processing(
                 original_response=result,
                 model=model,
-                optional_params=kwargs,
+                optional_params=call_kwargs,
                 original_function=original_function,
                 rules_obj=rules_obj,
             )
@@ -2070,7 +2076,7 @@ def client(original_function):
             _call_type_enum: Final = _CALL_TYPE_ENUM_MAP.get(call_type)
             if _call_type_enum is not None:
                 result = await async_post_call_success_deployment_hook(
-                    request_data=kwargs,
+                    request_data=call_kwargs,
                     response=result,
                     call_type=_call_type_enum,
                 )
@@ -2079,7 +2085,7 @@ def client(original_function):
             await _llm_caching_handler.async_set_cache(
                 result=result,
                 original_function=original_function,
-                kwargs=kwargs,
+                kwargs=call_kwargs,
                 args=args,
             )
 

tests/test_litellm/caching/test_caching.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import re
 from unittest.mock import MagicMock
@@ -6,8 +7,9 @@ import pytest
 
 import litellm.caching.redis_cache as redis_cache_module
 from litellm.caching.caching import Cache
+from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
 from litellm.caching.redis_cache import RedisCache, _RedisTimeoutLogThrottle
-from litellm.types.caching import LiteLLMCacheType, SemanticCacheScope
+from litellm.types.caching import EMBEDDING_CACHE_FORMAT_VERSION, LiteLLMCacheType, SemanticCacheScope
 from litellm.types.utils import Embedding, EmbeddingResponse, Usage
 
 
@@ -278,3 +280,112 @@ def test_exact_cache_key_includes_anthropic_messages_params(anthropic_param):
     assert baseline != cache.get_cache_key(
         model="claude-sonnet-4-5", messages=messages, **anthropic_param
     )
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_skips_write_when_one_input_yields_many_embeddings(monkeypatch):
+    """A cross-encoder behind /embeddings returns one score per document for a single
+    input string; caching data[0] per input would make the second call return 1 score."""
+    import litellm
+    from litellm import CustomLLM
+
+    class ScoreEveryDocument(CustomLLM):
+        provider_calls: int = 0
+
+        async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse:
+            self.provider_calls += 1
+            return EmbeddingResponse(
+                model=model,
+                data=[Embedding(embedding=[float(i)], index=i, object="embedding") for i in range(5)],
+            )
+
+    scorer = ScoreEveryDocument()
+    monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "score-every-doc", "custom_handler": scorer}])
+    monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "score-every-doc"])
+    monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "score-every-doc"])
+    monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
+
+    batch = '{"query": "q", "documents": ["a", "b", "c", "d", "e"]}'
+    first = await litellm.aembedding(model="score-every-doc/m", input=[batch])
+    await asyncio.gather(*_PENDING_CACHE_WRITES)
+    second = await litellm.aembedding(model="score-every-doc/m", input=[batch])
+
+    assert scorer.provider_calls == 2
+    assert [len(first.data), len(second.data)] == [5, 5]
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_refetches_entries_written_without_format_version(monkeypatch):
+    import litellm
+    from litellm import CustomLLM
+
+    class EmbedLength(CustomLLM):
+        provider_calls: int = 0
+
+        async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse:
+            self.provider_calls += 1
+            return EmbeddingResponse(
+                model=model,
+                data=[
+                    Embedding(embedding=[float(len(text))], index=idx, object="embedding")
+                    for idx, text in enumerate(input)
+                ],
+            )
+
+    embedder = EmbedLength()
+    monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "embed-length", "custom_handler": embedder}])
+    monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "embed-length"])
+    monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "embed-length"])
+    monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
+
+    await litellm.aembedding(model="embed-length/m", input=["abcd"])
+    await asyncio.gather(*_PENDING_CACHE_WRITES)
+    store = litellm.cache.cache.cache_dict
+    stored = [entry["response"] for entry in store.values()]
+    assert [entry["format_version"] for entry in stored] == [EMBEDDING_CACHE_FORMAT_VERSION], stored
+    legacy_store = {
+        key: {
+            **entry,
+            "response": {
+                field: value
+                for field, value in {**entry["response"], "embedding": [-1.0]}.items()
+                if field != "format_version"
+            },
+        }
+        for key, entry in store.items()
+    }
+    monkeypatch.setattr(litellm.cache.cache, "cache_dict", legacy_store)
+
+    refetched = await litellm.aembedding(model="embed-length/m", input=["abcd"])
+
+    assert embedder.provider_calls == 2, "an entry written without format_version must be a cache miss"
+    assert [item["embedding"] for item in refetched.data] == [[4.0]]
+
+
+@pytest.mark.asyncio
+async def test_embedding_cache_serves_base64_string_embeddings_on_repeat(monkeypatch):
+    import litellm
+    from litellm import CustomLLM
+
+    class Base64Embedder(CustomLLM):
+        provider_calls: int = 0
+
+        async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse:
+            self.provider_calls += 1
+            return EmbeddingResponse(
+                model=model,
+                data=[Embedding(embedding="AACAPwAAAEA=", index=idx, object="embedding") for idx, _ in enumerate(input)],
+            )
+
+    embedder = Base64Embedder()
+    monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "embed-b64", "custom_handler": embedder}])
+    monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "embed-b64"])
+    monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "embed-b64"])
+    monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
+
+    first = await litellm.aembedding(model="embed-b64/m", input=["abcd"])
+    await asyncio.gather(*_PENDING_CACHE_WRITES)
+    second = await litellm.aembedding(model="embed-b64/m", input=["abcd"])
+
+    assert embedder.provider_calls == 1, "a string embedding written to the cache must be served on repeat"
+    assert [item["embedding"] for item in second.data] == [item["embedding"] for item in first.data] == ["AACAPwAAAEA="]

tests/test_litellm/caching/test_caching_handler.py

@@ -11,7 +11,7 @@ from fastapi.testclient import TestClient
 from datetime import datetime
 from unittest.mock import AsyncMock
 
-from litellm.caching.caching_handler import LLMCachingHandler
+from litellm.caching.caching_handler import _PENDING_CACHE_WRITES, LLMCachingHandler
 
 
 @pytest.mark.asyncio
@@ -780,3 +780,46 @@ async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_repl
     assert hit.cached_result.choices[0].message.content == "done"
     logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
     assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
+
+
+@pytest.mark.asyncio
+async def test_partial_embedding_cache_hit_sends_only_misses_and_keeps_input_order(monkeypatch):
+    import litellm
+    from litellm import CustomLLM
+    from litellm.caching.caching import Cache
+    from litellm.types.utils import Embedding, EmbeddingResponse
+
+    class RecordingEmbedder(CustomLLM):
+        provider_inputs: tuple[tuple[str, ...], ...] = ()
+
+        async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse:
+            self.provider_inputs = (*self.provider_inputs, tuple(input))
+            return EmbeddingResponse(
+                model=model,
+                data=[
+                    Embedding(embedding=[float(len(text))], index=idx, object="embedding")
+                    for idx, text in enumerate(input)
+                ],
+            )
+
+    embedder = RecordingEmbedder()
+    monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "recording-embedder", "custom_handler": embedder}])
+    monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "recording-embedder"])
+    monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "recording-embedder"])
+    monkeypatch.setattr(litellm, "cache", Cache(type="local"))
+
+    await litellm.aembedding(model="recording-embedder/m", input=["aa", "bbbb"])
+    await asyncio.gather(*_PENDING_CACHE_WRITES)
+    mixed_input = ["c", "aa", "ddd", "bbbb", "eeeee"]
+    response = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
+    await asyncio.gather(*_PENDING_CACHE_WRITES)
+
+    assert embedder.provider_inputs == (("aa", "bbbb"), ("c", "ddd", "eeeee")), embedder.provider_inputs
+    assert [item["index"] for item in response.data] == [0, 1, 2, 3, 4]
+    assert [item["embedding"] for item in response.data] == [[float(len(text))] for text in mixed_input]
+    assert response._hidden_params["cache_hit"] is True, "a partial hit must still be reported as a cache hit"
+
+    repeat = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
+
+    assert len(embedder.provider_inputs) == 2, embedder.provider_inputs
+    assert [item["embedding"] for item in repeat.data] == [[float(len(text))] for text in mixed_input]