| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | from unittest.mock import MagicMock |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 10 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 11 | from vllm.parser.parser_manager import ParserManager |
| 12 | from vllm.tool_parsers import ToolParserManager |
| 13 | from vllm.tool_parsers.k2_horizon_tool_parser import K2HorizonToolParser |
| 14 | |
| 15 | pytestmark = pytest.mark.skip_global_cleanup |
| 16 | |
| 17 | TOOLS = [ |
| 18 | { |
| 19 | "type": "function", |
| 20 | "function": { |
| 21 | "name": "lookup", |
| 22 | "parameters": { |
| 23 | "type": "object", |
| 24 | "properties": { |
| 25 | "query": {"type": "string"}, |
| 26 | "limit": {"type": "integer"}, |
| 27 | "enabled": {"type": "boolean"}, |
| 28 | }, |
| 29 | }, |
| 30 | }, |
| 31 | }, |
| 32 | { |
| 33 | "type": "function", |
| 34 | "function": { |
| 35 | "name": "ping", |
| 36 | "parameters": {"type": "object", "properties": {}}, |
| 37 | }, |
| 38 | }, |
| 39 | ] |
| 40 | |
| 41 | |
| 42 | @pytest.fixture |
| 43 | def tokenizer(): |
| 44 | tokenizer = MagicMock() |
| 45 | tokenizer.get_vocab.return_value = {} |
| 46 | return tokenizer |
| 47 | |
| 48 | |
| 49 | def _request( |
| 50 | tool_format: str | None = None, |
| 51 | tool_choice="auto", |
| 52 | ) -> ChatCompletionRequest: |
| 53 | chat_template_kwargs = ( |
| 54 | {"tool_call_format": tool_format} if tool_format is not None else None |
| 55 | ) |
| 56 | return ChatCompletionRequest( |
| 57 | model="k2-horizon", |
| 58 | messages=[{"role": "user", "content": "test"}], |
| 59 | tools=TOOLS, |
| 60 | tool_choice=tool_choice, |
| 61 | chat_template_kwargs=chat_template_kwargs, |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | def _group(*calls: str, prefix: str = "", suffix: str = "") -> str: |
| 66 | return prefix + "<ifm|tool_calls>" + "\n".join(calls) + "</ifm|tool_calls>" + suffix |
| 67 | |
| 68 | |
| 69 | def _json_call(name: str, arguments: dict) -> str: |
| 70 | payload = json.dumps({"name": name, "arguments": arguments}) |
| 71 | return f"<ifm|tool_call>{payload}</ifm|tool_call>" |
| 72 | |
| 73 | |
| 74 | def _xml_call( |
| 75 | name: str, |
| 76 | arguments: list[tuple[str, str]], |
| 77 | *, |
| 78 | typed: bool = False, |
| 79 | ) -> str: |
| 80 | parts = [f"<ifm|tool_call>{name}"] |
| 81 | for key, value in arguments: |
| 82 | parts.append(f"<ifm|arg_key>{key}</ifm|arg_key>") |
| 83 | if typed: |
| 84 | arg_type = "integer" if key == "limit" else "string" |
| 85 | parts.append(f"<ifm|arg_type>{arg_type}</ifm|arg_type>") |
| 86 | parts.append(f"<ifm|arg_value>{value}</ifm|arg_value>") |
| 87 | parts.append("</ifm|tool_call>") |
| 88 | return "".join(parts) |
| 89 | |
| 90 | |
| 91 | def _collect_stream( |
| 92 | parser: K2HorizonToolParser, |
| 93 | request: ChatCompletionRequest, |
| 94 | output: str, |
| 95 | ) -> tuple[str, list[tuple[str, dict]]]: |
| 96 | messages: list[DeltaMessage | None] = [] |
| 97 | for char in output: |
| 98 | messages.append( |
| 99 | parser.extract_tool_calls_streaming( |
| 100 | previous_text="", |
| 101 | current_text="", |
| 102 | delta_text=char, |
| 103 | previous_token_ids=[], |
| 104 | current_token_ids=[], |
| 105 | delta_token_ids=[], |
| 106 | request=request, |
| 107 | ) |
| 108 | ) |
| 109 | messages.append(parser.finish_streaming()) |
| 110 | |
| 111 | content: list[str] = [] |
| 112 | calls: list[tuple[str, dict]] = [] |
| 113 | for message in messages: |
| 114 | if message is None: |
| 115 | continue |
| 116 | if message.content is not None: |
| 117 | content.append(message.content) |
| 118 | for tool_call in message.tool_calls: |
| 119 | assert tool_call.function is not None |
| 120 | assert tool_call.function.name is not None |
| 121 | assert tool_call.function.arguments is not None |
| 122 | calls.append( |
| 123 | ( |
| 124 | tool_call.function.name, |
| 125 | json.loads(tool_call.function.arguments), |
| 126 | ) |
| 127 | ) |
| 128 | return "".join(content), calls |
| 129 | |
| 130 | |
| 131 | def test_xml_default_parses_multiple_calls_and_coerces_schema(tokenizer): |
| 132 | parser = K2HorizonToolParser(tokenizer, _request().tools) |
| 133 | output = _group( |
| 134 | _xml_call("lookup", [("query", "123"), ("limit", "3")]), |
| 135 | _xml_call("ping", []), |
| 136 | prefix="Before ", |
| 137 | suffix=" after.", |
| 138 | ) |
| 139 | |
| 140 | result = parser.extract_tool_calls(output, _request()) |
| 141 | |
| 142 | assert parser.tool_format == "xml" |
| 143 | assert result.tools_called |
| 144 | assert result.content == "Before after." |
| 145 | assert [call.function.name for call in result.tool_calls] == ["lookup", "ping"] |
| 146 | assert json.loads(result.tool_calls[0].function.arguments) == { |
| 147 | "query": "123", |
| 148 | "limit": 3, |
| 149 | } |
| 150 | |
| 151 | |
| 152 | @pytest.mark.parametrize( |
| 153 | ("tool_format", "call", "expected"), |
| 154 | [ |
| 155 | ("json", _json_call("lookup", {"limit": "2"}), {"limit": 2}), |
| 156 | ( |
| 157 | "xml_typed", |
| 158 | _xml_call("lookup", [("query", "weather"), ("limit", "2")], typed=True), |
| 159 | {"query": "weather", "limit": 2}, |
| 160 | ), |
| 161 | ], |
| 162 | ) |
| 163 | def test_request_level_format_overrides(tokenizer, tool_format, call, expected): |
| 164 | request = _request(tool_format) |
| 165 | parser = K2HorizonToolParser(tokenizer, request.tools) |
| 166 | |
| 167 | result = parser.extract_tool_calls(_group(call), request) |
| 168 | |
| 169 | assert result.tools_called |
| 170 | assert json.loads(result.tool_calls[0].function.arguments) == expected |
| 171 | |
| 172 | |
| 173 | @pytest.mark.parametrize( |
| 174 | "tool_format", |
| 175 | ["yaml", "python", "xml_untyped", "", None, 1], |
| 176 | ) |
| 177 | def test_invalid_request_format_rejected_before_generation(tokenizer, tool_format): |
| 178 | parser = K2HorizonToolParser(tokenizer, _request().tools) |
| 179 | request = _request().model_copy( |
| 180 | update={"chat_template_kwargs": {"tool_call_format": tool_format}} |
| 181 | ) |
| 182 | |
| 183 | with pytest.raises(ValueError, match="Unsupported tool_call_format"): |
| 184 | parser.adjust_request(request) |
| 185 | |
| 186 | |
| 187 | @pytest.mark.parametrize( |
| 188 | ("tool_request", "call"), |
| 189 | [ |
| 190 | (_request(), _json_call("ping", {})), |
| 191 | (_request("json"), _xml_call("ping", [])), |
| 192 | (_request("xml"), _xml_call("lookup", [("query", "x")], typed=True)), |
| 193 | (_request("xml_typed"), _xml_call("lookup", [("query", "x")])), |
| 194 | ], |
| 195 | ) |
| 196 | def test_mismatched_formats_fail_closed(tokenizer, tool_request, call): |
| 197 | parser = K2HorizonToolParser(tokenizer, tool_request.tools) |
| 198 | output = _group(call) |
| 199 | |
| 200 | result = parser.extract_tool_calls(output, tool_request) |
| 201 | |
| 202 | assert not result.tools_called |
| 203 | assert result.content == output |
| 204 | |
| 205 | |
| 206 | @pytest.mark.parametrize( |
| 207 | "tool_choice", |
| 208 | [ |
| 209 | "auto", |
| 210 | "required", |
| 211 | {"type": "function", "function": {"name": "lookup"}}, |
| 212 | ], |
| 213 | ) |
| 214 | def test_auto_required_and_named_use_ifm_parser(tokenizer, tool_choice): |
| 215 | request = _request("json", tool_choice) |
| 216 | parser_cls = ParserManager.get_parser( |
| 217 | tool_parser_name="k2_horizon", |
| 218 | enable_auto_tools=True, |
| 219 | ) |
| 220 | assert parser_cls is not None |
| 221 | parser = parser_cls(tokenizer, request.tools) |
| 222 | |
| 223 | assert parser.adjust_request(request) is request |
| 224 | assert request.structured_outputs is None |
| 225 | _, content, tool_calls = parser.parse( |
| 226 | _group(_json_call("lookup", {"enabled": "true"})), |
| 227 | request, |
| 228 | enable_auto_tools=True, |
| 229 | ) |
| 230 | |
| 231 | assert content is None |
| 232 | assert tool_calls is not None |
| 233 | assert json.loads(tool_calls[0].arguments) == {"enabled": True} |
| 234 | |
| 235 | |
| 236 | def test_unknown_or_wrong_named_tool_fails_closed(tokenizer): |
| 237 | named_request = _request( |
| 238 | "json", {"type": "function", "function": {"name": "lookup"}} |
| 239 | ) |
| 240 | parser = K2HorizonToolParser(tokenizer, named_request.tools) |
| 241 | |
| 242 | unknown = parser.extract_tool_calls( |
| 243 | _group(_json_call("missing", {})), _request("json") |
| 244 | ) |
| 245 | wrong_named = parser.extract_tool_calls( |
| 246 | _group(_json_call("ping", {})), named_request |
| 247 | ) |
| 248 | |
| 249 | assert not unknown.tools_called |
| 250 | assert not wrong_named.tools_called |
| 251 | |
| 252 | |
| 253 | @pytest.mark.parametrize( |
| 254 | "output", |
| 255 | [ |
| 256 | "<ifm|tool_calls><ifm|tool_call>ping", |
| 257 | "<ifm|tool_calls><ifm|tool_call>ping</ifm|tool_call>", |
| 258 | _group("junk"), |
| 259 | ], |
| 260 | ) |
| 261 | def test_malformed_and_incomplete_input_is_content(tokenizer, output): |
| 262 | request = _request() |
| 263 | parser = K2HorizonToolParser(tokenizer, request.tools) |
| 264 | |
| 265 | result = parser.extract_tool_calls(output, request) |
| 266 | |
| 267 | assert not result.tools_called |
| 268 | assert result.content == output |
| 269 | |
| 270 | |
| 271 | @pytest.mark.parametrize( |
| 272 | ("tool_request", "call", "expected"), |
| 273 | [ |
| 274 | (_request(), _xml_call("lookup", [("limit", "2")]), {"limit": 2}), |
| 275 | (_request("json"), _json_call("lookup", {"limit": "2"}), {"limit": 2}), |
| 276 | ( |
| 277 | _request("xml_typed"), |
| 278 | _xml_call("lookup", [("limit", "2")], typed=True), |
| 279 | {"limit": 2}, |
| 280 | ), |
| 281 | ], |
| 282 | ) |
| 283 | def test_character_split_streaming(tokenizer, tool_request, call, expected): |
| 284 | parser = K2HorizonToolParser(tokenizer, tool_request.tools) |
| 285 | output = _group(call, prefix="Before", suffix="after") |
| 286 | |
| 287 | content, calls = _collect_stream(parser, tool_request, output) |
| 288 | |
| 289 | assert content == "Beforeafter" |
| 290 | assert calls == [("lookup", expected)] |
| 291 | |
| 292 | |
| 293 | def test_streaming_flushes_incomplete_markup_as_content(tokenizer): |
| 294 | request = _request() |
| 295 | parser = K2HorizonToolParser(tokenizer, request.tools) |
| 296 | output = "Before<ifm|tool_calls><ifm|tool_call>ping" |
| 297 | |
| 298 | content, calls = _collect_stream(parser, request, output) |
| 299 | |
| 300 | assert content == output |
| 301 | assert calls == [] |
| 302 | |
| 303 | |
| 304 | def test_whitespace_only_surrounding_content_is_not_preserved(tokenizer): |
| 305 | request = _request() |
| 306 | parser = K2HorizonToolParser(tokenizer, request.tools) |
| 307 | |
| 308 | result = parser.extract_tool_calls( |
| 309 | _group(_xml_call("ping", []), prefix=" \n", suffix="\t "), request |
| 310 | ) |
| 311 | |
| 312 | assert result.tools_called |
| 313 | assert result.content is None |
| 314 | |
| 315 | |
| 316 | def test_tool_parser_registered(): |
| 317 | assert ToolParserManager.get_tool_parser("k2_horizon") is K2HorizonToolParser |
| 318 | assert K2HorizonToolParser.supports_required_and_named is False |