@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 269 lines 8.07 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import MagicMock, patch
5
6 import pytest
7
8 from tests.tool_parsers.utils import (
9 run_tool_extraction,
10 run_tool_extraction_streaming,
11 )
12 from vllm.entrypoints.generate.base.protocol import FunctionCall
13 from vllm.tokenizers import TokenizerLike
14 from vllm.tool_parsers import ToolParser, ToolParserManager
15
16 # Test cases similar to pythonic parser but with Llama4 specific format
17 SIMPLE_FUNCTION_OUTPUT = "[get_weather(city='LA', metric='C')]"
18 SIMPLE_FUNCTION_CALL = FunctionCall(
19 name="get_weather",
20 arguments='{"city": "LA", "metric": "C"}',
21 )
22 MORE_TYPES_FUNCTION_OUTPUT = (
23 "[register_user(name='Doe', "
24 "age=9, "
25 "address={'city': 'LA', 'state': 'CA'}, "
26 "role=None, "
27 "passed_test=True, "
28 "aliases=['John', 'Johnny'])]"
29 )
30 MORE_TYPES_FUNCTION_CALL = FunctionCall(
31 name="register_user",
32 arguments='{"name": "Doe", '
33 '"age": 9, '
34 '"address": {"city": "LA", "state": "CA"}, '
35 '"role": null, '
36 '"passed_test": true, '
37 '"aliases": ["John", "Johnny"]}',
38 )
39 PARAMETERLESS_FUNCTION_OUTPUT = "[get_weather()]"
40 PARAMETERLESS_FUNCTION_CALL = FunctionCall(
41 name="get_weather",
42 arguments="{}",
43 )
44 EMPTY_DICT_FUNCTION_OUTPUT = "[do_something_cool(additional_data={})]"
45 EMPTY_DICT_FUNCTION_CALL = FunctionCall(
46 name="do_something_cool",
47 arguments='{"additional_data": {}}',
48 )
49 EMPTY_LIST_FUNCTION_OUTPUT = "[do_something_cool(steps=[])]"
50 EMPTY_LIST_FUNCTION_CALL = FunctionCall(
51 name="do_something_cool",
52 arguments='{"steps": []}',
53 )
54 ESCAPED_STRING_FUNCTION_OUTPUT = (
55 r"[get_weather(city='Martha\'s Vineyard', metric='\"cool units\"')]"
56 )
57 ESCAPED_STRING_FUNCTION_CALL = FunctionCall(
58 name="get_weather",
59 arguments='{"city": "Martha\'s Vineyard", "metric": "\\"cool units\\""}',
60 )
61 PYTHON_TAG_FUNCTION_OUTPUT = (
62 "<|python_start|>[get_weather(city='LA', metric='C')]<|python_end|>"
63 )
64
65
66 @pytest.mark.parametrize("streaming", [True, False])
67 def test_no_tool_call(streaming: bool, default_tokenizer: TokenizerLike):
68 tool_parser: ToolParser = ToolParserManager.get_tool_parser("llama4_pythonic")(
69 default_tokenizer
70 )
71 model_output = "How can I help you today?"
72
73 content, tool_calls = run_tool_extraction(
74 tool_parser, model_output, streaming=streaming
75 )
76
77 assert content == model_output
78 assert len(tool_calls) == 0
79
80
81 test_str = "<|python_start|>"
82 test_str += "[get_weather(city='LA', metric='C'),"
83 test_str += "register_user(name='Doe', age=9)]"
84 TEST_CASES = [
85 pytest.param(
86 True,
87 ESCAPED_STRING_FUNCTION_OUTPUT,
88 [ESCAPED_STRING_FUNCTION_CALL],
89 id="simple_streaming",
90 ),
91 pytest.param(
92 False, SIMPLE_FUNCTION_OUTPUT, [SIMPLE_FUNCTION_CALL], id="simple_nonstreaming"
93 ),
94 pytest.param(
95 True,
96 MORE_TYPES_FUNCTION_OUTPUT,
97 [MORE_TYPES_FUNCTION_CALL],
98 id="more_types_streaming",
99 ),
100 pytest.param(
101 False,
102 MORE_TYPES_FUNCTION_OUTPUT,
103 [MORE_TYPES_FUNCTION_CALL],
104 id="more_types_nonstreaming",
105 ),
106 pytest.param(
107 True,
108 PARAMETERLESS_FUNCTION_OUTPUT,
109 [PARAMETERLESS_FUNCTION_CALL],
110 id="parameterless_streaming",
111 ),
112 pytest.param(
113 False,
114 PARAMETERLESS_FUNCTION_OUTPUT,
115 [PARAMETERLESS_FUNCTION_CALL],
116 id="parameterless_nonstreaming",
117 ),
118 pytest.param(
119 True,
120 EMPTY_DICT_FUNCTION_OUTPUT,
121 [EMPTY_DICT_FUNCTION_CALL],
122 id="empty_dict_streaming",
123 ),
124 pytest.param(
125 False,
126 EMPTY_DICT_FUNCTION_OUTPUT,
127 [EMPTY_DICT_FUNCTION_CALL],
128 id="empty_dict_nonstreaming",
129 ),
130 pytest.param(
131 True,
132 EMPTY_LIST_FUNCTION_OUTPUT,
133 [EMPTY_LIST_FUNCTION_CALL],
134 id="empty_list_streaming",
135 ),
136 pytest.param(
137 False,
138 EMPTY_LIST_FUNCTION_OUTPUT,
139 [EMPTY_LIST_FUNCTION_CALL],
140 id="empty_list_nonstreaming",
141 ),
142 pytest.param(
143 True,
144 ESCAPED_STRING_FUNCTION_OUTPUT,
145 [ESCAPED_STRING_FUNCTION_CALL],
146 id="escaped_string_streaming",
147 ),
148 pytest.param(
149 False,
150 ESCAPED_STRING_FUNCTION_OUTPUT,
151 [ESCAPED_STRING_FUNCTION_CALL],
152 id="escaped_string_nonstreaming",
153 ),
154 pytest.param(
155 True,
156 "[get_weather(city='LA',metric='C'),register_user(name='Doe',age=9)]",
157 [
158 SIMPLE_FUNCTION_CALL,
159 FunctionCall(name="register_user", arguments='{"name": "Doe", "age": 9}'),
160 ],
161 id="parallel_calls_streaming",
162 ),
163 pytest.param(
164 False,
165 "[get_weather(city='LA',metric='C'),register_user(name='Doe',age=9)]",
166 [
167 SIMPLE_FUNCTION_CALL,
168 FunctionCall(name="register_user", arguments='{"name": "Doe", "age": 9}'),
169 ],
170 id="parallel_calls_nonstreaming",
171 ),
172 pytest.param(
173 True,
174 PYTHON_TAG_FUNCTION_OUTPUT,
175 [SIMPLE_FUNCTION_CALL],
176 id="python_tag_streaming",
177 ),
178 pytest.param(
179 False,
180 PYTHON_TAG_FUNCTION_OUTPUT,
181 [SIMPLE_FUNCTION_CALL],
182 id="python_tag_nonstreaming",
183 ),
184 pytest.param(
185 True,
186 test_str,
187 [
188 SIMPLE_FUNCTION_CALL,
189 FunctionCall(name="register_user", arguments='{"name": "Doe", "age": 9}'),
190 ],
191 id="parallel_calls_streaming",
192 ),
193 pytest.param(
194 False,
195 "<|python_start|>[get_weather(city='LA', metric='C'), "
196 "register_user(name='Doe', age=9)]",
197 [
198 SIMPLE_FUNCTION_CALL,
199 FunctionCall(name="register_user", arguments='{"name": "Doe", "age": 9}'),
200 ],
201 id="parallel_calls_nonstreaming",
202 ),
203 ]
204
205
206 @pytest.mark.parametrize("streaming, model_output, expected_tool_calls", TEST_CASES)
207 def test_tool_call(
208 streaming: bool,
209 model_output: str,
210 expected_tool_calls: list[FunctionCall],
211 default_tokenizer: TokenizerLike,
212 ):
213 tool_parser: ToolParser = ToolParserManager.get_tool_parser("llama4_pythonic")(
214 default_tokenizer
215 )
216
217 content, tool_calls = run_tool_extraction(
218 tool_parser, model_output, streaming=streaming
219 )
220
221 assert len(tool_calls) == len(expected_tool_calls)
222 for actual, expected in zip(tool_calls, expected_tool_calls):
223 assert actual.type == "function"
224 assert actual.function == expected
225
226
227 def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike):
228 tool_parser: ToolParser = ToolParserManager.get_tool_parser("llama4_pythonic")(
229 default_tokenizer
230 )
231 model_output_deltas = [
232 "<|python_start|>[get_weather(city='LA', metric='C'), "
233 "get_weather(), "
234 "do_something_cool(steps=[])]<|python_end|>",
235 ]
236
237 reconstructor = run_tool_extraction_streaming(
238 tool_parser, model_output_deltas, assert_one_tool_per_delta=False
239 )
240
241 assert reconstructor.other_content == ""
242 assert len(reconstructor.tool_calls) == 3
243 assert reconstructor.tool_calls[0].function == SIMPLE_FUNCTION_CALL
244 assert reconstructor.tool_calls[1].function == PARAMETERLESS_FUNCTION_CALL
245 assert reconstructor.tool_calls[2].function == EMPTY_LIST_FUNCTION_CALL
246
247
248 @pytest.mark.parametrize("streaming", [False])
249 def test_regex_timeout_handling(streaming: bool, default_tokenizer: TokenizerLike):
250 """Test regex timeout is handled gracefully."""
251 tool_parser: ToolParser = ToolParserManager.get_tool_parser("llama4_pythonic")(
252 default_tokenizer
253 )
254
255 fake_problematic_input = "hello world[A(A=" + "\t)A(A=,\t" * 2
256
257 # create a mock regex that raises TimeoutError
258 mock_regex = MagicMock()
259 mock_regex.match.side_effect = TimeoutError("Regex timeout")
260
261 with patch.object(tool_parser, "TOOL_CALL_REGEX", mock_regex):
262 content, tool_calls = run_tool_extraction(
263 tool_parser, fake_problematic_input, streaming=streaming
264 )
265
266 # should treat as regular text when regex times out
267 assert content == fake_problematic_input
268 assert len(tool_calls) == 0
269 mock_regex.match.assert_called_once()