From 6e17aba3ebadcb4d42b877f46a26197e2e739587 Mon Sep 17 00:00:00 2001 From: 4ktluffy Date: Tue, 6 Oct 2026 07:26:55 +0300 Subject: [PATCH 1/2] 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 --- sentry_sdk/integrations/litellm.py | 51 +++++++++++- tests/integrations/litellm/test_litellm.py | 97 ++++++++++++++++++++++ 2 files changed, 147 insertions(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index ce94441a4d..196cac3370 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -22,7 +22,7 @@ if TYPE_CHECKING: from datetime import datetime - from typing import Any, Dict, List + from typing import Any, Dict, List, Optional try: import litellm # type: ignore[import-not-found] @@ -36,6 +36,35 @@ _SPAN_KEY = "_sentry_span" +def _int_or_none(value: "Any") -> "Optional[int]": + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +def _get_cached_tokens(usage: "Any") -> "Optional[int]": + # LiteLLM normalizes every provider to the OpenAI shape, so cache reads land in + # prompt_tokens_details.cached_tokens (Anthropic's cache_read_input_tokens included). + # prompt_tokens already includes them, matching gen_ai.usage.input_tokens. + details = getattr(usage, "prompt_tokens_details", None) + cached = _int_or_none(getattr(details, "cached_tokens", None)) + if cached is None: + cached = _int_or_none(getattr(usage, "_cache_read_input_tokens", None)) + return cached + + +def _get_cache_write_tokens(usage: "Any") -> "Optional[int]": + details = getattr(usage, "prompt_tokens_details", None) + for attr in ("cache_write_tokens", "cache_creation_tokens"): + written = _int_or_none(getattr(details, attr, None)) + if written is not None: + return written + return _int_or_none(getattr(usage, "_cache_creation_input_tokens", None)) + + +def _get_reasoning_tokens(usage: "Any") -> "Optional[int]": + details = getattr(usage, "completion_tokens_details", None) + return _int_or_none(getattr(details, "reasoning_tokens", None)) + + def _store_span(kwargs: "Dict[str, Any]", span: "Any") -> None: kwargs[_SPAN_KEY] = span @@ -266,10 +295,30 @@ def _success_callback( if input_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) + input_tokens_cached = _get_cached_tokens(usage) + if input_tokens_cached is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, input_tokens_cached + ) + + input_tokens_cache_write = _get_cache_write_tokens(usage) + if input_tokens_cache_write is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, + input_tokens_cache_write, + ) + output_tokens = getattr(usage, "completion_tokens", None) if output_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + output_tokens_reasoning = _get_reasoning_tokens(usage) + if output_tokens_reasoning is not None: + set_on_span( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, + output_tokens_reasoning, + ) + total_tokens = getattr(usage, "total_tokens", None) if ( total_tokens is None diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index e475eaf346..c36bdd9b2d 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -35,6 +35,7 @@ async def __call__(self, *args, **kwargs): from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from openai import AsyncOpenAI, OpenAI from openai.types import CompletionUsage +from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails from sentry_sdk import start_transaction from sentry_sdk._types import BLOB_DATA_SUBSTITUTE @@ -3861,3 +3862,99 @@ def test_embeddings_data_collection( assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-ada-002" assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 + + +def test_cached_and_reasoning_token_usage( + reset_litellm_executor, + sentry_init, + capture_events, + get_model_response, + nonstreaming_chat_completions_model_response, +): + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + client = OpenAI(api_key="test-key") + model_response = get_model_response( + nonstreaming_chat_completions_model_response( + response_id="chatcmpl-test", + response_model="gpt-3.5-turbo", + message_content="Test response", + created=1234567890, + usage=CompletionUsage( + prompt_tokens=1200, + completion_tokens=300, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetails(cached_tokens=1024), + completion_tokens_details=CompletionTokensDetails(reasoning_tokens=256), + ), + ), + serialize_pydantic=True, + request_headers={"X-Stainless-Raw-Response": "true"}, + ) + + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + client=client, + ) + + litellm_utils.executor.shutdown(wait=True) + + (event,) = events + (span,) = ( + x + for x in event["spans"] + if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" + ) + + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 1200 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 1024 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 300 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 256 + assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 1500 + + +def test_cache_write_token_usage_from_anthropic_shaped_usage( + sentry_init, capture_events +): + """LiteLLM maps Anthropic's cache_read/cache_creation counts into prompt_tokens_details.""" + from litellm.types.utils import Usage + + sentry_init( + integrations=[LiteLLMIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + response = mock.Mock(spec=["usage", "model"]) + response.model = "claude-sonnet-4-5" + response.usage = Usage( + prompt_tokens=2600, + completion_tokens=120, + total_tokens=2720, + cache_read_input_tokens=2048, + cache_creation_input_tokens=512, + ) + kwargs = {"model": "claude-sonnet-4-5", "messages": [], "litellm_params": {}} + + with start_transaction(name="litellm test"): + _input_callback(kwargs) + _success_callback(kwargs, response, datetime.now(), datetime.now()) + + (event,) = events + (span,) = event["spans"] + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 2600 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 2048 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 512 From 10dbd45197013cd8ca12a326c8229b808c672352 Mon Sep 17 00:00:00 2001 From: 4ktluffy Date: Tue, 6 Oct 2026 13:50:28 +0300 Subject: [PATCH 2/2] test(litellm): Cover streaming, async and usage without details --- tests/integrations/litellm/test_litellm.py | 187 +++++++++++++++++++++ 1 file changed, 187 insertions(+) diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index c36bdd9b2d..420381e17b 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -3958,3 +3958,190 @@ def test_cache_write_token_usage_from_anthropic_shaped_usage( assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 2600 assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 2048 assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 512 + + +def test_streaming_cached_and_reasoning_token_usage( + reset_litellm_executor, + sentry_init, + capture_events, + get_model_response, + server_side_event_chunks, + streaming_chat_completions_model_response, +): + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + chunks = streaming_chat_completions_model_response( + usage=CompletionUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + message_contents=("Tes", "t r", "esp", "ons", "e"), + ) + + # With stream_options={"include_usage": True}, OpenAI sends usage in a final + # chunk that has no choices. + usage_chunk = chunks[-1].model_copy( + update={ + "choices": [], + "usage": CompletionUsage( + prompt_tokens=1200, + completion_tokens=300, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetails( + cached_tokens=1024, cache_write_tokens=128 + ), + completion_tokens_details=CompletionTokensDetails(reasoning_tokens=256), + ), + } + ) + + client = OpenAI(api_key="test-key") + model_response = get_model_response( + server_side_event_chunks( + [*chunks, usage_chunk], + include_event_type=False, + ), + request_headers={"X-Stainless-Raw-Response": "true"}, + ) + + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + client=client, + stream=True, + stream_options={"include_usage": True}, + ) + for _ in response: + pass + + streaming_handler.executor.shutdown(wait=True) + + (event,) = events + (span,) = ( + x + for x in event["spans"] + if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" + ) + + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 1200 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 1024 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 128 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 300 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 256 + assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 1500 + + +@pytest.mark.asyncio(loop_scope="session") +async def test_async_cached_and_reasoning_token_usage( + sentry_init, + capture_events, + get_model_response, + nonstreaming_chat_completions_model_response, +): + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + client = AsyncOpenAI(api_key="test-key") + model_response = get_model_response( + nonstreaming_chat_completions_model_response( + response_id="chatcmpl-test", + response_model="gpt-3.5-turbo", + message_content="Test response", + created=1234567890, + usage=CompletionUsage( + prompt_tokens=1200, + completion_tokens=300, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetails(cached_tokens=1024), + completion_tokens_details=CompletionTokensDetails(reasoning_tokens=256), + ), + ), + serialize_pydantic=True, + request_headers={"X-Stainless-Raw-Response": "true"}, + ) + + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello!"}], + client=client, + ) + + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) + + (event,) = events + (span,) = ( + x + for x in event["spans"] + if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" + ) + + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 1200 + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 1024 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 300 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 256 + assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 1500 + + +@pytest.mark.parametrize( + "usage", + [ + MockUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + CompletionUsage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + prompt_tokens_details=None, + completion_tokens_details=None, + ), + CompletionUsage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + prompt_tokens_details=PromptTokensDetails(cached_tokens=None), + completion_tokens_details=CompletionTokensDetails(reasoning_tokens=None), + ), + ], +) +def test_token_usage_without_details(sentry_init, capture_events, usage): + sentry_init( + integrations=[LiteLLMIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + response = MockCompletionResponse(usage=usage) + kwargs = {"model": "gpt-3.5-turbo", "messages": [], "litellm_params": {}} + + with start_transaction(name="litellm test"): + _input_callback(kwargs) + _success_callback(kwargs, response, datetime.now(), datetime.now()) + + (event,) = events + (span,) = event["spans"] + assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED not in span["data"] + assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE not in span["data"] + assert SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING not in span["data"]