From d4efaf8cf9b52d5d5c9ec1a29858df6d30513647 Mon Sep 17 00:00:00 2001 From: Junsu Lee Date: Mon, 27 Apr 2026 11:04:50 +0900 Subject: [PATCH] feat(arize): emit prompt-caching tokens as OpenInference span attributes `_set_usage_outputs` emits only total / completion / prompt / reasoning tokens, leaving `LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ` and `LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE` defined-but-unused. As a result, observability backends (Langfuse, Arize Phoenix) cannot display the prompt-cache breakdown nor apply correct cost calculation for Anthropic / Bedrock prompt caching (cache reads at 0.1x, cache writes at 1.25x of standard input). Read from `usage.prompt_tokens_details` so the change is provider-agnostic. LiteLLM's `Usage.__init__` already normalizes provider-specific cache fields onto this object: - Anthropic `cache_read_input_tokens` -> prompt_tokens_details.cached_tokens - DeepSeek `prompt_cache_hit_tokens` -> prompt_tokens_details.cached_tokens - OpenAI native cached_tokens -> prompt_tokens_details.cached_tokens - Anthropic `cache_creation_input_tokens` -> prompt_tokens_details.cache_creation_tokens Mapping: - prompt_tokens_details.cached_tokens -> llm.token_count.prompt_details.cache_read - prompt_tokens_details.cache_creation_tokens -> llm.token_count.prompt_details.cache_write Note: PR #24112 introduced a similar `cached_tokens` block in litellm_oss_staging_03_21_2026 but has not yet landed on main; PR #26506 on litellm-oss-staging-04-25-2026 refactored the function without including it. This change brings the emission directly to main with extended cache_creation coverage. Tests: - Anthropic: both cache_read and cache_write emitted - OpenAI: only cache_read emitted (no cache_write concept) - DeepSeek: prompt_cache_hit_tokens normalized -> cache_read emitted - No-cache: neither attribute emitted --- litellm/integrations/arize/_utils.py | 30 +++ .../integrations/arize/test_arize_utils.py | 179 ++++++++++++++++++ 2 files changed, 209 insertions(+) diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 8dfaa8b1425..d717c7f9b20 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -244,6 +244,36 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): reasoning_tokens, ) + # Prompt-caching tokens. LiteLLM normalizes provider-specific fields + # (Anthropic `cache_read_input_tokens`, DeepSeek `prompt_cache_hit_tokens`, + # OpenAI native `cached_tokens`) onto `prompt_tokens_details.cached_tokens`, + # so reading from there is provider-agnostic. `cache_creation_tokens` is + # Anthropic-specific but lives on the same object (None for others). + # The OpenInference SpanAttributes constants for these were defined but + # never populated, leaving observability backends (e.g. Langfuse) unable + # to display the cache breakdown. + prompt_tokens_details = usage.get("prompt_tokens_details") or {} + if hasattr(prompt_tokens_details, "get"): + cached_tokens = prompt_tokens_details.get("cached_tokens") + cache_creation_tokens = prompt_tokens_details.get("cache_creation_tokens") + else: + cached_tokens = getattr(prompt_tokens_details, "cached_tokens", None) + cache_creation_tokens = getattr( + prompt_tokens_details, "cache_creation_tokens", None + ) + if cached_tokens: + safe_set_attribute( + span, + span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, + cached_tokens, + ) + if cache_creation_tokens: + safe_set_attribute( + span, + span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE, + cache_creation_tokens, + ) + def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: """ diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/test_litellm/integrations/arize/test_arize_utils.py index a87a4167899..368bf1e787c 100644 --- a/tests/test_litellm/integrations/arize/test_arize_utils.py +++ b/tests/test_litellm/integrations/arize/test_arize_utils.py @@ -273,6 +273,185 @@ def test_arize_set_attributes_responses_api(): ) +def test_arize_set_attributes_anthropic_cache_tokens(): + """ + Anthropic prompt-caching populates both cache_read and cache_creation; + LiteLLM normalizes them onto `prompt_tokens_details`, and they should be + emitted as OpenInference cache_read / cache_write span attributes so + observability backends can display the cache breakdown and apply correct + cost calculation (cache reads at 0.1x, cache writes at 1.25x). + """ + from unittest.mock import MagicMock + + from litellm.types.utils import Choices, ModelResponse, Usage + + span = MagicMock() + kwargs = { + "model": "claude-sonnet-4", + "messages": [{"role": "user", "content": "Hi"}], + "standard_logging_object": { + "model_parameters": {"user": "test_user"}, + "metadata": {}, + "call_type": "completion", + }, + "optional_params": {"stream": False}, + "litellm_params": {"custom_llm_provider": "anthropic"}, + } + response_obj = ModelResponse( + usage=Usage( + prompt_tokens=4276, + completion_tokens=50, + total_tokens=4326, + cache_read_input_tokens=4000, + cache_creation_input_tokens=261, + ), + choices=[Choices(message={"role": "assistant", "content": "Hello"})], + model="claude-sonnet-4", + id="msg-cache-1", + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 4326) + span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 50) + span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 4276) + span.set_attribute.assert_any_call( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, 4000 + ) + span.set_attribute.assert_any_call( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE, 261 + ) + + +def test_arize_set_attributes_openai_cached_tokens(): + """ + OpenAI's native cache surfaces only cached_tokens (no creation). The + cache_read attribute must still be emitted; cache_write should be omitted. + """ + from unittest.mock import MagicMock + + from litellm.types.utils import ( + Choices, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, + ) + + span = MagicMock() + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hi"}], + "standard_logging_object": { + "model_parameters": {"user": "test_user"}, + "metadata": {}, + "call_type": "completion", + }, + "optional_params": {"stream": False}, + "litellm_params": {"custom_llm_provider": "openai"}, + } + response_obj = ModelResponse( + usage=Usage( + prompt_tokens=1000, + completion_tokens=100, + total_tokens=1100, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), + ), + choices=[Choices(message={"role": "assistant", "content": "Hello"})], + model="gpt-4o", + id="msg-openai-cache", + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + span.set_attribute.assert_any_call( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, 500 + ) + attribute_keys = [c.args[0] for c in span.set_attribute.call_args_list] + assert ( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE not in attribute_keys + ) + + +def test_arize_set_attributes_deepseek_cache_hit_tokens(): + """ + DeepSeek surfaces cache hits via `prompt_cache_hit_tokens`; LiteLLM + normalizes this onto `prompt_tokens_details.cached_tokens`, which the + OTEL emitter must pick up. + """ + from unittest.mock import MagicMock + + from litellm.types.utils import Choices, ModelResponse, Usage + + span = MagicMock() + kwargs = { + "model": "deepseek-chat", + "messages": [{"role": "user", "content": "Hi"}], + "standard_logging_object": { + "model_parameters": {"user": "test_user"}, + "metadata": {}, + "call_type": "completion", + }, + "optional_params": {"stream": False}, + "litellm_params": {"custom_llm_provider": "deepseek"}, + } + response_obj = ModelResponse( + usage=Usage( + prompt_tokens=2000, + completion_tokens=100, + total_tokens=2100, + prompt_cache_hit_tokens=1500, + ), + choices=[Choices(message={"role": "assistant", "content": "Hello"})], + model="deepseek-chat", + id="msg-deepseek-cache", + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + span.set_attribute.assert_any_call( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, 1500 + ) + + +def test_arize_set_attributes_no_cache_tokens_omits_attributes(): + """ + Without cache tokens (no caching, or first-time prompt below threshold), + cache_read / cache_write attributes must not be emitted. + """ + from unittest.mock import MagicMock + + from litellm.types.utils import Choices, ModelResponse, Usage + + span = MagicMock() + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hi"}], + "standard_logging_object": { + "model_parameters": {"user": "test_user"}, + "metadata": {}, + "call_type": "completion", + }, + "optional_params": {"stream": False}, + "litellm_params": {"custom_llm_provider": "openai"}, + } + response_obj = ModelResponse( + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + choices=[Choices(message={"role": "assistant", "content": "Hi"})], + model="gpt-4o", + id="resp-no-cache", + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + attribute_keys = [c.args[0] for c in span.set_attribute.call_args_list] + assert ( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ not in attribute_keys + ) + assert ( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE not in attribute_keys + ) + + class TestArizeLogger(CustomLogger): """ Custom logger implementation to capture standard_callback_dynamic_params.