copilot-swe-agent[bot]andhsliuustc0106
f227f28d19
Added comprehensive unit tests that verify: - Text-only modality filters out audio outputs - Audio-only modality filters out text but preserves usage info - Multiple modalities includes all requested outputs - No modalities parameter includes all outputs These tests use mocks to run quickly without requiring full model loading. Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>
261 lines
9.2 KiB
Python
261 lines
9.2 KiB
Python
# 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
|
|
|
|
import pytest
|
|
from vllm.entrypoints.openai.protocol import UsageInfo
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_request():
|
|
"""Create a mock chat completion request."""
|
|
request = MagicMock()
|
|
request.return_token_ids = False
|
|
return request
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_omni_output_text():
|
|
"""Create a mock text output."""
|
|
output = MagicMock()
|
|
output.final_output_type = "text"
|
|
output.request_output = MagicMock()
|
|
output.request_output.finished = True
|
|
output.request_output.prompt_token_ids = [1] * 10 # 10 prompt tokens
|
|
output.request_output.encoder_prompt_token_ids = None
|
|
output.request_output.num_cached_tokens = 0
|
|
output.request_output.prompt_logprobs = None
|
|
output.request_output.kv_transfer_params = None
|
|
|
|
# Mock a single output with token_ids
|
|
mock_output = MagicMock()
|
|
mock_output.token_ids = [1] * 20 # 20 completion tokens
|
|
output.request_output.outputs = [mock_output]
|
|
|
|
return output
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_omni_output_audio():
|
|
"""Create a mock audio output."""
|
|
output = MagicMock()
|
|
output.final_output_type = "audio"
|
|
output.request_output = MagicMock()
|
|
output.request_output.finished = True
|
|
output.request_output.multimodal_output = {"audio": MagicMock()}
|
|
|
|
# Mock outputs for audio choice
|
|
mock_output = MagicMock()
|
|
mock_output.index = 0
|
|
mock_output.token_ids = []
|
|
mock_output.stop_reason = None
|
|
output.request_output.outputs = [mock_output]
|
|
|
|
return output
|
|
|
|
|
|
@pytest.fixture
|
|
def serving_chat():
|
|
"""Create OmniOpenAIServingChat instance with mocked dependencies."""
|
|
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
|
|
|
instance = object.__new__(OmniOpenAIServingChat)
|
|
instance.enable_prompt_tokens_details = False
|
|
return instance
|
|
|
|
|
|
class TestModalitiesFiltering:
|
|
"""Tests for modalities filtering in chat_completion_full_generator."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_text_only_modality_filters_audio(
|
|
self, serving_chat, mock_request, mock_omni_output_text, mock_omni_output_audio
|
|
):
|
|
"""Test that modalities=["text"] filters out audio outputs."""
|
|
mock_request.modalities = ["text"]
|
|
|
|
# Create a mock result generator that yields both text and audio
|
|
async def mock_generator():
|
|
yield mock_omni_output_text
|
|
yield mock_omni_output_audio
|
|
|
|
# Mock the methods we need
|
|
serving_chat.get_chat_request_role = MagicMock(return_value="assistant")
|
|
serving_chat._create_text_choice = MagicMock(
|
|
return_value=(
|
|
[MagicMock()], # choices_data
|
|
UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
|
|
None, # prompt_logprobs
|
|
None, # prompt_token_ids
|
|
None, # kv_transfer_params
|
|
)
|
|
)
|
|
serving_chat._create_audio_choice = MagicMock(return_value=[MagicMock()])
|
|
|
|
# Call the method
|
|
result = await serving_chat.chat_completion_full_generator(
|
|
request=mock_request,
|
|
result_generator=mock_generator(),
|
|
request_id="test-123",
|
|
model_name="test-model",
|
|
conversation=[],
|
|
tokenizer=MagicMock(),
|
|
request_metadata=MagicMock(),
|
|
)
|
|
|
|
# Verify text choice was created
|
|
assert serving_chat._create_text_choice.called
|
|
# Verify audio choice was NOT created (filtered out)
|
|
assert not serving_chat._create_audio_choice.called
|
|
|
|
# Verify response has only 1 choice
|
|
assert len(result.choices) == 1
|
|
# Verify usage is populated
|
|
assert result.usage.prompt_tokens == 10
|
|
assert result.usage.completion_tokens == 20
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_audio_only_modality_filters_text(
|
|
self, serving_chat, mock_request, mock_omni_output_text, mock_omni_output_audio
|
|
):
|
|
"""Test that modalities=["audio"] filters out text but preserves usage."""
|
|
mock_request.modalities = ["audio"]
|
|
|
|
# Create a mock result generator that yields both text and audio
|
|
async def mock_generator():
|
|
yield mock_omni_output_text
|
|
yield mock_omni_output_audio
|
|
|
|
# Mock the methods we need
|
|
serving_chat.get_chat_request_role = MagicMock(return_value="assistant")
|
|
serving_chat._create_text_choice = MagicMock(
|
|
return_value=(
|
|
[MagicMock()], # choices_data
|
|
UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
|
|
None, # prompt_logprobs
|
|
None, # prompt_token_ids
|
|
None, # kv_transfer_params
|
|
)
|
|
)
|
|
serving_chat._create_audio_choice = MagicMock(return_value=[MagicMock()])
|
|
|
|
# Call the method
|
|
result = await serving_chat.chat_completion_full_generator(
|
|
request=mock_request,
|
|
result_generator=mock_generator(),
|
|
request_id="test-123",
|
|
model_name="test-model",
|
|
conversation=[],
|
|
tokenizer=MagicMock(),
|
|
request_metadata=MagicMock(),
|
|
)
|
|
|
|
# Verify text choice was NOT created (filtered out)
|
|
assert not serving_chat._create_text_choice.called
|
|
# Verify audio choice was created
|
|
assert serving_chat._create_audio_choice.called
|
|
|
|
# Verify response has only 1 choice (audio)
|
|
assert len(result.choices) == 1
|
|
|
|
# CRITICAL: Verify usage is still populated even though text was filtered
|
|
# This is the regression that was fixed
|
|
assert result.usage.prompt_tokens == 10
|
|
assert result.usage.completion_tokens == 20
|
|
assert result.usage.total_tokens == 30
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_modalities_includes_all(
|
|
self, serving_chat, mock_request, mock_omni_output_text, mock_omni_output_audio
|
|
):
|
|
"""Test that modalities=["text", "audio"] includes both outputs."""
|
|
mock_request.modalities = ["text", "audio"]
|
|
|
|
# Create a mock result generator that yields both text and audio
|
|
async def mock_generator():
|
|
yield mock_omni_output_text
|
|
yield mock_omni_output_audio
|
|
|
|
# Mock the methods we need
|
|
serving_chat.get_chat_request_role = MagicMock(return_value="assistant")
|
|
serving_chat._create_text_choice = MagicMock(
|
|
return_value=(
|
|
[MagicMock()], # choices_data
|
|
UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
|
|
None, # prompt_logprobs
|
|
None, # prompt_token_ids
|
|
None, # kv_transfer_params
|
|
)
|
|
)
|
|
serving_chat._create_audio_choice = MagicMock(return_value=[MagicMock()])
|
|
|
|
# Call the method
|
|
result = await serving_chat.chat_completion_full_generator(
|
|
request=mock_request,
|
|
result_generator=mock_generator(),
|
|
request_id="test-123",
|
|
model_name="test-model",
|
|
conversation=[],
|
|
tokenizer=MagicMock(),
|
|
request_metadata=MagicMock(),
|
|
)
|
|
|
|
# Verify both choices were created
|
|
assert serving_chat._create_text_choice.called
|
|
assert serving_chat._create_audio_choice.called
|
|
|
|
# Verify response has 2 choices
|
|
assert len(result.choices) == 2
|
|
# Verify usage is populated
|
|
assert result.usage.prompt_tokens == 10
|
|
assert result.usage.completion_tokens == 20
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_modalities_includes_all(
|
|
self, serving_chat, mock_request, mock_omni_output_text, mock_omni_output_audio
|
|
):
|
|
"""Test that when modalities is None, all outputs are included."""
|
|
mock_request.modalities = None
|
|
|
|
# Create a mock result generator that yields both text and audio
|
|
async def mock_generator():
|
|
yield mock_omni_output_text
|
|
yield mock_omni_output_audio
|
|
|
|
# Mock the methods we need
|
|
serving_chat.get_chat_request_role = MagicMock(return_value="assistant")
|
|
serving_chat._create_text_choice = MagicMock(
|
|
return_value=(
|
|
[MagicMock()], # choices_data
|
|
UsageInfo(prompt_tokens=10, completion_tokens=20, total_tokens=30),
|
|
None, # prompt_logprobs
|
|
None, # prompt_token_ids
|
|
None, # kv_transfer_params
|
|
)
|
|
)
|
|
serving_chat._create_audio_choice = MagicMock(return_value=[MagicMock()])
|
|
|
|
# Call the method
|
|
result = await serving_chat.chat_completion_full_generator(
|
|
request=mock_request,
|
|
result_generator=mock_generator(),
|
|
request_id="test-123",
|
|
model_name="test-model",
|
|
conversation=[],
|
|
tokenizer=MagicMock(),
|
|
request_metadata=MagicMock(),
|
|
)
|
|
|
|
# Verify both choices were created
|
|
assert serving_chat._create_text_choice.called
|
|
assert serving_chat._create_audio_choice.called
|
|
|
|
# Verify response has 2 choices
|
|
assert len(result.choices) == 2
|