A high-throughput and memory-efficient inference and serving engine for LLMs
[Misc] Added validation for Cohere /v2/embed input field exclusivity (#45640)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jun 16, 2026 at 00:42 UTC
81d8f4ebacaf4b0bf85ec559d5a0db1bcf5ade87
2 files changed
+65
-1
tests/entrypoints/pooling/embed/test_io_processor.py
+54
-1
@@ -3,7 +3,7 @@
3
"""Unit tests for EmbedIOProcessor."""
4
5
import pytest
6
-from pydantic import TypeAdapter
6
+from pydantic import TypeAdapter, ValidationError
7
8
from vllm import PoolingParams
9
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
@@ -105,6 +105,59 @@ class TestEmbeddingRequestParsing:
105
assert request.chat_template_kwargs == {"instruction": "Represent the query: "}
106
107
108
+class TestCohereEmbedRequestParsing:
109
+ """Unit tests for Cohere embed request parsing."""
110
+
111
+ @pytest.mark.parametrize(
112
+ "request_body",
113
+ [
114
+ {"model": "test"},
115
+ {"model": "test", "texts": ["hello"], "images": ["image-uri"]},
116
+ {
117
+ "model": "test",
118
+ "texts": ["hello"],
119
+ "inputs": [
120
+ {"content": [{"type": "text", "text": "hello"}]},
121
+ ],
122
+ },
123
+ {
124
+ "model": "test",
125
+ "images": ["image-uri"],
126
+ "inputs": [
127
+ {"content": [{"type": "text", "text": "hello"}]},
128
+ ],
129
+ },
130
+ {"model": "test", "texts": []},
131
+ {"model": "test", "images": []},
132
+ {"model": "test", "inputs": []},
133
+ ],
134
+ )
135
+ def test_rejects_invalid_input_field_combinations(self, request_body):
136
+ with pytest.raises(
137
+ ValidationError,
138
+ match="Exactly one of texts, images, or inputs must be provided",
139
+ ):
140
+ CohereEmbedRequest(**request_body)
141
+
142
+ @pytest.mark.parametrize(
143
+ "request_body",
144
+ [
145
+ {"model": "test", "texts": ["hello"]},
146
+ {"model": "test", "images": ["image-uri"]},
147
+ {
148
+ "model": "test",
149
+ "inputs": [
150
+ {"content": [{"type": "text", "text": "hello"}]},
151
+ ],
152
+ },
153
+ ],
154
+ )
155
+ def test_accepts_exactly_one_non_empty_input_field(self, request_body):
156
+ request = CohereEmbedRequest(**request_body)
157
+
158
+ assert request.model == "test"
159
+
160
+
161
class TestResolveTruncation:
162
"""Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
163
vllm/entrypoints/pooling/embed/protocol.py
+11
@@ -224,6 +224,17 @@ class CohereEmbedRequest(BaseModel):
224
max_tokens: int | None = None
225
priority: int = 0
226
227
+ @model_validator(mode="after")
228
+ def validate_input_fields(self):
229
+ input_fields = (self.texts, self.images, self.inputs)
230
+ provided_fields = [field for field in input_fields if field is not None]
231
+ if len(provided_fields) != 1 or not provided_fields[0]:
232
+ raise ValueError(
233
+ "Exactly one of texts, images, or inputs must be provided, "
234
+ "and it must be non-empty"
235
+ )
236
+ return self
237
+
238
239
# ---------------------------------------------------------------------------
240
# Cohere /v2/embed — response models