fix(litellm): Record cached, cache-write and reasoning tokens

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

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,286 @@ 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
+
+
+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"]