Skip to content

Commit 6e17aba

Browse files
committed
fix(litellm): Record cached, cache-write and reasoning tokens
LiteLLM normalizes every provider's usage into prompt_tokens_details and completion_tokens_details, but the integration only read prompt, completion and total tokens. Read cached, cache-write and reasoning token counts from those details, falling back to LiteLLM's private cache fields. Fixes #5455
1 parent ab069a5 commit 6e17aba

2 files changed

Lines changed: 147 additions & 1 deletion

File tree

‎sentry_sdk/integrations/litellm.py‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222

2323
if TYPE_CHECKING:
2424
from datetime import datetime
25-
from typing import Any, Dict, List
25+
from typing import Any, Dict, List, Optional
2626

2727
try:
2828
import litellm # type: ignore[import-not-found]
@@ -36,6 +36,35 @@
3636
_SPAN_KEY = "_sentry_span"
3737

3838

39+
def _int_or_none(value: "Any") -> "Optional[int]":
40+
return value if isinstance(value, int) and not isinstance(value, bool) else None
41+
42+
43+
def _get_cached_tokens(usage: "Any") -> "Optional[int]":
44+
# LiteLLM normalizes every provider to the OpenAI shape, so cache reads land in
45+
# prompt_tokens_details.cached_tokens (Anthropic's cache_read_input_tokens included).
46+
# prompt_tokens already includes them, matching gen_ai.usage.input_tokens.
47+
details = getattr(usage, "prompt_tokens_details", None)
48+
cached = _int_or_none(getattr(details, "cached_tokens", None))
49+
if cached is None:
50+
cached = _int_or_none(getattr(usage, "_cache_read_input_tokens", None))
51+
return cached
52+
53+
54+
def _get_cache_write_tokens(usage: "Any") -> "Optional[int]":
55+
details = getattr(usage, "prompt_tokens_details", None)
56+
for attr in ("cache_write_tokens", "cache_creation_tokens"):
57+
written = _int_or_none(getattr(details, attr, None))
58+
if written is not None:
59+
return written
60+
return _int_or_none(getattr(usage, "_cache_creation_input_tokens", None))
61+
62+
63+
def _get_reasoning_tokens(usage: "Any") -> "Optional[int]":
64+
details = getattr(usage, "completion_tokens_details", None)
65+
return _int_or_none(getattr(details, "reasoning_tokens", None))
66+
67+
3968
def _store_span(kwargs: "Dict[str, Any]", span: "Any") -> None:
4069
kwargs[_SPAN_KEY] = span
4170

@@ -266,10 +295,30 @@ def _success_callback(
266295
if input_tokens is not None:
267296
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
268297

298+
input_tokens_cached = _get_cached_tokens(usage)
299+
if input_tokens_cached is not None:
300+
set_on_span(
301+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, input_tokens_cached
302+
)
303+
304+
input_tokens_cache_write = _get_cache_write_tokens(usage)
305+
if input_tokens_cache_write is not None:
306+
set_on_span(
307+
SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE,
308+
input_tokens_cache_write,
309+
)
310+
269311
output_tokens = getattr(usage, "completion_tokens", None)
270312
if output_tokens is not None:
271313
set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
272314

315+
output_tokens_reasoning = _get_reasoning_tokens(usage)
316+
if output_tokens_reasoning is not None:
317+
set_on_span(
318+
SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING,
319+
output_tokens_reasoning,
320+
)
321+
273322
total_tokens = getattr(usage, "total_tokens", None)
274323
if (
275324
total_tokens is None

‎tests/integrations/litellm/test_litellm.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ async def __call__(self, *args, **kwargs):
3535
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
3636
from openai import AsyncOpenAI, OpenAI
3737
from openai.types import CompletionUsage
38+
from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails
3839

3940
from sentry_sdk import start_transaction
4041
from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
@@ -3861,3 +3862,99 @@ def test_embeddings_data_collection(
38613862
assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings"
38623863
assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-ada-002"
38633864
assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5
3865+
3866+
3867+
def test_cached_and_reasoning_token_usage(
3868+
reset_litellm_executor,
3869+
sentry_init,
3870+
capture_events,
3871+
get_model_response,
3872+
nonstreaming_chat_completions_model_response,
3873+
):
3874+
sentry_init(
3875+
integrations=[LiteLLMIntegration()],
3876+
disabled_integrations=[StdlibIntegration],
3877+
traces_sample_rate=1.0,
3878+
stream_gen_ai_spans=False,
3879+
)
3880+
events = capture_events()
3881+
3882+
client = OpenAI(api_key="test-key")
3883+
model_response = get_model_response(
3884+
nonstreaming_chat_completions_model_response(
3885+
response_id="chatcmpl-test",
3886+
response_model="gpt-3.5-turbo",
3887+
message_content="Test response",
3888+
created=1234567890,
3889+
usage=CompletionUsage(
3890+
prompt_tokens=1200,
3891+
completion_tokens=300,
3892+
total_tokens=1500,
3893+
prompt_tokens_details=PromptTokensDetails(cached_tokens=1024),
3894+
completion_tokens_details=CompletionTokensDetails(reasoning_tokens=256),
3895+
),
3896+
),
3897+
serialize_pydantic=True,
3898+
request_headers={"X-Stainless-Raw-Response": "true"},
3899+
)
3900+
3901+
with mock.patch.object(
3902+
client.completions._client._client,
3903+
"send",
3904+
return_value=model_response,
3905+
), start_transaction(name="litellm test"):
3906+
litellm.completion(
3907+
model="gpt-3.5-turbo",
3908+
messages=[{"role": "user", "content": "Hello!"}],
3909+
client=client,
3910+
)
3911+
3912+
litellm_utils.executor.shutdown(wait=True)
3913+
3914+
(event,) = events
3915+
(span,) = (
3916+
x
3917+
for x in event["spans"]
3918+
if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm"
3919+
)
3920+
3921+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 1200
3922+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 1024
3923+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 300
3924+
assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 256
3925+
assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 1500
3926+
3927+
3928+
def test_cache_write_token_usage_from_anthropic_shaped_usage(
3929+
sentry_init, capture_events
3930+
):
3931+
"""LiteLLM maps Anthropic's cache_read/cache_creation counts into prompt_tokens_details."""
3932+
from litellm.types.utils import Usage
3933+
3934+
sentry_init(
3935+
integrations=[LiteLLMIntegration()],
3936+
traces_sample_rate=1.0,
3937+
stream_gen_ai_spans=False,
3938+
)
3939+
events = capture_events()
3940+
3941+
response = mock.Mock(spec=["usage", "model"])
3942+
response.model = "claude-sonnet-4-5"
3943+
response.usage = Usage(
3944+
prompt_tokens=2600,
3945+
completion_tokens=120,
3946+
total_tokens=2720,
3947+
cache_read_input_tokens=2048,
3948+
cache_creation_input_tokens=512,
3949+
)
3950+
kwargs = {"model": "claude-sonnet-4-5", "messages": [], "litellm_params": {}}
3951+
3952+
with start_transaction(name="litellm test"):
3953+
_input_callback(kwargs)
3954+
_success_callback(kwargs, response, datetime.now(), datetime.now())
3955+
3956+
(event,) = events
3957+
(span,) = event["spans"]
3958+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 2600
3959+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 2048
3960+
assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 512

0 commit comments

Comments
 (0)