| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | |
| 6 | import pytest |
| 7 | from transformers import AutoTokenizer |
| 8 | |
| 9 | from tests.tool_parsers.utils import ( |
| 10 | run_tool_extraction, |
| 11 | run_tool_extraction_streaming, |
| 12 | ) |
| 13 | from vllm.entrypoints.generate.base.protocol import FunctionCall |
| 14 | from vllm.tokenizers import TokenizerLike |
| 15 | from vllm.tool_parsers import ToolParser, ToolParserManager |
| 16 | |
| 17 | |
| 18 | @pytest.fixture(scope="function") |
| 19 | def default_tokenizer() -> TokenizerLike: |
| 20 | """Override module-scoped default_tokenizer because gigachat tests |
| 21 | mutate the tokenizer via ``add_tokens``.""" |
| 22 | return AutoTokenizer.from_pretrained("openai-community/gpt2") |
| 23 | |
| 24 | |
| 25 | MSG_SEP_TOKEN = "<|message_sep|>\n\n" |
| 26 | ROLE_SEP_TOKEN = "<|role_sep|>\n" |
| 27 | EOS_TOKEN = "</s>" |
| 28 | TOOL_HEADER_GIGACHAT3 = f"function call{ROLE_SEP_TOKEN}" |
| 29 | TOOL_HEADER_GIGACHAT31 = "<|function_call|>" |
| 30 | |
| 31 | |
| 32 | SIMPLE_ARGS_DICT = { |
| 33 | "action": "create", |
| 34 | "id": "preferences", |
| 35 | } |
| 36 | SIMPLE_FUNCTION_JSON = json.dumps( |
| 37 | { |
| 38 | "name": "manage_user_memory", |
| 39 | "arguments": SIMPLE_ARGS_DICT, |
| 40 | }, |
| 41 | ensure_ascii=False, |
| 42 | ) |
| 43 | SIMPLE_FUNCTION_OUTPUT_GIGACHAT3 = ( |
| 44 | f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{SIMPLE_FUNCTION_JSON}" |
| 45 | ) |
| 46 | SIMPLE_FUNCTION_OUTPUT_GIGACHAT31 = f"{TOOL_HEADER_GIGACHAT31}{SIMPLE_FUNCTION_JSON}" |
| 47 | SIMPLE_FUNCTION_CALL = FunctionCall( |
| 48 | name="manage_user_memory", |
| 49 | arguments=json.dumps(SIMPLE_ARGS_DICT, ensure_ascii=False), |
| 50 | ) |
| 51 | |
| 52 | |
| 53 | PARAMETERLESS_FUNCTION_JSON = json.dumps( |
| 54 | { |
| 55 | "name": "manage_user_memory", |
| 56 | "arguments": {}, |
| 57 | }, |
| 58 | ensure_ascii=False, |
| 59 | ) |
| 60 | PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3 = ( |
| 61 | f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{PARAMETERLESS_FUNCTION_JSON}" |
| 62 | ) |
| 63 | PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31 = ( |
| 64 | f"{TOOL_HEADER_GIGACHAT31}{PARAMETERLESS_FUNCTION_JSON}" |
| 65 | ) |
| 66 | PARAMETERLESS_FUNCTION_CALL = FunctionCall( |
| 67 | name="manage_user_memory", |
| 68 | arguments=json.dumps({}, ensure_ascii=False), |
| 69 | ) |
| 70 | |
| 71 | |
| 72 | COMPLEX_ARGS_DICT = { |
| 73 | "action": "create", |
| 74 | "id": "preferences", |
| 75 | "content": { |
| 76 | "short_answers": True, |
| 77 | "hate_emojis": True, |
| 78 | "english_ui": False, |
| 79 | "russian_math_explanations": True, |
| 80 | }, |
| 81 | } |
| 82 | COMPLEX_FUNCTION_JSON = json.dumps( |
| 83 | { |
| 84 | "name": "manage_user_memory", |
| 85 | "arguments": COMPLEX_ARGS_DICT, |
| 86 | }, |
| 87 | ensure_ascii=False, |
| 88 | ) |
| 89 | COMPLEX_FUNCTION_OUTPUT_GIGACHAT3 = ( |
| 90 | f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{COMPLEX_FUNCTION_JSON}" |
| 91 | ) |
| 92 | COMPLEX_FUNCTION_OUTPUT_GIGACHAT31 = f"{TOOL_HEADER_GIGACHAT31}{COMPLEX_FUNCTION_JSON}" |
| 93 | COMPLEX_FUNCTION_CALL = FunctionCall( |
| 94 | name="manage_user_memory", |
| 95 | arguments=json.dumps(COMPLEX_ARGS_DICT, ensure_ascii=False), |
| 96 | ) |
| 97 | |
| 98 | |
| 99 | CONTENT_TEXT = "I'll check that for you." |
| 100 | MIXED_OUTPUT_GIGACHAT3 = f"{CONTENT_TEXT}{SIMPLE_FUNCTION_OUTPUT_GIGACHAT3}" |
| 101 | MIXED_OUTPUT_GIGACHAT31 = f"{CONTENT_TEXT}{SIMPLE_FUNCTION_OUTPUT_GIGACHAT31}" |
| 102 | |
| 103 | |
| 104 | @pytest.fixture(name="gigachat_tokenizer") |
| 105 | def fixture_gigachat_tokenizer(default_tokenizer: TokenizerLike): |
| 106 | default_tokenizer.add_tokens( |
| 107 | [ |
| 108 | MSG_SEP_TOKEN, |
| 109 | ROLE_SEP_TOKEN, |
| 110 | TOOL_HEADER_GIGACHAT31, |
| 111 | EOS_TOKEN, |
| 112 | ] |
| 113 | ) |
| 114 | return default_tokenizer |
| 115 | |
| 116 | |
| 117 | @pytest.mark.parametrize("streaming", [True, False]) |
| 118 | def test_no_tool_call(streaming: bool, gigachat_tokenizer: TokenizerLike): |
| 119 | tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")( |
| 120 | gigachat_tokenizer |
| 121 | ) |
| 122 | model_output = "How can I help you today?" |
| 123 | content, tool_calls = run_tool_extraction( |
| 124 | tool_parser, model_output, streaming=streaming |
| 125 | ) |
| 126 | assert content == model_output |
| 127 | assert len(tool_calls) == 0 |
| 128 | |
| 129 | |
| 130 | TEST_CASES = [ |
| 131 | pytest.param( |
| 132 | True, |
| 133 | SIMPLE_FUNCTION_OUTPUT_GIGACHAT3, |
| 134 | [SIMPLE_FUNCTION_CALL], |
| 135 | None, |
| 136 | id="simple_streaming_gigachat3", |
| 137 | ), |
| 138 | pytest.param( |
| 139 | False, |
| 140 | SIMPLE_FUNCTION_OUTPUT_GIGACHAT3, |
| 141 | [SIMPLE_FUNCTION_CALL], |
| 142 | None, |
| 143 | id="simple_nonstreaming_gigachat3", |
| 144 | ), |
| 145 | pytest.param( |
| 146 | True, |
| 147 | PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3, |
| 148 | [PARAMETERLESS_FUNCTION_CALL], |
| 149 | None, |
| 150 | id="parameterless_streaming_gigachat3", |
| 151 | ), |
| 152 | pytest.param( |
| 153 | False, |
| 154 | PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3, |
| 155 | [PARAMETERLESS_FUNCTION_CALL], |
| 156 | None, |
| 157 | id="parameterless_nonstreaming_gigachat3", |
| 158 | ), |
| 159 | pytest.param( |
| 160 | True, |
| 161 | COMPLEX_FUNCTION_OUTPUT_GIGACHAT3, |
| 162 | [COMPLEX_FUNCTION_CALL], |
| 163 | None, |
| 164 | id="complex_streaming_gigachat3", |
| 165 | ), |
| 166 | pytest.param( |
| 167 | False, |
| 168 | COMPLEX_FUNCTION_OUTPUT_GIGACHAT3, |
| 169 | [COMPLEX_FUNCTION_CALL], |
| 170 | None, |
| 171 | id="complex_nonstreaming_gigachat3", |
| 172 | ), |
| 173 | pytest.param( |
| 174 | True, |
| 175 | MIXED_OUTPUT_GIGACHAT3, |
| 176 | [SIMPLE_FUNCTION_CALL], |
| 177 | CONTENT_TEXT, |
| 178 | id="mixed_content_streaming_gigachat3", |
| 179 | ), |
| 180 | pytest.param( |
| 181 | False, |
| 182 | MIXED_OUTPUT_GIGACHAT3, |
| 183 | [SIMPLE_FUNCTION_CALL], |
| 184 | CONTENT_TEXT, |
| 185 | id="mixed_content_nonstreaming_gigachat3", |
| 186 | ), |
| 187 | pytest.param( |
| 188 | True, |
| 189 | MIXED_OUTPUT_GIGACHAT3 + EOS_TOKEN, |
| 190 | [SIMPLE_FUNCTION_CALL], |
| 191 | CONTENT_TEXT, |
| 192 | id="mixed_content_streaming_with_eos_gigachat3", |
| 193 | ), |
| 194 | pytest.param( |
| 195 | False, |
| 196 | MIXED_OUTPUT_GIGACHAT3 + EOS_TOKEN, |
| 197 | [SIMPLE_FUNCTION_CALL], |
| 198 | CONTENT_TEXT, |
| 199 | id="mixed_content_nonstreaming_with_eos_gigachat3", |
| 200 | ), |
| 201 | pytest.param( |
| 202 | True, |
| 203 | SIMPLE_FUNCTION_OUTPUT_GIGACHAT31, |
| 204 | [SIMPLE_FUNCTION_CALL], |
| 205 | None, |
| 206 | id="simple_streaming_gigachat31", |
| 207 | ), |
| 208 | pytest.param( |
| 209 | False, |
| 210 | SIMPLE_FUNCTION_OUTPUT_GIGACHAT31, |
| 211 | [SIMPLE_FUNCTION_CALL], |
| 212 | None, |
| 213 | id="simple_nonstreaming_gigachat31", |
| 214 | ), |
| 215 | pytest.param( |
| 216 | True, |
| 217 | PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31, |
| 218 | [PARAMETERLESS_FUNCTION_CALL], |
| 219 | None, |
| 220 | id="parameterless_streaming_gigachat31", |
| 221 | ), |
| 222 | pytest.param( |
| 223 | False, |
| 224 | PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31, |
| 225 | [PARAMETERLESS_FUNCTION_CALL], |
| 226 | None, |
| 227 | id="parameterless_nonstreaming_gigachat31", |
| 228 | ), |
| 229 | pytest.param( |
| 230 | True, |
| 231 | COMPLEX_FUNCTION_OUTPUT_GIGACHAT31, |
| 232 | [COMPLEX_FUNCTION_CALL], |
| 233 | None, |
| 234 | id="complex_streaming_gigachat31", |
| 235 | ), |
| 236 | pytest.param( |
| 237 | False, |
| 238 | COMPLEX_FUNCTION_OUTPUT_GIGACHAT31, |
| 239 | [COMPLEX_FUNCTION_CALL], |
| 240 | None, |
| 241 | id="complex_nonstreaming_gigachat31", |
| 242 | ), |
| 243 | pytest.param( |
| 244 | True, |
| 245 | MIXED_OUTPUT_GIGACHAT31, |
| 246 | [SIMPLE_FUNCTION_CALL], |
| 247 | CONTENT_TEXT, |
| 248 | id="mixed_content_streaming_gigachat31", |
| 249 | ), |
| 250 | pytest.param( |
| 251 | False, |
| 252 | MIXED_OUTPUT_GIGACHAT31, |
| 253 | [SIMPLE_FUNCTION_CALL], |
| 254 | CONTENT_TEXT, |
| 255 | id="mixed_content_nonstreaming_gigachat31", |
| 256 | ), |
| 257 | pytest.param( |
| 258 | True, |
| 259 | MIXED_OUTPUT_GIGACHAT31 + EOS_TOKEN, |
| 260 | [SIMPLE_FUNCTION_CALL], |
| 261 | CONTENT_TEXT, |
| 262 | id="mixed_content_streaming_with_eos_gigachat31", |
| 263 | ), |
| 264 | pytest.param( |
| 265 | False, |
| 266 | MIXED_OUTPUT_GIGACHAT31 + EOS_TOKEN, |
| 267 | [SIMPLE_FUNCTION_CALL], |
| 268 | CONTENT_TEXT, |
| 269 | id="mixed_content_nonstreaming_with_eos_gigachat31", |
| 270 | ), |
| 271 | ] |
| 272 | |
| 273 | |
| 274 | @pytest.mark.parametrize( |
| 275 | "streaming, model_output, expected_tool_calls, expected_content", TEST_CASES |
| 276 | ) |
| 277 | def test_tool_call( |
| 278 | streaming: bool, |
| 279 | model_output: str, |
| 280 | expected_tool_calls: list[FunctionCall], |
| 281 | expected_content: str | None, |
| 282 | gigachat_tokenizer: TokenizerLike, |
| 283 | ): |
| 284 | tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")( |
| 285 | gigachat_tokenizer |
| 286 | ) |
| 287 | content, tool_calls = run_tool_extraction( |
| 288 | tool_parser, model_output, streaming=streaming |
| 289 | ) |
| 290 | if content == "": |
| 291 | content = None |
| 292 | assert content == expected_content |
| 293 | assert len(tool_calls) == len(expected_tool_calls) |
| 294 | for actual, expected in zip(tool_calls, expected_tool_calls): |
| 295 | assert actual.type == "function" |
| 296 | assert actual.function.name == expected.name |
| 297 | actual_args = json.loads(actual.function.arguments) |
| 298 | expected_args = json.loads(expected.arguments) |
| 299 | assert actual_args == expected_args |
| 300 | |
| 301 | |
| 302 | @pytest.mark.parametrize( |
| 303 | "model_output_deltas", |
| 304 | [ |
| 305 | pytest.param( |
| 306 | [ |
| 307 | CONTENT_TEXT[:3], |
| 308 | CONTENT_TEXT[3:5], |
| 309 | CONTENT_TEXT[5:], |
| 310 | MSG_SEP_TOKEN, |
| 311 | TOOL_HEADER_GIGACHAT3, |
| 312 | COMPLEX_FUNCTION_JSON[:40], |
| 313 | COMPLEX_FUNCTION_JSON[40:-1], |
| 314 | COMPLEX_FUNCTION_JSON[-1], |
| 315 | ], |
| 316 | id="gigachat3", |
| 317 | ), |
| 318 | pytest.param( |
| 319 | [ |
| 320 | CONTENT_TEXT[:3], |
| 321 | CONTENT_TEXT[3:5], |
| 322 | CONTENT_TEXT[5:], |
| 323 | TOOL_HEADER_GIGACHAT31, |
| 324 | COMPLEX_FUNCTION_JSON[:40], |
| 325 | COMPLEX_FUNCTION_JSON[40:-1], |
| 326 | COMPLEX_FUNCTION_JSON[-1], |
| 327 | ], |
| 328 | id="gigachat31", |
| 329 | ), |
| 330 | ], |
| 331 | ) |
| 332 | def test_streaming_tool_call_with_large_steps( |
| 333 | model_output_deltas: list[str], |
| 334 | gigachat_tokenizer: TokenizerLike, |
| 335 | ): |
| 336 | """Test that the closing braces are streamed correctly.""" |
| 337 | tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")( |
| 338 | gigachat_tokenizer |
| 339 | ) |
| 340 | reconstructor = run_tool_extraction_streaming( |
| 341 | tool_parser, |
| 342 | model_output_deltas, |
| 343 | assert_one_tool_per_delta=False, |
| 344 | ) |
| 345 | assert len(reconstructor.tool_calls) == 1 |
| 346 | call = reconstructor.tool_calls[0] |
| 347 | assert call.type == "function" |
| 348 | assert call.function.name == "manage_user_memory" |
| 349 | args_dict = json.loads(call.function.arguments) |
| 350 | assert args_dict == COMPLEX_ARGS_DICT |