From 43fa24e832971da51c97fc434eac2c2de870a446 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Wed, 17 Jun 2026 01:57:12 -0500 Subject: [PATCH] [Misc] Validate Cohere Embed Mixed Content Payloads (#45873) Signed-off-by: Taneem Ibrahim --- .../pooling/embed/test_io_processor.py | 42 +++++++++++++++++++ vllm/entrypoints/pooling/embed/protocol.py | 11 +++++ 2 files changed, 53 insertions(+) diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index f0dea74044..8f8f8faa8a 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -150,6 +150,16 @@ class TestCohereEmbedRequestParsing: {"content": [{"type": "text", "text": "hello"}]}, ], }, + { + "model": "test", + "inputs": [ + { + "content": [ + {"type": "image_url", "image_url": {"url": "image-uri"}} + ] + }, + ], + }, ], ) def test_accepts_exactly_one_non_empty_input_field(self, request_body): @@ -157,6 +167,38 @@ class TestCohereEmbedRequestParsing: assert request.model == "test" + @pytest.mark.parametrize( + ("content", "error"), + [ + ( + {"type": "text"}, + "CohereEmbedContent with type='text' requires text", + ), + ( + {"type": "image_url"}, + "CohereEmbedContent with type='image_url' requires image_url.url", + ), + ( + {"type": "image_url", "image_url": {}}, + "CohereEmbedContent with type='image_url' requires image_url.url", + ), + ( + {"type": "image_url", "image_url": {"url": ""}}, + "CohereEmbedContent with type='image_url' requires image_url.url", + ), + ], + ) + def test_rejects_invalid_mixed_content_payloads(self, content, error): + with pytest.raises(ValidationError, match=error): + CohereEmbedRequest( + model="test", + inputs=[ + { + "content": [content], + }, + ], + ) + class TestResolveTruncation: """Unit tests for EmbedIOProcessor._resolve_cohere_truncation.""" diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index 8ec908f451..2dcc848c8c 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -207,6 +207,17 @@ class CohereEmbedContent(BaseModel): text: str | None = None image_url: dict[str, str] | None = None + @model_validator(mode="after") + def validate_content_payload(self): + if self.type == "text": + if self.text is None: + raise ValueError("CohereEmbedContent with type='text' requires text") + elif not self.image_url or not self.image_url.get("url"): + raise ValueError( + "CohereEmbedContent with type='image_url' requires image_url.url" + ) + return self + class CohereEmbedInput(BaseModel): content: list[CohereEmbedContent]