Refactor tests to match project style guidelines

- Removed class-based test structure
- Moved imports to top of file
- Added section dividers (============) between test groups
- Removed inline imports
- Followed same structure as test_serving_chat_sampling_params.py
- All tests are now standalone functions with proper docstrings

This aligns with the test style guide at docs/contributing/ci/tests_style.md

Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>

tests/entrypoints/openai_api/test_serving_chat_modalities.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Unit tests for OmniOpenAIServingChat modalities filtering.
+
+Tests that modalities parameter correctly filters outputs in non-streaming
+chat completions while preserving usage/billing information.
+"""
+
+from unittest.mock import MagicMock
+
+# =============================================================================
+# Tests for modalities set creation logic
+# =============================================================================
+
+
+def test_requested_modalities_set_creation_with_list():
+    """Test that requested_modalities set is created from list."""
+    # Simulate what happens in the code
+    request = MagicMock()
+    request.modalities = ["text", "audio"]
+
+    # This is the logic from the fix
+    requested_modalities = set(request.modalities) if hasattr(request, "modalities") and request.modalities else None
+
+    assert requested_modalities == {"text", "audio"}
+
+
+def test_requested_modalities_set_creation_with_none():
+    """Test that requested_modalities is None when modalities is None."""
+    request = MagicMock()
+    request.modalities = None
+
+    requested_modalities = set(request.modalities) if hasattr(request, "modalities") and request.modalities else None
+
+    assert requested_modalities is None
+
+
+def test_requested_modalities_set_creation_when_missing():
+    """Test that requested_modalities is None when attribute missing."""
+    request = MagicMock(spec=[])  # No modalities attribute
+
+    requested_modalities = set(request.modalities) if hasattr(request, "modalities") and request.modalities else None
+
+    assert requested_modalities is None
+
+
+# =============================================================================
+# Tests for modality filtering logic
+# =============================================================================
+
+
+def test_modality_filtering_text_in_set():
+    """Test that text output passes filter when text is requested."""
+    requested_modalities = {"text", "audio"}
+    final_output_type = "text"
+
+    # This is the filtering logic from the fix
+    should_filter = requested_modalities is not None and final_output_type not in requested_modalities
+
+    assert not should_filter  # text should NOT be filtered
+
+
+def test_modality_filtering_audio_not_in_set():
+    """Test that audio output is filtered when only text requested."""
+    requested_modalities = {"text"}
+    final_output_type = "audio"
+
+    should_filter = requested_modalities is not None and final_output_type not in requested_modalities
+
+    assert should_filter  # audio SHOULD be filtered
+
+
+def test_modality_filtering_with_none_modalities():
+    """Test that nothing is filtered when modalities is None."""
+    requested_modalities = None
+    final_output_type = "text"
+
+    should_filter = requested_modalities is not None and final_output_type not in requested_modalities
+
+    assert not should_filter  # Nothing should be filtered
+
+    final_output_type = "audio"
+    should_filter = requested_modalities is not None and final_output_type not in requested_modalities
+
+    assert not should_filter  # Nothing should be filtered
+
+
+# =============================================================================
+# Tests for usage extraction logic
+# =============================================================================
+
+
+def test_usage_extraction_logic():
+    """Test that usage info is extracted correctly from request_output."""
+    # Mock the omni_outputs structure
+    omni_outputs = MagicMock()
+    omni_outputs.final_output_type = "text"
+    omni_outputs.request_output = MagicMock()
+    omni_outputs.request_output.prompt_token_ids = [1] * 10
+    omni_outputs.request_output.encoder_prompt_token_ids = None
+
+    mock_output = MagicMock()
+    mock_output.token_ids = [1] * 20
+    omni_outputs.request_output.outputs = [mock_output]
+
+    # Simulate the usage extraction logic
+    final_output_type = omni_outputs.final_output_type
+    if final_output_type == "text" and omni_outputs.request_output is not None:
+        final_res = omni_outputs.request_output
+        if final_res.prompt_token_ids is not None:
+            num_prompt_tokens = len(final_res.prompt_token_ids)
+            if final_res.encoder_prompt_token_ids is not None:
+                num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
+            num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs)
+
+            assert num_prompt_tokens == 10
+            assert num_generated_tokens == 20

vllm_omni/entrypoints/openai/serving_chat.py

@@ -1297,12 +1297,44 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         prompt_token_ids = None
         kv_transfer_params = None
 
+        # Build requested modalities set for filtering
+        requested_modalities = (
+            set(request.modalities) if hasattr(request, "modalities") and request.modalities else None
+        )
+
         for omni_outputs in final_outputs:
             choices_data = []
             if omni_outputs.request_output is not None and not getattr(omni_outputs.request_output, "finished", False):
                 continue
 
-            if omni_outputs.final_output_type == "text":
+            final_output_type = omni_outputs.final_output_type
+
+            # Extract usage info from text output regardless of modality filtering
+            # This ensures usage/billing data is captured even for audio-only requests
+            if final_output_type == "text" and omni_outputs.request_output is not None:
+                final_res = omni_outputs.request_output
+                if final_res.prompt_token_ids is not None:
+                    num_prompt_tokens = len(final_res.prompt_token_ids)
+                    if final_res.encoder_prompt_token_ids is not None:
+                        num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
+                    num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs)
+                    usage = UsageInfo(
+                        prompt_tokens=num_prompt_tokens,
+                        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)
+                    prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs)
+                    prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None
+                    kv_transfer_params = final_res.kv_transfer_params
+
+            # Filter outputs based on requested modalities
+            if requested_modalities is not None and final_output_type not in requested_modalities:
+                logger.warning(f"final output type: {final_output_type} is not needed by the request")
+                continue
+
+            if final_output_type == "text":
                 (
                     choices_data,
                     usage,
@@ -1310,12 +1342,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     prompt_token_ids,
                     kv_transfer_params,
                 ) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role)
-            elif omni_outputs.final_output_type == "audio":
+            elif final_output_type == "audio":
                 choices_data = self._create_audio_choice(omni_outputs, role, request, stream=False)
-            elif omni_outputs.final_output_type == "image":
+            elif final_output_type == "image":
                 choices_data = self._create_image_choice(omni_outputs, role, request, stream=False)
             else:
-                logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}")
+                logger.warning(f"Unsupported final output type: {final_output_type}")
                 continue
             choices.extend(choices_data)