From 5730af0a95299387c233e12dfbdbe728aace0e58 Mon Sep 17 00:00:00 2001 From: Ian Eaves Date: Thu, 20 Aug 2026 15:05:05 -0500 Subject: [PATCH] [Bugfix] Preserve prompt token usage details (#5181) Signed-off-by: Ian Eaves --- .buildkite/cuda/test-merge.yml | 2 + .buildkite/cuda/test-ready.yml | 2 + tests/e2e/online_serving/test_qwen3_omni.py | 18 +- .../openai_api/test_serving_chat_metrics.py | 68 ++++++ .../test_serving_chat_sampling_params.py | 160 ++++++------ .../entrypoints/test_stream_finish_reason.py | 228 +++++------------- tests/helpers/runtime.py | 3 + tests/helpers/serving_chat.py | 160 ++++++++++++ vllm_omni/entrypoints/openai/serving_chat.py | 34 ++- 9 files changed, 418 insertions(+), 257 deletions(-) create mode 100644 tests/helpers/serving_chat.py diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 69913bf55..6a402254b 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -119,6 +119,8 @@ steps: source_file_dependencies: - tests/e2e/offline_inference/test_qwen3_omni.py - tests/e2e/online_serving/test_qwen3_omni.py + - tests/helpers/runtime.py + - vllm_omni/entrypoints/openai/serving_chat.py - vllm_omni/model_executor/models/qwen3_omni/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py - vllm_omni/model_executor/models/common/snake_activation.py diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 7f7e44cd2..ec3d7485a 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -168,6 +168,8 @@ steps: - label: "Omni ยท Qwen3-Omni Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen3_omni.py + - tests/helpers/runtime.py + - vllm_omni/entrypoints/openai/serving_chat.py - vllm_omni/model_executor/models/qwen3_omni/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py - vllm_omni/model_executor/models/common/snake_activation.py diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index a54d2fe8a..036740ead 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -149,10 +149,19 @@ def _run_prefix_cache_check(openai_client, request_config: dict): 2. The number of cached tokens is divisible by the block size used in test_params, because currently upstream vLLM does not cache partial blocks. + 3. The first request also reports prompt token details; details must be + present whenever --enable-prompt-tokens-details is on, even when zero + tokens are cached. (An exact cached_tokens == 0 check would be + order-dependent here: the module-scoped server may have already cached + the shared system-prompt blocks from earlier tests.) + + Returns the (first, second) responses for additional per-test assertions. """ - openai_client.send_omni_request(request_config, request_num=1)[0] + first_response = openai_client.send_omni_request(request_config, request_num=1)[0] cached_response = openai_client.send_omni_request(request_config, request_num=1)[0] + assert first_response.cached_tokens is not None + # Ensure that we have a prefix cache hit on the second request and that only the last # partial block is uncached (since currently we don't cache partial blocks). num_cached_tokens = cached_response.cached_tokens @@ -163,6 +172,8 @@ def _run_prefix_cache_check(openai_client, request_config: dict): assert num_cached_tokens % BLOCK_SIZE == 0 assert (num_cached_tokens + num_uncached_tokens) == num_prompt_tokens + return first_response, cached_response + @pytest.mark.advanced_model @pytest.mark.core_model @@ -195,7 +206,10 @@ def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None: "stream": False, "modalities": ["text"], } - _run_prefix_cache_check(openai_client, request_config) + first_response, _ = _run_prefix_cache_check(openai_client, request_config) + + assert first_response.multimodal_tokens is not None + assert first_response.multimodal_tokens.get("image", 0) > 0 @pytest.mark.advanced_model diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py index 0647c40a3..2fdf9f646 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_chat_metrics.py @@ -2,9 +2,18 @@ """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field.""" from types import SimpleNamespace +from unittest.mock import MagicMock import pytest +from tests.helpers.serving_chat import ( + build_serving_chat, + collect_stream, + make_request, + make_text_omni_output, + parse_sse_chunks, +) + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -54,6 +63,65 @@ def test_omni_chat_completion_stream_response_metrics(): assert response.metrics == {"stage_latency": 0.5} +@pytest.mark.asyncio +async def test_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens(): + serving_chat = build_serving_chat() + serving_chat.enable_prompt_tokens_details = True + request = make_request(modalities=["text"], include_usage=True) + + async def result_generator(): + yield make_text_omni_output( + token_ids=[10], + finish_reason="stop", + num_cached_tokens=0, + ) + + raw_lines = await collect_stream( + serving_chat.chat_completion_stream_generator( + request=request, + result_generator=result_generator(), + request_id="test-req", + model_name="test-model", + conversation=[], + tokenizer=MagicMock(), + request_metadata=MagicMock(), + mm_token_counts={"image": 7}, + ) + ) + + chunks = parse_sse_chunks(raw_lines) + usage = next(chunk["usage"] for chunk in chunks if chunk.get("usage")) + + assert usage["prompt_tokens_details"] == { + "cached_tokens": 0, + "multimodal_tokens": {"image": 7}, + } + + +def test_non_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens(): + serving_chat = build_serving_chat() + serving_chat.enable_prompt_tokens_details = True + request = make_request(modalities=["text"], stream=False) + omni_output = make_text_omni_output( + token_ids=[10], + finish_reason="stop", + num_cached_tokens=0, + ) + + _, usage, _, _, _ = serving_chat._create_text_choice( + request=request, + omni_outputs=omni_output, + tokenizer=MagicMock(), + conversation=[], + role="assistant", + mm_token_counts={"audio": 5}, + ) + + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 0 + assert usage.prompt_tokens_details.multimodal_tokens == {"audio": 5} + + def test_create_image_choice_exposes_diffusion_metrics(): """Ensure image chat content exposes profiler metrics for clients.""" from PIL import Image diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index e7f014a3b..f81e61cf1 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -14,6 +14,8 @@ from pytest_mock import MockerFixture from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.sampling_params import SamplingParams +from tests.helpers.serving_chat import build_serving_chat + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -77,18 +79,8 @@ def mock_engine_client( return engine_client -@pytest.fixture -def serving_chat(mock_engine_client): - """Create OmniOpenAIServingChat instance with mocked dependencies.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - # Create instance without calling __init__ - instance = object.__new__(OmniOpenAIServingChat) - instance.engine_client = mock_engine_client - return instance - - -def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_chat): +def test_serving_boundary_normalizes_declared_root_and_nested_extras(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset({"cfg_text_scale", "negative_prompt"}) request = ChatCompletionRequest( model="test", @@ -118,7 +110,8 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha assert diffusion_request_args["modalities"] == ["image"] -def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat): +def test_unknown_root_extra_does_not_claim_canonical_extra(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset() request = ChatCompletionRequest( model="test", @@ -135,7 +128,8 @@ def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat): assert "pipeline_option" not in diffusion_request_args -def test_serving_boundary_rejects_invalid_quality(serving_chat): +def test_serving_boundary_rejects_invalid_quality(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset() request = ChatCompletionRequest( model="test", @@ -148,7 +142,8 @@ def test_serving_boundary_rejects_invalid_quality(serving_chat): serving_chat._normalize_diffusion_request_args(request) -def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat): +def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset() request = ChatCompletionRequest( model="test", @@ -170,7 +165,8 @@ def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat): ], ids=["registered-model-extra", "unregistered-common-alias"], ) -def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, request_kwargs): +def test_cfg_scale_owner_conflicts_are_rejected(mock_engine_client, registered, request_kwargs): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset(registered) request = ChatCompletionRequest(model="test", messages=[], **request_kwargs) @@ -180,10 +176,11 @@ def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, reques @pytest.mark.parametrize("diffusion_mode", [True, False], ids=["pure", "mixed"]) def test_duplicate_extras_return_the_same_bad_request_before_dispatch( - serving_chat, + mock_engine_client, mocker: MockerFixture, diffusion_mode: bool, ): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_mode = diffusion_mode serving_chat._diffusion_extra_body_params = frozenset({"sample_solver"}) serving_chat.engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] @@ -211,9 +208,10 @@ def test_duplicate_extras_return_the_same_bad_request_before_dispatch( pure_dispatch.assert_not_awaited() -def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, mocker: MockerFixture): +def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(mock_engine_client, mocker: MockerFixture): from vllm_omni.inputs.data import OmniDiffusionSamplingParams + serving_chat = build_serving_chat(engine_client=mock_engine_client) captured: dict[str, object] = {} async def generate(**kwargs): @@ -261,19 +259,20 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, assert response.error.message == "No output generated from AsyncOmni" -def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, mocker: MockerFixture): +def test_mixed_consumer_keeps_root_common_args_with_nested_extras(mock_engine_client, mocker: MockerFixture): from vllm_omni.inputs.data import OmniDiffusionSamplingParams - serving_chat.engine_client.stage_configs = [ + mock_engine_client.stage_configs = [ SimpleNamespace(stage_type="llm", is_comprehension=True), SimpleNamespace(stage_type="diffusion", is_comprehension=False), ] - serving_chat.engine_client.default_sampling_params_list = [ + mock_engine_client.default_sampling_params_list = [ SamplingParams(), OmniDiffusionSamplingParams(), ] - serving_chat.engine_client.output_modalities = ["image"] - serving_chat.engine_client.errored = False + mock_engine_client.output_modalities = ["image"] + mock_engine_client.errored = False + mock_engine_client.renderer = SimpleNamespace(get_tokenizer=lambda: object()) captured: dict[str, object] = {} @@ -285,21 +284,15 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, captured.update(kwargs) return results() - serving_chat.engine_client.generate = generate - serving_chat.__dict__.update( - _diffusion_mode=False, - _diffusion_extra_body_params=frozenset(), + mock_engine_client.generate = generate + serving_chat = build_serving_chat( + engine_client=mock_engine_client, models=SimpleNamespace(model_name=lambda _: "test"), - renderer=SimpleNamespace(get_tokenizer=lambda: object()), online_renderer=SimpleNamespace(validate_chat_template=lambda **_: None), - parser_cls=None, - use_harmony=False, - enable_auto_tools=False, - exclude_tools_when_tool_choice_none=False, trust_request_chat_template=True, - chat_template=None, - chat_template_content_format="auto", ) + serving_chat._diffusion_mode = False + serving_chat._diffusion_extra_body_params = frozenset() mocker.patch.multiple( serving_chat, _check_model=mocker.AsyncMock(return_value=None), @@ -386,8 +379,9 @@ def test_openai_sampling_fields_contains_expected_fields(): # ============================================================================= -def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_request): +def test_preserves_yaml_defaults_when_no_request_params(mock_engine_client, mock_request): """Test that YAML defaults are preserved when request has no params.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._build_sampling_params_list_from_request(mock_request) assert len(result) == 2 @@ -400,8 +394,9 @@ def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_reque assert comprehension_params.repetition_penalty == 1.05 # YAML custom param preserved -def test_request_temperature_overrides_yaml_default(serving_chat, mock_request): +def test_request_temperature_overrides_yaml_default(mock_engine_client, mock_request): """Test that request temperature overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.temperature = 0.8 mock_request.model_fields_set = {"temperature"} @@ -413,8 +408,9 @@ def test_request_temperature_overrides_yaml_default(serving_chat, mock_request): assert comprehension_params.top_k == 1 # YAML custom param preserved -def test_request_top_p_overrides_yaml_default(serving_chat, mock_request): +def test_request_top_p_overrides_yaml_default(mock_engine_client, mock_request): """Test that request top_p overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.top_p = 0.95 mock_request.model_fields_set = {"top_p"} @@ -425,8 +421,9 @@ def test_request_top_p_overrides_yaml_default(serving_chat, mock_request): assert comprehension_params.temperature == 0.4 # Preserved from YAML -def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request): +def test_request_max_tokens_overrides_yaml_default(mock_engine_client, mock_request): """Test that request max_tokens overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.max_tokens = 100 mock_request.model_fields_set = {"max_tokens"} @@ -435,15 +432,17 @@ def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request): assert result[0].max_tokens == 100 -def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_request): +def test_max_tokens_uses_yaml_default_when_not_specified(mock_engine_client, mock_request): """Test that max_tokens falls back to YAML default when not in request.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._build_sampling_params_list_from_request(mock_request) assert result[0].max_tokens == 4353 -def test_request_seed_overrides_yaml_default(serving_chat, mock_request): +def test_request_seed_overrides_yaml_default(mock_engine_client, mock_request): """Test that request seed overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.seed = 123 mock_request.model_fields_set = {"seed"} @@ -454,8 +453,9 @@ def test_request_seed_overrides_yaml_default(serving_chat, mock_request): assert comprehension_params.temperature == 0.4 # Preserved from YAML -def test_request_frequency_penalty_overrides(serving_chat, mock_request): +def test_request_frequency_penalty_overrides(mock_engine_client, mock_request): """Test that request frequency_penalty is applied.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.frequency_penalty = 0.5 mock_request.model_fields_set = {"frequency_penalty"} @@ -464,8 +464,9 @@ def test_request_frequency_penalty_overrides(serving_chat, mock_request): assert result[0].frequency_penalty == 0.5 -def test_request_presence_penalty_overrides(serving_chat, mock_request): +def test_request_presence_penalty_overrides(mock_engine_client, mock_request): """Test that request presence_penalty is applied.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.presence_penalty = 0.3 mock_request.model_fields_set = {"presence_penalty"} @@ -474,8 +475,9 @@ def test_request_presence_penalty_overrides(serving_chat, mock_request): assert result[0].presence_penalty == 0.3 -def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request): +def test_non_comprehension_stages_use_cloned_defaults(mock_engine_client, mock_request): """Test that non-comprehension stages always use cloned YAML defaults.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.max_tokens = 50 mock_request.temperature = 0.1 @@ -488,8 +490,9 @@ def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request assert other_params.seed == 42 # YAML default -def test_multiple_params_override_together(serving_chat, mock_request): +def test_multiple_params_override_together(mock_engine_client, mock_request): """Test that multiple request params can override together.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.max_tokens = 200 mock_request.temperature = 0.7 mock_request.top_p = 0.85 @@ -514,15 +517,17 @@ def test_multiple_params_override_together(serving_chat, mock_request): # ============================================================================= -def test_apply_request_overrides_clones_params(serving_chat, mock_request, default_comprehension_params): +def test_apply_request_overrides_clones_params(mock_engine_client, mock_request, default_comprehension_params): """Test that _apply_request_overrides returns a cloned object.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request) assert result is not default_comprehension_params # Different object -def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request, default_comprehension_params): +def test_apply_request_overrides_preserves_defaults(mock_engine_client, mock_request, default_comprehension_params): """Test that _apply_request_overrides preserves defaults when request has None.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request) assert result.temperature == 0.4 @@ -531,8 +536,9 @@ def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request, assert result.top_k == 1 # YAML custom param -def test_apply_request_overrides_applies_values(serving_chat, mock_request, default_comprehension_params): +def test_apply_request_overrides_applies_values(mock_engine_client, mock_request, default_comprehension_params): """Test that _apply_request_overrides applies non-None request values.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.temperature = 0.8 mock_request.seed = 123 mock_request.model_fields_set = {"temperature", "seed"} @@ -550,8 +556,9 @@ def test_apply_request_overrides_applies_values(serving_chat, mock_request, defa # ============================================================================= -def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker): +def test_apply_overrides_empty_stop_list_preserves_default(mock_engine_client, mocker): """Test that request.stop=[] does NOT override YAML default stop words.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"]) request = mocker.MagicMock() request.temperature = None @@ -573,8 +580,9 @@ def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker) assert result.stop == ["<|im_end|>"] # YAML default preserved -def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mocker): +def test_apply_overrides_nonempty_stop_list_overrides_default(mock_engine_client, mocker): """Test that request.stop=["\\n"] overrides YAML default stop words.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"]) request = mocker.MagicMock() request.temperature = None @@ -596,8 +604,9 @@ def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mock assert result.stop == ["\n"] # Overridden by request -def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mocker): +def test_apply_overrides_empty_stop_token_ids_preserves_default(mock_engine_client, mocker): """Test that request.stop_token_ids=[] does NOT override YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3]) request = mocker.MagicMock() request.temperature = None @@ -617,8 +626,9 @@ def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mo assert result.stop_token_ids == [2, 3] # YAML default preserved -def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat, mocker): +def test_apply_overrides_nonempty_stop_token_ids_overrides_default(mock_engine_client, mocker): """Test that request.stop_token_ids=[100] overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3]) request = mocker.MagicMock() request.temperature = None @@ -640,8 +650,9 @@ def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat, assert result.stop_token_ids == [100] # Overridden by request -def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker): +def test_apply_overrides_mixed_empty_and_nonempty_lists(mock_engine_client, mocker): """Test mixing empty and non-empty list fields with scalar fields.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams( temperature=0.4, stop=["<|end|>"], @@ -669,8 +680,9 @@ def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker): assert result.stop_token_ids == [100, 200] # Non-empty list DID override -def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocker): +def test_apply_overrides_none_scalar_still_preserves_default(mock_engine_client, mocker): """Regression: ensure None scalar values still don't override defaults.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, max_tokens=100, seed=42) request = mocker.MagicMock() request.temperature = None @@ -694,8 +706,9 @@ def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocke assert result.seed == 42 -def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocker): +def test_apply_overrides_both_lists_empty_preserves_defaults(mock_engine_client, mocker): """Test that both stop=[] and stop_token_ids=[] preserve YAML defaults.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams( temperature=0.5, stop=["<|end|>", "\\n"], @@ -722,9 +735,10 @@ def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocke assert result.stop_token_ids == [2, 32000] -def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock_request): +def test_build_sampling_params_list_empty_stop_preserves_yaml(mock_engine_client, mock_request): """Test that empty stop list in request preserves YAML defaults via _build_sampling_params_list_from_request.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.stop = [] mock_request.stop_token_ids = [] @@ -736,21 +750,18 @@ def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock assert comprehension_params.stop_token_ids == [] -def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(): +def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(mocker: MockerFixture): """AURA callers may pass 3 semantic model params for a 4-stage engine pipeline.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) default_params = [ SamplingParams(max_tokens=10), SamplingParams(max_tokens=20), SamplingParams(max_tokens=30), SamplingParams(max_tokens=40), ] - instance.engine_client = SimpleNamespace( - stage_configs=[SimpleNamespace(stage_type="llm") for _ in range(4)], - default_sampling_params_list=default_params, - ) + engine_client = mocker.MagicMock() + engine_client.stage_configs = [SimpleNamespace(stage_type="llm") for _ in range(4)] + engine_client.default_sampling_params_list = default_params + instance = build_serving_chat(engine_client=engine_client) result = instance._to_sampling_params_list( [ @@ -772,44 +783,35 @@ def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(): def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client): """Test finding comprehension stage when it's at index 0.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) - instance.engine_client = mock_engine_client + instance = build_serving_chat(engine_client=mock_engine_client) assert instance._get_comprehension_stage_index() == 0 def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture): """Test finding comprehension stage when it's at index 1.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) - other = mocker.MagicMock() other.is_comprehension = False comprehension = mocker.MagicMock() comprehension.is_comprehension = True - instance.engine_client = mocker.MagicMock() - instance.engine_client.stage_configs = [other, comprehension] + engine_client = mocker.MagicMock() + engine_client.stage_configs = [other, comprehension] + instance = build_serving_chat(engine_client=engine_client) assert instance._get_comprehension_stage_index() == 1 def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixture): """Test that ValueError is raised when no comprehension stage exists.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) - stage1 = mocker.MagicMock() stage1.is_comprehension = False stage2 = mocker.MagicMock() stage2.is_comprehension = False - instance.engine_client = mocker.MagicMock() - instance.engine_client.stage_configs = [stage1, stage2] + engine_client = mocker.MagicMock() + engine_client.stage_configs = [stage1, stage2] + instance = build_serving_chat(engine_client=engine_client) with pytest.raises(ValueError, match="No comprehension stage"): instance._get_comprehension_stage_index() diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index b02936838..97ed2370b 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -16,19 +16,20 @@ Key invariants tested: - voice/speaker parameter compatibility in chat completions """ -import json import time from unittest.mock import MagicMock import pytest -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ChatCompletionResponseStreamChoice, -) -from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.outputs import CompletionOutput, RequestOutput +from tests.helpers.serving_chat import ( + build_serving_chat, + collect_stream, + make_request, + make_text_omni_output, + parse_sse_chunks, +) from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.outputs import OmniRequestOutput @@ -40,43 +41,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # --------------------------------------------------------------------------- -def _make_text_omni_output( - request_id: str = "test-req", - text: str = "hello", - token_ids: list[int] | None = None, - finish_reason: str | None = None, - index: int = 0, - num_prompt_tokens: int = 3, -) -> OmniRequestOutput: - """Build an OmniRequestOutput wrapping a text RequestOutput.""" - if token_ids is None: - token_ids = [10, 11, 12] - res = RequestOutput( - request_id=request_id, - prompt="test", - prompt_token_ids=list(range(num_prompt_tokens)), - prompt_logprobs=None, - outputs=[ - CompletionOutput( - index=index, - text=text, - token_ids=token_ids, - cumulative_logprob=0.0, - logprobs=None, - finish_reason=finish_reason, - stop_reason=None, - ) - ], - finished=finish_reason is not None, - ) - return OmniRequestOutput.from_stage_output( - res, - request_id=request_id, - final_output_type="text", - finished=finish_reason is not None, - ) - - def _make_audio_omni_output( request_id: str = "test-req", index: int = 0, @@ -118,82 +82,6 @@ def _make_audio_omni_output( ) -def _mock_audio_choices(index: int = 0, role: str = "assistant"): - return [ - ChatCompletionResponseStreamChoice( - index=index, - delta=DeltaMessage(role=role, content="dGVzdA=="), - logprobs=None, - finish_reason="stop", - ) - ] - - -def _build_serving_chat(): - """Create a minimal OmniOpenAIServingChat for testing.""" - mock_engine = MagicMock() - mock_engine.errored = False - - models = OpenAIServingModels( - engine_client=mock_engine, - base_model_paths=[], - ) - mock_render = MagicMock() - - instance = OmniOpenAIServingChat( - engine_client=mock_engine, - models=models, - response_role="assistant", - online_renderer=mock_render, - request_logger=None, - chat_template=None, - chat_template_content_format="auto", - ) - instance._create_audio_choice = MagicMock( - side_effect=lambda omni_res, role, request, stream=False: _mock_audio_choices( - index=omni_res.outputs[0].index, - role=role, - ) - ) - return instance - - -def _make_request(modalities: list[str], n: int = 1) -> ChatCompletionRequest: - req = ChatCompletionRequest( - model="test-model", - messages=[{"role": "user", "content": "hello"}], - n=n, - stream=True, - ) - req.modalities = modalities # type: ignore[attr-defined] - return req - - -def _parse_sse_chunks(lines: list[str]) -> list[dict]: - """Parse SSE lines into JSON dicts.""" - prefix = "data: " - chunks = [] - for line in lines: - line = line.strip() - if not line.startswith(prefix): - continue - payload = line[len(prefix) :].strip() - if payload == "[DONE]": - continue - try: - chunks.append(json.loads(payload)) - except json.JSONDecodeError: - pass - return chunks - - -async def _collect_stream(gen): - result = [] - async for item in gen: - result.append(item) - return result - - # --------------------------------------------------------------------------- # Tests: finish_reason correctness # --------------------------------------------------------------------------- @@ -202,14 +90,14 @@ async def _collect_stream(gen): @pytest.mark.asyncio async def test_single_modality_text_only_one_stop(): """Text-only streaming: exactly one chunk has finish_reason='stop'.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text"]) async def result_generator(): - yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) - yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") + yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) + yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -221,7 +109,7 @@ async def test_single_modality_text_only_one_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [c["choices"][0]["finish_reason"] for c in chunks if c.get("choices")] assert finish_reasons[-1] == "stop" @@ -233,15 +121,15 @@ async def test_single_modality_text_only_one_stop(): @pytest.mark.asyncio async def test_multi_modal_text_audio_only_last_stop(): """text+audio: text finish sends finish_reason=null, audio sends stop.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) async def result_generator(): - yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) - yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") + yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) + yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") yield _make_audio_omni_output() - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -253,7 +141,7 @@ async def test_multi_modal_text_audio_only_last_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1 @@ -269,18 +157,18 @@ async def test_multi_modal_text_audio_only_last_stop(): @pytest.mark.asyncio async def test_multi_modal_n2_independent_per_choice(): """n=2 with text+audio: each choice gets exactly one stop, at the end.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"], n=2) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"], n=2) async def result_generator(): - yield _make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0) - yield _make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1) - yield _make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0) - yield _make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1) + yield make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0) + yield make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1) + yield make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0) + yield make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1) yield _make_audio_omni_output(index=0) yield _make_audio_omni_output(index=1) - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -292,7 +180,7 @@ async def test_multi_modal_n2_independent_per_choice(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) per_choice: dict[int, list] = {} for c in chunks: for ch in c.get("choices", []): @@ -306,13 +194,13 @@ async def test_multi_modal_n2_independent_per_choice(): @pytest.mark.asyncio async def test_single_modality_audio_only_one_stop(): """Audio-only streaming: the audio chunk carries finish_reason='stop'.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["audio"]) async def result_generator(): yield _make_audio_omni_output() - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -324,7 +212,7 @@ async def test_single_modality_audio_only_one_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1 @@ -349,8 +237,8 @@ async def test_streaming_audio_metrics_resolve_replica_id( import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod from vllm_omni.entrypoints.client_request_state import ClientRequestState - serving_chat = _build_serving_chat() - request = _make_request(modalities=["audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["audio"]) req_state = ClientRequestState( request_id="internal-req", external_request_id="test-req", @@ -377,7 +265,7 @@ async def test_streaming_audio_metrics_resolve_replica_id( async def result_generator(): yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400) - await _collect_stream( + await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -408,15 +296,15 @@ async def test_streaming_audio_metrics_resolve_replica_id( async def test_declared_modality_not_produced_emits_fallback_stop(): """If request.modalities declares ["text","audio"] but engine only produces text, a fallback stop chunk is emitted at stream end.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) async def result_generator(): # Engine only produces text, no audio output at all - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) - yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) + yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -428,7 +316,7 @@ async def test_declared_modality_not_produced_emits_fallback_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] # Text finish is suppressed (audio not seen yet), but fallback stop @@ -441,15 +329,15 @@ async def test_declared_modality_not_produced_emits_fallback_stop(): async def test_declared_modality_not_produced_text_finish_suppressed(): """When text finishes but audio (declared in modalities) never appears, the text finish chunk has finish_reason=null (suppressed).""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) async def result_generator(): - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) - yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) + yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") # No audio output โ€” stream ends - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -461,7 +349,7 @@ async def test_declared_modality_not_produced_text_finish_suppressed(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) # Find the text finish chunk (content "!") for c in chunks: @@ -485,8 +373,8 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): yields (field, value) tuples, so the loop raised ``AttributeError: 'tuple' object has no attribute 'finish_reason'``. """ - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) empty_audio = _make_audio_omni_output() empty_audio.outputs[0].multimodal_output = {"audio": []} @@ -497,11 +385,11 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): serving_chat._create_audio_choice = create_audio_choice async def result_generator(): - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) - yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) + yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") yield empty_audio - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -514,7 +402,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): ) assert not any("Error in chat completion stream generator" in line for line in raw_lines) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}" assert finish_reasons[-1] == "stop" @@ -525,18 +413,18 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices(): """Any ErrorResponse from the audio path is skipped, not iterated.""" from vllm.entrypoints.openai.engine.protocol import ErrorResponse - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) serving_chat._create_audio_choice = MagicMock( side_effect=lambda omni_res, role, request, stream=False: serving_chat._create_error_response("boom") ) async def result_generator(): - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop") yield _make_audio_omni_output() - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -550,7 +438,7 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices(): assert isinstance(serving_chat._create_error_response("boom"), ErrorResponse) assert not any("AttributeError" in line for line in raw_lines) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}" diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 56f7aaccc..51df9ef84 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -770,6 +770,7 @@ class OmniResponse: success: bool = False prompt_tokens: int | None = None cached_tokens: int | None = None + multimodal_tokens: dict[str, int] | None = None logprobs: list | None = None #: HTTP status + error text for the error-handling path (e.g. validator #: rejections); populated when the OpenAI client raises an APIError. @@ -1089,6 +1090,7 @@ class OpenAIClientHandler: result.prompt_tokens = chunk.usage.prompt_tokens if details := getattr(chunk.usage, "prompt_tokens_details", None): result.cached_tokens = details.cached_tokens + result.multimodal_tokens = getattr(details, "multimodal_tokens", None) if audio_data: merged_seg = _merge_base64_audio_to_segment(audio_data) @@ -1121,6 +1123,7 @@ class OpenAIClientHandler: result.prompt_tokens = usage.prompt_tokens if details := getattr(usage, "prompt_tokens_details", None): result.cached_tokens = details.cached_tokens + result.multimodal_tokens = getattr(details, "multimodal_tokens", None) if audio_data: result.audio_bytes = base64.b64decode(audio_data) result.text_content = text_content diff --git a/tests/helpers/serving_chat.py b/tests/helpers/serving_chat.py new file mode 100644 index 000000000..f1f2710b5 --- /dev/null +++ b/tests/helpers/serving_chat.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared scaffolding for OmniOpenAIServingChat unit tests. + +Provides factories for a minimal serving-chat instance, chat requests, and +OmniRequestOutput objects, plus SSE stream collection/parsing utilities. +""" + +import json +from unittest.mock import MagicMock + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponseStreamChoice, +) +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.outputs import CompletionOutput, RequestOutput + +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.outputs import OmniRequestOutput + + +def make_text_omni_output( + request_id: str = "test-req", + text: str = "hello", + token_ids: list[int] | None = None, + finish_reason: str | None = None, + index: int = 0, + num_prompt_tokens: int = 3, + num_cached_tokens: int | None = None, +) -> OmniRequestOutput: + """Build an OmniRequestOutput wrapping a text RequestOutput.""" + if token_ids is None: + token_ids = [10, 11, 12] + res = RequestOutput( + request_id=request_id, + prompt="test", + prompt_token_ids=list(range(num_prompt_tokens)), + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=index, + text=text, + token_ids=token_ids, + cumulative_logprob=0.0, + logprobs=None, + finish_reason=finish_reason, + stop_reason=None, + ) + ], + finished=finish_reason is not None, + num_cached_tokens=num_cached_tokens, + ) + return OmniRequestOutput.from_stage_output( + res, + request_id=request_id, + final_output_type="text", + finished=finish_reason is not None, + ) + + +def mock_audio_choices(index: int = 0, role: str = "assistant"): + return [ + ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage(role=role, content="dGVzdA=="), + logprobs=None, + finish_reason="stop", + ) + ] + + +def build_serving_chat( + *, + engine_client=None, + models=None, + online_renderer=None, + response_role: str = "assistant", + request_logger=None, + chat_template=None, + chat_template_content_format="auto", + **init_overrides, +) -> OmniOpenAIServingChat: + """Create a real OmniOpenAIServingChat via its actual __init__. + + engine_client/models/online_renderer default to minimal mocks; pass your + own to wire in test-specific behavior (e.g. engine_client.stage_configs). + Any other OmniOpenAIServingChat.__init__ keyword โ€” enable_auto_tools, + enable_prompt_tokens_details, trust_request_chat_template, etc. โ€” can be + passed through **init_overrides instead of set on the instance after + construction. + """ + if engine_client is None: + engine_client = MagicMock() + engine_client.errored = False + if models is None: + models = OpenAIServingModels( + engine_client=engine_client, + base_model_paths=[], + ) + if online_renderer is None: + online_renderer = MagicMock() + + instance = OmniOpenAIServingChat( + engine_client=engine_client, + models=models, + response_role=response_role, + online_renderer=online_renderer, + request_logger=request_logger, + chat_template=chat_template, + chat_template_content_format=chat_template_content_format, + **init_overrides, + ) + instance._create_audio_choice = MagicMock( + side_effect=lambda omni_res, role, request, stream=False: mock_audio_choices( + index=omni_res.outputs[0].index, + role=role, + ) + ) + return instance + + +def make_request( + modalities: list[str], + n: int = 1, + *, + stream: bool = True, + include_usage: bool = False, +) -> ChatCompletionRequest: + req = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + n=n, + stream=stream, + stream_options={"include_usage": True} if include_usage else None, + ) + req.modalities = modalities # type: ignore[attr-defined] + return req + + +def parse_sse_chunks(lines: list[str]) -> list[dict]: + """Parse SSE lines into JSON dicts.""" + prefix = "data: " + chunks = [] + for line in lines: + line = line.strip() + if not line.startswith(prefix): + continue + payload = line[len(prefix) :].strip() + if payload == "[DONE]": + continue + chunks.append(json.loads(payload)) + return chunks + + +async def collect_stream(gen): + result = [] + async for item in gen: + result.append(item) + return result diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 752d93577..2bcbecc33 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -58,7 +58,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionResponseStreamChoice, ChatMessage, ) -from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.chat_completion.serving import ( + OpenAIServingChat, + _get_mm_token_counts, + _make_prompt_tokens_details, +) from vllm.entrypoints.openai.engine.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -67,7 +71,6 @@ from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, FunctionCall, FunctionDefinition, - PromptTokenUsageInfo, RequestResponseMetadata, ToolCall, UsageInfo, @@ -730,8 +733,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Schedule the request and get the result generator. generators: list[AsyncGenerator[RequestOutput, None]] = [] + mm_token_counts: dict[str, int] | None = None try: for i, engine_prompt in enumerate(engine_prompts): + if self.enable_prompt_tokens_details: + mm_token_counts = _get_mm_token_counts(engine_prompt) if hasattr(request, "sampling_params_list"): sampling_params_list = self._to_sampling_params_list(request.sampling_params_list) else: @@ -814,6 +820,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer, request_metadata, reasoning_parser, + mm_token_counts=mm_token_counts, raw_request=raw_request, ) @@ -827,6 +834,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer, request_metadata, reasoning_parser, + mm_token_counts=mm_token_counts, ) except ValueError as e: return self.create_error_response(e) @@ -1387,6 +1395,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer: AnyTokenizer, request_metadata: RequestResponseMetadata, reasoning_parser: ReasoningParser | None = None, + mm_token_counts: dict[str, int] | None = None, raw_request: Request | None = None, ): created_time = int(time.time()) @@ -1407,6 +1416,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): stop_reason_emitted: list[bool] = [False] * num_choices num_prompt_tokens = 0 num_cached_tokens = None + num_cache_creation_tokens = None if self.use_harmony: harmony_parsers = [get_streamable_parser_for_assistant() for _ in range(num_choices)] harmony_tools_streamed = [False] * num_choices @@ -1503,6 +1513,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): num_prompt_tokens += len(res.encoder_prompt_token_ids) num_cached_tokens = res.num_cached_tokens + num_cache_creation_tokens = res.num_cache_creation_tokens # Send first response for each choice with role # NOTE: num_choices defaults to 1 so this usually executes once per request for i in range(num_choices): @@ -2234,8 +2245,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): completion_tokens=completion_tokens, total_tokens=num_prompt_tokens + completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: - final_usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=num_cached_tokens) + final_usage.prompt_tokens_details = _make_prompt_tokens_details( + self.enable_prompt_tokens_details, + num_cached_tokens, + num_cache_creation_tokens, + mm_token_counts, + ) final_usage_chunk = OmniChatCompletionStreamResponse( id=request_id, @@ -2307,6 +2322,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, reasoning_parser: ReasoningParser | None = None, + mm_token_counts: dict[str, int] | None = None, ) -> ErrorResponse | OmniChatCompletionResponse: created_time = int(time.time()) final_res: RequestOutput | None = None @@ -2361,6 +2377,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation, role, reasoning_parser, + mm_token_counts=mm_token_counts, ) final_res = omni_outputs else: @@ -2462,6 +2479,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation: list[ConversationMessage], role: str, reasoning_parser: ReasoningParser | None = None, + mm_token_counts: dict[str, int] | None = None, ): final_res = omni_outputs if self.tool_call_id_type == "kimi_k2": @@ -2698,8 +2716,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): 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: - usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens) + usage.prompt_tokens_details = _make_prompt_tokens_details( + self.enable_prompt_tokens_details, + final_res.num_cached_tokens, + final_res.num_cache_creation_tokens, + mm_token_counts, + ) prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs) prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None