@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 200 lines 6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 from openai.types.chat.chat_completion import ChatCompletion as OpenAIChatCompletion
6 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
7
8 from vllm.entrypoints.generate.base.protocol import (
9 DeltaFunctionCall,
10 DeltaMessage,
11 DeltaToolCall,
12 FunctionCall,
13 ToolCall,
14 )
15 from vllm.entrypoints.openai.chat_completion.protocol import (
16 ChatCompletionRequest,
17 ChatCompletionResponse,
18 ChatCompletionResponseChoice,
19 ChatCompletionResponseStreamChoice,
20 ChatCompletionStreamResponse,
21 ChatMessage,
22 )
23 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
24 from vllm.entrypoints.serve.engine.protocol import UsageInfo
25 from vllm.parser.abstract_parser import DelegatingParser
26
27 pytestmark = pytest.mark.skip_global_cleanup
28
29
30 class _DummyDelegatingParser(DelegatingParser):
31 def is_reasoning_end(self, input_ids: list[int]) -> bool:
32 return False
33
34 def extract_reasoning(self, model_output: str, request):
35 return None, model_output
36
37 def extract_reasoning_streaming(
38 self,
39 previous_text: str,
40 current_text: str,
41 delta_text: str,
42 previous_token_ids: list[int],
43 current_token_ids: list[int],
44 delta_token_ids: list[int],
45 ):
46 return None
47
48
49 def test_chat_completion_named_tool_choice_with_none_content():
50 request = ChatCompletionRequest.model_validate(
51 {
52 "model": "test-model",
53 "messages": [{"role": "user", "content": "test"}],
54 "tools": [
55 {
56 "type": "function",
57 "function": {
58 "name": "get_weather",
59 "parameters": {"type": "object", "properties": {}},
60 },
61 }
62 ],
63 "tool_choice": {"type": "function", "function": {"name": "get_weather"}},
64 }
65 )
66 parser = _DummyDelegatingParser(tokenizer=None)
67
68 tool_calls, content = parser._extract_tool_calls(
69 content=None,
70 request=request,
71 enable_auto_tools=True,
72 )
73
74 assert content is None
75 assert tool_calls == []
76
77
78 def test_responses_parser_allows_named_tool_choice_with_none_content():
79 request = ResponsesRequest.model_validate(
80 {
81 "model": "test-model",
82 "input": "test",
83 "tools": [
84 {
85 "type": "function",
86 "name": "get_weather",
87 "parameters": {"type": "object", "properties": {}},
88 }
89 ],
90 "tool_choice": {"type": "function", "name": "get_weather"},
91 }
92 )
93 parser = _DummyDelegatingParser(tokenizer=None)
94
95 tool_calls, content = parser._extract_tool_calls(
96 content=None,
97 request=request,
98 enable_auto_tools=False,
99 )
100
101 assert content is None
102 assert tool_calls == []
103
104
105 def _chat_response(message: ChatMessage) -> ChatCompletionResponse:
106 return ChatCompletionResponse(
107 model="test-model",
108 choices=[
109 ChatCompletionResponseChoice(
110 index=0,
111 message=message,
112 finish_reason="stop",
113 )
114 ],
115 usage=UsageInfo(prompt_tokens=1, completion_tokens=1, total_tokens=2),
116 )
117
118
119 def test_chat_completion_response_omits_empty_tool_calls_payload():
120 response = _chat_response(ChatMessage(role="assistant", content="done"))
121
122 payload = response.model_dump()
123 payload_exclude_unset = response.model_dump(exclude_unset=True)
124
125 assert "tool_calls" not in payload["choices"][0]["message"]
126 assert "tool_calls" not in payload_exclude_unset["choices"][0]["message"]
127 parsed = OpenAIChatCompletion.model_validate(payload)
128 assert parsed.choices[0].message.tool_calls is None
129
130
131 def test_chat_completion_response_keeps_non_empty_tool_calls_payload():
132 response = _chat_response(
133 ChatMessage(
134 role="assistant",
135 content="",
136 tool_calls=[
137 ToolCall(
138 function=FunctionCall(
139 name="get_weather",
140 arguments='{"city": "Beijing"}',
141 )
142 )
143 ],
144 )
145 )
146
147 message = response.model_dump()["choices"][0]["message"]
148
149 assert len(message["tool_calls"]) == 1
150 assert message["tool_calls"][0]["function"]["name"] == "get_weather"
151
152
153 def _stream_response(delta: DeltaMessage) -> ChatCompletionStreamResponse:
154 return ChatCompletionStreamResponse(
155 id="chatcmpl-test",
156 object="chat.completion.chunk",
157 created=1,
158 model="test-model",
159 choices=[
160 ChatCompletionResponseStreamChoice(
161 index=0,
162 delta=delta,
163 finish_reason=None,
164 )
165 ],
166 )
167
168
169 def test_chat_completion_stream_response_omits_empty_tool_calls_payload():
170 response = _stream_response(DeltaMessage(content="done"))
171
172 payload = response.model_dump(exclude_unset=True)
173 payload_json = response.model_dump_json(exclude_unset=True)
174
175 assert "tool_calls" not in payload["choices"][0]["delta"]
176 parsed = ChatCompletionChunk.model_validate_json(payload_json)
177 assert parsed.choices[0].delta.tool_calls is None
178
179
180 def test_chat_completion_stream_response_keeps_non_empty_tool_calls_payload():
181 response = _stream_response(
182 DeltaMessage(
183 tool_calls=[
184 DeltaToolCall(
185 index=0,
186 id="call-test",
187 type="function",
188 function=DeltaFunctionCall(
189 name="get_weather",
190 arguments='{"city": "Beijing"}',
191 ),
192 )
193 ]
194 )
195 )
196
197 delta = response.model_dump(exclude_unset=True)["choices"][0]["delta"]
198
199 assert len(delta["tool_calls"]) == 1
200 assert delta["tool_calls"][0]["function"]["name"] == "get_weather"