| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the Cohere /v2/embed API with generic (non-Cohere) models. |
| 4 | |
| 5 | Validates that the Cohere v2 embed endpoint works correctly with standard |
| 6 | embedding models, covering text embedding, embedding type conversions, |
| 7 | response structure, batching, normalisation, and semantic similarity. |
| 8 | """ |
| 9 | |
| 10 | import struct |
| 11 | |
| 12 | import numpy as np |
| 13 | import pybase64 as base64 |
| 14 | import pytest |
| 15 | import requests |
| 16 | |
| 17 | from tests.utils import RemoteOpenAIServer |
| 18 | |
| 19 | DTYPE = "bfloat16" |
| 20 | |
| 21 | MODELS: list[tuple[str, list[str]]] = [ |
| 22 | ("intfloat/multilingual-e5-small", []), |
| 23 | ( |
| 24 | "Snowflake/snowflake-arctic-embed-m-v1.5", |
| 25 | [ |
| 26 | "--trust_remote_code", |
| 27 | "--hf_overrides", |
| 28 | '{"matryoshka_dimensions":[256]}', |
| 29 | ], |
| 30 | ), |
| 31 | ] |
| 32 | |
| 33 | |
| 34 | @pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0]) |
| 35 | def model_config(request): |
| 36 | return request.param |
| 37 | |
| 38 | |
| 39 | @pytest.fixture(scope="module") |
| 40 | def model_name(model_config): |
| 41 | return model_config[0] |
| 42 | |
| 43 | |
| 44 | @pytest.fixture(scope="module") |
| 45 | def server(model_config): |
| 46 | name, extra_args = model_config |
| 47 | args = [ |
| 48 | "--runner", |
| 49 | "pooling", |
| 50 | "--dtype", |
| 51 | DTYPE, |
| 52 | "--enforce-eager", |
| 53 | "--max-model-len", |
| 54 | "512", |
| 55 | "--gpu-memory-utilization", |
| 56 | "0.02", |
| 57 | ] + extra_args |
| 58 | with RemoteOpenAIServer(name, args) as remote_server: |
| 59 | yield remote_server |
| 60 | |
| 61 | |
| 62 | def _cohere_embed( |
| 63 | server: RemoteOpenAIServer, |
| 64 | model_name: str, |
| 65 | texts: list[str] | None = None, |
| 66 | images: list[str] | None = None, |
| 67 | input_type: str | None = None, |
| 68 | embedding_types: list[str] | None = None, |
| 69 | ) -> dict: |
| 70 | body: dict = {"model": model_name} |
| 71 | if input_type is not None: |
| 72 | body["input_type"] = input_type |
| 73 | if texts is not None: |
| 74 | body["texts"] = texts |
| 75 | if images is not None: |
| 76 | body["images"] = images |
| 77 | if embedding_types is not None: |
| 78 | body["embedding_types"] = embedding_types |
| 79 | resp = requests.post(server.url_for("/v2/embed"), json=body) |
| 80 | resp.raise_for_status() |
| 81 | return resp.json() |
| 82 | |
| 83 | |
| 84 | def _openai_embed( |
| 85 | server: RemoteOpenAIServer, model_name: str, texts: list[str] |
| 86 | ) -> dict: |
| 87 | body = {"model": model_name, "input": texts, "encoding_format": "float"} |
| 88 | resp = requests.post(server.url_for("/v1/embeddings"), json=body) |
| 89 | resp.raise_for_status() |
| 90 | return resp.json() |
| 91 | |
| 92 | |
| 93 | def _cosine_sim(a: list[float], b: list[float]) -> float: |
| 94 | va, vb = np.array(a), np.array(b) |
| 95 | return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb))) |
| 96 | |
| 97 | |
| 98 | # ----------------------------------------------------------- |
| 99 | # Text embedding tests |
| 100 | # ----------------------------------------------------------- |
| 101 | |
| 102 | |
| 103 | def test_basic_embed(server: RemoteOpenAIServer, model_name: str): |
| 104 | r = _cohere_embed( |
| 105 | server, model_name, texts=["hello world"], embedding_types=["float"] |
| 106 | ) |
| 107 | assert "embeddings" in r |
| 108 | assert len(r["embeddings"]["float"]) == 1 |
| 109 | assert len(r["embeddings"]["float"][0]) > 0 |
| 110 | |
| 111 | |
| 112 | def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str): |
| 113 | """An input_type not defined in the model's prompt config should be |
| 114 | rejected with a 400 error.""" |
| 115 | body = { |
| 116 | "model": model_name, |
| 117 | "input_type": "nonexistent_type", |
| 118 | "texts": ["hello world"], |
| 119 | "embedding_types": ["float"], |
| 120 | } |
| 121 | resp = requests.post(server.url_for("/v2/embed"), json=body) |
| 122 | assert resp.status_code == 400 |
| 123 | assert "Unsupported input_type" in resp.json()["error"]["message"] |
| 124 | |
| 125 | |
| 126 | def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str): |
| 127 | """Omitting input_type should always work (no prompt prefix applied).""" |
| 128 | body = { |
| 129 | "model": model_name, |
| 130 | "texts": ["hello world"], |
| 131 | "embedding_types": ["float"], |
| 132 | } |
| 133 | resp = requests.post(server.url_for("/v2/embed"), json=body) |
| 134 | assert resp.status_code == 200 |
| 135 | data = resp.json() |
| 136 | assert len(data["embeddings"]["float"]) == 1 |
| 137 | |
| 138 | |
| 139 | def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str): |
| 140 | """v1 (OpenAI) and v2 (Cohere) endpoints should produce the same |
| 141 | float embeddings for a generic model.""" |
| 142 | texts = ["hello world"] |
| 143 | v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"]) |
| 144 | v1 = _openai_embed(server, model_name, texts) |
| 145 | cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"]) |
| 146 | assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}" |
| 147 | |
| 148 | |
| 149 | def test_embedding_types(server: RemoteOpenAIServer, model_name: str): |
| 150 | r = _cohere_embed( |
| 151 | server, |
| 152 | model_name, |
| 153 | texts=["test"], |
| 154 | embedding_types=["float", "binary", "ubinary"], |
| 155 | ) |
| 156 | dim = len(r["embeddings"]["float"][0]) |
| 157 | assert len(r["embeddings"]["binary"][0]) == dim // 8 |
| 158 | assert len(r["embeddings"]["ubinary"][0]) == dim // 8 |
| 159 | |
| 160 | |
| 161 | def test_response_structure(server: RemoteOpenAIServer, model_name: str): |
| 162 | r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"]) |
| 163 | assert "id" in r |
| 164 | assert "embeddings" in r |
| 165 | assert "texts" in r |
| 166 | assert r["texts"] == ["test"] |
| 167 | assert "meta" in r |
| 168 | assert r["meta"]["api_version"]["version"] == "2" |
| 169 | assert "billed_units" in r["meta"] |
| 170 | assert r["meta"]["billed_units"]["input_tokens"] > 0 |
| 171 | assert r["meta"]["billed_units"]["image_tokens"] == 0 |
| 172 | |
| 173 | |
| 174 | def test_batch(server: RemoteOpenAIServer, model_name: str): |
| 175 | texts = ["apple", "banana", "cherry"] |
| 176 | r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"]) |
| 177 | assert len(r["embeddings"]["float"]) == 3 |
| 178 | dim = len(r["embeddings"]["float"][0]) |
| 179 | for emb in r["embeddings"]["float"]: |
| 180 | assert len(emb) == dim |
| 181 | |
| 182 | |
| 183 | def test_l2_normalized(server: RemoteOpenAIServer, model_name: str): |
| 184 | r = _cohere_embed( |
| 185 | server, model_name, texts=["hello world"], embedding_types=["float"] |
| 186 | ) |
| 187 | emb = np.array(r["embeddings"]["float"][0]) |
| 188 | assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01 |
| 189 | |
| 190 | |
| 191 | def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str): |
| 192 | r = _cohere_embed( |
| 193 | server, |
| 194 | model_name, |
| 195 | texts=["machine learning", "deep learning", "chocolate cake recipe"], |
| 196 | embedding_types=["float"], |
| 197 | ) |
| 198 | embs = r["embeddings"]["float"] |
| 199 | cos_related = _cosine_sim(embs[0], embs[1]) |
| 200 | cos_unrelated = _cosine_sim(embs[0], embs[2]) |
| 201 | assert cos_related > cos_unrelated |
| 202 | |
| 203 | |
| 204 | def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str): |
| 205 | body = {"model": model_name} |
| 206 | resp = requests.post(server.url_for("/v2/embed"), json=body) |
| 207 | assert resp.status_code == 400 |
| 208 | |
| 209 | |
| 210 | def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str): |
| 211 | r = _cohere_embed( |
| 212 | server, |
| 213 | model_name, |
| 214 | texts=["test encoding"], |
| 215 | embedding_types=["float", "base64"], |
| 216 | ) |
| 217 | float_emb = r["embeddings"]["float"][0] |
| 218 | b64_str = r["embeddings"]["base64"][0] |
| 219 | decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str)) |
| 220 | np.testing.assert_allclose(float_emb, decoded, rtol=1e-5) |
| 221 | |
| 222 | |
| 223 | # ----------------------------------------------------------- |
| 224 | # Truncation tests |
| 225 | # ----------------------------------------------------------- |
| 226 | |
| 227 | |
| 228 | def _cohere_embed_raw( |
| 229 | server: RemoteOpenAIServer, |
| 230 | body: dict, |
| 231 | ) -> requests.Response: |
| 232 | return requests.post(server.url_for("/v2/embed"), json=body) |
| 233 | |
| 234 | |
| 235 | def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str): |
| 236 | """truncate=END should silently truncate long input.""" |
| 237 | long_text = " ".join(["word"] * 2000) |
| 238 | body = { |
| 239 | "model": model_name, |
| 240 | "texts": [long_text], |
| 241 | "embedding_types": ["float"], |
| 242 | "truncate": "END", |
| 243 | } |
| 244 | resp = _cohere_embed_raw(server, body) |
| 245 | assert resp.status_code == 200 |
| 246 | data = resp.json() |
| 247 | assert len(data["embeddings"]["float"]) == 1 |
| 248 | |
| 249 | |
| 250 | def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str): |
| 251 | """truncate=START should silently truncate long input from the start.""" |
| 252 | long_text = " ".join(["word"] * 2000) |
| 253 | body = { |
| 254 | "model": model_name, |
| 255 | "texts": [long_text], |
| 256 | "embedding_types": ["float"], |
| 257 | "truncate": "START", |
| 258 | } |
| 259 | resp = _cohere_embed_raw(server, body) |
| 260 | assert resp.status_code == 200 |
| 261 | data = resp.json() |
| 262 | assert len(data["embeddings"]["float"]) == 1 |
| 263 | |
| 264 | |
| 265 | def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str): |
| 266 | """truncate=NONE should error when input exceeds model context.""" |
| 267 | long_text = " ".join(["word"] * 2000) |
| 268 | body = { |
| 269 | "model": model_name, |
| 270 | "texts": [long_text], |
| 271 | "embedding_types": ["float"], |
| 272 | "truncate": "NONE", |
| 273 | } |
| 274 | resp = _cohere_embed_raw(server, body) |
| 275 | assert resp.status_code == 400 |
| 276 | |
| 277 | |
| 278 | def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str): |
| 279 | """START and END truncation should produce different embeddings |
| 280 | when the input is long enough to actually be truncated. |
| 281 | |
| 282 | We construct input with distinct tokens at the start vs end |
| 283 | so that keeping different halves produces different embeddings. |
| 284 | """ |
| 285 | start_words = " ".join([f"alpha{i}" for i in range(300)]) |
| 286 | end_words = " ".join([f"omega{i}" for i in range(300)]) |
| 287 | long_text = start_words + " " + end_words |
| 288 | |
| 289 | body_end = { |
| 290 | "model": model_name, |
| 291 | "texts": [long_text], |
| 292 | "embedding_types": ["float"], |
| 293 | "truncate": "END", |
| 294 | } |
| 295 | body_start = { |
| 296 | "model": model_name, |
| 297 | "texts": [long_text], |
| 298 | "embedding_types": ["float"], |
| 299 | "truncate": "START", |
| 300 | } |
| 301 | r_end = _cohere_embed_raw(server, body_end).json() |
| 302 | r_start = _cohere_embed_raw(server, body_start).json() |
| 303 | |
| 304 | emb_end = r_end["embeddings"]["float"][0] |
| 305 | emb_start = r_start["embeddings"]["float"][0] |
| 306 | cos = _cosine_sim(emb_end, emb_start) |
| 307 | assert cos < 0.99, ( |
| 308 | f"START and END truncation should produce different embeddings " |
| 309 | f"for long input, but cosine similarity was {cos}" |
| 310 | ) |