scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(caching): preserve prompt_tokens_details through embedding cache round-trip

Browse source 
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.
This commit is contained in:
michelligabriele 2026-03-30 20:08:30 +02:00
parent a278933f18
commit 56a587d874
No known key found for this signature in database
4 changed files with 305 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
litellm/caching/caching.py
View file

@ -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(

64
litellm/caching/caching_handler.py
View file

@ -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,

1
litellm/types/caching.py
View file

@ -118,3 +118,4 @@ class CachedEmbedding(TypedDict):
index: Optional[int]
object: Optional[str]
model: Optional[str]
prompt_tokens_details: Optional[dict]

180
tests/test_litellm/caching/test_caching_handler.py
View file

@ -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