fix(caching): preserve prompt_tokens_details through embedding cache round-trip - #26653
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.
…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.
Greptile SummaryThis PR fixes a bug where Confidence Score: 4/5Safe to merge; no correctness-breaking issues found, only P2 observability and edge-case concerns. All findings are P2 (style/best-practice level). The core logic is sound, backward compatibility is preserved via litellm/caching/caching_handler.py (silent exception swallowing in two places) and litellm/caching/caching.py (batch distribution assumption).
|
| Filename | Overview |
|---|---|
| litellm/types/caching.py | Adds prompt_tokens_details: Optional[dict] to CachedEmbedding TypedDict; follows the same convention as existing optional fields, and .get() is used everywhere so old cached entries without this key are handled correctly. |
| litellm/caching/caching.py | Adds _get_per_item_prompt_tokens_details() and threads details through _convert_to_cached_embedding() and add_embedding_response_to_cache(). The integer-remainder distribution for multi-item batches is correct for uniform batches but can mis-assign counts in mixed text/image batches. |
| litellm/caching/caching_handler.py | Aggregates per-item prompt_tokens_details on cache retrieval and merges them in combine_usage via the new _merge_prompt_tokens_details helper. Main concern is the bare except Exception: return None in two places which silently drops details without any log output. |
| tests/test_litellm/caching/test_caching_handler.py | Adds 5 well-scoped unit tests covering single-item preservation, backward-compat (old cached entries), multi-item aggregation, combine_usage merge, and None-handling. All tests use mocks and make no real network calls. |
Sequence Diagram
sequenceDiagram
participant Caller
participant CachingHandler
participant Cache
participant Provider
Note over Caller,Provider: Cache MISS (first call)
Caller->>CachingHandler: aembedding(input=[img1, img2])
CachingHandler->>Cache: get(k1), get(k2) - both miss
CachingHandler->>Provider: embedding(input=[img1, img2])
Provider-->>CachingHandler: EmbeddingResponse with prompt_tokens_details image_count 2
CachingHandler->>Cache: set k1 with CachedEmbedding prompt_tokens_details image_count 1
CachingHandler->>Cache: set k2 with CachedEmbedding prompt_tokens_details image_count 1
CachingHandler-->>Caller: EmbeddingResponse usage image_count 2
Note over Caller,Provider: Full Cache HIT (second call)
Caller->>CachingHandler: aembedding(input=[img1, img2])
CachingHandler->>Cache: get(k1), get(k2) - both hit
CachingHandler->>CachingHandler: aggregate details, image_count 1 plus 1
CachingHandler-->>Caller: EmbeddingResponse usage image_count 2
Note over Caller,Provider: Partial Cache HIT
Caller->>CachingHandler: aembedding(input=[img1, img3])
CachingHandler->>Cache: get(k1) hit, get(k3) miss
CachingHandler->>Provider: embedding(input=[img3])
Provider-->>CachingHandler: EmbeddingResponse image_count 1
CachingHandler->>CachingHandler: combine_usage, merge prompt_tokens_details
CachingHandler-->>Caller: EmbeddingResponse usage image_count 2
Reviews (1): 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: | ||
| from litellm.types.utils import PromptTokensDetailsWrapper | ||
|
|
||
| try: | ||
| prompt_tokens_details = PromptTokensDetailsWrapper( | ||
| **aggregated_details | ||
| ) | ||
| except Exception: | ||
| prompt_tokens_details = None |
There was a problem hiding this comment.
prompt_tokens_details without any diagnostic
PromptTokensDetailsWrapper(**aggregated_details) is swallowed silently on exception — if aggregated_details contains a key that Pydantic rejects (e.g. a field from a provider that isn't yet modelled in PromptTokensDetailsWrapper), the entire details object is discarded and callers see None with no log entry. The same pattern is repeated in _merge_prompt_tokens_details (line 566-569). Adding a verbose_logger.warning before the return None would make diagnosing this far easier.
| if aggregated_details: | |
| from litellm.types.utils import PromptTokensDetailsWrapper | |
| try: | |
| prompt_tokens_details = PromptTokensDetailsWrapper( | |
| **aggregated_details | |
| ) | |
| except Exception: | |
| prompt_tokens_details = None | |
| try: | |
| prompt_tokens_details = PromptTokensDetailsWrapper( | |
| **aggregated_details | |
| ) | |
| except Exception as e: | |
| verbose_logger.warning( | |
| f"LiteLLM Cache: failed to reconstruct prompt_tokens_details from cached embedding: {e}. Dropping details." | |
| ) | |
| prompt_tokens_details = 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 |
There was a problem hiding this comment.
When a batch contains a mix of text and image inputs (e.g. ["hello", "<base64_img>", "world"]) and the provider returns image_count=1, the distribution assigns image_count=1 to item 0 (the text input) instead of item 1 (the image), because the remainder goes to the lowest indices. Upon retrieval, the aggregated image_count is still correct, but the per-item association is wrong — a cached text-only re-run of item 0 would incorrectly carry an image_count. This is an inherent limitation of distributing aggregate counts without per-item provenance data; a code comment clarifying this caveat would help future maintainers.
0dd64ba
into
litellm_internal_staging
…round-trip (BerriAI#26653) * fix(caching): preserve prompt_tokens_details through embedding cache 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. * style: apply Black formatting to caching files * fix(caching): address Greptile review — cyclic import, guarded construction, 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.
…round-trip (BerriAI#26653) * fix(caching): preserve prompt_tokens_details through embedding cache 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. * style: apply Black formatting to caching files * fix(caching): address Greptile review — cyclic import, guarded construction, 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.
2 participants
The embedding caching layer was dropping
prompt_tokens_details(includingimage_count) becauseCachedEmbeddinghad no field for usage metadata and the cache retrieval code reconstructedUsagewithout it. This caused inconsistent responses where the first call returnedimage_countbut cached responses did not, breaking cost tracking for multimodal embeddings.Add
prompt_tokens_detailstoCachedEmbedding, persist per-item details during cache storage, aggregate them on retrieval, and merge them incombine_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 reviewType
🐛 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