fix(caching): preserve prompt_tokens_details through embedding cache round-trip - #24806
Conversation
…round-trip The embedding caching layer was dropping prompt_tokens_details (including image_count) because CachedEmbedding had no field for usage metadata and the cache retrieval code reconstructed Usage without it. This caused inconsistent responses where the first call returned image_count but cached responses did not, breaking cost tracking for multimodal embeddings. Add prompt_tokens_details to CachedEmbedding, persist per-item details during cache storage, aggregate them on retrieval, and merge them in combine_usage() for partial cache hits.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes a bug where Key changes:
Confidence Score: 4/5Safe to merge after the result.data None guard is addressed; remaining findings are low-risk style issues. The core fix is correct and well-tested with 5 new mock-only unit tests. A P1 issue exists in _get_per_item_prompt_tokens_details where len(result.data) raises TypeError if result.data is None. A P2 edge case exists in _merge_prompt_tokens_details where zero-valued falsy fields could silently be dropped. The TypedDict total=True concern is a type-checking nit with no runtime impact. litellm/caching/caching.py — _get_per_item_prompt_tokens_details needs a None guard on result.data
|
| Filename | Overview |
|---|---|
| litellm/caching/caching.py | Added _get_per_item_prompt_tokens_details helper that distributes batch-level token details per item, and threaded those details through _convert_to_cached_embedding and add_embedding_response_to_cache. Minor defensive gap: result.data is not guarded against None before calling len(). |
| litellm/caching/caching_handler.py | Updated _process_async_embedding_cached_response to aggregate prompt_tokens_details from cached items; added _merge_prompt_tokens_details helper and wired it into combine_usage. Logic is sound; minor falsy-zero edge case in the non-numeric/non-dict branch of _merge_prompt_tokens_details. |
| litellm/types/caching.py | Added prompt_tokens_details: Optional[dict] to CachedEmbedding; because TypedDict defaults to total=True this new key is required for mypy but backward-compatible at runtime (all read sites use .get()). |
| tests/test_litellm/caching/test_caching_handler.py | 5 new mock-only tests covering single-item preservation, backward compatibility, multi-item aggregation, combine_usage merge, and None handling. All use MagicMock/AsyncMock, no real network calls. Pre-existing test unchanged. |
Sequence Diagram
sequenceDiagram
participant Caller
participant CachingHandler
participant Cache
participant Backend
Note over Caller,Backend: Write path
Caller->>CachingHandler: aembedding(inputs)
CachingHandler->>Cache: add_embedding_response_to_cache
Cache->>Cache: _get_per_item_prompt_tokens_details
Cache->>Backend: store CachedEmbedding with prompt_tokens_details
Note over Caller,Backend: Full cache hit
Caller->>CachingHandler: aembedding(inputs)
CachingHandler->>Backend: fetch cached items
Backend-->>CachingHandler: items with prompt_tokens_details
CachingHandler->>CachingHandler: aggregate details per item
CachingHandler-->>Caller: EmbeddingResponse with usage
Note over Caller,Backend: Partial cache hit
CachingHandler->>Backend: some hits, some misses
CachingHandler->>CachingHandler: live call for misses
CachingHandler->>CachingHandler: combine_usage merges details
CachingHandler-->>Caller: merged EmbeddingResponse
Reviews (3): Last reviewed commit: "fix(caching): address Greptile review — ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if aggregated_details: | ||
| prompt_tokens_details = PromptTokensDetailsWrapper( | ||
| **aggregated_details | ||
| ) |
There was a problem hiding this comment.
PromptTokensDetailsWrapper construction may fail on unexpected cached keys
aggregated_details is built from cr.get("prompt_tokens_details") which may come from a cache entry written by a different version of LiteLLM (either older or newer). If that stored dict contains a key that is not a known field of PromptTokensDetailsWrapper and the model is configured with extra = "forbid", PromptTokensDetailsWrapper(**aggregated_details) will raise a ValidationError and crash the entire cache-retrieval path for all inputs in the batch.
A simple try/except is sufficient to keep backward and forward compatibility:
if aggregated_details:
try:
prompt_tokens_details = PromptTokensDetailsWrapper(
**aggregated_details
)
except Exception:
verbose_logger.debug(
"LiteLLM Cache: Could not reconstruct prompt_tokens_details "
f"from cached data: {aggregated_details}"
)
prompt_tokens_details = NoneThe same risk applies to PromptTokensDetailsWrapper(**merged) in _merge_prompt_tokens_details (line 543).
| def _get_per_item_prompt_tokens_details( | ||
| self, | ||
| result: EmbeddingResponse, | ||
| idx_in_result_data: int, | ||
| ) -> Optional[dict]: | ||
| """ | ||
| Extract per-item prompt_tokens_details from a response for caching. | ||
|
|
||
| For single-item responses (common for multimodal providers like Bedrock Titan, | ||
| Nova, Vertex AI), returns the full prompt_tokens_details. | ||
| For multi-item responses, distributes integer fields evenly across items | ||
| so that summing all per-item details reconstructs the original totals. | ||
| """ | ||
| if ( | ||
| result.usage is None | ||
| or result.usage.prompt_tokens_details is None | ||
| ): | ||
| return None | ||
|
|
||
| details = result.usage.prompt_tokens_details | ||
| if hasattr(details, "model_dump"): | ||
| details_dict = details.model_dump(exclude_none=True) | ||
| elif isinstance(details, dict): | ||
| details_dict = {k: v for k, v in details.items() if v is not None} | ||
| else: | ||
| return None | ||
|
|
||
| if not details_dict: | ||
| return None | ||
|
|
||
| num_items = len(result.data) | ||
| if num_items <= 1: | ||
| return details_dict | ||
|
|
||
| # Distribute integer/float fields evenly across items | ||
| per_item: dict = {} | ||
| for key, value in details_dict.items(): | ||
| if isinstance(value, int): | ||
| quotient, remainder = divmod(value, num_items) | ||
| per_item[key] = quotient + (1 if idx_in_result_data < remainder else 0) | ||
| elif isinstance(value, float): | ||
| per_item[key] = value / num_items | ||
| else: | ||
| per_item[key] = value | ||
| return per_item if per_item else None |
There was a problem hiding this comment.
_get_per_item_prompt_tokens_details contains non-trivial remainder-distributing arithmetic (e.g. 3 images across 2 items → [2, 1]), yet all five new tests exercise only the cache-retrieval/aggregation side by pre-populating mock cached_result dicts that already contain per-item values. A regression in the distribution logic (wrong idx_in_result_data, off-by-one in the remainder check) would not be caught.
Consider adding a unit test that:
- Constructs a real
EmbeddingResponsewithusage.prompt_tokens_details.image_count = Nandlen(data) = K - Calls
cache._get_per_item_prompt_tokens_details(result, idx)for everyidxinrange(K) - Asserts that the per-item values sum back to
N
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| Embedding, | ||
| EmbeddingResponse, | ||
| ModelResponse, | ||
| PromptTokensDetailsWrapper, |
Check failure
Code scanning / CodeQL
Module-level cyclic import
…uction, nested dict merge Move PromptTokensDetailsWrapper to inline import to resolve CodeQL cyclic import warning. Guard PromptTokensDetailsWrapper construction with try/except to handle unexpected cached keys. Add recursive dict merging in _merge_prompt_tokens_details for nested fields like cache_creation_token_details.
|
Closing in favour of #26653, which is the same change refiled against |
2 participants
The embedding caching layer was dropping prompt_tokens_details (including image_count) because CachedEmbedding had no field for usage metadata and the cache retrieval code reconstructed Usage without it. This caused inconsistent responses where the first call returned image_count but cached responses did not, breaking cost tracking for multimodal embeddings.
Add prompt_tokens_details to CachedEmbedding, persist per-item details during cache storage, aggregate them on retrieval, and merge them in combine_usage() for partial cache hits.
Relevant issues
Related to #9029 (image_count support for multimodal embeddings)
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🐛 Bug Fix
Changes
litellm/types/caching.py— Addedprompt_tokens_details: Optional[dict]field toCachedEmbeddingTypedDict so usage metadata survives serialization to any cache backend (Redis, S3, disk)litellm/caching/caching.py— Updated_convert_to_cached_embedding()to accept and store per-item details; added_get_per_item_prompt_tokens_details()helper that extracts details from the response (distributing evenly for multi-item batches); updatedadd_embedding_response_to_cache()to pass details throughlitellm/caching/caching_handler.py— Updated_process_async_embedding_cached_response()to aggregateprompt_tokens_detailsfrom cached items back intoPromptTokensDetailsWrapper; updatedcombine_usage()to merge details during partial cache hits; added_merge_prompt_tokens_details()helpertests/test_litellm/caching/test_caching_handler.py— 5 new tests: single-item preservation, backward compatibility (old cached entries without the field), multi-item aggregation,combine_usagemerge, and None handling