diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index 8ca0d1604b14..a16fa83fe326 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -58,9 +58,12 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str) -> choice = completion.choices[0] assert len(choice.text) >= 5 assert choice.finish_reason == "length" - assert completion.usage == openai.types.CompletionUsage( - completion_tokens=5, prompt_tokens=6, total_tokens=11 - ) + assert completion.usage is not None + assert completion.usage.completion_tokens == 5 + assert completion.usage.prompt_tokens == 6 + assert completion.usage.total_tokens == 11 + assert completion.usage.prompt_tokens_details is not None + assert completion.usage.prompt_tokens_details.cached_tokens == 0 # test using token IDs completion = await client.completions.create( diff --git a/tests/entrypoints/serve/disagg/test_generate_stream.py b/tests/entrypoints/serve/disagg/test_generate_stream.py index ac5b8bcd9158..bd52863342d7 100644 --- a/tests/entrypoints/serve/disagg/test_generate_stream.py +++ b/tests/entrypoints/serve/disagg/test_generate_stream.py @@ -512,3 +512,46 @@ async def mock_generate(*args, **kwargs): usage_chunk = parsed[-2] assert usage_chunk["choices"] == [] assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2 + + +@pytest.mark.asyncio +async def test_stream_prompt_tokens_details_zero_cached(): + """enable_prompt_tokens_details includes cached_tokens=0 in final usage. + + Regression test for https://github.com/vllm-project/vllm/issues/44377: + zero cached tokens must not be treated as falsy and omitted. + """ + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + num_cached_tokens=0, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine, enable_prompt_tokens_details=True) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions(include_usage=True), + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + parsed = _parse_sse_chunks(chunks) + # Usage-only chunk (before [DONE]) + usage_chunk = parsed[-2] + assert usage_chunk["choices"] == [] + # Zero cached tokens must be present, not omitted + assert usage_chunk["usage"]["prompt_tokens_details"] is not None + assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 0 diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 9dd9a34162e8..ee8ad51889db 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -881,7 +881,7 @@ async def chat_completion_stream_generator( completion_tokens=completion_tokens, total_tokens=num_prompt_tokens + completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) @@ -1305,7 +1305,10 @@ async def chat_completion_full_generator( completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) - if self.enable_prompt_tokens_details and final_res.num_cached_tokens: + if ( + self.enable_prompt_tokens_details + and final_res.num_cached_tokens is not None + ): usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=final_res.num_cached_tokens ) diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index ed85323d8067..bd7e26b2b164 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -443,7 +443,7 @@ async def completion_stream_generator( total_tokens=total_prompt_tokens + total_completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage_info.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) @@ -583,7 +583,7 @@ def request_output_to_completion_response( if ( self.enable_prompt_tokens_details and last_final_res - and last_final_res.num_cached_tokens + and last_final_res.num_cached_tokens is not None ): usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=last_final_res.num_cached_tokens diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/serve/disagg/serving.py index 72aeb8437739..0bb29c68d01e 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/serve/disagg/serving.py @@ -307,7 +307,10 @@ async def serve_tokens_full_generator( completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) - if self.enable_prompt_tokens_details and final_res.num_cached_tokens: + if ( + self.enable_prompt_tokens_details + and final_res.num_cached_tokens is not None + ): # This info is not available at the /coordinator level usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=final_res.num_cached_tokens @@ -424,7 +427,7 @@ async def serve_tokens_stream_generator( total_tokens=num_prompt_tokens + total_completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage_info.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens )