A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 16:30 UTC
[Test] Add tests for n parameter in chat completions API (#35283)
Signed-off-by: KrxGu <krishom70@gmail.com>
Krish Gupta committed
Feb 26, 2026 at 14:44 UTC
3827c8c55aaa6622fd96b0c846a38b94444ebb80
1 file changed
+206
tests/entrypoints/openai/test_chat.py
+206
@@ -3,6 +3,7 @@
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
@@ -13,6 +14,11 @@ import requests
14
import torch
15
from openai import BadRequestError
16
17
+from vllm.entrypoints.openai.chat_completion.protocol import (
18
+ ChatCompletionRequest,
19
+)
20
+from vllm.sampling_params import SamplingParams
21
+
22
from ...utils import RemoteOpenAIServer
23
24
# any model with a chat template should work here
@@ -815,3 +821,203 @@ async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenA
821
822
assert chat_output.keys() == invocation_output.keys()
823
assert chat_output["choices"] == invocation_output["choices"]
824
+
825
+
826
+# Test n parameter for chat completions
827
+# Tests that the n parameter works correctly for regular sampling
828
+# (non-beam search) in chat completions, addressing issue #34305.
829
+
830
+
831
+@pytest.mark.asyncio
832
+@pytest.mark.parametrize(
833
+ "model_name",
834
+ [MODEL_NAME],
835
+)
836
+async def test_chat_completion_n_parameter_non_streaming(
837
+ client: openai.AsyncOpenAI, model_name: str
838
+):
839
+ """Test that n parameter returns multiple choices for non-streaming requests."""
840
+ messages = [
841
+ {"role": "system", "content": "You are a helpful assistant."},
842
+ {"role": "user", "content": "What is the opposite of big?"},
843
+ ]
844
+
845
+ # Test with n=3
846
+ chat_completion = await client.chat.completions.create(
847
+ model=model_name,
848
+ messages=messages,
849
+ max_completion_tokens=20,
850
+ temperature=0.7,
851
+ n=3,
852
+ stream=False,
853
+ )
854
+
855
+ assert len(chat_completion.choices) == 3
856
+
857
+ # Verify each choice has content and correct index
858
+ for i, choice in enumerate(chat_completion.choices):
859
+ assert choice.index == i
860
+ assert choice.message.content is not None
861
+ assert len(choice.message.content) > 0
862
+
863
+ # Verify all responses are different (highly likely with temperature > 0)
864
+ contents = [choice.message.content for choice in chat_completion.choices]
865
+ assert len(set(contents)) > 1, "Expected different responses with n=3"
866
+
867
+
868
+@pytest.mark.asyncio
869
+@pytest.mark.parametrize(
870
+ "model_name",
871
+ [MODEL_NAME],
872
+)
873
+async def test_chat_completion_n_parameter_streaming(
874
+ client: openai.AsyncOpenAI, model_name: str
875
+):
876
+ """Test that n parameter returns multiple choices for streaming requests."""
877
+ messages = [
878
+ {"role": "system", "content": "You are a helpful assistant."},
879
+ {"role": "user", "content": "What is the capital of France?"},
880
+ ]
881
+
882
+ stream = await client.chat.completions.create(
883
+ model=model_name,
884
+ messages=messages,
885
+ max_completion_tokens=15,
886
+ temperature=0.7,
887
+ n=2,
888
+ stream=True,
889
+ )
890
+
891
+ # Collect all chunks using defaultdict for dynamic handling
892
+ chunks_by_index = defaultdict(list)
893
+ async for chunk in stream:
894
+ for choice in chunk.choices:
895
+ if choice.delta.content:
896
+ chunks_by_index[choice.index].append(choice.delta.content)
897
+
898
+ # Verify both choices received content
899
+ assert len(chunks_by_index[0]) > 0, "Choice 0 received no content chunks"
900
+ assert len(chunks_by_index[1]) > 0, "Choice 1 received no content chunks"
901
+
902
+ # Reconstruct full responses
903
+ response_0 = "".join(chunks_by_index[0])
904
+ response_1 = "".join(chunks_by_index[1])
905
+
906
+ assert len(response_0) > 0, "Choice 0 has empty response"
907
+ assert len(response_1) > 0, "Choice 1 has empty response"
908
+
909
+
910
+@pytest.mark.asyncio
911
+@pytest.mark.parametrize(
912
+ "model_name",
913
+ [MODEL_NAME],
914
+)
915
+async def test_chat_completion_n_with_seed(client: openai.AsyncOpenAI, model_name: str):
916
+ """Test that n parameter works correctly with seed parameter."""
917
+ messages = [
918
+ {"role": "user", "content": "Say hello."},
919
+ ]
920
+
921
+ # Test that seed parameter is accepted and works with n > 1
922
+ chat_completion = await client.chat.completions.create(
923
+ model=model_name,
924
+ messages=messages,
925
+ max_completion_tokens=10,
926
+ temperature=0.8,
927
+ n=2,
928
+ seed=42,
929
+ stream=False,
930
+ )
931
+
932
+ # Verify we get n=2 choices
933
+ assert len(chat_completion.choices) == 2
934
+
935
+ # Verify both choices have valid content
936
+ for i, choice in enumerate(chat_completion.choices):
937
+ assert choice.index == i
938
+ assert choice.message.content is not None
939
+ assert len(choice.message.content) > 0
940
+
941
+
942
+@pytest.mark.asyncio
943
+@pytest.mark.parametrize(
944
+ "model_name",
945
+ [MODEL_NAME],
946
+)
947
+async def test_chat_completion_n_equals_1(client: openai.AsyncOpenAI, model_name: str):
948
+ """Test that n=1 (default) still works correctly."""
949
+ messages = [
950
+ {"role": "user", "content": "Hello!"},
951
+ ]
952
+
953
+ chat_completion = await client.chat.completions.create(
954
+ model=model_name,
955
+ messages=messages,
956
+ max_completion_tokens=10,
957
+ temperature=0.7,
958
+ n=1,
959
+ stream=False,
960
+ )
961
+
962
+ assert len(chat_completion.choices) == 1
963
+ assert chat_completion.choices[0].index == 0
964
+ assert chat_completion.choices[0].message.content is not None
965
+
966
+
967
+# Unit tests for n parameter in ChatCompletionRequest.to_sampling_params()
968
+def test_chat_completion_request_n_parameter_to_sampling_params():
969
+ """Test that n parameter is correctly passed to SamplingParams."""
970
+ # Test with n=3
971
+ request = ChatCompletionRequest(
972
+ model="test-model",
973
+ messages=[{"role": "user", "content": "Hello"}],
974
+ n=3,
975
+ max_tokens=10,
976
+ )
977
+
978
+ sampling_params = request.to_sampling_params(
979
+ max_tokens=10,
980
+ default_sampling_params={},
981
+ )
982
+
983
+ assert isinstance(sampling_params, SamplingParams)
984
+ assert sampling_params.n == 3, f"Expected n=3, got n={sampling_params.n}"
985
+
986
+
987
+def test_chat_completion_request_n_parameter_default():
988
+ """Test that n parameter defaults to 1."""
989
+ request = ChatCompletionRequest(
990
+ model="test-model",
991
+ messages=[{"role": "user", "content": "Hello"}],
992
+ # n not specified, should default to 1
993
+ max_tokens=10,
994
+ )
995
+
996
+ assert request.n == 1, "n should default to 1"
997
+ sampling_params = request.to_sampling_params(
998
+ max_tokens=10,
999
+ default_sampling_params={},
1000
+ )
1001
+
1002
+ # SamplingParams.from_optional converts None to 1
1003
+ assert sampling_params.n == 1, f"Expected n=1 (default), got n={sampling_params.n}"
1004
+
1005
+
1006
+def test_chat_completion_request_n_parameter_various_values():
1007
+ """Test n parameter with various values."""
1008
+ for n_value in [1, 2, 5, 10]:
1009
+ request = ChatCompletionRequest(
1010
+ model="test-model",
1011
+ messages=[{"role": "user", "content": "Test"}],
1012
+ n=n_value,
1013
+ max_tokens=10,
1014
+ )
1015
+
1016
+ sampling_params = request.to_sampling_params(
1017
+ max_tokens=10,
1018
+ default_sampling_params={},
1019
+ )
1020
+
1021
+ assert sampling_params.n == n_value, (
1022
+ f"Expected n={n_value}, got n={sampling_params.n}"
1023
+ )