@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 155 lines 5.51 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
5
6 import pytest
7
8 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
9 from vllm.tool_parsers.functiongemma_tool_parser import FunctionGemmaToolParser
10
11
12 @pytest.fixture
13 def mock_tokenizer():
14 tokenizer = MagicMock()
15 tokenizer.encode.return_value = [1, 2, 3]
16 tokenizer.get_vocab.return_value = {}
17 return tokenizer
18
19
20 @pytest.fixture
21 def parser(mock_tokenizer):
22 return FunctionGemmaToolParser(mock_tokenizer)
23
24
25 @pytest.fixture
26 def mock_request():
27 request = MagicMock(spec=ChatCompletionRequest)
28 request.tools = []
29 request.tool_choice = "auto"
30 request.parallel_tool_calls = True
31 return request
32
33
34 class TestExtractToolCalls:
35 def test_no_tool_calls(self, parser, mock_request):
36 model_output = "Hello, how can I help you today?"
37 result = parser.extract_tool_calls(model_output, mock_request)
38
39 assert result.tools_called is False
40 assert result.tool_calls == []
41 assert result.content == model_output
42
43 def test_single_tool_call(self, parser, mock_request):
44 model_output = (
45 "<start_function_call>call:get_weather{location:<escape>London<escape>}"
46 "<end_function_call>"
47 )
48 result = parser.extract_tool_calls(model_output, mock_request)
49
50 assert result.tools_called is True
51 assert len(result.tool_calls) == 1
52 assert result.tool_calls[0].function.name == "get_weather"
53 assert '"location": "London"' in result.tool_calls[0].function.arguments
54
55 def test_multiple_arguments(self, parser, mock_request):
56 model_output = (
57 "<start_function_call>call:get_weather{"
58 "location:<escape>San Francisco<escape>,"
59 "unit:<escape>celsius<escape>}"
60 "<end_function_call>"
61 )
62 result = parser.extract_tool_calls(model_output, mock_request)
63
64 assert result.tools_called is True
65 assert len(result.tool_calls) == 1
66 assert result.tool_calls[0].function.name == "get_weather"
67 args = result.tool_calls[0].function.arguments
68 assert "San Francisco" in args
69 assert "celsius" in args
70
71 def test_text_before_tool_call(self, parser, mock_request):
72 model_output = (
73 "Let me check the weather for you. "
74 "<start_function_call>call:get_weather{location:<escape>Paris<escape>}"
75 "<end_function_call>"
76 )
77 result = parser.extract_tool_calls(model_output, mock_request)
78
79 assert result.tools_called is True
80 assert result.content == "Let me check the weather for you."
81
82 def test_multiple_tool_calls(self, parser, mock_request):
83 model_output = (
84 "<start_function_call>call:get_weather{location:<escape>London<escape>}"
85 "<end_function_call>"
86 "<start_function_call>call:get_time{timezone:<escape>UTC<escape>}"
87 "<end_function_call>"
88 )
89 result = parser.extract_tool_calls(model_output, mock_request)
90
91 assert result.tools_called is True
92 assert len(result.tool_calls) == 2
93 assert result.tool_calls[0].function.name == "get_weather"
94 assert result.tool_calls[1].function.name == "get_time"
95
96
97 class TestParseArguments:
98 def test_empty_arguments(self, parser):
99 result = parser._parse_arguments("")
100 assert result == {}
101
102 def test_single_string_argument(self, parser):
103 result = parser._parse_arguments("city:<escape>Tokyo<escape>")
104 assert result == {"city": "Tokyo"}
105
106 def test_multiple_arguments(self, parser):
107 args_str = "city:<escape>Tokyo<escape>,country:<escape>Japan<escape>"
108 result = parser._parse_arguments(args_str)
109 assert result == {"city": "Tokyo", "country": "Japan"}
110
111 def test_numeric_argument(self, parser):
112 result = parser._parse_arguments("count:<escape>42<escape>")
113 assert result == {"count": 42}
114
115 def test_boolean_argument(self, parser):
116 result = parser._parse_arguments("enabled:<escape>true<escape>")
117 assert result == {"enabled": True}
118
119 def test_argument_with_spaces(self, parser):
120 result = parser._parse_arguments("message:<escape>Hello World<escape>")
121 assert result == {"message": "Hello World"}
122
123
124 class TestAdjustRequest:
125 def test_skip_special_tokens_disabled(self, parser, mock_request):
126 mock_request.tools = [{"type": "function", "function": {"name": "test"}}]
127 mock_request.tool_choice = "auto"
128 mock_request.skip_special_tokens = True
129
130 result = parser.adjust_request(mock_request)
131 assert result.skip_special_tokens is False
132
133 def test_skip_special_tokens_when_tool_choice_none(self, parser, mock_request):
134 mock_request.tools = [{"type": "function", "function": {"name": "test"}}]
135 mock_request.tool_choice = "none"
136 mock_request.skip_special_tokens = True
137
138 result = parser.adjust_request(mock_request)
139 assert result.skip_special_tokens is True
140
141
142 class TestBufferDeltaText:
143 def test_regular_text_not_buffered(self, parser):
144 result = parser._buffer_delta_text("hello")
145 assert result == "hello"
146 assert parser.buffered_delta_text == ""
147
148 def test_complete_tag_flushed(self, parser):
149 parser.buffered_delta_text = "<start_function_"
150 result = parser._buffer_delta_text("call>")
151 assert "<start_function_call>" in result
152
153
154 if __name__ == "__main__":
155 pytest.main([__file__, "-v"])