@@ -35,6 +35,7 @@ async def __call__(self, *args, **kwargs):
3535from litellm .llms .custom_httpx .http_handler import AsyncHTTPHandler , HTTPHandler
3636from openai import AsyncOpenAI , OpenAI
3737from openai .types import CompletionUsage
38+ from openai .types .completion_usage import CompletionTokensDetails , PromptTokensDetails
3839
3940from sentry_sdk import start_transaction
4041from 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