| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copy |
| 4 | import json |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.parser.mistral import MistralToolCall |
| 9 | from vllm.sampling_params import SamplingParams |
| 10 | from vllm.tokenizers.mistral import MistralTokenizer |
| 11 | from vllm.tool_parsers.mistral_tool_parser import MistralToolParser |
| 12 | |
| 13 | from ...utils import check_logprobs_close |
| 14 | |
| 15 | MODELS = [ |
| 16 | "mistralai/Mistral-7B-Instruct-v0.3", |
| 17 | ] |
| 18 | |
| 19 | MISTRAL_FORMAT_MODELS = [ |
| 20 | "mistralai/Mistral-7B-Instruct-v0.3", |
| 21 | # uses the v3-Tekken tokenizer |
| 22 | "mistralai/Ministral-8B-Instruct-2410", |
| 23 | # Mistral-Nemo is too big for CI, but passes locally |
| 24 | # "mistralai/Mistral-Nemo-Instruct-2407" |
| 25 | ] |
| 26 | |
| 27 | SAMPLING_PARAMS = SamplingParams(max_tokens=512, temperature=0.0, logprobs=5) |
| 28 | SYMBOLIC_LANG_PROMPTS = [ |
| 29 | "勇敢な船乗りについての詩を書く", # japanese |
| 30 | "寫一首關於勇敢的水手的詩", # chinese |
| 31 | "ပုံပြင်လေးပြောပြပါ်:\n", # burmese |
| 32 | "Repeat the phrase 'URGENCY🌶️':\nURGENCY🌶️\nURGENCY🌶️\n", # see https://github.com/vllm-project/vllm/pull/9625 |
| 33 | ] |
| 34 | |
| 35 | # for function calling |
| 36 | TOOLS = [ |
| 37 | { |
| 38 | "type": "function", |
| 39 | "function": { |
| 40 | "name": "get_current_weather", |
| 41 | "description": "Get the current weather in a given location", |
| 42 | "parameters": { |
| 43 | "type": "object", |
| 44 | "properties": { |
| 45 | "city": { |
| 46 | "type": "string", |
| 47 | "description": "The city to find the weather for, e.g. " |
| 48 | "'San Francisco'", |
| 49 | }, |
| 50 | "state": { |
| 51 | "type": "string", |
| 52 | "description": "the two-letter abbreviation for the state that " |
| 53 | "the city is in, e.g. 'CA' which would mean 'California'", |
| 54 | }, |
| 55 | "unit": { |
| 56 | "type": "string", |
| 57 | "description": "The unit to fetch the temperature in", |
| 58 | "enum": ["celsius", "fahrenheit"], |
| 59 | }, |
| 60 | }, |
| 61 | "required": ["city", "state", "unit"], |
| 62 | }, |
| 63 | }, |
| 64 | }, |
| 65 | { |
| 66 | "type": "function", |
| 67 | "function": { |
| 68 | "name": "rewrite", |
| 69 | "description": "Rewrites text", |
| 70 | "parameters": { |
| 71 | "type": "object", |
| 72 | "required": [], |
| 73 | "properties": { |
| 74 | "text": { |
| 75 | "type": "string", |
| 76 | "description": "The input text to rewrite.", |
| 77 | } |
| 78 | }, |
| 79 | }, |
| 80 | }, |
| 81 | }, |
| 82 | ] |
| 83 | MSGS = [ |
| 84 | {"role": "system", "content": "You are an assistant."}, |
| 85 | { |
| 86 | "role": "user", |
| 87 | "content": "Could you please rewrite the below article? \n\n My English needs " |
| 88 | "improving, maybe I make errors.", |
| 89 | }, |
| 90 | { |
| 91 | "role": "assistant", |
| 92 | "content": "", |
| 93 | "tool_calls": [ |
| 94 | { |
| 95 | "id": "bbc5b7ede", |
| 96 | "type": "function", |
| 97 | "function": { |
| 98 | "name": "rewrite", |
| 99 | "arguments": '{"text":"My English needs improving, maybe ' |
| 100 | 'I make errors."}', |
| 101 | }, |
| 102 | } |
| 103 | ], |
| 104 | }, |
| 105 | { |
| 106 | "role": "tool", |
| 107 | "content": '{"action":"rewrite","outcome":"My English needs improving, maybe ' |
| 108 | 'I make errors."}', |
| 109 | "tool_call_id": "bbc5b7ede", |
| 110 | "name": "rewrite", |
| 111 | }, |
| 112 | { |
| 113 | "role": "assistant", |
| 114 | "content": "---\n\nMy English needs improving, maybe I make errors", |
| 115 | }, |
| 116 | { |
| 117 | "role": "user", |
| 118 | "content": ( |
| 119 | "Can you tell me what the temperate will be in Dallas, in fahrenheit?" |
| 120 | ), |
| 121 | }, |
| 122 | ] |
| 123 | |
| 124 | SAMPLE_JSON_SCHEMA = { |
| 125 | "type": "object", |
| 126 | "properties": { |
| 127 | "name": {"type": "string"}, |
| 128 | "age": {"type": "integer"}, |
| 129 | "skills": { |
| 130 | "type": "array", |
| 131 | "items": {"type": "string", "maxLength": 10}, |
| 132 | "minItems": 3, |
| 133 | }, |
| 134 | "work_history": { |
| 135 | "type": "array", |
| 136 | "items": { |
| 137 | "type": "object", |
| 138 | "properties": { |
| 139 | "company": {"type": "string"}, |
| 140 | "duration": {"type": "number"}, |
| 141 | "position": {"type": "string"}, |
| 142 | }, |
| 143 | "required": ["company", "position"], |
| 144 | }, |
| 145 | }, |
| 146 | }, |
| 147 | "required": ["name", "age", "skills", "work_history"], |
| 148 | } |
| 149 | |
| 150 | |
| 151 | @pytest.mark.parametrize("model", MODELS) |
| 152 | @pytest.mark.parametrize("dtype", ["bfloat16"]) |
| 153 | @pytest.mark.parametrize("max_tokens", [64]) |
| 154 | @pytest.mark.parametrize("num_logprobs", [5]) |
| 155 | def test_models( |
| 156 | hf_runner, |
| 157 | vllm_runner, |
| 158 | example_prompts, |
| 159 | model: str, |
| 160 | dtype: str, |
| 161 | max_tokens: int, |
| 162 | num_logprobs: int, |
| 163 | ) -> None: |
| 164 | # TODO(sang): Sliding window should be tested separately. |
| 165 | with hf_runner(model, dtype=dtype) as hf_model: |
| 166 | hf_outputs = hf_model.generate_greedy_logprobs_limit( |
| 167 | example_prompts, max_tokens, num_logprobs |
| 168 | ) |
| 169 | |
| 170 | with vllm_runner(model, dtype=dtype, tokenizer_mode="mistral") as vllm_model: |
| 171 | vllm_outputs = vllm_model.generate_greedy_logprobs( |
| 172 | example_prompts, max_tokens, num_logprobs |
| 173 | ) |
| 174 | |
| 175 | check_logprobs_close( |
| 176 | outputs_0_lst=hf_outputs, |
| 177 | outputs_1_lst=vllm_outputs, |
| 178 | name_0="hf", |
| 179 | name_1="vllm", |
| 180 | ) |
| 181 | |
| 182 | |
| 183 | @pytest.mark.parametrize("model", MISTRAL_FORMAT_MODELS) |
| 184 | @pytest.mark.parametrize("dtype", ["bfloat16"]) |
| 185 | @pytest.mark.parametrize("max_tokens", [64]) |
| 186 | @pytest.mark.parametrize("num_logprobs", [5]) |
| 187 | def test_mistral_format( |
| 188 | vllm_runner, |
| 189 | example_prompts, |
| 190 | model: str, |
| 191 | dtype: str, |
| 192 | max_tokens: int, |
| 193 | num_logprobs: int, |
| 194 | ) -> None: |
| 195 | with vllm_runner( |
| 196 | model, |
| 197 | dtype=dtype, |
| 198 | tokenizer_mode="mistral", |
| 199 | load_format="mistral", |
| 200 | config_format="mistral", |
| 201 | ) as mistral_format_model: |
| 202 | mistral_format_outputs = mistral_format_model.generate_greedy_logprobs( |
| 203 | example_prompts, max_tokens, num_logprobs |
| 204 | ) |
| 205 | |
| 206 | with vllm_runner( |
| 207 | model, |
| 208 | dtype=dtype, |
| 209 | tokenizer_mode="hf", |
| 210 | load_format="safetensors", |
| 211 | config_format="hf", |
| 212 | ) as hf_format_model: |
| 213 | hf_format_outputs = hf_format_model.generate_greedy_logprobs( |
| 214 | example_prompts, max_tokens, num_logprobs |
| 215 | ) |
| 216 | |
| 217 | check_logprobs_close( |
| 218 | outputs_0_lst=hf_format_outputs, |
| 219 | outputs_1_lst=mistral_format_outputs, |
| 220 | name_0="hf", |
| 221 | name_1="mistral", |
| 222 | ) |
| 223 | |
| 224 | |
| 225 | @pytest.mark.parametrize("model", MISTRAL_FORMAT_MODELS) |
| 226 | @pytest.mark.parametrize("dtype", ["bfloat16"]) |
| 227 | def test_mistral_symbolic_languages(vllm_runner, model: str, dtype: str) -> None: |
| 228 | with vllm_runner( |
| 229 | model, |
| 230 | dtype=dtype, |
| 231 | max_model_len=8192, |
| 232 | tokenizer_mode="mistral", |
| 233 | config_format="mistral", |
| 234 | load_format="mistral", |
| 235 | ) as vllm_model: |
| 236 | for prompt in SYMBOLIC_LANG_PROMPTS: |
| 237 | msg = {"role": "user", "content": prompt} |
| 238 | outputs = vllm_model.llm.chat([msg], sampling_params=SAMPLING_PARAMS) |
| 239 | assert "�" not in outputs[0].outputs[0].text.strip() |
| 240 | |
| 241 | |
| 242 | @pytest.mark.parametrize("model", MISTRAL_FORMAT_MODELS) |
| 243 | @pytest.mark.parametrize("dtype", ["bfloat16"]) |
| 244 | def test_mistral_function_calling(vllm_runner, model: str, dtype: str) -> None: |
| 245 | with vllm_runner( |
| 246 | model, |
| 247 | dtype=dtype, |
| 248 | tokenizer_mode="mistral", |
| 249 | config_format="mistral", |
| 250 | load_format="mistral", |
| 251 | ) as vllm_model: |
| 252 | msgs = copy.deepcopy(MSGS) |
| 253 | outputs = vllm_model.llm.chat( |
| 254 | msgs, tools=TOOLS, sampling_params=SAMPLING_PARAMS |
| 255 | ) |
| 256 | |
| 257 | tokenizer = vllm_model.llm.get_tokenizer() |
| 258 | tool_parser = MistralToolParser(tokenizer) |
| 259 | |
| 260 | model_output = outputs[0].outputs[0].text.strip() |
| 261 | assert model_output.startswith(tool_parser.bot_token), model_output |
| 262 | parsed_message = tool_parser.extract_tool_calls(model_output, None) |
| 263 | |
| 264 | assert parsed_message.tools_called |
| 265 | |
| 266 | assert MistralToolCall.is_valid_id(parsed_message.tool_calls[0].id) |
| 267 | assert parsed_message.tool_calls[0].function.name == "get_current_weather" |
| 268 | assert ( |
| 269 | parsed_message.tool_calls[0].function.arguments |
| 270 | == '{"city": "Dallas", "state": "TX", "unit": "fahrenheit"}' |
| 271 | ) # noqa |
| 272 | assert parsed_message.content is None |
| 273 | |
| 274 | |
| 275 | def test_mistral_function_call_nested_json(): |
| 276 | """Ensure that the function-name regex captures the entire outermost |
| 277 | JSON block, including nested braces.""" |
| 278 | |
| 279 | # Create a minimal stub tokenizer that provides the few attributes the |
| 280 | # parser accesses (`version` and `get_vocab`). |
| 281 | class _StubMistralTokenizer(MistralTokenizer): |
| 282 | version = 11 # Satisfy the version check |
| 283 | |
| 284 | def __init__(self): |
| 285 | pass |
| 286 | |
| 287 | @staticmethod |
| 288 | def get_vocab(): |
| 289 | # Provide the special TOOL_CALLS token expected by the parser. |
| 290 | return {"[TOOL_CALLS]": 0} |
| 291 | |
| 292 | tokenizer = _StubMistralTokenizer() |
| 293 | parser = MistralToolParser(tokenizer) |
| 294 | |
| 295 | # Craft a model output featuring nested JSON inside the arguments. |
| 296 | args_dict = { |
| 297 | "city": "Dallas", |
| 298 | "state": "TX", |
| 299 | "unit": "fahrenheit", |
| 300 | "sub_dict": {"foo": "bar", "inner": {"x": 1, "y": 2}}, |
| 301 | } |
| 302 | |
| 303 | model_output = f"{parser.bot_token}get_current_weather{json.dumps(args_dict)}" |
| 304 | |
| 305 | parsed = parser.extract_tool_calls(model_output, None) |
| 306 | |
| 307 | # Assertions: the tool call is detected and the full nested JSON is parsed |
| 308 | # without truncation. |
| 309 | assert parsed.tools_called |
| 310 | |
| 311 | assert MistralToolCall.is_valid_id(parsed.tool_calls[0].id) |
| 312 | assert parsed.tool_calls[0].function.name == "get_current_weather" |
| 313 | assert json.loads(parsed.tool_calls[0].function.arguments) == args_dict |
| 314 | # No additional content outside the tool call should be returned. |
| 315 | assert parsed.content is None |
| 316 | |
| 317 | # multiple calls |
| 318 | multiple_args_dict = [ |
| 319 | { |
| 320 | "city": "Dallas", |
| 321 | "state": "TX", |
| 322 | "unit": "fahrenheit", |
| 323 | "sub_dict": {"foo": "bar", "inner": {"x": 1, "y": 2}}, |
| 324 | }, |
| 325 | {}, |
| 326 | {"a": 0}, |
| 327 | {"a": 1, "b": "c"}, |
| 328 | ] |
| 329 | names = ["get_current_weather", "get_current_weather_2", "random", "random_2"] |
| 330 | |
| 331 | model_output = "".join( |
| 332 | [ |
| 333 | f"{parser.bot_token}{name}{json.dumps(args)}" |
| 334 | for name, args in zip(names, multiple_args_dict) |
| 335 | ] |
| 336 | ) |
| 337 | |
| 338 | parsed = parser.extract_tool_calls(model_output, None) |
| 339 | |
| 340 | # Assertions: the tool call is detected and the full nested JSON is parsed |
| 341 | # without truncation. |
| 342 | assert parsed.tools_called |
| 343 | assert len(parsed.tool_calls) == len(multiple_args_dict) |
| 344 | |
| 345 | for i, tool_call in enumerate(parsed.tool_calls): |
| 346 | assert MistralToolCall.is_valid_id(tool_call.id) |
| 347 | assert tool_call.function.name == names[i] |
| 348 | assert json.loads(tool_call.function.arguments) == multiple_args_dict[i] |
| 349 | # No additional content outside the tool call should be returned. |
| 350 | assert parsed.content is None |