From 56a587d874c618810447428e556dfe1671cda28c Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Mon, 30 Mar 2026 20:08:30 +0200 Subject: [PATCH] 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. --- litellm/caching/caching.py | 62 +++++- litellm/caching/caching_handler.py | 64 +++++++ litellm/types/caching.py | 1 + .../caching/test_caching_handler.py | 180 ++++++++++++++++++ 4 files changed, 305 insertions(+), 2 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 406a4f8c98a..97e5fa330cb 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -647,7 +647,10 @@ class Cache: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") def _convert_to_cached_embedding( - self, embedding_response: Any, model: Optional[str] + self, + embedding_response: Any, + model: Optional[str], + prompt_tokens_details: Optional[dict] = None, ) -> CachedEmbedding: """ Convert any embedding response into the standardized CachedEmbedding TypedDict format. @@ -659,6 +662,7 @@ class Cache: "index": embedding_response.get("index"), "object": embedding_response.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } elif hasattr(embedding_response, "model_dump"): data = embedding_response.model_dump() @@ -667,6 +671,7 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } else: data = vars(embedding_response) @@ -675,10 +680,57 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } except KeyError as e: raise ValueError(f"Missing expected key in embedding response: {e}") + 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 + def add_embedding_response_to_cache( self, result: EmbeddingResponse, @@ -690,10 +742,16 @@ class Cache: kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx_in_result_data] + # Extract per-item prompt_tokens_details from response usage + prompt_tokens_details = self._get_per_item_prompt_tokens_details( + result=result, + idx_in_result_data=idx_in_result_data, + ) + # Always convert to properly typed CachedEmbedding model_name = result.model embedding_dict: CachedEmbedding = self._convert_to_cached_embedding( - embedding_response, model_name + embedding_response, model_name, prompt_tokens_details=prompt_tokens_details, ) cache_key, cached_data, kwargs = self._add_cache_logic( diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 7cdbd3fc03d..ab292bc6f33 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -52,6 +52,7 @@ from litellm.types.utils import ( Embedding, EmbeddingResponse, ModelResponse, + PromptTokensDetailsWrapper, TextCompletionResponse, TranscriptionResponse, Usage, @@ -413,6 +414,7 @@ class LLMCachingHandler: final_embedding_cached_response._hidden_params["cache_hit"] = True prompt_tokens = 0 + aggregated_details: Optional[dict] = None for val in non_null_list: idx, cr = val # (idx, cr) tuple if cr is not None: @@ -429,11 +431,30 @@ class LLMCachingHandler: prompt_tokens += token_counter( text=kwargs_input_as_list[idx], count_response_tokens=True ) + # Aggregate prompt_tokens_details from cached items + item_details = cr.get("prompt_tokens_details") + if item_details: + if aggregated_details is None: + aggregated_details = {} + for key, value in item_details.items(): + if isinstance(value, (int, float)): + aggregated_details[key] = ( + aggregated_details.get(key, 0) + value + ) + else: + aggregated_details[key] = value + ## USAGE + prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None + if aggregated_details: + prompt_tokens_details = PromptTokensDetailsWrapper( + **aggregated_details + ) usage = Usage( prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens, + prompt_tokens_details=prompt_tokens_details, ) final_embedding_cached_response.usage = usage if len(remaining_list) == 0: @@ -476,8 +497,51 @@ class LLMCachingHandler: prompt_tokens=usage1.prompt_tokens + usage2.prompt_tokens, completion_tokens=usage1.completion_tokens + usage2.completion_tokens, total_tokens=usage1.total_tokens + usage2.total_tokens, + prompt_tokens_details=self._merge_prompt_tokens_details( + usage1.prompt_tokens_details, + usage2.prompt_tokens_details, + ), ) + def _merge_prompt_tokens_details( + self, + details1: Optional[PromptTokensDetailsWrapper], + details2: Optional[PromptTokensDetailsWrapper], + ) -> Optional[PromptTokensDetailsWrapper]: + """Merge two PromptTokensDetailsWrapper objects by summing numeric fields.""" + if details1 is None and details2 is None: + return None + if details1 is None: + return details2 + if details2 is None: + return details1 + + dict1 = ( + details1.model_dump(exclude_none=True) + if hasattr(details1, "model_dump") + else {} + ) + dict2 = ( + details2.model_dump(exclude_none=True) + if hasattr(details2, "model_dump") + else {} + ) + + merged: dict = {} + for key in set(dict1.keys()) | set(dict2.keys()): + v1 = dict1.get(key, 0) + v2 = dict2.get(key, 0) + if isinstance(v1, (int, float)) and isinstance(v2, (int, float)): + merged[key] = v1 + v2 + elif v1: + merged[key] = v1 + else: + merged[key] = v2 + + if not merged: + return None + return PromptTokensDetailsWrapper(**merged) + def _combine_cached_embedding_response_with_api_result( self, _caching_handler_response: CachingHandlerResponse, diff --git a/litellm/types/caching.py b/litellm/types/caching.py index c8194ce2e7d..f8050b292c7 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -118,3 +118,4 @@ class CachedEmbedding(TypedDict): index: Optional[int] object: Optional[str] model: Optional[str] + prompt_tokens_details: Optional[dict] diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 837ce7d405d..742a4f410d4 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -52,3 +52,183 @@ async def test_process_async_embedding_cached_response(): print(f"response: {response}") assert len(response.data) == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_preserves_prompt_tokens_details(): + """Test that prompt_tokens_details (including image_count) survives a full cache hit.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_backward_compat_no_prompt_tokens_details(): + """Test that old cached items without prompt_tokens_details still work.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Old-format cached item — no prompt_tokens_details field + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-ada-002", + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-ada-002", "input": "test"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is None + + +@pytest.mark.asyncio +async def test_embedding_cache_aggregates_multiple_image_counts(): + """Test that image_count is summed correctly across multiple cached items.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + { + "embedding": [0.031, 0.042], + "index": 1, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={ + "model": "amazon.titan-embed-image-v1", + "input": ["img1", "img2"], + }, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 2 + + +def test_combine_usage_merges_prompt_tokens_details(): + """Test that combine_usage merges prompt_tokens_details from both Usage objects.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + usage1 = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + usage2 = Usage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2), + ) + + combined = llm_caching_handler.combine_usage(usage1, usage2) + + assert combined.prompt_tokens == 30 + assert combined.total_tokens == 30 + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 3 + + +def test_combine_usage_handles_none_details(): + """Test that combine_usage works when one or both sides have null prompt_tokens_details.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Both null + usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20) + combined = llm_caching_handler.combine_usage(usage_a, usage_b) + assert combined.prompt_tokens_details is None + + # Only first has details + usage_c = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + combined = llm_caching_handler.combine_usage(usage_c, usage_b) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1 + + # Only second has details + combined = llm_caching_handler.combine_usage(usage_a, usage_c) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1