Skip to content

Commit 562ad43

Browse files
jiayin-nvidiajiayin-nvidia
committed
Refactor test_openai_embeddings
1 parent 1d4d263 commit 562ad43

File tree

1 file changed

+62
-175
lines changed

1 file changed

+62
-175
lines changed

‎tests/integration/inference/test_openai_embeddings.py‎

Lines changed: 62 additions & 175 deletions
Original file line numberDiff line numberDiff line change
@@ -149,27 +149,13 @@ def test_openai_embeddings_single_string(compat_client, client_with_models, embe
149149
skip_if_model_doesnt_support_openai_embeddings(client_with_models, embedding_model_id)
150150

151151
input_text = "Hello, world!"
152-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
153152

154-
# For asymmetric models, verify that calling without extra_body raises an error
155-
if is_asymmetric_model(client_with_models, embedding_model_id):
156-
kwargs_without_extra = {
157-
"model": embedding_model_id,
158-
"input": input_text,
159-
"encoding_format": "float",
160-
}
161-
with pytest.raises(Exception): # noqa: B017
162-
compat_client.embeddings.create(**kwargs_without_extra)
163-
164-
kwargs = {
165-
"model": embedding_model_id,
166-
"input": input_text,
167-
"encoding_format": "float",
168-
}
169-
if is_asymmetric_model(client_with_models, embedding_model_id):
170-
kwargs["extra_body"] = extra_body
171-
172-
response = compat_client.embeddings.create(**kwargs)
153+
response = compat_client.embeddings.create(
154+
model=embedding_model_id,
155+
input=input_text,
156+
encoding_format="float",
157+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
158+
)
173159

174160
assert response.object == "list"
175161

@@ -188,26 +174,13 @@ def test_openai_embeddings_multiple_strings(compat_client, client_with_models, e
188174
skip_if_model_doesnt_support_openai_embeddings(client_with_models, embedding_model_id)
189175

190176
input_texts = ["Hello, world!", "How are you today?", "This is a test."]
191-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
192177

193-
if is_asymmetric_model(client_with_models, embedding_model_id):
194-
kwargs_without_extra = {
195-
"model": embedding_model_id,
196-
"input": input_texts,
197-
"encoding_format": "float",
198-
}
199-
with pytest.raises(Exception): # noqa: B017
200-
compat_client.embeddings.create(**kwargs_without_extra)
201-
202-
kwargs = {
203-
"model": embedding_model_id,
204-
"input": input_texts,
205-
"encoding_format": "float",
206-
}
207-
if is_asymmetric_model(client_with_models, embedding_model_id):
208-
kwargs["extra_body"] = extra_body
209-
210-
response = compat_client.embeddings.create(**kwargs)
178+
response = compat_client.embeddings.create(
179+
model=embedding_model_id,
180+
input=input_texts,
181+
encoding_format="float",
182+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
183+
)
211184

212185
assert response.object == "list"
213186

@@ -228,26 +201,13 @@ def test_openai_embeddings_with_encoding_format_float(compat_client, client_with
228201
skip_if_model_doesnt_support_openai_embeddings(client_with_models, embedding_model_id)
229202

230203
input_text = "Test encoding format"
231-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
232204

233-
if is_asymmetric_model(client_with_models, embedding_model_id):
234-
kwargs_without_extra = {
235-
"model": embedding_model_id,
236-
"input": input_text,
237-
"encoding_format": "float",
238-
}
239-
with pytest.raises(Exception): # noqa: B017
240-
compat_client.embeddings.create(**kwargs_without_extra)
241-
242-
kwargs = {
243-
"model": embedding_model_id,
244-
"input": input_text,
245-
"encoding_format": "float",
246-
}
247-
if is_asymmetric_model(client_with_models, embedding_model_id):
248-
kwargs["extra_body"] = extra_body
249-
250-
response = compat_client.embeddings.create(**kwargs)
205+
response = compat_client.embeddings.create(
206+
model=embedding_model_id,
207+
input=input_text,
208+
encoding_format="float",
209+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
210+
)
251211

252212
assert response.object == "list"
253213
assert len(response.data) == 1
@@ -262,26 +222,13 @@ def test_openai_embeddings_with_dimensions(compat_client, client_with_models, em
262222

263223
input_text = "Test dimensions parameter"
264224
dimensions = 16
265-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
266225

267-
if is_asymmetric_model(client_with_models, embedding_model_id):
268-
kwargs_without_extra = {
269-
"model": embedding_model_id,
270-
"input": input_text,
271-
"dimensions": dimensions,
272-
}
273-
with pytest.raises(Exception): # noqa: B017
274-
compat_client.embeddings.create(**kwargs_without_extra)
275-
276-
kwargs = {
277-
"model": embedding_model_id,
278-
"input": input_text,
279-
"dimensions": dimensions,
280-
}
281-
if is_asymmetric_model(client_with_models, embedding_model_id):
282-
kwargs["extra_body"] = extra_body
283-
284-
response = compat_client.embeddings.create(**kwargs)
226+
response = compat_client.embeddings.create(
227+
model=embedding_model_id,
228+
input=input_text,
229+
dimensions=dimensions,
230+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
231+
)
285232

286233
assert response.object == "list"
287234
assert len(response.data) == 1
@@ -297,26 +244,13 @@ def test_openai_embeddings_with_user_parameter(compat_client, client_with_models
297244

298245
input_text = "Test user parameter"
299246
user_id = "test-user-123"
300-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
301247

302-
if is_asymmetric_model(client_with_models, embedding_model_id):
303-
kwargs_without_extra = {
304-
"model": embedding_model_id,
305-
"input": input_text,
306-
"user": user_id,
307-
}
308-
with pytest.raises(Exception): # noqa: B017
309-
compat_client.embeddings.create(**kwargs_without_extra)
310-
311-
kwargs = {
312-
"model": embedding_model_id,
313-
"input": input_text,
314-
"user": user_id,
315-
}
316-
if is_asymmetric_model(client_with_models, embedding_model_id):
317-
kwargs["extra_body"] = extra_body
318-
319-
response = compat_client.embeddings.create(**kwargs)
248+
response = compat_client.embeddings.create(
249+
model=embedding_model_id,
250+
input=input_text,
251+
user=user_id,
252+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
253+
)
320254

321255
assert response.object == "list"
322256
assert len(response.data) == 1
@@ -328,17 +262,12 @@ def test_openai_embeddings_empty_list_error(compat_client, client_with_models, e
328262
"""Test that empty list input raises an appropriate error."""
329263
skip_if_model_doesnt_support_openai_embeddings(client_with_models, embedding_model_id)
330264

331-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
332-
333-
kwargs = {
334-
"model": embedding_model_id,
335-
"input": [],
336-
}
337-
if is_asymmetric_model(client_with_models, embedding_model_id):
338-
kwargs["extra_body"] = extra_body
339-
340265
with pytest.raises(Exception): # noqa: B017
341-
compat_client.embeddings.create(**kwargs)
266+
compat_client.embeddings.create(
267+
model=embedding_model_id,
268+
input=[],
269+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
270+
)
342271

343272

344273
def test_openai_embeddings_invalid_model_error(compat_client, client_with_models, embedding_model_id):
@@ -349,6 +278,7 @@ def test_openai_embeddings_invalid_model_error(compat_client, client_with_models
349278
compat_client.embeddings.create(
350279
model="invalid-model-id",
351280
input="Test text",
281+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
352282
)
353283

354284

@@ -358,35 +288,20 @@ def test_openai_embeddings_different_inputs_different_outputs(compat_client, cli
358288

359289
input_text1 = "This is the first text"
360290
input_text2 = "This is completely different content"
361-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
362291

363-
if is_asymmetric_model(client_with_models, embedding_model_id):
364-
kwargs_without_extra = {
365-
"model": embedding_model_id,
366-
"input": input_text1,
367-
"encoding_format": "float",
368-
}
369-
with pytest.raises(Exception): # noqa: B017
370-
compat_client.embeddings.create(**kwargs_without_extra)
371-
372-
kwargs1 = {
373-
"model": embedding_model_id,
374-
"input": input_text1,
375-
"encoding_format": "float",
376-
}
377-
if is_asymmetric_model(client_with_models, embedding_model_id):
378-
kwargs1["extra_body"] = extra_body
379-
380-
kwargs2 = {
381-
"model": embedding_model_id,
382-
"input": input_text2,
383-
"encoding_format": "float",
384-
}
385-
if is_asymmetric_model(client_with_models, embedding_model_id):
386-
kwargs2["extra_body"] = extra_body
387-
388-
response1 = compat_client.embeddings.create(**kwargs1)
389-
response2 = compat_client.embeddings.create(**kwargs2)
292+
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
293+
response1 = compat_client.embeddings.create(
294+
model=embedding_model_id,
295+
input=input_text1,
296+
encoding_format="float",
297+
extra_body=extra_body,
298+
)
299+
response2 = compat_client.embeddings.create(
300+
model=embedding_model_id,
301+
input=input_text2,
302+
encoding_format="float",
303+
extra_body=extra_body,
304+
)
390305

391306
embedding1 = response1.data[0].embedding
392307
embedding2 = response2.data[0].embedding
@@ -404,28 +319,14 @@ def test_openai_embeddings_with_encoding_format_base64(compat_client, client_wit
404319

405320
input_text = "Test base64 encoding format"
406321
dimensions = 12
407-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
408322

409-
if is_asymmetric_model(client_with_models, embedding_model_id):
410-
kwargs_without_extra = {
411-
"model": embedding_model_id,
412-
"input": input_text,
413-
"encoding_format": "base64",
414-
"dimensions": dimensions,
415-
}
416-
with pytest.raises(Exception): # noqa: B017
417-
compat_client.embeddings.create(**kwargs_without_extra)
418-
419-
kwargs = {
420-
"model": embedding_model_id,
421-
"input": input_text,
422-
"encoding_format": "base64",
423-
"dimensions": dimensions,
424-
}
425-
if is_asymmetric_model(client_with_models, embedding_model_id):
426-
kwargs["extra_body"] = extra_body
427-
428-
response = compat_client.embeddings.create(**kwargs)
323+
response = compat_client.embeddings.create(
324+
model=embedding_model_id,
325+
input=input_text,
326+
encoding_format="base64",
327+
dimensions=dimensions,
328+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
329+
)
429330

430331
# Validate response structure
431332
assert response.object == "list"
@@ -451,27 +352,13 @@ def test_openai_embeddings_base64_batch_processing(compat_client, client_with_mo
451352
skip_if_model_doesnt_support_encoding_format_base64(client_with_models, embedding_model_id)
452353

453354
input_texts = ["First text for base64", "Second text for base64", "Third text for base64"]
454-
extra_body = get_extra_body_for_model(client_with_models, embedding_model_id)
455-
456-
if is_asymmetric_model(client_with_models, embedding_model_id):
457-
kwargs_without_extra = {
458-
"model": embedding_model_id,
459-
"input": input_texts,
460-
"encoding_format": "base64",
461-
}
462-
with pytest.raises(Exception): # noqa: B017
463-
compat_client.embeddings.create(**kwargs_without_extra)
464-
465-
kwargs = {
466-
"model": embedding_model_id,
467-
"input": input_texts,
468-
"encoding_format": "base64",
469-
}
470-
if is_asymmetric_model(client_with_models, embedding_model_id):
471-
kwargs["extra_body"] = extra_body
472-
473-
response = compat_client.embeddings.create(**kwargs)
474355

356+
response = compat_client.embeddings.create(
357+
model=embedding_model_id,
358+
input=input_texts,
359+
encoding_format="base64",
360+
extra_body=get_extra_body_for_model(client_with_models, embedding_model_id),
361+
)
475362
# Validate response structure
476363
assert response.object == "list"
477364

0 commit comments

Comments
 (0)