main
py 135 lines 4.06 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import openai
5 import pytest
6 import pytest_asyncio
7
8 from tests.utils import RemoteOpenAIServer
9 from vllm.config import ModelConfig
10
11 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
12
13
14 def get_vocab_size(model_name):
15 config = ModelConfig(
16 model=model_name,
17 seed=0,
18 dtype="bfloat16",
19 )
20 return config.get_vocab_size()
21
22
23 @pytest.fixture(scope="module")
24 def server():
25 args = [
26 "--dtype",
27 "bfloat16",
28 "--max-model-len",
29 "1024",
30 "--enforce-eager",
31 ]
32
33 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
34 yield remote_server
35
36
37 @pytest_asyncio.fixture
38 async def client(server):
39 async with server.get_async_client() as async_client:
40 yield async_client
41
42
43 @pytest.mark.asyncio
44 async def test_chat_logit_bias_valid(client):
45 """Test that valid logit_bias values are accepted in chat completions."""
46 vocab_size = get_vocab_size(MODEL_NAME)
47 valid_token_id = vocab_size - 1
48
49 completion = await client.chat.completions.create(
50 model=MODEL_NAME,
51 messages=[{"role": "user", "content": "Testing valid logit bias"}],
52 max_tokens=5,
53 logit_bias={str(valid_token_id): 1.0},
54 )
55
56 assert completion.choices[0].message.content is not None
57
58
59 @pytest.mark.asyncio
60 async def test_chat_logit_bias_invalid(client):
61 """Test that invalid logit_bias values are rejected in chat completions."""
62 vocab_size = get_vocab_size(MODEL_NAME)
63 invalid_token_id = vocab_size + 1
64
65 with pytest.raises(openai.BadRequestError) as excinfo:
66 await client.chat.completions.create(
67 model=MODEL_NAME,
68 messages=[{"role": "user", "content": "Testing invalid logit bias"}],
69 max_tokens=5,
70 logit_bias={str(invalid_token_id): 1.0},
71 )
72
73 error = excinfo.value
74 error_message = str(error)
75
76 assert error.status_code == 400
77 assert str(invalid_token_id) in error_message
78 assert str(vocab_size) in error_message
79
80
81 @pytest.mark.asyncio
82 async def test_chat_logit_bias_non_integer_key(client):
83 """Test that a non-integer logit_bias key is rejected with a clean,
84 informative error instead of a raw 'invalid literal for int()' message."""
85 with pytest.raises(openai.BadRequestError) as excinfo:
86 await client.chat.completions.create(
87 model=MODEL_NAME,
88 messages=[{"role": "user", "content": "Testing invalid logit bias key"}],
89 max_tokens=5,
90 logit_bias={"not_a_token_id": 50},
91 )
92
93 error = excinfo.value
94 error_message = str(error)
95
96 assert error.status_code == 400
97 assert "not_a_token_id" in error_message
98 assert "logit_bias" in error_message
99
100
101 @pytest.mark.asyncio
102 async def test_chat_logit_bias_non_numeric_value(client):
103 """Test that a non-numeric logit_bias value is rejected with a message
104 that names the specific offending token, not just a generic TypeError."""
105 with pytest.raises(openai.BadRequestError) as excinfo:
106 await client.chat.completions.create(
107 model=MODEL_NAME,
108 messages=[{"role": "user", "content": "Testing invalid logit bias value"}],
109 max_tokens=5,
110 logit_bias={"1": "not_a_number"},
111 )
112
113 error = excinfo.value
114 error_message = str(error)
115
116 assert error.status_code == 400
117 assert "logit_bias" in error_message
118
119
120 @pytest.mark.asyncio
121 async def test_chat_logit_bias_multiple_non_integer_keys(client):
122 """Test that ALL invalid logit_bias keys are reported together,
123 not just the first one encountered."""
124 with pytest.raises(openai.BadRequestError) as excinfo:
125 await client.chat.completions.create(
126 model=MODEL_NAME,
127 messages=[{"role": "user", "content": "Testing multiple bad keys"}],
128 max_tokens=5,
129 logit_bias={"bad1": 50.0, "bad2": 20.0},
130 )
131
132 error_message = str(excinfo.value)
133 assert excinfo.value.status_code == 400
134 assert "bad1" in error_message
135 assert "bad2" in error_message