VertexAI Anthropic - streaming cost tracking w/ prompt caching fixes (#12188)

* fix(rebuild-usage-object---ensure-cache_tokens-is-set): Ensures cache tokens is correctly set

Fixes https://github.com/BerriAI/litellm/issues/12149

* test(test_stream_chunk_builder_utils.py): add unit test to ensure cached tokens is part of stream chunk builder

Ensures standardized values are used
This commit is contained in:
Krish Dholakia 2025-06-30 22:00:10 -07:00 • committed by GitHub
parent 847c4514c4
commit 49ed3065f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 38 additions and 3 deletions

View file

@ -107,9 +107,9 @@ class ChunkProcessor:
self, tool_call_chunks: List[Dict[str, Any]]
) -> List[ChatCompletionMessageToolCall]:
tool_calls_list: List[ChatCompletionMessageToolCall] = []
tool_call_map: Dict[
int, Dict[str, Any]
] = {} # Map to store tool calls by index
tool_call_map: Dict[int, Dict[str, Any]] = (
{}
) # Map to store tool calls by index
for chunk in tool_call_chunks:
choices = chunk["choices"]
@ -415,6 +415,8 @@ class ChunkProcessor:
if prompt_tokens_details is not None:
returned_usage.prompt_tokens_details = prompt_tokens_details
# Return a new usage object with the new values
returned_usage = Usage(**returned_usage.model_dump())
return returned_usage

View file

@ -140,3 +140,35 @@ def test_generic_cost_per_token_above_200k_tokens():
* usage.completion_tokens,
10,
)
def test_generic_cost_per_token_anthropic_prompt_caching():
model = "claude-sonnet-4@20250514"
usage = Usage(
completion_tokens=90,
prompt_tokens=28436,
total_tokens=28526,
completion_tokens_details=CompletionTokensDetailsWrapper(
accepted_prediction_tokens=None,
audio_tokens=None,
reasoning_tokens=0,
rejected_prediction_tokens=None,
text_tokens=None,
),
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None
),
cache_creation_input_tokens=118,
cache_read_input_tokens=28432,
)
custom_llm_provider = "vertex_ai"
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider,
)
print(f"prompt_cost: {prompt_cost}")
assert prompt_cost < 0.085

View file

@ -242,3 +242,4 @@ def test_cache_read_input_tokens_retained():
assert usage.cache_creation_input_tokens == 4
assert usage.cache_read_input_tokens == 11775
assert usage.prompt_tokens_details.cached_tokens == 11775