@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 191 lines 6.35 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # ruff: noqa: E501
4
5 import json
6 from unittest.mock import MagicMock
7
8 import pytest
9
10 from tests.tool_parsers.utils import (
11 run_tool_extraction,
12 run_tool_extraction_streaming,
13 )
14 from vllm.entrypoints.generate.base.protocol import FunctionCall, ToolCall
15 from vllm.tool_parsers import ToolParser, ToolParserManager
16
17
18 def make_tool_call(name, arguments):
19 return ToolCall(
20 type="function",
21 function=FunctionCall(name=name, arguments=json.dumps(arguments)),
22 )
23
24
25 # TODO: add reason prefix and suffix.
26
27
28 @pytest.mark.parametrize(
29 "model_output,expected_tool_calls,expected_content",
30 [
31 # No tool call
32 ("How can I help you today?", [], "How can I help you today?"),
33 # Single tool call, no content
34 (
35 '<tool_calls>[{"name": "get_weather", "arguments": {"city": "San Francisco", "metric": "celsius"}}]</tool_calls>', # noqa: E501
36 [
37 make_tool_call(
38 "get_weather", {"city": "San Francisco", "metric": "celsius"}
39 )
40 ],
41 None,
42 ),
43 # Multiple tool calls
44 (
45 '<tool_calls>[{"name": "get_weather", "arguments": {"city": "San Francisco", "metric": "celsius"}}, {"name": "register_user", "arguments": {"name": "John Doe", "age": 37, "address": {"city": "San Francisco", "state": "CA"}, "role": null, "passed_test": true, "aliases": ["John", "Johnny"]}}]</tool_calls>', # noqa: E501
46 [
47 make_tool_call(
48 "get_weather", {"city": "San Francisco", "metric": "celsius"}
49 ),
50 make_tool_call(
51 "register_user",
52 {
53 "name": "John Doe",
54 "age": 37,
55 "address": {"city": "San Francisco", "state": "CA"},
56 "role": None,
57 "passed_test": True,
58 "aliases": ["John", "Johnny"],
59 },
60 ),
61 ],
62 None,
63 ),
64 # Content before tool call
65 (
66 'I will call the tool now. <tool_calls>[{"name": "get_weather", "arguments": {"city": "Boston"}}]</tool_calls>', # noqa: E501
67 [make_tool_call("get_weather", {"city": "Boston"})],
68 "I will call the tool now. ",
69 ),
70 # Content after tool call (should be stripped)
71 (
72 '<tool_calls>[{"name": "get_weather", "arguments": {"city": "Seattle"}}]</tool_calls>\nThank you!', # noqa: E501
73 [make_tool_call("get_weather", {"city": "Seattle"})],
74 None,
75 ),
76 (
77 '<tool_calls>[{"name": "complex_tool", "arguments": {"level1": {"level2": {"level3": {"value": 123}}}}}]</tool_calls>',
78 [
79 make_tool_call(
80 "complex_tool", {"level1": {"level2": {"level3": {"value": 123}}}}
81 )
82 ],
83 None,
84 ),
85 ],
86 )
87 def test_hunyuan_a13b_tool_parser_extract(
88 model_output, expected_tool_calls, expected_content
89 ):
90 mock_tokenizer = MagicMock()
91 tool_parser: ToolParser = ToolParserManager.get_tool_parser("hunyuan_a13b")(
92 mock_tokenizer
93 )
94 content, tool_calls = run_tool_extraction(
95 tool_parser, model_output, streaming=False
96 )
97
98 # align the random id.
99 for idx in range(len(tool_calls)):
100 tool_calls[idx].id = expected_tool_calls[idx].id
101 assert tool_calls == expected_tool_calls
102 assert content == expected_content
103
104
105 # Streaming test: simulate incremental output
106 @pytest.mark.parametrize(
107 "model_deltas,expected_tool_calls",
108 [
109 (
110 [
111 '<tool_calls>[{"name": "get_weather", ',
112 '"arguments": {"city": "San Francisco", ',
113 '"metric": "celsius"}}]',
114 "</tool_calls>",
115 ],
116 [
117 make_tool_call(
118 "get_weather", {"city": "San Francisco", "metric": "celsius"}
119 )
120 ],
121 ),
122 (
123 [
124 '<tool_calls>[{"name":',
125 ' "get_weather",',
126 ' "arguments":',
127 ' {"city": "Boston"}',
128 "}]",
129 "</tool_calls>",
130 ],
131 [make_tool_call("get_weather", {"city": "Boston"})],
132 ),
133 (
134 [
135 "",
136 '<tool_calls>[{"name":',
137 ' "get_weather",',
138 ' "arguments":',
139 ' {"city": "Boston"}',
140 "}]",
141 "</tool_calls>",
142 "\n</answer>",
143 ],
144 [make_tool_call("get_weather", {"city": "Boston"})],
145 ),
146 pytest.param(
147 [
148 '<tool_calls>[{"name": "complex_tool",',
149 ' "arguments": ',
150 ' {"level1": {"level2": ',
151 '{"level3": {"value": 123}}}}}',
152 "]</tool_calls>",
153 ],
154 [
155 make_tool_call(
156 "complex_tool", {"level1": {"level2": {"level3": {"value": 123}}}}
157 )
158 ],
159 marks=pytest.mark.xfail(
160 reason="stream parsing not support nested json yet."
161 ),
162 ),
163 ],
164 )
165 def test_hunyuan_a13b_tool_parser_streaming(model_deltas, expected_tool_calls):
166 mock_tokenizer = MagicMock()
167
168 tool_parser: ToolParser = ToolParserManager.get_tool_parser("hunyuan_a13b")(
169 mock_tokenizer
170 )
171 reconstructor = run_tool_extraction_streaming(
172 tool_parser, model_deltas, assert_one_tool_per_delta=False
173 )
174
175 # align the random id.
176 for idx in range(len(reconstructor.tool_calls)):
177 reconstructor.tool_calls[idx].id = expected_tool_calls[idx].id
178
179 assert reconstructor.tool_calls == expected_tool_calls
180
181
182 def test_hunyuan_a13b_tool_parser_non_ascii():
183 mock_tokenizer = MagicMock()
184 tool_parser: ToolParser = ToolParserManager.get_tool_parser("hunyuan_a13b")(
185 mock_tokenizer
186 )
187 model_output = '<tool_calls>[{"name": "get_weather", "arguments": {"city": "北京"}}]</tool_calls>'
188 _, tool_calls = run_tool_extraction(tool_parser, model_output, streaming=False)
189 args = tool_calls[0].function.arguments
190 assert "北京" in args
191 assert "\\u" not in args