A high-throughput and memory-efficient inference and serving engine for LLMs
[Misc] Validate Cohere Embed Mixed Content Payloads (#45873)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Jun 17, 2026 at 01:57 UTC
43fa24e832971da51c97fc434eac2c2de870a446
2 files changed
+53
tests/entrypoints/pooling/embed/test_io_processor.py
+42
@@ -150,6 +150,16 @@ class TestCohereEmbedRequestParsing:
150
{"content": [{"type": "text", "text": "hello"}]},
151
],
152
},
153
+ {
154
+ "model": "test",
155
+ "inputs": [
156
+ {
157
+ "content": [
158
+ {"type": "image_url", "image_url": {"url": "image-uri"}}
159
+ ]
160
+ },
161
+ ],
162
+ },
163
],
164
)
165
def test_accepts_exactly_one_non_empty_input_field(self, request_body):
@@ -157,6 +167,38 @@ class TestCohereEmbedRequestParsing:
167
168
assert request.model == "test"
169
170
+ @pytest.mark.parametrize(
171
+ ("content", "error"),
172
+ [
173
+ (
174
+ {"type": "text"},
175
+ "CohereEmbedContent with type='text' requires text",
176
+ ),
177
+ (
178
+ {"type": "image_url"},
179
+ "CohereEmbedContent with type='image_url' requires image_url.url",
180
+ ),
181
+ (
182
+ {"type": "image_url", "image_url": {}},
183
+ "CohereEmbedContent with type='image_url' requires image_url.url",
184
+ ),
185
+ (
186
+ {"type": "image_url", "image_url": {"url": ""}},
187
+ "CohereEmbedContent with type='image_url' requires image_url.url",
188
+ ),
189
+ ],
190
+ )
191
+ def test_rejects_invalid_mixed_content_payloads(self, content, error):
192
+ with pytest.raises(ValidationError, match=error):
193
+ CohereEmbedRequest(
194
+ model="test",
195
+ inputs=[
196
+ {
197
+ "content": [content],
198
+ },
199
+ ],
200
+ )
201
+
202
203
class TestResolveTruncation:
204
"""Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
vllm/entrypoints/pooling/embed/protocol.py
+11
@@ -207,6 +207,17 @@ class CohereEmbedContent(BaseModel):
207
text: str | None = None
208
image_url: dict[str, str] | None = None
209
210
+ @model_validator(mode="after")
211
+ def validate_content_payload(self):
212
+ if self.type == "text":
213
+ if self.text is None:
214
+ raise ValueError("CohereEmbedContent with type='text' requires text")
215
+ elif not self.image_url or not self.image_url.get("url"):
216
+ raise ValueError(
217
+ "CohereEmbedContent with type='image_url' requires image_url.url"
218
+ )
219
+ return self
220
+
221
222
class CohereEmbedInput(BaseModel):
223
content: list[CohereEmbedContent]