mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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:
parent
3d2b8fed32
commit
d4efaf8cf9
2 changed files with 209 additions and 0 deletions
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue