scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

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

Browse source 
* 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.
This commit is contained in:
michelligabriele 2026-04-28 17:25:11 +02:00 • committed by GitHub
parent 62920a0cb2
commit 0dd64baa66
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 328 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
litellm/caching/caching.py
View file

@ -650,7 +650,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.
@ -662,6 +665,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()
@ -670,6 +674,7 @@ class Cache:
"index": data.get("index"),
"object": data.get("object"),
"model": model,
"prompt_tokens_details": prompt_tokens_details,
}
else:
data = vars(embedding_response)
@ -678,10 +683,54 @@ 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,
@ -693,10 +742,18 @@ 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(

88
litellm/caching/caching_handler.py
View file

@ -59,6 +59,7 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import PromptTokensDetailsWrapper
else:
LiteLLMLoggingObj = Any
@ -415,6 +416,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:
@ -431,11 +433,35 @@ 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:
from litellm.types.utils import PromptTokensDetailsWrapper
try:
prompt_tokens_details = PromptTokensDetailsWrapper(
**aggregated_details
)
except Exception:
prompt_tokens_details = None
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:
@ -478,8 +504,70 @@ 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 isinstance(v1, dict) and isinstance(v2, dict):
# Recursively merge nested dicts (e.g. cache_creation_token_details)
nested: dict = {}
for nk in set(v1.keys()) | set(v2.keys()):
nv1 = v1.get(nk, 0)
nv2 = v2.get(nk, 0)
if isinstance(nv1, (int, float)) and isinstance(nv2, (int, float)):
nested[nk] = nv1 + nv2
elif nv1:
nested[nk] = nv1
else:
nested[nk] = nv2
merged[key] = nested
elif v1:
merged[key] = v1
else:
merged[key] = v2
if not merged:
return None
from litellm.types.utils import PromptTokensDetailsWrapper
try:
return PromptTokensDetailsWrapper(**merged)
except Exception:
return None
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