| 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 | from transformers import AutoTokenizer |
| 8 | |
| 9 | from vllm.entrypoints.generate.base.protocol import ExtractedToolCallInformation |
| 10 | from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser |
| 11 | |
| 12 | LLAMA_MODEL = "meta-llama/Llama-3.2-1B-Instruct" |
| 13 | |
| 14 | |
| 15 | @pytest.fixture(scope="module") |
| 16 | def llama_tokenizer(): |
| 17 | return AutoTokenizer.from_pretrained(LLAMA_MODEL) |
| 18 | |
| 19 | |
| 20 | @pytest.fixture |
| 21 | def parser(llama_tokenizer): |
| 22 | return Llama3JsonToolParser(llama_tokenizer) |
| 23 | |
| 24 | |
| 25 | def test_extract_tool_calls_simple(parser): |
| 26 | # Test with a simple tool call |
| 27 | model_output = ( |
| 28 | 'Here is the result: {"name": "getOpenIncidentsTool", ' |
| 29 | '"parameters": {}} Would you like to know more?' |
| 30 | ) |
| 31 | result = parser.extract_tool_calls(model_output, None) |
| 32 | |
| 33 | assert isinstance(result, ExtractedToolCallInformation) |
| 34 | assert result.tools_called is True |
| 35 | assert len(result.tool_calls) == 1 |
| 36 | assert result.tool_calls[0].type == "function" |
| 37 | assert result.tool_calls[0].function.name == "getOpenIncidentsTool" |
| 38 | assert result.tool_calls[0].function.arguments == "{}" |
| 39 | assert result.content is None |
| 40 | |
| 41 | |
| 42 | def test_extract_tool_calls_with_arguments(parser): |
| 43 | # Test with a tool call that has arguments |
| 44 | model_output = ( |
| 45 | '{"name": "searchTool", "parameters": {"query": "test query", "limit": 10}}' |
| 46 | ) |
| 47 | result = parser.extract_tool_calls(model_output, None) |
| 48 | |
| 49 | assert result.tools_called is True |
| 50 | assert len(result.tool_calls) == 1 |
| 51 | assert result.tool_calls[0].function.name == "searchTool" |
| 52 | assert '"query": "test query"' in result.tool_calls[0].function.arguments |
| 53 | assert '"limit": 10' in result.tool_calls[0].function.arguments |
| 54 | |
| 55 | |
| 56 | def test_extract_tool_calls_no_json(parser): |
| 57 | # Test with text that doesn't contain a JSON object |
| 58 | model_output = "This is just some text without any tool calls" |
| 59 | result = parser.extract_tool_calls(model_output, None) |
| 60 | |
| 61 | assert result.tools_called is False |
| 62 | assert len(result.tool_calls) == 0 |
| 63 | assert result.content == model_output |
| 64 | |
| 65 | |
| 66 | def test_extract_tool_calls_invalid_json(parser): |
| 67 | # Test with invalid JSON |
| 68 | model_output = '{"name": "invalidTool", "parameters": {invalid json}' |
| 69 | result = parser.extract_tool_calls(model_output, None) |
| 70 | |
| 71 | assert result.tools_called is False |
| 72 | assert len(result.tool_calls) == 0 |
| 73 | assert result.content == model_output |
| 74 | |
| 75 | |
| 76 | def test_extract_tool_calls_with_arguments_key(parser): |
| 77 | # Test with a tool call that uses "arguments" instead of "parameters" |
| 78 | model_output = '{"name": "searchTool", "arguments": {"query": "test"}}' |
| 79 | result = parser.extract_tool_calls(model_output, None) |
| 80 | |
| 81 | assert result.tools_called is True |
| 82 | assert len(result.tool_calls) == 1 |
| 83 | assert result.tool_calls[0].function.name == "searchTool" |
| 84 | assert '"query": "test"' in result.tool_calls[0].function.arguments |
| 85 | |
| 86 | |
| 87 | def test_extract_tool_calls_multiple_json(parser): |
| 88 | # Test with multiple JSONs separated by semicolons |
| 89 | model_output = ( |
| 90 | '{"name": "searchTool", "parameters": {"query": "test1"}}; ' |
| 91 | '{"name": "getOpenIncidentsTool", "parameters": {}}; ' |
| 92 | '{"name": "searchTool", "parameters": {"query": "test2"}}' |
| 93 | ) |
| 94 | result = parser.extract_tool_calls(model_output, None) |
| 95 | |
| 96 | assert result.tools_called is True |
| 97 | assert len(result.tool_calls) == 3 |
| 98 | |
| 99 | # Check first tool call |
| 100 | assert result.tool_calls[0].function.name == "searchTool" |
| 101 | assert '"query": "test1"' in result.tool_calls[0].function.arguments |
| 102 | |
| 103 | # Check second tool call |
| 104 | assert result.tool_calls[1].function.name == "getOpenIncidentsTool" |
| 105 | assert result.tool_calls[1].function.arguments == "{}" |
| 106 | |
| 107 | # Check third tool call |
| 108 | assert result.tool_calls[2].function.name == "searchTool" |
| 109 | assert '"query": "test2"' in result.tool_calls[2].function.arguments |
| 110 | |
| 111 | |
| 112 | def test_extract_tool_calls_multiple_json_with_whitespace(parser): |
| 113 | # Test with multiple JSONs separated by semicolons and extra whitespace |
| 114 | model_output = ( |
| 115 | '{"name": "searchTool", "parameters": {"query": "test1"}} ; ' |
| 116 | '{"name": "getOpenIncidentsTool", "parameters": {}} ; ' |
| 117 | '{"name": "searchTool", "parameters": {"query": "test2"}}' |
| 118 | ) |
| 119 | result = parser.extract_tool_calls(model_output, None) |
| 120 | |
| 121 | assert result.tools_called is True |
| 122 | assert len(result.tool_calls) == 3 |
| 123 | assert result.tool_calls[0].function.name == "searchTool" |
| 124 | assert result.tool_calls[1].function.name == "getOpenIncidentsTool" |
| 125 | assert result.tool_calls[2].function.name == "searchTool" |
| 126 | |
| 127 | |
| 128 | def test_extract_tool_calls_multiple_json_with_surrounding_text(parser): |
| 129 | # Test with multiple JSONs and surrounding text |
| 130 | model_output = ( |
| 131 | "Here are the results: " |
| 132 | '{"name": "searchTool", "parameters": {"query": "test1"}}; ' |
| 133 | '{"name": "getOpenIncidentsTool", "parameters": {}}; ' |
| 134 | '{"name": "searchTool", "parameters": {"query": "test2"}} ' |
| 135 | "Would you like to know more?" |
| 136 | ) |
| 137 | result = parser.extract_tool_calls(model_output, None) |
| 138 | |
| 139 | assert result.tools_called is True |
| 140 | assert len(result.tool_calls) == 3 |
| 141 | assert result.tool_calls[0].function.name == "searchTool" |
| 142 | assert result.tool_calls[1].function.name == "getOpenIncidentsTool" |
| 143 | assert result.tool_calls[2].function.name == "searchTool" |
| 144 | |
| 145 | |
| 146 | def test_extract_tool_calls_deeply_nested_json(parser): |
| 147 | # Test with deeply nested JSON parameters (5 levels) |
| 148 | model_output = ( |
| 149 | '{"name": "complexTool", ' |
| 150 | '"parameters": {' |
| 151 | '"level1": {' |
| 152 | '"level2": {' |
| 153 | '"level3": {' |
| 154 | '"level4": {' |
| 155 | '"value": "deep"' |
| 156 | "}}}}}}" |
| 157 | ) |
| 158 | result = parser.extract_tool_calls(model_output, None) |
| 159 | |
| 160 | assert result.tools_called is True |
| 161 | assert len(result.tool_calls) == 1 |
| 162 | assert result.tool_calls[0].function.name == "complexTool" |
| 163 | # Verify the nested structure is preserved in the arguments |
| 164 | import json |
| 165 | |
| 166 | args = json.loads(result.tool_calls[0].function.arguments) |
| 167 | assert args["level1"]["level2"]["level3"]["level4"]["value"] == "deep" |
| 168 | |
| 169 | |
| 170 | def test_extract_tool_calls_multiple_with_deep_nesting(parser): |
| 171 | # Test with multiple tool calls where some have deeply nested parameters |
| 172 | model_output = ( |
| 173 | '{"name": "simpleTool", "parameters": {"value": "test"}}; ' |
| 174 | '{"name": "complexTool", "parameters": ' |
| 175 | '{"config": {"database": {"connection": {"pool": {"size": 10}}}}}}' |
| 176 | ) |
| 177 | result = parser.extract_tool_calls(model_output, None) |
| 178 | |
| 179 | assert result.tools_called is True |
| 180 | assert len(result.tool_calls) == 2 |
| 181 | |
| 182 | # Check first tool call |
| 183 | assert result.tool_calls[0].function.name == "simpleTool" |
| 184 | import json |
| 185 | |
| 186 | args0 = json.loads(result.tool_calls[0].function.arguments) |
| 187 | assert args0["value"] == "test" |
| 188 | |
| 189 | # Check second tool call with deep nesting |
| 190 | assert result.tool_calls[1].function.name == "complexTool" |
| 191 | args1 = json.loads(result.tool_calls[1].function.arguments) |
| 192 | assert args1["config"]["database"]["connection"]["pool"]["size"] == 10 |
| 193 | |
| 194 | |
| 195 | def test_extract_tool_calls_with_quotes_and_brackets_in_string(parser): |
| 196 | # Test with quotes and brackets inside quoted string values |
| 197 | model_output = ( |
| 198 | '{"name": "searchTool", ' |
| 199 | '"parameters": {' |
| 200 | '"query": "test {value} [complex]",' |
| 201 | '"nested": {"inner": "more {brackets}"}' |
| 202 | "}}" |
| 203 | ) |
| 204 | result = parser.extract_tool_calls(model_output, None) |
| 205 | |
| 206 | assert result.tools_called is True |
| 207 | assert len(result.tool_calls) == 1 |
| 208 | assert result.tool_calls[0].function.name == "searchTool" |
| 209 | # Verify the string values are preserved including brackets and quotes |
| 210 | import json |
| 211 | |
| 212 | args = json.loads(result.tool_calls[0].function.arguments) |
| 213 | assert args["query"] == "test {value} [complex]" |
| 214 | assert args["nested"]["inner"] == "more {brackets}" |
| 215 | |
| 216 | |
| 217 | def test_extract_tool_calls_with_escaped_quotes_in_nested_json(parser): |
| 218 | # Test with escaped quotes in deeply nested JSON |
| 219 | model_output = ( |
| 220 | '{"name": "parserTool", "parameters": {"text": "He said \\"Hello {world}\\""}}' |
| 221 | ) |
| 222 | result = parser.extract_tool_calls(model_output, None) |
| 223 | |
| 224 | assert result.tools_called is True |
| 225 | assert len(result.tool_calls) == 1 |
| 226 | assert result.tool_calls[0].function.name == "parserTool" |
| 227 | # Verify escaped quotes are preserved |
| 228 | import json |
| 229 | |
| 230 | args = json.loads(result.tool_calls[0].function.arguments) |
| 231 | assert args["text"] == 'He said "Hello {world}"' |
| 232 | |
| 233 | |
| 234 | def test_extract_tool_calls_missing_name_key(parser): |
| 235 | # Test that missing "name" key returns content |
| 236 | model_output = '{"parameters": {}}' |
| 237 | result = parser.extract_tool_calls(model_output, None) |
| 238 | |
| 239 | assert result.tools_called is False |
| 240 | assert len(result.tool_calls) == 0 |
| 241 | assert result.content == model_output |
| 242 | |
| 243 | |
| 244 | def test_extract_tool_calls_missing_parameters_and_arguments_key(parser): |
| 245 | # Test that missing both "parameters" and "arguments" keys returns content |
| 246 | model_output = '{"name": "toolWithoutParams"}' |
| 247 | result = parser.extract_tool_calls(model_output, None) |
| 248 | |
| 249 | assert result.tools_called is False |
| 250 | assert len(result.tool_calls) == 0 |
| 251 | assert result.content == model_output |
| 252 | |
| 253 | |
| 254 | def test_regex_timeout_handling(parser): |
| 255 | """Test regex timeout is handled gracefully.""" |
| 256 | fake_problematic_input = "{hello world[A(A=" + "\t)A(A=,\t" * 2 |
| 257 | |
| 258 | # create a mock regex that raises TimeoutError |
| 259 | mock_regex = MagicMock() |
| 260 | mock_regex.finditer.side_effect = TimeoutError("Regex timeout") |
| 261 | |
| 262 | with patch.object(parser, "tool_call_start_regex", mock_regex): |
| 263 | result = parser.extract_tool_calls(fake_problematic_input, None) |
| 264 | |
| 265 | # should treat as regular text when regex times out |
| 266 | assert result.content == fake_problematic_input |
| 267 | assert result.tools_called is False |
| 268 | assert len(result.tool_calls) == 0 |
| 269 | mock_regex.finditer.assert_called_once() |