main
py 1,168 lines 35.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # imports for structured outputs tests
5 import json
6 from collections import defaultdict
7
8 import jsonschema
9 import openai # use the official client for correctness check
10 import pytest
11 import pytest_asyncio
12 import regex as re
13 import requests
14 import torch
15 from openai import BadRequestError
16
17 from tests.utils import RemoteOpenAIServer
18 from vllm.entrypoints.openai.chat_completion.protocol import (
19 ChatCompletionRequest,
20 )
21 from vllm.exceptions import VLLMValidationError
22 from vllm.sampling_params import SamplingParams
23
24 # any model with a chat template should work here
25 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
26
27
28 @pytest.fixture(scope="module")
29 def zephyr_lora_files():
30 """Download zephyr LoRA files once per test session."""
31 from vllm.transformers_utils.repo_utils import hf_api
32
33 return hf_api().snapshot_download(repo_id="typeof/zephyr-7b-beta-lora")
34
35
36 @pytest.fixture(scope="module")
37 def server(zephyr_lora_files):
38 args = [
39 # use half precision for speed and memory savings in CI environment
40 "--dtype",
41 "bfloat16",
42 "--max-model-len",
43 "8192",
44 "--enforce-eager",
45 # lora config below
46 "--enable-lora",
47 "--lora-modules",
48 f"zephyr-lora={zephyr_lora_files}",
49 "--max-lora-rank",
50 "64",
51 "--max-cpu-loras",
52 "2",
53 "--max-num-seqs",
54 "128",
55 ]
56
57 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
58 yield remote_server
59
60
61 @pytest_asyncio.fixture
62 async def client(server):
63 async with server.get_async_client() as async_client:
64 yield async_client
65
66
67 @pytest.mark.asyncio
68 @pytest.mark.parametrize(
69 # first test base model, then test loras
70 "model_name",
71 [MODEL_NAME, "zephyr-lora"],
72 )
73 async def test_no_logprobs_chat(client: openai.AsyncOpenAI, model_name: str):
74 messages = [
75 {"role": "system", "content": "you are a helpful assistant"},
76 {"role": "user", "content": "what is 1+1?"},
77 ]
78
79 chat_completion = await client.chat.completions.create(
80 model=model_name,
81 messages=messages,
82 max_completion_tokens=5,
83 temperature=0.0,
84 logprobs=False,
85 )
86
87 choice = chat_completion.choices[0]
88 assert choice.logprobs is None
89
90
91 @pytest.mark.asyncio
92 @pytest.mark.parametrize(
93 # just test 1 lora hereafter
94 "model_name",
95 [MODEL_NAME, "zephyr-lora"],
96 )
97 async def test_zero_logprobs_chat(client: openai.AsyncOpenAI, model_name: str):
98 messages = [
99 {"role": "system", "content": "you are a helpful assistant"},
100 {"role": "user", "content": "what is 1+1?"},
101 ]
102
103 chat_completion = await client.chat.completions.create(
104 model=model_name,
105 messages=messages,
106 max_completion_tokens=5,
107 temperature=0.0,
108 logprobs=True,
109 top_logprobs=0,
110 )
111
112 choice = chat_completion.choices[0]
113 assert choice.logprobs is not None
114 assert choice.logprobs.content is not None
115 assert len(choice.logprobs.content[0].top_logprobs) == 0
116
117
118 @pytest.mark.asyncio
119 @pytest.mark.parametrize(
120 "model_name",
121 [MODEL_NAME, "zephyr-lora"],
122 )
123 async def test_some_logprobs_chat(client: openai.AsyncOpenAI, model_name: str):
124 messages = [
125 {"role": "system", "content": "you are a helpful assistant"},
126 {"role": "user", "content": "what is 1+1?"},
127 ]
128
129 chat_completion = await client.chat.completions.create(
130 model=model_name,
131 messages=messages,
132 max_completion_tokens=5,
133 temperature=0.0,
134 logprobs=True,
135 top_logprobs=5,
136 )
137
138 choice = chat_completion.choices[0]
139 assert choice.logprobs is not None
140 assert choice.logprobs.content is not None
141 assert len(choice.logprobs.content[0].top_logprobs) == 5
142
143
144 @pytest.mark.asyncio
145 @pytest.mark.parametrize(
146 "model_name",
147 [MODEL_NAME, "zephyr-lora"],
148 )
149 async def test_too_many_chat_logprobs(client: openai.AsyncOpenAI, model_name: str):
150 messages = [
151 {"role": "system", "content": "you are a helpful assistant"},
152 {"role": "user", "content": "what is 1+1?"},
153 ]
154
155 # Default max_logprobs is 20, so this should raise an error
156 with pytest.raises((openai.BadRequestError, openai.APIError)):
157 stream = await client.chat.completions.create(
158 model=model_name,
159 messages=messages,
160 max_completion_tokens=10,
161 logprobs=True,
162 top_logprobs=21,
163 stream=True,
164 )
165 async for chunk in stream:
166 ...
167
168 with pytest.raises(openai.BadRequestError):
169 await client.chat.completions.create(
170 model=model_name,
171 messages=messages,
172 max_completion_tokens=10,
173 logprobs=True,
174 top_logprobs=30,
175 stream=False,
176 )
177
178 # the server should still work afterwards
179 chat_completion = await client.chat.completions.create(
180 model=model_name, messages=messages, max_completion_tokens=10, stream=False
181 )
182 message = chat_completion.choices[0].message
183 assert message.content is not None and len(message.content) >= 0
184
185
186 @pytest.mark.asyncio
187 @pytest.mark.parametrize(
188 "model_name, prompt_logprobs",
189 [(MODEL_NAME, 1), (MODEL_NAME, 0), (MODEL_NAME, -1), (MODEL_NAME, None)],
190 )
191 async def test_prompt_logprobs_chat(
192 client: openai.AsyncOpenAI, model_name: str, prompt_logprobs: int | None
193 ):
194 params: dict = {
195 "messages": [
196 {"role": "system", "content": "You are a helpful assistant."},
197 {"role": "user", "content": "Who won the world series in 2020?"},
198 {
199 "role": "assistant",
200 "content": "The Los Angeles Dodgers won the World Series in 2020.",
201 },
202 {"role": "user", "content": "Where was it played?"},
203 ],
204 "model": model_name,
205 }
206
207 if prompt_logprobs is not None:
208 params["extra_body"] = {"prompt_logprobs": prompt_logprobs}
209
210 if prompt_logprobs is not None and prompt_logprobs < 0:
211 with pytest.raises(BadRequestError):
212 await client.chat.completions.create(**params)
213 else:
214 completion = await client.chat.completions.create(**params)
215 if prompt_logprobs is not None:
216 assert completion.prompt_logprobs is not None
217 assert len(completion.prompt_logprobs) > 0
218 else:
219 assert completion.prompt_logprobs is None
220
221
222 @pytest.mark.asyncio
223 @pytest.mark.parametrize(
224 "model_name",
225 [MODEL_NAME],
226 )
227 async def test_more_than_one_prompt_logprobs_chat(
228 client: openai.AsyncOpenAI, model_name: str
229 ):
230 params: dict = {
231 "messages": [
232 {"role": "system", "content": "You are a helpful assistant."},
233 {"role": "user", "content": "Who won the world series in 2020?"},
234 {
235 "role": "assistant",
236 "content": "The Los Angeles Dodgers won the World Series in 2020.",
237 },
238 {"role": "user", "content": "Where was it played?"},
239 ],
240 "model": model_name,
241 "extra_body": {"prompt_logprobs": 1},
242 }
243
244 completion_1 = await client.chat.completions.create(**params)
245
246 params["extra_body"] = {"prompt_logprobs": 2}
247 completion_2 = await client.chat.completions.create(**params)
248
249 assert len(completion_1.prompt_logprobs[3]) == 1
250 assert len(completion_2.prompt_logprobs[3]) == 2
251
252
253 @pytest.mark.asyncio
254 @pytest.mark.parametrize(
255 "model_name",
256 [MODEL_NAME, "zephyr-lora"],
257 )
258 async def test_single_chat_session(client: openai.AsyncOpenAI, model_name: str):
259 messages = [
260 {"role": "system", "content": "you are a helpful assistant"},
261 {"role": "user", "content": "what is 1+1?"},
262 ]
263 # test single completion
264 chat_completion = await client.chat.completions.create(
265 model=model_name,
266 messages=messages,
267 max_completion_tokens=5,
268 logprobs=True,
269 top_logprobs=5,
270 )
271 assert chat_completion.id is not None
272 assert len(chat_completion.choices) == 1
273
274 choice = chat_completion.choices[0]
275
276 assert choice.finish_reason == "length"
277 assert chat_completion.usage == openai.types.CompletionUsage(
278 completion_tokens=5, prompt_tokens=37, total_tokens=42
279 )
280
281 message = choice.message
282 assert message.content is not None and len(message.content) >= 5
283 assert message.role == "assistant"
284 messages.append({"role": "assistant", "content": message.content})
285
286 # test multi-turn dialogue
287 messages.append({"role": "user", "content": "express your result in json"})
288 chat_completion = await client.chat.completions.create(
289 model=model_name,
290 messages=messages,
291 max_completion_tokens=5,
292 )
293 message = chat_completion.choices[0].message
294 assert message.content is not None and len(message.content) >= 0
295
296
297 @pytest.mark.asyncio
298 @pytest.mark.parametrize(
299 # just test 1 lora hereafter
300 "model_name",
301 [MODEL_NAME, "zephyr-lora"],
302 )
303 async def test_chat_streaming(client: openai.AsyncOpenAI, model_name: str):
304 messages = [
305 {"role": "system", "content": "you are a helpful assistant"},
306 {"role": "user", "content": "what is 1+1?"},
307 ]
308
309 # test single completion
310 chat_completion = await client.chat.completions.create(
311 model=model_name,
312 messages=messages,
313 max_completion_tokens=10,
314 temperature=0.0,
315 )
316 output = chat_completion.choices[0].message.content
317 stop_reason = chat_completion.choices[0].finish_reason
318
319 # test streaming
320 stream = await client.chat.completions.create(
321 model=model_name,
322 messages=messages,
323 max_completion_tokens=10,
324 temperature=0.0,
325 stream=True,
326 )
327 chunks: list[str] = []
328 finish_reason_count = 0
329 async for chunk in stream:
330 delta = chunk.choices[0].delta
331 if delta.role:
332 assert delta.role == "assistant"
333 if delta.content:
334 chunks.append(delta.content)
335 if chunk.choices[0].finish_reason is not None:
336 finish_reason_count += 1
337 # finish reason should only return in last block
338 assert finish_reason_count == 1
339 assert chunk.choices[0].finish_reason == stop_reason
340 assert delta.content
341 assert "".join(chunks) == output
342
343
344 @pytest.mark.asyncio
345 @pytest.mark.parametrize(
346 "model_name",
347 ["HuggingFaceH4/zephyr-7b-beta", "zephyr-lora"],
348 )
349 async def test_chat_completion_stream_options(
350 client: openai.AsyncOpenAI, model_name: str
351 ):
352 messages = [
353 {"role": "system", "content": "You are a helpful assistant."},
354 {"role": "user", "content": "What is the capital of France?"},
355 ]
356
357 # Test stream=True, stream_options={"include_usage": False}
358 stream = await client.chat.completions.create(
359 model=model_name,
360 messages=messages,
361 max_completion_tokens=10,
362 temperature=0.0,
363 stream=True,
364 stream_options={"include_usage": False},
365 )
366 async for chunk in stream:
367 assert chunk.usage is None
368
369 # Test stream=True, stream_options={"include_usage": True,
370 # "continuous_usage_stats": False}}
371 stream = await client.chat.completions.create(
372 model=model_name,
373 messages=messages,
374 max_completion_tokens=10,
375 temperature=0.0,
376 stream=True,
377 stream_options={"include_usage": True, "continuous_usage_stats": False},
378 )
379
380 async for chunk in stream:
381 if chunk.choices[0].finish_reason is None:
382 assert chunk.usage is None
383 else:
384 assert chunk.usage is None
385 final_chunk = await anext(stream)
386 assert final_chunk.usage is not None
387 assert final_chunk.usage.prompt_tokens > 0
388 assert final_chunk.usage.completion_tokens > 0
389 assert final_chunk.usage.total_tokens == (
390 final_chunk.usage.prompt_tokens + final_chunk.usage.completion_tokens
391 )
392 assert final_chunk.choices == []
393
394 # Test stream=False, stream_options={"include_usage": None}
395 with pytest.raises(BadRequestError):
396 await client.chat.completions.create(
397 model=model_name,
398 messages=messages,
399 max_completion_tokens=10,
400 temperature=0.0,
401 stream=False,
402 stream_options={"include_usage": None},
403 )
404
405 # Test stream=False, stream_options={"include_usage": True}
406 with pytest.raises(BadRequestError):
407 await client.chat.completions.create(
408 model=model_name,
409 messages=messages,
410 max_completion_tokens=10,
411 temperature=0.0,
412 stream=False,
413 stream_options={"include_usage": True},
414 )
415
416 # Test stream=True, stream_options={"include_usage": True,
417 # "continuous_usage_stats": True}
418 stream = await client.chat.completions.create(
419 model=model_name,
420 messages=messages,
421 max_completion_tokens=10,
422 extra_body=dict(min_tokens=10),
423 temperature=0.0,
424 stream=True,
425 stream_options={
426 "include_usage": True,
427 "continuous_usage_stats": True,
428 },
429 )
430 last_completion_tokens = 0
431 async for chunk in stream:
432 assert chunk.usage.prompt_tokens >= 0
433 assert (
434 last_completion_tokens == 0
435 or chunk.usage.completion_tokens > last_completion_tokens
436 or (
437 not chunk.choices
438 and chunk.usage.completion_tokens == last_completion_tokens
439 )
440 )
441 assert chunk.usage.total_tokens == (
442 chunk.usage.prompt_tokens + chunk.usage.completion_tokens
443 )
444 last_completion_tokens = chunk.usage.completion_tokens
445
446 assert last_completion_tokens == 10
447
448
449 @pytest.mark.asyncio
450 async def test_structured_outputs_choice_chat(
451 client: openai.AsyncOpenAI,
452 sample_structured_outputs_choices,
453 ):
454 messages = [
455 {"role": "system", "content": "you are a helpful assistant"},
456 {
457 "role": "user",
458 "content": "The best language for type-safe systems programming is ",
459 },
460 ]
461 chat_completion = await client.chat.completions.create(
462 model=MODEL_NAME,
463 messages=messages,
464 max_completion_tokens=10,
465 temperature=0.7,
466 extra_body=dict(
467 structured_outputs={"choice": sample_structured_outputs_choices}
468 ),
469 )
470 choice1 = chat_completion.choices[0].message.content
471 assert choice1 in sample_structured_outputs_choices
472
473 messages.append({"role": "assistant", "content": choice1})
474 messages.append({"role": "user", "content": "I disagree, pick another one"})
475 chat_completion = await client.chat.completions.create(
476 model=MODEL_NAME,
477 messages=messages,
478 max_completion_tokens=10,
479 temperature=0.7,
480 extra_body=dict(
481 structured_outputs={"choice": sample_structured_outputs_choices}
482 ),
483 )
484 choice2 = chat_completion.choices[0].message.content
485 assert choice2 in sample_structured_outputs_choices
486 assert choice1 != choice2
487
488
489 @pytest.mark.asyncio
490 async def test_structured_outputs_json_chat(
491 client: openai.AsyncOpenAI,
492 sample_json_schema,
493 ):
494 messages = [
495 {"role": "system", "content": "you are a helpful assistant"},
496 {
497 "role": "user",
498 "content": f"Give an example JSON for an employee profile that "
499 f"fits this schema: {sample_json_schema}",
500 },
501 ]
502 chat_completion = await client.chat.completions.create(
503 model=MODEL_NAME,
504 messages=messages,
505 max_completion_tokens=1000,
506 extra_body=dict(structured_outputs={"json": sample_json_schema}),
507 )
508 message = chat_completion.choices[0].message
509 assert message.content is not None
510 json1 = json.loads(message.content)
511 jsonschema.validate(instance=json1, schema=sample_json_schema)
512
513 messages.append({"role": "assistant", "content": message.content})
514 messages.append(
515 {"role": "user", "content": "Give me another one with a different name and age"}
516 )
517 chat_completion = await client.chat.completions.create(
518 model=MODEL_NAME,
519 messages=messages,
520 max_completion_tokens=1000,
521 extra_body=dict(structured_outputs={"json": sample_json_schema}),
522 )
523 message = chat_completion.choices[0].message
524 assert message.content is not None
525 json2 = json.loads(message.content)
526 jsonschema.validate(instance=json2, schema=sample_json_schema)
527 assert json1["name"] != json2["name"]
528 assert json1["age"] != json2["age"]
529
530
531 @pytest.mark.asyncio
532 async def test_structured_outputs_regex_chat(
533 client: openai.AsyncOpenAI,
534 sample_regex,
535 ):
536 messages = [
537 {"role": "system", "content": "you are a helpful assistant"},
538 {
539 "role": "user",
540 "content": f"Give an example IP address with this regex: {sample_regex}",
541 },
542 ]
543 chat_completion = await client.chat.completions.create(
544 model=MODEL_NAME,
545 messages=messages,
546 max_completion_tokens=20,
547 extra_body=dict(structured_outputs={"regex": sample_regex}),
548 )
549 ip1 = chat_completion.choices[0].message.content
550 assert ip1 is not None
551 assert re.fullmatch(sample_regex, ip1) is not None
552
553 messages.append({"role": "assistant", "content": ip1})
554 messages.append({"role": "user", "content": "Give me a different one"})
555 chat_completion = await client.chat.completions.create(
556 model=MODEL_NAME,
557 messages=messages,
558 max_completion_tokens=20,
559 extra_body=dict(structured_outputs={"regex": sample_regex}),
560 )
561 ip2 = chat_completion.choices[0].message.content
562 assert ip2 is not None
563 assert re.fullmatch(sample_regex, ip2) is not None
564 assert ip1 != ip2
565
566
567 @pytest.mark.asyncio
568 async def test_structured_outputs_type_error(client: openai.AsyncOpenAI):
569 messages = [
570 {"role": "system", "content": "you are a helpful assistant"},
571 {
572 "role": "user",
573 "content": "The best language for type-safe systems programming is ",
574 },
575 ]
576
577 with pytest.raises(openai.BadRequestError):
578 _ = await client.chat.completions.create(
579 model=MODEL_NAME,
580 messages=messages,
581 extra_body=dict(structured_outputs={"regex": {1: "Python", 2: "C++"}}),
582 )
583
584
585 @pytest.mark.asyncio
586 async def test_structured_outputs_choice_chat_logprobs(
587 client: openai.AsyncOpenAI, sample_structured_outputs_choices
588 ):
589 messages = [
590 {"role": "system", "content": "you are a helpful assistant"},
591 {
592 "role": "user",
593 "content": "The best language for type-safe systems programming is ",
594 },
595 ]
596 chat_completion = await client.chat.completions.create(
597 model=MODEL_NAME,
598 messages=messages,
599 max_completion_tokens=10,
600 logprobs=True,
601 top_logprobs=5,
602 extra_body=dict(
603 structured_outputs={"choice": sample_structured_outputs_choices}
604 ),
605 )
606
607 assert chat_completion.choices[0].logprobs is not None
608 assert chat_completion.choices[0].logprobs.content is not None
609 top_logprobs = chat_completion.choices[0].logprobs.content[0].top_logprobs
610
611 # -9999.0 is the minimum logprob returned by OpenAI
612 for item in top_logprobs:
613 assert item.logprob >= -9999.0, f"Failed (top_logprobs={top_logprobs})"
614
615
616 @pytest.mark.asyncio
617 async def test_response_format_json_object(client: openai.AsyncOpenAI):
618 for _ in range(2):
619 resp = await client.chat.completions.create(
620 model=MODEL_NAME,
621 messages=[
622 {
623 "role": "user",
624 "content": (
625 "what is 1+1? please respond with a JSON object, "
626 'the format is {"result": 2}'
627 ),
628 }
629 ],
630 response_format={"type": "json_object"},
631 )
632
633 content = resp.choices[0].message.content
634 assert content is not None
635
636 loaded = json.loads(content)
637 assert loaded == {"result": 2}, loaded
638
639
640 @pytest.mark.asyncio
641 async def test_response_format_json_schema(client: openai.AsyncOpenAI):
642 prompt = 'what is 1+1? The format is "result": 2'
643 # Check that this prompt cannot lead to a valid JSON without json_schema
644 for _ in range(2):
645 resp = await client.chat.completions.create(
646 model=MODEL_NAME,
647 messages=[{"role": "user", "content": prompt}],
648 )
649 content = resp.choices[0].message.content
650 assert content is not None
651 with pytest.raises((json.JSONDecodeError, AssertionError)):
652 loaded = json.loads(content)
653 assert loaded == {"result": 2}, loaded
654
655 for _ in range(2):
656 resp = await client.chat.completions.create(
657 model=MODEL_NAME,
658 messages=[{"role": "user", "content": prompt}],
659 response_format={
660 "type": "json_schema",
661 "json_schema": {
662 "name": "foo_test",
663 "schema": {
664 "type": "object",
665 "properties": {
666 "result": {"type": "integer"},
667 },
668 },
669 },
670 },
671 )
672
673 content = resp.choices[0].message.content
674 assert content is not None
675
676 loaded = json.loads(content)
677 assert loaded == {"result": 2}, loaded
678
679
680 @pytest.mark.asyncio
681 async def test_response_format_text(client: openai.AsyncOpenAI):
682 for _ in range(2):
683 resp = await client.chat.completions.create(
684 model=MODEL_NAME,
685 messages=[
686 {
687 "role": "user",
688 "content": "what is 1+1?",
689 }
690 ],
691 max_completion_tokens=10,
692 response_format={"type": "text"},
693 )
694
695 content = resp.choices[0].message.content
696 assert content is not None
697
698
699 @pytest.mark.asyncio
700 async def test_extra_fields_allowed(client: openai.AsyncOpenAI):
701 resp = await client.chat.completions.create(
702 model=MODEL_NAME,
703 messages=[
704 {
705 "role": "user",
706 "content": "what is 1+1?",
707 "extra_field": "0",
708 }
709 ], # type: ignore
710 temperature=0,
711 seed=0,
712 )
713
714 content = resp.choices[0].message.content
715 assert content is not None
716
717
718 @pytest.mark.asyncio
719 async def test_complex_message_content(client: openai.AsyncOpenAI):
720 content = [
721 {
722 "type": "text",
723 "text": "what is 1+1? please provide the result without any other text.",
724 }
725 ]
726 resp = await client.chat.completions.create(
727 model=MODEL_NAME,
728 messages=[
729 {
730 "role": "user",
731 "content": content,
732 }
733 ],
734 temperature=0,
735 seed=0,
736 )
737 content = resp.choices[0].message.content
738 assert content == "2"
739
740
741 @pytest.mark.asyncio
742 async def test_custom_role(client: openai.AsyncOpenAI):
743 # Not sure how the model handles custom roles so we just check that
744 # both string and complex message content are handled in the same way
745
746 resp1 = await client.chat.completions.create(
747 model=MODEL_NAME,
748 messages=[
749 {
750 "role": "my-custom-role",
751 "content": "what is 1+1?",
752 }
753 ], # type: ignore
754 temperature=0,
755 seed=0,
756 )
757
758 resp2 = await client.chat.completions.create(
759 model=MODEL_NAME,
760 messages=[
761 {
762 "role": "my-custom-role",
763 "content": [{"type": "text", "text": "what is 1+1?"}],
764 }
765 ], # type: ignore
766 temperature=0,
767 seed=0,
768 )
769
770 content1 = resp1.choices[0].message.content
771 content2 = resp2.choices[0].message.content
772 assert content1 == content2
773
774
775 @pytest.mark.asyncio
776 async def test_long_seed(client: openai.AsyncOpenAI):
777 for seed in [torch.iinfo(torch.long).min - 1, torch.iinfo(torch.long).max + 1]:
778 with pytest.raises(BadRequestError) as exc_info:
779 await client.chat.completions.create(
780 model=MODEL_NAME,
781 messages=[
782 {
783 "role": "system",
784 "content": "You are a helpful assistant.",
785 }
786 ],
787 temperature=0,
788 seed=seed,
789 )
790
791 assert (
792 "greater_than_equal" in exc_info.value.message
793 or "less_than_equal" in exc_info.value.message
794 )
795
796
797 @pytest.mark.asyncio
798 async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenAI):
799 messages = [
800 {"role": "system", "content": "you are a helpful assistant"},
801 {"role": "user", "content": "what is 1+1?"},
802 ]
803
804 request_args = {
805 "model": MODEL_NAME,
806 "messages": messages,
807 "max_completion_tokens": 5,
808 "temperature": 0.0,
809 "logprobs": False,
810 }
811
812 # Use raw HTTP for both endpoints so we compare server responses
813 # directly, without the openai SDK injecting extra fields
814 # (e.g. `moderation` added in newer SDK versions).
815 chat_response = requests.post(
816 server.url_for("v1/chat/completions"), json=request_args
817 )
818 chat_response.raise_for_status()
819
820 invocation_response = requests.post(
821 server.url_for("invocations"), json=request_args
822 )
823 invocation_response.raise_for_status()
824
825 chat_output = chat_response.json()
826 invocation_output = invocation_response.json()
827
828 extra_keys = invocation_output.keys() - chat_output.keys()
829 missing_keys = chat_output.keys() - invocation_output.keys()
830 assert missing_keys == set()
831 assert extra_keys <= {"moderation"}
832 assert chat_output["choices"] == invocation_output["choices"]
833
834
835 # Test n parameter for chat completions
836 # Tests that the n parameter works correctly for regular sampling
837 # (non-beam search) in chat completions, addressing issue #34305.
838
839
840 @pytest.mark.asyncio
841 @pytest.mark.parametrize(
842 "model_name",
843 [MODEL_NAME],
844 )
845 async def test_chat_completion_n_parameter_non_streaming(
846 client: openai.AsyncOpenAI, model_name: str
847 ):
848 """Test that n parameter returns multiple choices for non-streaming requests."""
849 messages = [
850 {"role": "system", "content": "You are a helpful assistant."},
851 {"role": "user", "content": "What is the opposite of big?"},
852 ]
853
854 # Test with n=3
855 chat_completion = await client.chat.completions.create(
856 model=model_name,
857 messages=messages,
858 max_completion_tokens=50,
859 temperature=1.0,
860 n=3,
861 seed=42,
862 stream=False,
863 )
864
865 assert len(chat_completion.choices) == 3
866
867 # Verify each choice has content and correct index
868 for i, choice in enumerate(chat_completion.choices):
869 assert choice.index == i
870 assert choice.message.content is not None
871 assert len(choice.message.content) > 0
872
873 contents = [choice.message.content for choice in chat_completion.choices]
874 assert len(set(contents)) > 1, "Expected different responses with n=3"
875
876
877 @pytest.mark.asyncio
878 @pytest.mark.parametrize(
879 "model_name",
880 [MODEL_NAME],
881 )
882 async def test_chat_completion_n_parameter_streaming(
883 client: openai.AsyncOpenAI, model_name: str
884 ):
885 """Test that n parameter returns multiple choices for streaming requests."""
886 messages = [
887 {"role": "system", "content": "You are a helpful assistant."},
888 {"role": "user", "content": "What is the capital of France?"},
889 ]
890
891 stream = await client.chat.completions.create(
892 model=model_name,
893 messages=messages,
894 max_completion_tokens=15,
895 temperature=0.7,
896 n=2,
897 stream=True,
898 )
899
900 # Collect all chunks using defaultdict for dynamic handling
901 chunks_by_index = defaultdict(list)
902 async for chunk in stream:
903 for choice in chunk.choices:
904 if choice.delta.content:
905 chunks_by_index[choice.index].append(choice.delta.content)
906
907 # Verify both choices received content
908 assert len(chunks_by_index[0]) > 0, "Choice 0 received no content chunks"
909 assert len(chunks_by_index[1]) > 0, "Choice 1 received no content chunks"
910
911 # Reconstruct full responses
912 response_0 = "".join(chunks_by_index[0])
913 response_1 = "".join(chunks_by_index[1])
914
915 assert len(response_0) > 0, "Choice 0 has empty response"
916 assert len(response_1) > 0, "Choice 1 has empty response"
917
918
919 @pytest.mark.asyncio
920 @pytest.mark.parametrize(
921 "model_name",
922 [MODEL_NAME],
923 )
924 async def test_chat_completion_n_with_seed(client: openai.AsyncOpenAI, model_name: str):
925 """Test that n parameter works correctly with seed parameter."""
926 messages = [
927 {"role": "user", "content": "Say hello."},
928 ]
929
930 # Test that seed parameter is accepted and works with n > 1
931 chat_completion = await client.chat.completions.create(
932 model=model_name,
933 messages=messages,
934 max_completion_tokens=10,
935 temperature=0.8,
936 n=2,
937 seed=42,
938 stream=False,
939 )
940
941 # Verify we get n=2 choices
942 assert len(chat_completion.choices) == 2
943
944 # Verify both choices have valid content
945 for i, choice in enumerate(chat_completion.choices):
946 assert choice.index == i
947 assert choice.message.content is not None
948 assert len(choice.message.content) > 0
949
950
951 @pytest.mark.asyncio
952 @pytest.mark.parametrize(
953 "model_name",
954 [MODEL_NAME],
955 )
956 async def test_chat_completion_n_equals_1(client: openai.AsyncOpenAI, model_name: str):
957 """Test that n=1 (default) still works correctly."""
958 messages = [
959 {"role": "user", "content": "Hello!"},
960 ]
961
962 chat_completion = await client.chat.completions.create(
963 model=model_name,
964 messages=messages,
965 max_completion_tokens=10,
966 temperature=0.7,
967 n=1,
968 stream=False,
969 )
970
971 assert len(chat_completion.choices) == 1
972 assert chat_completion.choices[0].index == 0
973 assert chat_completion.choices[0].message.content is not None
974
975
976 # Unit tests for n parameter in ChatCompletionRequest.to_sampling_params()
977 def test_chat_completion_request_n_parameter_to_sampling_params():
978 """Test that n parameter is correctly passed to SamplingParams."""
979 # Test with n=3
980 request = ChatCompletionRequest(
981 model="test-model",
982 messages=[{"role": "user", "content": "Hello"}],
983 n=3,
984 max_tokens=10,
985 )
986
987 sampling_params = request.to_sampling_params(
988 max_tokens=10,
989 default_sampling_params={},
990 )
991
992 assert isinstance(sampling_params, SamplingParams)
993 assert sampling_params.n == 3, f"Expected n=3, got n={sampling_params.n}"
994
995
996 def test_chat_completion_request_n_parameter_default():
997 """Test that n parameter defaults to 1."""
998 request = ChatCompletionRequest(
999 model="test-model",
1000 messages=[{"role": "user", "content": "Hello"}],
1001 # n not specified, should default to 1
1002 max_tokens=10,
1003 )
1004
1005 assert request.n == 1, "n should default to 1"
1006 sampling_params = request.to_sampling_params(
1007 max_tokens=10,
1008 default_sampling_params={},
1009 )
1010
1011 # SamplingParams.from_optional converts None to 1
1012 assert sampling_params.n == 1, f"Expected n=1 (default), got n={sampling_params.n}"
1013
1014
1015 def test_chat_completion_request_accepts_model_specific_reasoning_effort():
1016 request = ChatCompletionRequest(
1017 model="test-model",
1018 messages=[{"role": "user", "content": "Hello"}],
1019 reasoning_effort="max",
1020 )
1021
1022 chat_params = request.build_chat_params(
1023 default_template=None,
1024 default_template_content_format="auto",
1025 )
1026
1027 assert request.reasoning_effort == "max"
1028 assert chat_params.chat_template_kwargs["reasoning_effort"] == "max"
1029
1030
1031 def test_chat_completion_request_rejects_unknown_reasoning_effort():
1032 with pytest.raises(ValueError, match="Input should be"):
1033 ChatCompletionRequest(
1034 model="test-model",
1035 messages=[{"role": "user", "content": "Hello"}],
1036 reasoning_effort="extra_high",
1037 )
1038
1039
1040 def test_chat_completion_request_n_parameter_various_values():
1041 """Test n parameter with various values."""
1042 for n_value in [1, 2, 5, 10]:
1043 request = ChatCompletionRequest(
1044 model="test-model",
1045 messages=[{"role": "user", "content": "Test"}],
1046 n=n_value,
1047 max_tokens=10,
1048 )
1049
1050 sampling_params = request.to_sampling_params(
1051 max_tokens=10,
1052 default_sampling_params={},
1053 )
1054
1055 assert sampling_params.n == n_value, (
1056 f"Expected n={n_value}, got n={sampling_params.n}"
1057 )
1058
1059
1060 def test_chat_completion_request_n_parameter_exceeds_default_limit(
1061 monkeypatch: pytest.MonkeyPatch,
1062 ):
1063 """Test that n values exceeding the default limit are rejected."""
1064 import vllm.envs as envs
1065
1066 monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
1067 if hasattr(envs.__getattr__, "cache_clear"):
1068 envs.__getattr__.cache_clear()
1069
1070 max_n = envs.VLLM_MAX_N_SEQUENCES
1071 request = ChatCompletionRequest(
1072 model="test-model",
1073 messages=[{"role": "user", "content": "Test"}],
1074 n=max_n + 1,
1075 max_tokens=10,
1076 )
1077
1078 with pytest.raises(VLLMValidationError, match="n must be at most"):
1079 request.to_sampling_params(
1080 max_tokens=10,
1081 default_sampling_params={},
1082 )
1083
1084
1085 def test_chat_completion_request_n_parameter_at_limit(
1086 monkeypatch: pytest.MonkeyPatch,
1087 ):
1088 """Test that n at exactly the limit is accepted."""
1089 import vllm.envs as envs
1090
1091 monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
1092 if hasattr(envs.__getattr__, "cache_clear"):
1093 envs.__getattr__.cache_clear()
1094
1095 max_n = envs.VLLM_MAX_N_SEQUENCES
1096 request = ChatCompletionRequest(
1097 model="test-model",
1098 messages=[{"role": "user", "content": "Test"}],
1099 n=max_n,
1100 max_tokens=10,
1101 )
1102
1103 sampling_params = request.to_sampling_params(
1104 max_tokens=10,
1105 default_sampling_params={},
1106 )
1107 assert sampling_params.n == max_n
1108
1109
1110 def test_chat_completion_request_n_parameter_custom_limit(
1111 monkeypatch: pytest.MonkeyPatch,
1112 ):
1113 """Test that VLLM_MAX_N_SEQUENCES env var overrides the default limit."""
1114 import vllm.envs as envs
1115
1116 monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", "128")
1117 if hasattr(envs.__getattr__, "cache_clear"):
1118 envs.__getattr__.cache_clear()
1119
1120 request = ChatCompletionRequest(
1121 model="test-model",
1122 messages=[{"role": "user", "content": "Test"}],
1123 n=128,
1124 max_tokens=10,
1125 )
1126
1127 sampling_params = request.to_sampling_params(
1128 max_tokens=10,
1129 default_sampling_params={},
1130 )
1131 assert sampling_params.n == 128
1132
1133 request_over = ChatCompletionRequest(
1134 model="test-model",
1135 messages=[{"role": "user", "content": "Test"}],
1136 n=129,
1137 max_tokens=10,
1138 )
1139
1140 with pytest.raises(VLLMValidationError, match="n must be at most 128"):
1141 request_over.to_sampling_params(
1142 max_tokens=10,
1143 default_sampling_params={},
1144 )
1145
1146
1147 def test_chat_completion_request_n_parameter_massive_value(
1148 monkeypatch: pytest.MonkeyPatch,
1149 ):
1150 """Test that astronomically large n values are rejected (CVE fix)."""
1151 import vllm.envs as envs
1152
1153 monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
1154 if hasattr(envs.__getattr__, "cache_clear"):
1155 envs.__getattr__.cache_clear()
1156
1157 request = ChatCompletionRequest(
1158 model="test-model",
1159 messages=[{"role": "user", "content": "Test"}],
1160 n=100_000_000,
1161 max_tokens=1,
1162 )
1163
1164 with pytest.raises(VLLMValidationError, match="n must be at most"):
1165 request.to_sampling_params(
1166 max_tokens=1,
1167 default_sampling_params={},
1168 )