| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from unittest.mock import MagicMock |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.entrypoints.pooling import factories |
| 9 | from vllm.entrypoints.pooling.factories import init_pooling_io_processors |
| 10 | from vllm.entrypoints.pooling.pooling.io_processor import ( |
| 11 | PluginWithIOProcessorPlugins, |
| 12 | UnsupportedCombinedTaskIOProcessor, |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | def _bge_m3_config(io_processor_plugin=None): |
| 17 | model_config = MagicMock() |
| 18 | model_config.get_pooling_task.return_value = "embed&token_classify" |
| 19 | model_config.io_processor_plugin = io_processor_plugin |
| 20 | model_config.hf_config.to_dict.return_value = {} |
| 21 | model_config.architecture = "BgeM3EmbeddingModel" |
| 22 | |
| 23 | vllm_config = MagicMock(model_config=model_config) |
| 24 | renderer = MagicMock() |
| 25 | renderer._executor = MagicMock() |
| 26 | chat_template_config = MagicMock( |
| 27 | chat_template=None, |
| 28 | chat_template_content_format="auto", |
| 29 | trust_request_chat_template=False, |
| 30 | ) |
| 31 | return vllm_config, renderer, chat_template_config |
| 32 | |
| 33 | |
| 34 | def test_combined_task_without_plugin_uses_rejection_processor(): |
| 35 | vllm_config, renderer, chat_template_config = _bge_m3_config() |
| 36 | |
| 37 | processors = init_pooling_io_processors( |
| 38 | supported_tasks=("embed", "embed&token_classify"), |
| 39 | vllm_config=vllm_config, |
| 40 | renderer=renderer, |
| 41 | chat_template_config=chat_template_config, |
| 42 | ) |
| 43 | |
| 44 | assert processors.keys() == {"embed&token_classify"} |
| 45 | assert isinstance( |
| 46 | processors["embed&token_classify"], UnsupportedCombinedTaskIOProcessor |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | def test_combined_task_with_plugin_uses_plugin_processor(monkeypatch): |
| 51 | vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin") |
| 52 | monkeypatch.setattr(factories, "has_io_processor", lambda *_: True) |
| 53 | monkeypatch.setattr( |
| 54 | "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor", |
| 55 | lambda *_: MagicMock(), |
| 56 | ) |
| 57 | |
| 58 | processors = init_pooling_io_processors( |
| 59 | supported_tasks=("embed", "embed&token_classify"), |
| 60 | vllm_config=vllm_config, |
| 61 | renderer=renderer, |
| 62 | chat_template_config=chat_template_config, |
| 63 | ) |
| 64 | |
| 65 | assert processors.keys() == {"embed&token_classify", "plugin"} |
| 66 | assert isinstance(processors["plugin"], PluginWithIOProcessorPlugins) |
| 67 | |
| 68 | |
| 69 | def test_combined_task_plain_pooling_request_has_actionable_error(monkeypatch): |
| 70 | from vllm.entrypoints.pooling.pooling.protocol import PoolingCompletionRequest |
| 71 | from vllm.entrypoints.pooling.pooling.serving import ServingPooling |
| 72 | |
| 73 | vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin") |
| 74 | monkeypatch.setattr(factories, "has_io_processor", lambda *_: True) |
| 75 | monkeypatch.setattr( |
| 76 | "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor", |
| 77 | lambda *_: MagicMock(), |
| 78 | ) |
| 79 | |
| 80 | engine_client = MagicMock(renderer=renderer, vllm_config=vllm_config) |
| 81 | models = MagicMock(model_config=vllm_config.model_config) |
| 82 | serving = ServingPooling( |
| 83 | engine_client, |
| 84 | models, |
| 85 | supported_tasks=("embed", "embed&token_classify"), |
| 86 | request_logger=None, |
| 87 | chat_template_config=chat_template_config, |
| 88 | ) |
| 89 | request = PoolingCompletionRequest(model="BAAI/bge-m3", input=["hola"]) |
| 90 | |
| 91 | assert serving.io_processors.keys() == {"embed&token_classify", "plugin"} |
| 92 | io_processor = serving.get_io_processor(request) |
| 93 | with pytest.raises(ValueError, match="plugin request with a 'data' field"): |
| 94 | io_processor.create_pooling_params(request) |