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
This commit is contained in:
Junsu Lee 2026-04-27 11:04:50 +09:00
parent 3d2b8fed32
commit d4efaf8cf9
2 changed files with 209 additions and 0 deletions

View file

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

View file

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