fix(caching): restore stored prompt_tokens on embedding cache hits instead of recomputing

[michelligabriele]

Relevant issues

N/A

Linear ticket

Resolves LIT-2323

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays 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:

Screenshots / Proof of Fix

On an /embeddings cache hit, usage was rebuilt by re-tokenizing the raw input with token_counter instead of restoring what the provider returned. For multimodal image embeddings (e.g. amazon.titan-embed-image-v1) the provider reports prompt_tokens=0 and carries cost via prompt_tokens_details.image_count, so the recompute tokenizes the base64 blob and reports thousands of phantom tokens — the same input returns different usage depending on whether it hits the cache.

Reproduction: identical embedding request issued twice (miss → hit) through the caching path, with a provider response of prompt_tokens=0, image_count=1 for a base64 image input.

Before (cache hit recomputes prompt_tokens over the base64 input):

MISS usage: {"prompt_tokens": 0, "total_tokens": 0, "image_count": 1}
HIT  usage: {"prompt_tokens": 6601, "total_tokens": 6601, "image_count": 1}
AssertionError: usage mismatch: miss={'prompt_tokens': 0, ...} hit={'prompt_tokens': 6601, ...}

After (cache hit restores the stored provider value verbatim):

MISS usage: {"prompt_tokens": 0, "total_tokens": 0, "image_count": 1}
HIT  usage: {"prompt_tokens": 0, "total_tokens": 0, "image_count": 1}
HIT cache_hit flag: True
PASS: cache hit restored the provider usage verbatim (prompt_tokens=0), no token_counter recompute.

Unit tests: 17/17 passing in tests/test_litellm/caching/ (5 new + the full #26653 regression set for prompt_tokens_details).

Type

🐛 Bug Fix

Changes

• litellm/types/caching.py — add prompt_tokens: Optional[int] to CachedEmbedding (additive, no migration needed).
• litellm/caching/caching.py — new _get_per_item_prompt_tokens helper: single-item responses store the full usage.prompt_tokens; multi-item responses distribute it evenly (with remainder) so per-item values sum back to the original total. Wired into _convert_to_cached_embedding / add_embedding_response_to_cache, mirroring how prompt_tokens_details is stored since fix(caching): preserve prompt_tokens_details through embedding cache round-trip #26653.
• litellm/caching/caching_handler.py — on cache hit, prefer the stored prompt_tokens over the token_counter recompute. The check is is not None so a legitimate stored 0 is honored. Legacy cache entries without the field fall back to the existing token_counter behavior — fully backward compatible.
• tests/test_litellm/caching/test_caching_handler.py — 3 new tests: image-input hit restores prompt_tokens=0; multi-item hit sums stored per-item values; legacy entries still use the token_counter fallback.
• tests/test_litellm/caching/test_caching.py — 2 new tests for the store-side distribution helper (single item incl. 0, remainder split [4,3,3]).

[codecov]

Codecov Report

❌ Patch coverage is 93.75000% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/caching/caching.py	90.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR fixes a cache hit regression for embedding responses where prompt_tokens was recomputed by token_counter instead of restoring the provider's original value — causing multimodal image embeddings (e.g. amazon.titan-embed-image-v1, which reports prompt_tokens=0) to return thousands of phantom tokens on cache hits.

• Store side (caching.py): a new _get_per_item_prompt_tokens helper distributes the response's usage.prompt_tokens evenly across cached items, so per-item values always sum back to the original total. The value is stored in CachedEmbedding.prompt_tokens.
• Retrieval side (caching_handler.py): on cache hit, the handler checks cached_prompt_tokens is not None — correctly honoring a stored 0 — and sums the stored values; legacy entries without the field fall back to the old token_counter path for full backward compatibility.
• Three new handler tests and two new unit tests (mocks only) cover the image-input zero case, multi-item summation, and the legacy fallback.

Confidence Score: 5/5

Narrowly scoped to embedding cache hit handling, fully backward compatible via the is not None guard for legacy entries, and well-tested with mock-only unit tests.

The store-side distribution logic mirrors an identical pre-existing helper, the retrieval-side guard correctly distinguishes a stored 0 from a missing field, and every changed code path is covered by new tests.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/types/caching.py	Adds prompt_tokens: Optional[int] to CachedEmbedding TypedDict — additive, no migration needed.
litellm/caching/caching.py	Adds _get_per_item_prompt_tokens helper mirroring the existing _get_per_item_prompt_tokens_details distribution logic, and wires it into _convert_to_cached_embedding / add_embedding_response_to_cache.
litellm/caching/caching_handler.py	On cache hit, prefers stored prompt_tokens (with correct is not None guard for zero) over token_counter recompute; legacy entries fall back to the old path — backward compatible.
tests/test_litellm/caching/test_caching_handler.py	Adds 3 new async tests: image-input hit restores prompt_tokens=0; multi-item hit sums stored per-item values; legacy entries still use the token_counter fallback. All use mocks only.
tests/test_litellm/caching/test_caching.py	Adds 2 unit tests for _get_per_item_prompt_tokens: single-item returns full value (incl. 0), multi-item distributes with remainder [4, 3, 3].

_{Reviews (1): Last reviewed commit: "fix(caching): restore stored prompt_toke..." | Re-trigger Greptile}

[greptile-apps]

Greptile Summary

This PR fixes a bug where embedding cache hits recomputed prompt_tokens via token_counter instead of restoring the provider-returned value, causing multimodal image embeddings (e.g. amazon.titan-embed-image-v1) to report thousands of phantom tokens on cache hits (the base64 payload was tokenized rather than the original prompt_tokens=0 being honored).

• Adds prompt_tokens: Optional[int] to CachedEmbedding and a _get_per_item_prompt_tokens helper that stores the full value for single-item responses and distributes it evenly (with remainder on leading items) for multi-item responses — mirroring the existing _get_per_item_prompt_tokens_details pattern.
• In _process_async_embedding_cached_response, prefers the stored value (is not None so a legitimate 0 is honored) and falls through to token_counter only for legacy cache entries that predate this field — fully backward-compatible.

Confidence Score: 5/5

Safe to merge — the change is strictly additive to the cache schema and the new code path is gated behind an is not None check, leaving all existing behavior intact for legacy cache entries.

The fix is narrow and well-tested: the distribution helper is a direct analogue of the already-proven _get_per_item_prompt_tokens_details method, the is not None guard correctly distinguishes a stored 0 from a missing field, and all three new handler tests plus two new unit tests cover the main scenarios including backward compatibility.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/caching/caching_handler.py	On embedding cache hit, prefers stored prompt_tokens (including a legitimate 0) over token_counter recompute; falls through to token_counter only for legacy entries with no stored value. is not None guard correctly honors a stored 0.
litellm/caching/caching.py	Adds _get_per_item_prompt_tokens helper (mirroring existing _get_per_item_prompt_tokens_details) and threads prompt_tokens through _convert_to_cached_embedding / add_embedding_response_to_cache. Logic handles single-item, multi-item (with even distribution + remainder on leading items), and None-usage cases correctly.
litellm/types/caching.py	Adds prompt_tokens: Optional[int] to CachedEmbedding TypedDict — purely additive, no migration needed, backward-compatible with existing cached entries.
tests/test_litellm/caching/test_caching_handler.py	Three new async unit tests: image-input hit restores prompt_tokens=0, multi-item hit sums stored per-item values, and legacy entries trigger token_counter fallback. All tests use mocks with no real network calls.
tests/test_litellm/caching/test_caching.py	Two new tests for _get_per_item_prompt_tokens: single-item with 0, and remainder-aware distribution [4, 3, 3] for 10 tokens across 3 items. An existing test got minor whitespace-only reformatting with no semantic change.

_{Reviews (2): Last reviewed commit: "fix(caching): restore stored prompt_toke..." | Re-trigger Greptile}

[Sameerlite]

LGTM