From c3a8f1a569005d6bb833dc66136c3134ee6e89be Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Wed, 10 Jun 2026 11:50:02 -0400 Subject: [PATCH 1/3] adding proper /v1/embeddings support for messages + chat_template_kwargs and unit tests along with it Signed-off-by: Taneem Ibrahim --- .../pooling/embed/test_io_processor.py | 150 ++++++++++++++++++ .../entrypoints/pooling/embed/io_processor.py | 50 ++++++ vllm/entrypoints/pooling/embed/protocol.py | 61 ++++++- 3 files changed, 259 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index 341ccbd5f0c5..18cb64954f8d 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -3,6 +3,7 @@ """Unit tests for EmbedIOProcessor.""" import pytest +from pydantic import TypeAdapter from vllm import PoolingParams from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor @@ -10,10 +11,61 @@ CohereEmbedContent, CohereEmbedInput, CohereEmbedRequest, + EmbeddingChatRequest, + EmbeddingCompletionRequest, + EmbeddingRequest, ) from vllm.entrypoints.pooling.typing import PoolingServeContext +class TestEmbeddingRequestParsing: + """Unit tests for OpenAI embedding request parsing.""" + + def test_input_messages_parses_as_chat_request(self): + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "input": [{"role": "user", "content": "hello"}], + "chat_template_kwargs": {"instruction": "Represent the query: "}, + } + ) + + assert isinstance(request, EmbeddingChatRequest) + assert request.messages == [{"role": "user", "content": "hello"}] + assert request.chat_template_kwargs == {"instruction": "Represent the query: "} + + def test_batched_input_messages_parses_as_chat_request(self): + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "input": [ + [{"role": "user", "content": "hello"}], + [{"role": "user", "content": "goodbye"}], + ], + "chat_template_kwargs": {"instruction": "Represent the query: "}, + } + ) + + assert isinstance(request, EmbeddingChatRequest) + assert request.messages == [{"role": "user", "content": "hello"}] + assert request.messages_batch == [ + [{"role": "user", "content": "hello"}], + [{"role": "user", "content": "goodbye"}], + ] + assert request.chat_template_kwargs == {"instruction": "Represent the query: "} + + def test_token_ids_still_parse_as_completion_request(self): + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "input": [[1, 2, 3], [4, 5]], + } + ) + + assert isinstance(request, EmbeddingCompletionRequest) + assert request.input == [[1, 2, 3], [4, 5]] + + class TestResolveTruncation: """Unit tests for EmbedIOProcessor._resolve_cohere_truncation.""" @@ -324,3 +376,101 @@ def batch_render_chat( }, ) ] + + +class TestPreProcessOpenAIEmbeddingChatOnline: + """Unit tests for OpenAI embedding chat preprocessing.""" + + class _FakeModelConfig: + max_model_len = 128 + encoder_config: dict[str, object] = {} + pooler_config = None + multimodal_config = None + is_encoder_decoder = False + + class _FakeRenderer: + tokenizer = object() + + def __init__(self): + self.calls = [] + + def render_chat( + self, + all_messages, + chat_params, + tok_params, + prompt_extras=None, + ): + self.calls.append( + { + "all_messages": all_messages, + "chat_params": chat_params, + "tok_params": tok_params, + "prompt_extras": prompt_extras, + } + ) + return all_messages, [ + {"prompt_token_ids": [index]} for index, _ in enumerate(all_messages) + ] + + @classmethod + def _make_handler(cls, renderer): + handler = object.__new__(EmbedIOProcessor) + handler.renderer = renderer + handler.model_config = cls._FakeModelConfig() + handler.chat_template = "template" + handler.chat_template_content_format = "auto" + handler.trust_request_chat_template = False + handler.enable_chunked_processing = False + return handler + + @staticmethod + def _make_context(request) -> PoolingServeContext[EmbeddingChatRequest]: + return PoolingServeContext( + request=request, + pooling_params=PoolingParams(), + model_name="test", + request_id="embd-test", + ) + + def test_chat_template_kwargs_forwarded_for_batched_input_messages(self): + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "input": [ + [{"role": "user", "content": "hello"}], + [{"role": "user", "content": "goodbye"}], + ], + "add_generation_prompt": True, + "chat_template_kwargs": {"instruction": "Represent the query: "}, + "mm_processor_kwargs": {"max_pixels": 1}, + "cache_salt": "salt", + } + ) + assert isinstance(request, EmbeddingChatRequest) + + renderer = self._FakeRenderer() + handler = self._make_handler(renderer) + ctx = self._make_context(request) + + handler.pre_process_online(ctx) + + assert ctx.engine_inputs == [ + {"prompt_token_ids": [0]}, + {"prompt_token_ids": [1]}, + ] + assert len(renderer.calls) == 1 + + call = renderer.calls[0] + assert call["all_messages"] == request.messages_batch + assert call["prompt_extras"] == { + "mm_processor_kwargs": {"max_pixels": 1}, + "cache_salt": "salt", + } + + chat_template_kwargs = call["chat_params"].chat_template_kwargs + assert chat_template_kwargs["instruction"] == "Represent the query: " + assert chat_template_kwargs["add_generation_prompt"] is True + assert chat_template_kwargs["continue_final_message"] is False + assert "tools" not in chat_template_kwargs + assert chat_template_kwargs["tokenize"] is False diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 8c28f9f3d4e7..a40ceedbea49 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -66,6 +66,8 @@ def __init__(self, *args, **kwargs): def pre_process_online(self, ctx: PoolingServeContext): if isinstance(ctx.request, CohereEmbedRequest): self._pre_process_cohere_online(ctx) + elif isinstance(ctx.request, EmbeddingChatRequest): + self._pre_process_openai_chat_online(ctx) else: super().pre_process_online(ctx) @@ -367,6 +369,54 @@ def create_pooling_params(self, request): ) return super().create_pooling_params(request) + def _pre_process_openai_chat_online( + self, ctx: PoolingServeContext[EmbeddingChatRequest] + ) -> None: + request = ctx.request + self._validate_chat_template( + request_chat_template=request.chat_template, + chat_template_kwargs=request.chat_template_kwargs, + trust_request_chat_template=self.trust_request_chat_template, + ) + + all_messages = request.messages_batch or [request.messages] + ctx.engine_inputs = self._batch_render_openai_chat(request, all_messages) + + def _batch_render_openai_chat( + self, + request: EmbeddingChatRequest, + all_messages: Sequence[list[ChatCompletionMessageParam]], + ) -> list[EngineInput]: + renderer = self.renderer + mm_config = self.model_config.multimodal_config + + tok_params = request.build_tok_params(self.model_config) + chat_params = request.build_chat_params( + self.chat_template, + self.chat_template_content_format, + ).with_defaults( + merge_kwargs( + None, + dict( + tools=None, + tokenize=is_mistral_tokenizer(renderer.tokenizer), + ), + ), + default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None), + ) + + _, engine_inputs = renderer.render_chat( + all_messages, + chat_params, + tok_params, + prompt_extras={ + k: v + for k in ("mm_processor_kwargs", "cache_salt") + if (v := getattr(request, k, None)) is not None + }, + ) + return engine_inputs + def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None: """Convert a ``CohereEmbedRequest`` into engine prompts. diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index d886e3199f7c..0d45e6a9ac76 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -10,12 +10,13 @@ import struct import time from collections.abc import Sequence -from typing import Literal, TypeAlias +from typing import Any, Literal, TypeAlias, cast import pybase64 as base64 -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator from vllm import PoolingParams +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo from vllm.utils import random_uuid @@ -42,12 +43,68 @@ def to_pooling_params(self): ) +def _is_chat_message(value: Any) -> bool: + return isinstance(value, dict) and isinstance(value.get("role"), str) + + +def _is_chat_messages(value: Any) -> bool: + return ( + isinstance(value, list) + and bool(value) + and all(_is_chat_message(item) for item in value) + ) + + +def _is_batched_chat_messages(value: Any) -> bool: + return ( + isinstance(value, list) + and bool(value) + and all(_is_chat_messages(item) for item in value) + ) + + class EmbeddingChatRequest( PoolingBasicRequestMixin, ChatRequestMixin, EmbedRequestMixin, EmbeddingTokenizeParamsMixin, ): + messages_batch: list[list[ChatCompletionMessageParam]] | None = Field( + default=None, + exclude=True, + ) + + @model_validator(mode="before") + @classmethod + def normalize_input_messages(cls, data): + if not isinstance(data, dict): + return data + + messages_data = data.get("messages") + if _is_batched_chat_messages(messages_data): + messages_batch = cast(list[list[ChatCompletionMessageParam]], messages_data) + normalized = dict(data) + normalized["messages"] = messages_batch[0] + normalized["messages_batch"] = messages_batch + return normalized + + if "messages" in data or "input" not in data: + return data + + normalized = dict(data) + input_data = data["input"] + if _is_chat_messages(input_data): + normalized["messages"] = input_data + elif _is_batched_chat_messages(input_data): + messages_batch = cast(list[list[ChatCompletionMessageParam]], input_data) + normalized["messages"] = messages_batch[0] + normalized["messages_batch"] = messages_batch + else: + return data + + normalized.pop("input") + return normalized + def to_pooling_params(self): return PoolingParams( task="embed", From 05efea7142737efa2e088e62b5061701ce460621 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Fri, 12 Jun 2026 11:08:36 -0500 Subject: [PATCH 2/3] Split the message-shaped input / batched chat embedding shape into a dedicated request class so EmbeddingChatRequest stays focused on the existing top-level messages API Signed-off-by: Taneem Ibrahim --- .../pooling/embed/test_io_processor.py | 67 ++++++++++-- vllm/entrypoints/pooling/base/protocol.py | 12 +- .../entrypoints/pooling/embed/io_processor.py | 35 +++++- vllm/entrypoints/pooling/embed/protocol.py | 103 +++++++++++++----- vllm/entrypoints/pooling/typing.py | 10 +- 5 files changed, 181 insertions(+), 46 deletions(-) diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index 18cb64954f8d..f4f1f4aa4005 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -11,6 +11,9 @@ CohereEmbedContent, CohereEmbedInput, CohereEmbedRequest, + EmbeddingBatchChatInputRequest, + EmbeddingBatchChatRequest, + EmbeddingChatInputRequest, EmbeddingChatRequest, EmbeddingCompletionRequest, EmbeddingRequest, @@ -30,11 +33,12 @@ def test_input_messages_parses_as_chat_request(self): } ) - assert isinstance(request, EmbeddingChatRequest) + assert isinstance(request, EmbeddingChatInputRequest) + assert request.input == [{"role": "user", "content": "hello"}] assert request.messages == [{"role": "user", "content": "hello"}] assert request.chat_template_kwargs == {"instruction": "Represent the query: "} - def test_batched_input_messages_parses_as_chat_request(self): + def test_batched_input_messages_parses_as_batch_chat_input_request(self): request = TypeAdapter(EmbeddingRequest).validate_python( { "model": "test", @@ -46,9 +50,12 @@ def test_batched_input_messages_parses_as_chat_request(self): } ) - assert isinstance(request, EmbeddingChatRequest) - assert request.messages == [{"role": "user", "content": "hello"}] - assert request.messages_batch == [ + assert isinstance(request, EmbeddingBatchChatInputRequest) + assert request.input == [ + [{"role": "user", "content": "hello"}], + [{"role": "user", "content": "goodbye"}], + ] + assert request.messages == [ [{"role": "user", "content": "hello"}], [{"role": "user", "content": "goodbye"}], ] @@ -65,6 +72,38 @@ def test_token_ids_still_parse_as_completion_request(self): assert isinstance(request, EmbeddingCompletionRequest) assert request.input == [[1, 2, 3], [4, 5]] + def test_messages_still_parses_as_chat_request(self): + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "messages": [{"role": "user", "content": "hello"}], + "chat_template_kwargs": {"instruction": "Represent the query: "}, + } + ) + + assert isinstance(request, EmbeddingChatRequest) + assert request.messages == [{"role": "user", "content": "hello"}] + assert request.chat_template_kwargs == {"instruction": "Represent the query: "} + + def test_batched_messages_parses_as_batch_chat_request(self): + request = TypeAdapter(EmbeddingRequest).validate_python( + { + "model": "test", + "messages": [ + [{"role": "user", "content": "hello"}], + [{"role": "user", "content": "goodbye"}], + ], + "chat_template_kwargs": {"instruction": "Represent the query: "}, + } + ) + + assert isinstance(request, EmbeddingBatchChatRequest) + assert request.messages == [ + [{"role": "user", "content": "hello"}], + [{"role": "user", "content": "goodbye"}], + ] + assert request.chat_template_kwargs == {"instruction": "Represent the query: "} + class TestResolveTruncation: """Unit tests for EmbedIOProcessor._resolve_cohere_truncation.""" @@ -425,7 +464,19 @@ def _make_handler(cls, renderer): return handler @staticmethod - def _make_context(request) -> PoolingServeContext[EmbeddingChatRequest]: + def _make_context( + request: ( + EmbeddingChatRequest + | EmbeddingBatchChatRequest + | EmbeddingChatInputRequest + | EmbeddingBatchChatInputRequest + ), + ) -> PoolingServeContext[ + EmbeddingChatRequest + | EmbeddingBatchChatRequest + | EmbeddingChatInputRequest + | EmbeddingBatchChatInputRequest + ]: return PoolingServeContext( request=request, pooling_params=PoolingParams(), @@ -447,7 +498,7 @@ def test_chat_template_kwargs_forwarded_for_batched_input_messages(self): "cache_salt": "salt", } ) - assert isinstance(request, EmbeddingChatRequest) + assert isinstance(request, EmbeddingBatchChatInputRequest) renderer = self._FakeRenderer() handler = self._make_handler(renderer) @@ -462,7 +513,7 @@ def test_chat_template_kwargs_forwarded_for_batched_input_messages(self): assert len(renderer.calls) == 1 call = renderer.calls[0] - assert call["all_messages"] == request.messages_batch + assert call["all_messages"] == request.messages assert call["prompt_extras"] == { "mm_processor_kwargs": {"max_pixels": 1}, "cache_salt": "salt", diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index 9e410a2b540d..81ad303ad902 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -168,11 +168,7 @@ class CompletionRequestMixin(OpenAIBaseModel): # --8<-- [end:completion-extra-params] -class ChatRequestMixin(OpenAIBaseModel): - # --8<-- [start:chat-params] - messages: list[ChatCompletionMessageParam] - # --8<-- [end:chat-params] - +class ChatRequestOptionsMixin(OpenAIBaseModel): # --8<-- [start:chat-extra-params] add_generation_prompt: bool = Field( default=False, @@ -256,6 +252,12 @@ def build_chat_params( ) +class ChatRequestMixin(ChatRequestOptionsMixin): + # --8<-- [start:chat-params] + messages: list[ChatCompletionMessageParam] + # --8<-- [end:chat-params] + + class EncodingRequestMixin(OpenAIBaseModel): # --8<-- [start:encoding-params] encoding_format: EncodingFormat = "float" diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index a40ceedbea49..d2e6f23c149d 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -36,6 +36,9 @@ CohereEmbedContent, CohereEmbedInput, CohereEmbedRequest, + EmbeddingBatchChatInputRequest, + EmbeddingBatchChatRequest, + EmbeddingChatInputRequest, EmbeddingChatRequest, EmbeddingCompletionRequest, ) @@ -66,7 +69,15 @@ def __init__(self, *args, **kwargs): def pre_process_online(self, ctx: PoolingServeContext): if isinstance(ctx.request, CohereEmbedRequest): self._pre_process_cohere_online(ctx) - elif isinstance(ctx.request, EmbeddingChatRequest): + elif isinstance( + ctx.request, + ( + EmbeddingChatRequest, + EmbeddingBatchChatRequest, + EmbeddingChatInputRequest, + EmbeddingBatchChatInputRequest, + ), + ): self._pre_process_openai_chat_online(ctx) else: super().pre_process_online(ctx) @@ -370,7 +381,13 @@ def create_pooling_params(self, request): return super().create_pooling_params(request) def _pre_process_openai_chat_online( - self, ctx: PoolingServeContext[EmbeddingChatRequest] + self, + ctx: PoolingServeContext[ + EmbeddingChatRequest + | EmbeddingBatchChatRequest + | EmbeddingChatInputRequest + | EmbeddingBatchChatInputRequest + ], ) -> None: request = ctx.request self._validate_chat_template( @@ -379,12 +396,22 @@ def _pre_process_openai_chat_online( trust_request_chat_template=self.trust_request_chat_template, ) - all_messages = request.messages_batch or [request.messages] + if isinstance( + request, (EmbeddingBatchChatRequest, EmbeddingBatchChatInputRequest) + ): + all_messages = request.messages + else: + all_messages = [request.messages] ctx.engine_inputs = self._batch_render_openai_chat(request, all_messages) def _batch_render_openai_chat( self, - request: EmbeddingChatRequest, + request: ( + EmbeddingChatRequest + | EmbeddingBatchChatRequest + | EmbeddingChatInputRequest + | EmbeddingBatchChatInputRequest + ), all_messages: Sequence[list[ChatCompletionMessageParam]], ) -> list[EngineInput]: renderer = self.renderer diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index 0d45e6a9ac76..99a07e4d828e 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -10,7 +10,7 @@ import struct import time from collections.abc import Sequence -from typing import Any, Literal, TypeAlias, cast +from typing import Annotated, Any, Literal, TypeAlias import pybase64 as base64 from pydantic import BaseModel, Field, model_validator @@ -22,6 +22,7 @@ from ..base.protocol import ( ChatRequestMixin, + ChatRequestOptionsMixin, CompletionRequestMixin, EmbeddingTokenizeParamsMixin, EmbedRequestMixin, @@ -69,51 +70,97 @@ class EmbeddingChatRequest( EmbedRequestMixin, EmbeddingTokenizeParamsMixin, ): - messages_batch: list[list[ChatCompletionMessageParam]] | None = Field( - default=None, - exclude=True, + """OpenAI embeddings request with one top-level chat conversation.""" + + def to_pooling_params(self): + return PoolingParams( + task="embed", + dimensions=self.dimensions, + use_activation=self.use_activation, + ) + + +class EmbeddingBatchChatRequest( + PoolingBasicRequestMixin, + ChatRequestOptionsMixin, + EmbedRequestMixin, + EmbeddingTokenizeParamsMixin, +): + """OpenAI embeddings request with batched top-level chat conversations. + + Mirrors ``BatchChatCompletionRequest`` by keeping batched conversations in + ``messages`` instead of introducing a separate batch-specific field. + """ + + messages: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = ( + Field(..., min_length=1) ) + def to_pooling_params(self): + return PoolingParams( + task="embed", + dimensions=self.dimensions, + use_activation=self.use_activation, + ) + + +class EmbeddingChatInputRequest( + EmbeddingChatRequest, +): + """OpenAI embeddings request with one chat conversation in ``input``.""" + + input: list[ChatCompletionMessageParam] + @model_validator(mode="before") @classmethod def normalize_input_messages(cls, data): if not isinstance(data, dict): return data - messages_data = data.get("messages") - if _is_batched_chat_messages(messages_data): - messages_batch = cast(list[list[ChatCompletionMessageParam]], messages_data) - normalized = dict(data) - normalized["messages"] = messages_batch[0] - normalized["messages_batch"] = messages_batch - return normalized - if "messages" in data or "input" not in data: return data - normalized = dict(data) input_data = data["input"] - if _is_chat_messages(input_data): - normalized["messages"] = input_data - elif _is_batched_chat_messages(input_data): - messages_batch = cast(list[list[ChatCompletionMessageParam]], input_data) - normalized["messages"] = messages_batch[0] - normalized["messages_batch"] = messages_batch - else: + if not _is_chat_messages(input_data): return data - normalized.pop("input") + normalized = dict(data) + normalized["messages"] = input_data return normalized - def to_pooling_params(self): - return PoolingParams( - task="embed", - dimensions=self.dimensions, - use_activation=self.use_activation, - ) +class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest): + """OpenAI embeddings request with batched chat conversations in ``input``.""" + + input: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = ( + Field(..., min_length=1) + ) -EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest + @model_validator(mode="before") + @classmethod + def normalize_input_messages(cls, data): + if not isinstance(data, dict): + return data + + if "messages" in data or "input" not in data: + return data + + input_data = data["input"] + if not _is_batched_chat_messages(input_data): + return data + + normalized = dict(data) + normalized["messages"] = input_data + return normalized + + +EmbeddingRequest: TypeAlias = ( + EmbeddingCompletionRequest + | EmbeddingChatRequest + | EmbeddingBatchChatRequest + | EmbeddingChatInputRequest + | EmbeddingBatchChatInputRequest +) # --------------------------------------------------------------------------- diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index ffcd3e7be434..748a4a6f6dce 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -19,7 +19,10 @@ ) from .embed.protocol import ( CohereEmbedRequest, + EmbeddingBatchChatInputRequest, + EmbeddingBatchChatRequest, EmbeddingBytesResponse, + EmbeddingChatInputRequest, EmbeddingChatRequest, EmbeddingCompletionRequest, EmbeddingResponse, @@ -41,7 +44,12 @@ ) PoolingChatLikeRequest: TypeAlias = ( - EmbeddingChatRequest | ClassificationChatRequest | PoolingChatRequest + EmbeddingChatRequest + | EmbeddingBatchChatRequest + | EmbeddingChatInputRequest + | EmbeddingBatchChatInputRequest + | ClassificationChatRequest + | PoolingChatRequest ) AnyPoolingRequest: TypeAlias = ( From 41faab1797778c04ad5988dd7319ca2c3a8b4c58 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 14 Jun 2026 16:05:10 -0500 Subject: [PATCH 3/3] Fix pooling chat request typing for batched embeddings Signed-off-by: Taneem Ibrahim --- vllm/entrypoints/pooling/typing.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 748a4a6f6dce..2cf384900532 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -19,12 +19,11 @@ ) from .embed.protocol import ( CohereEmbedRequest, - EmbeddingBatchChatInputRequest, - EmbeddingBatchChatRequest, EmbeddingBytesResponse, EmbeddingChatInputRequest, EmbeddingChatRequest, EmbeddingCompletionRequest, + EmbeddingRequest, EmbeddingResponse, ) from .pooling.protocol import ( @@ -45,15 +44,14 @@ PoolingChatLikeRequest: TypeAlias = ( EmbeddingChatRequest - | EmbeddingBatchChatRequest | EmbeddingChatInputRequest - | EmbeddingBatchChatInputRequest | ClassificationChatRequest | PoolingChatRequest ) AnyPoolingRequest: TypeAlias = ( - PoolingCompletionLikeRequest + EmbeddingRequest + | PoolingCompletionLikeRequest | PoolingChatLikeRequest | IOProcessorRequest | ScoringRequest