@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dd5d1ef780
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 02:15 UTC

[Bugfix] Mistral tool parser streaming update (#19425)

Signed-off-by: avigny <47987522+avigny@users.noreply.github.com> Signed-off-by: Chauncey <chaunceyjiang@gmail.com> Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Jeff Cook <jeff@jeffcook.io> Co-authored-by: sfbemerk <benjaminmerkel@mail.de> Co-authored-by: Chauncey <chaunceyjiang@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

avigny committed Dec 3, 2025 at 18:45 UTC dd5d1ef780b5b73b9817e5dc8fe9b3e98a399e20
4 files changed +1272 -202
requirements/common.txt
+1
@@ -46,6 +46,7 @@ scipy # Required for phi-4-multimodal-instruct
46 ninja # Required for xgrammar, rocm, tpu, xpu
47 pybase64 # fast base64 implementation
48 cbor2 # Required for cross-language serialization of hashable objects
49 +ijson # Required for mistral streaming tool parser
50 setproctitle # Used to set process names for better debugging and monitoring
51 openai-harmony >= 0.0.3 # Required for gpt-oss
52 anthropic == 0.71.0
tests/tool_use/test_mistral_tool_parser.py new
+847
@@ -0,0 +1,847 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +from collections.abc import Generator
6 +
7 +import partial_json_parser
8 +import pytest
9 +from mistral_common.protocol.instruct.messages import AssistantMessage
10 +from mistral_common.protocol.instruct.request import InstructRequest
11 +from mistral_common.protocol.instruct.tool_calls import FunctionCall, ToolCall
12 +from partial_json_parser.core.options import Allow
13 +
14 +from vllm.entrypoints.openai.protocol import DeltaMessage, DeltaToolCall
15 +from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolParser
16 +from vllm.tokenizers import (
17 + MistralTokenizer,
18 + TokenizerLike,
19 + get_tokenizer,
20 +)
21 +from vllm.tokenizers.detokenizer_utils import detokenize_incrementally
22 +
23 +
24 +@pytest.fixture(scope="module")
25 +def mistral_pre_v11_tokenizer():
26 + MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
27 + return get_tokenizer(tokenizer_name=MODEL)
28 +
29 +
30 +@pytest.fixture(scope="module")
31 +def mistral_tokenizer():
32 + MODEL = "mistralai/Mistral-Small-3.2-24B-Instruct-2506"
33 + return get_tokenizer(tokenizer_name=MODEL, tokenizer_mode="mistral")
34 +
35 +
36 +@pytest.fixture
37 +def mistral_pre_v11_tool_parser(mistral_pre_v11_tokenizer):
38 + return MistralToolParser(mistral_pre_v11_tokenizer)
39 +
40 +
41 +@pytest.fixture
42 +def mistral_tool_parser(mistral_tokenizer):
43 + return MistralToolParser(mistral_tokenizer)
44 +
45 +
46 +def assert_tool_calls(
47 + actual_tool_calls: list[ToolCall] | list[DeltaToolCall],
48 + expected_tool_calls: list[ToolCall],
49 +):
50 + assert len(actual_tool_calls) == len(expected_tool_calls)
51 +
52 + for actual_tool_call, expected_tool_call in zip(
53 + actual_tool_calls, expected_tool_calls
54 + ):
55 + assert isinstance(actual_tool_call.id, str)
56 + assert len(actual_tool_call.id) == 9
57 +
58 + if isinstance(actual_tool_call, ToolCall):
59 + assert actual_tool_call.type == "function"
60 + elif isinstance(actual_tool_call, DeltaToolCall):
61 + assert actual_tool_call.function is not None
62 + assert actual_tool_call.function.name is not None
63 + assert actual_tool_call.function.arguments is not None
64 + assert actual_tool_call.function is not None
65 + assert actual_tool_call.function.name == expected_tool_call.function.name, (
66 + f"got wrong function name:${actual_tool_call.function.name}"
67 + )
68 + assert (
69 + actual_tool_call.function.arguments == expected_tool_call.function.arguments
70 + ), f"got wrong function argument:${actual_tool_call.function.arguments}"
71 +
72 +
73 +def fix_tool_call_tokenization(
74 + tokens: list[int],
75 + mistral_tool_parser: MistralToolParser,
76 + mistral_tokenizer: TokenizerLike,
77 +):
78 + """
79 + Replaces the textual token sequence for [TOOL_CALLS]
80 + with its single special token ID.
81 + """
82 + textual_tool_call_token_ids = mistral_tokenizer.encode(
83 + text=mistral_tool_parser.bot_token,
84 + add_special_tokens=False,
85 + )
86 + # textual_tool_call_token_ids must not contain special tokens like bos, eos etc
87 + special_tool_call_token_ids = [mistral_tool_parser.bot_token_id]
88 +
89 + # If the input is too short to contain the sequence, no replacement is possible
90 + if not tokens or len(tokens) < len(textual_tool_call_token_ids):
91 + return tokens
92 +
93 + result_tokens = []
94 + i = 0
95 + target_len = len(textual_tool_call_token_ids)
96 +
97 + while i < len(tokens):
98 + # Check if the slice from the current position matches the target sequence
99 + if tokens[i : i + target_len] == textual_tool_call_token_ids:
100 + # If it matches, add the replacement and jump the index forward
101 + result_tokens.extend(special_tool_call_token_ids)
102 + i += target_len
103 + else:
104 + # Otherwise, just add the current token and move to the next one
105 + result_tokens.append(tokens[i])
106 + i += 1
107 +
108 + return result_tokens
109 +
110 +
111 +def stream_delta_message_generator(
112 + mistral_tool_parser: MistralToolParser,
113 + mistral_tokenizer: TokenizerLike,
114 + model_output: str | None,
115 + tools: list[tuple[str, str]] | None,
116 +) -> Generator[DeltaMessage, None, None]:
117 + if (
118 + isinstance(mistral_tokenizer, MistralTokenizer)
119 + and mistral_tokenizer.version >= 11
120 + ):
121 + # With the newer versions of the tokenizer,
122 + # we cannot tokenize free text
123 + # so we need to create a list of messages to get tokenized
124 + assert tools is not None
125 + assistant_msg = AssistantMessage(
126 + tool_calls=[
127 + ToolCall(
128 + function=FunctionCall(
129 + name=name,
130 + arguments=arg,
131 + )
132 + )
133 + for (name, arg) in tools
134 + ],
135 + )
136 + request = InstructRequest(
137 + messages=[assistant_msg],
138 + )
139 + all_token_ids = mistral_tokenizer.instruct.encode_instruct(request).tokens
140 + else:
141 + # Older versions of the tokenizer are
142 + # able to encode directly the model's output (free text) into tokens
143 + assert model_output is not None
144 + all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False)
145 +
146 + all_token_ids = fix_tool_call_tokenization(
147 + all_token_ids, mistral_tool_parser, mistral_tokenizer
148 + )
149 +
150 + previous_text = ""
151 + previous_tokens = None
152 + prefix_offset = 0
153 + read_offset = 0
154 + for i, delta_token in enumerate(all_token_ids):
155 + delta_token_ids = [delta_token]
156 + previous_token_ids = all_token_ids[:i]
157 + current_token_ids = all_token_ids[: i + 1]
158 +
159 + (new_tokens, delta_text, new_prefix_offset, new_read_offset) = (
160 + detokenize_incrementally(
161 + tokenizer=mistral_tokenizer,
162 + all_input_ids=current_token_ids,
163 + prev_tokens=previous_tokens,
164 + prefix_offset=prefix_offset,
165 + read_offset=read_offset,
166 + skip_special_tokens=isinstance(mistral_tokenizer, MistralTokenizer),
167 + spaces_between_special_tokens=True,
168 + )
169 + )
170 +
171 + current_text = previous_text + delta_text
172 +
173 + delta_message = mistral_tool_parser.extract_tool_calls_streaming(
174 + previous_text,
175 + current_text,
176 + delta_text,
177 + previous_token_ids,
178 + current_token_ids,
179 + delta_token_ids,
180 + request=None, # type: ignore[arg-type]
181 + )
182 + if delta_message:
183 + yield delta_message
184 +
185 + previous_text = current_text
186 + previous_tokens = (
187 + previous_tokens + new_tokens if previous_tokens else new_tokens
188 + )
189 + prefix_offset = new_prefix_offset
190 + read_offset = new_read_offset
191 +
192 +
193 +def test_extract_tool_calls_no_tools(mistral_pre_v11_tool_parser):
194 + model_output = "This is a test"
195 + extracted_tool_calls = mistral_pre_v11_tool_parser.extract_tool_calls(
196 + model_output, request=None
197 + ) # type: ignore[arg-type]
198 + assert not extracted_tool_calls.tools_called
199 + assert extracted_tool_calls.tool_calls == []
200 + assert extracted_tool_calls.content == model_output
201 +
202 +
203 +@pytest.mark.parametrize(
204 + ids=[
205 + "single_tool_add",
206 + "single_tool_weather",
207 + "argument_before_name",
208 + "argument_before_name_and_name_in_argument",
209 + ],
210 + argnames=["model_output", "expected_tool_calls", "expected_content"],
211 + argvalues=[
212 + (
213 + """[TOOL_CALLS][{"name": "add", "arguments":{"a": 3.5, "b": 4}}]""", # noqa: E501
214 + [
215 + ToolCall(
216 + function=FunctionCall(
217 + name="add", arguments=json.dumps({"a": 3.5, "b": 4})
218 + )
219 + )
220 + ],
221 + None,
222 + ),
223 + (
224 + """[TOOL_CALLS] [{"name": "get_current_weather", "arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501
225 + [
226 + ToolCall(
227 + function=FunctionCall(
228 + name="get_current_weather",
229 + arguments=json.dumps(
230 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
231 + ),
232 + )
233 + )
234 + ],
235 + None,
236 + ),
237 + (
238 + """[TOOL_CALLS] [{"arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501
239 + [
240 + ToolCall(
241 + function=FunctionCall(
242 + name="get_current_weather",
243 + arguments=json.dumps(
244 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
245 + ),
246 + )
247 + )
248 + ],
249 + None,
250 + ),
251 + (
252 + """[TOOL_CALLS] [{"arguments":{"name": "John Doe"}, "name": "get_age"}]""", # noqa: E501
253 + [
254 + ToolCall(
255 + function=FunctionCall(
256 + name="get_age",
257 + arguments=json.dumps(
258 + {
259 + "name": "John Doe",
260 + }
261 + ),
262 + )
263 + )
264 + ],
265 + None,
266 + ),
267 + ],
268 +)
269 +def test_extract_tool_calls_pre_v11_tokenizer(
270 + mistral_pre_v11_tool_parser, model_output, expected_tool_calls, expected_content
271 +):
272 + extracted_tool_calls = mistral_pre_v11_tool_parser.extract_tool_calls(
273 + model_output, request=None
274 + ) # type: ignore[arg-type]
275 + assert extracted_tool_calls.tools_called
276 +
277 + assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)
278 +
279 + assert extracted_tool_calls.content == expected_content
280 +
281 +
282 +@pytest.mark.parametrize(
283 + ids=[
284 + "single_tool_add",
285 + "single_tool_weather",
286 + "multiple_tool_calls",
287 + ],
288 + argnames=["model_output", "expected_tool_calls", "expected_content"],
289 + argvalues=[
290 + (
291 + """[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501
292 + [
293 + ToolCall(
294 + function=FunctionCall(
295 + name="add_this_and_that",
296 + arguments=json.dumps({"a": 3.5, "b": 4}),
297 + )
298 + )
299 + ],
300 + None,
301 + ),
302 + (
303 + """[TOOL_CALLS]get_current_weather{"city": "San Francisco", "state": "CA", "unit": "celsius"}""", # noqa: E501
304 + [
305 + ToolCall(
306 + function=FunctionCall(
307 + name="get_current_weather",
308 + arguments=json.dumps(
309 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
310 + ),
311 + )
312 + )
313 + ],
314 + None,
315 + ),
316 + (
317 + """[TOOL_CALLS]add{"a": 3.5, "b": 4}[TOOL_CALLS]multiply{"a": 3, "b": 6}""", # noqa: E501
318 + [
319 + ToolCall(
320 + function=FunctionCall(
321 + name="add", arguments=json.dumps({"a": 3.5, "b": 4})
322 + )
323 + ),
324 + ToolCall(
325 + function=FunctionCall(
326 + name="multiply", arguments=json.dumps({"a": 3, "b": 6})
327 + )
328 + ),
329 + ],
330 + None,
331 + ),
332 + ],
333 +)
334 +def test_extract_tool_calls(
335 + mistral_tool_parser, model_output, expected_tool_calls, expected_content
336 +):
337 + extracted_tool_calls = mistral_tool_parser.extract_tool_calls(
338 + model_output, request=None
339 + ) # type: ignore[arg-type]
340 + assert extracted_tool_calls.tools_called
341 +
342 + assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)
343 +
344 + assert extracted_tool_calls.content == expected_content
345 +
346 +
347 +def _test_extract_tool_calls_streaming(
348 + tool_parser, tokenizer, model_output, tools, expected_tool_calls, expected_content
349 +):
350 + other_content: str = ""
351 + function_names: list[str] = []
352 + function_args_strs: list[str] = []
353 + tool_call_idx: int = -1
354 + tool_call_ids: list[str | None] = []
355 +
356 + for delta_message in stream_delta_message_generator(
357 + tool_parser, tokenizer, model_output, tools
358 + ):
359 + # role should never be streamed from tool parser
360 + assert not delta_message.role
361 +
362 + if delta_message.content:
363 + other_content += delta_message.content
364 +
365 + streamed_tool_calls = delta_message.tool_calls
366 +
367 + if streamed_tool_calls and len(streamed_tool_calls) > 0:
368 + # make sure only one diff is present - correct even for parallel
369 + assert len(streamed_tool_calls) == 1
370 + tool_call = streamed_tool_calls[0]
371 +
372 + assert len(tool_parser.prev_tool_call_arr) > 0
373 +
374 + # if a new tool is being called, set up empty arguments
375 + if tool_call.index != tool_call_idx:
376 + tool_call_idx = tool_call.index
377 + function_args_strs.append("")
378 + tool_call_ids.append(None)
379 +
380 + # if a tool call ID is streamed, make sure one hasn't been already
381 + if tool_call.id and not tool_call_ids[tool_call.index]:
382 + tool_call_ids[tool_call.index] = tool_call.id
383 +
384 + # if parts of the function start being streamed
385 + if tool_call.function:
386 + # if the function name is defined, set it. it should be streamed
387 + # IN ENTIRETY, exactly one time.
388 + if tool_call.function.name:
389 + assert isinstance(tool_call.function.name, str)
390 + function_names.append(tool_call.function.name)
391 +
392 + if tool_call.function.arguments:
393 + # make sure they're a string and then add them to the list
394 + assert isinstance(tool_call.function.arguments, str)
395 +
396 + function_args_strs[tool_call.index] += tool_call.function.arguments
397 +
398 + assert other_content == expected_content
399 +
400 + actual_tool_calls = [
401 + ToolCall(
402 + id=tool_call_id,
403 + function=FunctionCall(
404 + name=function_name,
405 + arguments=partial_json_parser.ensure_json(
406 + function_args_str, Allow.OBJ | Allow.STR
407 + ),
408 + ),
409 + )
410 + for tool_call_id, function_name, function_args_str in zip(
411 + tool_call_ids, function_names, function_args_strs
412 + )
413 + ]
414 + assert_tool_calls(actual_tool_calls, expected_tool_calls)
415 +
416 +
417 +@pytest.mark.parametrize(
418 + ids=[
419 + "no_tools",
420 + "single_tool_add",
421 + "single_tool_add_strings",
422 + "single_tool_weather",
423 + "argument_before_name",
424 + "argument_before_name_and_name_in_argument",
425 + "multiple_tools",
426 + ],
427 + argnames=["model_output", "expected_tool_calls", "expected_content"],
428 + argvalues=[
429 + ("""This is a test""", [], """This is a test"""),
430 + (
431 + """[TOOL_CALLS] [ {"name":"add" , "arguments" : {"a": 3, "b": 4} } ]""", # noqa: E501
432 + [
433 + ToolCall(
434 + function=FunctionCall(
435 + name="add", arguments=json.dumps({"a": 3, "b": 4})
436 + )
437 + )
438 + ],
439 + "",
440 + ),
441 + (
442 + """[TOOL_CALLS] [{"name": "add", "arguments":{"a": "3", "b": "4"}}]""", # noqa: E501
443 + [
444 + ToolCall(
445 + function=FunctionCall(
446 + name="add", arguments=json.dumps({"a": "3", "b": "4"})
447 + )
448 + )
449 + ],
450 + "",
451 + ),
452 + (
453 + """[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501
454 + [
455 + ToolCall(
456 + function=FunctionCall(
457 + name="get_current_weather",
458 + arguments=json.dumps(
459 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
460 + ),
461 + )
462 + )
463 + ],
464 + "",
465 + ),
466 + (
467 + """[TOOL_CALLS] [{"arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501
468 + [
469 + ToolCall(
470 + function=FunctionCall(
471 + name="get_current_weather",
472 + arguments=json.dumps(
473 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
474 + ),
475 + )
476 + )
477 + ],
478 + "",
479 + ),
480 + (
481 + """[TOOL_CALLS] [{"arguments": {"name": "John Doe"}, "name": "get_age"}]""", # noqa: E501
482 + [
483 + ToolCall(
484 + function=FunctionCall(
485 + name="get_age",
486 + arguments=json.dumps(
487 + {
488 + "name": "John Doe",
489 + }
490 + ),
491 + )
492 + )
493 + ],
494 + "",
495 + ),
496 + (
497 + """[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, {"name": "get_current_weather", "arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501
498 + [
499 + ToolCall(
500 + function=FunctionCall(
501 + name="add", arguments=json.dumps({"a": 3.5, "b": 4})
502 + )
503 + ),
504 + ToolCall(
505 + function=FunctionCall(
506 + name="get_current_weather",
507 + arguments=json.dumps(
508 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
509 + ),
510 + )
511 + ),
512 + ],
513 + "",
514 + ),
515 + ],
516 +)
517 +def test_extract_tool_calls_streaming_pre_v11_tokenizer(
518 + mistral_pre_v11_tool_parser,
519 + mistral_pre_v11_tokenizer,
520 + model_output,
521 + expected_tool_calls,
522 + expected_content,
523 +):
524 + _test_extract_tool_calls_streaming(
525 + mistral_pre_v11_tool_parser,
526 + mistral_pre_v11_tokenizer,
527 + model_output,
528 + None,
529 + expected_tool_calls,
530 + expected_content,
531 + )
532 +
533 +
534 +@pytest.mark.parametrize(
535 + ids=[
536 + "single_tool_add",
537 + "single_tool_add_strings",
538 + "multiple_tools",
539 + ],
540 + argnames=["tools", "expected_tool_calls", "expected_content"],
541 + argvalues=[
542 + (
543 + [("add", '{"a": 3, "b": 4}')],
544 + # [TOOL_CALLS]add{"a": 3, "b": 4}
545 + [
546 + ToolCall(
547 + function=FunctionCall(
548 + name="add", arguments=json.dumps({"a": 3, "b": 4})
549 + )
550 + )
551 + ],
552 + "",
553 + ),
554 + (
555 + [("add_two_strings", '{"a": "3", "b": "4"}')],
556 + # [TOOL_CALLS]add_two_strings{"a": "3", "b": "4"}
557 + [
558 + ToolCall(
559 + function=FunctionCall(
560 + name="add_two_strings",
561 + arguments=json.dumps({"a": "3", "b": "4"}),
562 + )
563 + )
564 + ],
565 + "",
566 + ),
567 + (
568 + [
569 + ("add", '{"a": 3.5, "b": 4}'),
570 + (
571 + "get_current_weather",
572 + '{"city": "San Francisco", "state": "CA", "unit": "celsius"}', # noqa: E501
573 + ),
574 + ],
575 + # [TOOL_CALLS]add{"a": 3.5, "b": 4}[TOOL_CALLS]get_current_weather{"city": "San Francisco", "state": "CA", "unit": "celsius"} # noqa: E501
576 + [
577 + ToolCall(
578 + function=FunctionCall(
579 + name="add", arguments=json.dumps({"a": 3.5, "b": 4})
580 + )
581 + ),
582 + ToolCall(
583 + function=FunctionCall(
584 + name="get_current_weather",
585 + arguments=json.dumps(
586 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
587 + ),
588 + )
589 + ),
590 + ],
591 + "",
592 + ),
593 + ],
594 +)
595 +def test_extract_tool_calls_streaming(
596 + mistral_tool_parser,
597 + mistral_tokenizer,
598 + tools,
599 + expected_tool_calls,
600 + expected_content,
601 +):
602 + _test_extract_tool_calls_streaming(
603 + mistral_tool_parser,
604 + mistral_tokenizer,
605 + None,
606 + tools,
607 + expected_tool_calls,
608 + expected_content,
609 + )
610 +
611 +
612 +@pytest.mark.parametrize(
613 + ids=[
614 + "single_tool_add",
615 + "single_tool_weather",
616 + "multiple_tool_calls",
617 + "content_before_tool",
618 + ],
619 + argnames=["model_output", "expected_tool_calls", "expected_content"],
620 + argvalues=[
621 + (
622 + """[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501
623 + [
624 + ToolCall(
625 + function=FunctionCall(
626 + name="add_this_and_that",
627 + arguments=json.dumps({"a": 3.5, "b": 4}),
628 + )
629 + )
630 + ],
631 + "",
632 + ),
633 + (
634 + """[TOOL_CALLS]get_current_weather{"city": "San Francisco", "state": "CA", "unit": "celsius"}""", # noqa: E501
635 + [
636 + ToolCall(
637 + function=FunctionCall(
638 + name="get_current_weather",
639 + arguments=json.dumps(
640 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
641 + ),
642 + )
643 + )
644 + ],
645 + "",
646 + ),
647 + (
648 + """[TOOL_CALLS]add{"a": 3.5, "b": 4}[TOOL_CALLS]multiply{"a": 3, "b": 6}""", # noqa: E501
649 + [
650 + ToolCall(
651 + function=FunctionCall(
652 + name="add", arguments=json.dumps({"a": 3.5, "b": 4})
653 + )
654 + ),
655 + ToolCall(
656 + function=FunctionCall(
657 + name="multiply", arguments=json.dumps({"a": 3, "b": 6})
658 + )
659 + ),
660 + ],
661 + "",
662 + ),
663 + (
664 + # Additional content should not be after the tool calls
665 + """bla[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501
666 + [
667 + ToolCall(
668 + function=FunctionCall(
669 + name="add_this_and_that",
670 + arguments=json.dumps({"a": 3.5, "b": 4}),
671 + )
672 + )
673 + ],
674 + "bla",
675 + ),
676 + ],
677 +)
678 +def test_extract_tool_calls_streaming_one_chunk(
679 + mistral_tool_parser,
680 + mistral_tokenizer,
681 + model_output,
682 + expected_tool_calls,
683 + expected_content,
684 +):
685 + if isinstance(mistral_tokenizer, MistralTokenizer):
686 + all_token_ids = mistral_tokenizer.encode(model_output)
687 + else:
688 + all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False)
689 + all_token_ids = fix_tool_call_tokenization(
690 + all_token_ids, mistral_tool_parser, mistral_tokenizer
691 + )
692 +
693 + delta_message = mistral_tool_parser.extract_tool_calls_streaming(
694 + previous_text="",
695 + current_text=model_output,
696 + delta_text=model_output,
697 + previous_token_ids=[],
698 + current_token_ids=all_token_ids,
699 + delta_token_ids=all_token_ids,
700 + request=None,
701 + ) # type: ignore[arg-type]
702 + assert isinstance(delta_message, DeltaMessage)
703 + assert len(delta_message.tool_calls) == len(expected_tool_calls)
704 +
705 + assert_tool_calls(delta_message.tool_calls, expected_tool_calls)
706 +
707 + if delta_message.content is None:
708 + assert expected_content == ""
709 + else:
710 + assert delta_message.content == expected_content
711 +
712 +
713 +@pytest.mark.parametrize(
714 + ids=[
715 + "no_tools",
716 + "single_tool_add",
717 + "single_tool_add_strings",
718 + "single_tool_weather",
719 + "argument_before_name",
720 + "argument_before_name_and_name_in_argument",
721 + "multiple_tools",
722 + ],
723 + argnames=["model_output", "expected_tool_calls", "expected_content"],
724 + argvalues=[
725 + ("""This is a test""", [], """This is a test"""),
726 + (
727 + """[TOOL_CALLS] [ {"name":"add" , "arguments" : {"a": 3, "b": 4} } ]""", # noqa: E501
728 + [
729 + ToolCall(
730 + function=FunctionCall(
731 + name="add", arguments=json.dumps({"a": 3, "b": 4})
732 + )
733 + )
734 + ],
735 + "",
736 + ),
737 + (
738 + """[TOOL_CALLS] [{"name": "add", "arguments":{"a": "3", "b": "4"}}]""", # noqa: E501
739 + [
740 + ToolCall(
741 + function=FunctionCall(
742 + name="add", arguments=json.dumps({"a": "3", "b": "4"})
743 + )
744 + )
745 + ],
746 + "",
747 + ),
748 + (
749 + """[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501
750 + [
751 + ToolCall(
752 + function=FunctionCall(
753 + name="get_current_weather",
754 + arguments=json.dumps(
755 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
756 + ),
757 + )
758 + )
759 + ],
760 + "",
761 + ),
762 + (
763 + """[TOOL_CALLS] [{"arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501
764 + [
765 + ToolCall(
766 + function=FunctionCall(
767 + name="get_current_weather",
768 + arguments=json.dumps(
769 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
770 + ),
771 + )
772 + )
773 + ],
774 + "",
775 + ),
776 + (
777 + """[TOOL_CALLS] [{"arguments": {"name": "John Doe"}, "name": "get_age"}]""", # noqa: E501
778 + [
779 + ToolCall(
780 + function=FunctionCall(
781 + name="get_age",
782 + arguments=json.dumps(
783 + {
784 + "name": "John Doe",
785 + }
786 + ),
787 + )
788 + )
789 + ],
790 + "",
791 + ),
792 + (
793 + """[TOOL_CALLS] [{"arguments": {"a": 3.5, "b": 4}, "name": "add"}, {"arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501
794 + [
795 + ToolCall(
796 + function=FunctionCall(
797 + name="add", arguments=json.dumps({"a": 3.5, "b": 4})
798 + )
799 + ),
800 + ToolCall(
801 + function=FunctionCall(
802 + name="get_current_weather",
803 + arguments=json.dumps(
804 + {"city": "San Francisco", "state": "CA", "unit": "celsius"}
805 + ),
806 + )
807 + ),
808 + ],
809 + "",
810 + ),
811 + ],
812 +)
813 +def test_extract_tool_calls_streaming_pre_v11_tokenizer_one_chunk(
814 + mistral_pre_v11_tool_parser,
815 + mistral_pre_v11_tokenizer,
816 + model_output,
817 + expected_tool_calls,
818 + expected_content,
819 +):
820 + if isinstance(mistral_pre_v11_tokenizer, MistralTokenizer):
821 + all_token_ids = mistral_pre_v11_tokenizer.encode(model_output)
822 + else:
823 + all_token_ids = mistral_pre_v11_tokenizer.encode(
824 + model_output, add_special_tokens=False
825 + )
826 + all_token_ids = fix_tool_call_tokenization(
827 + all_token_ids, mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer
828 + )
829 +
830 + delta_message = mistral_pre_v11_tool_parser.extract_tool_calls_streaming(
831 + previous_text="",
832 + current_text=model_output,
833 + delta_text=model_output,
834 + previous_token_ids=[],
835 + current_token_ids=all_token_ids,
836 + delta_token_ids=all_token_ids,
837 + request=None,
838 + ) # type: ignore[arg-type]
839 + assert isinstance(delta_message, DeltaMessage)
840 + assert len(delta_message.tool_calls) == len(expected_tool_calls)
841 +
842 + assert_tool_calls(delta_message.tool_calls, expected_tool_calls)
843 +
844 + if delta_message.content is None:
845 + assert expected_content == ""
846 + else:
847 + assert delta_message.content == expected_content
tests/tool_use/utils.py
+27 -1
@@ -123,7 +123,7 @@ CONFIGS: dict[str, ServerConfig] = {
123 "supports_parallel": True,
124 "extended": True,
125 },
126 - "mistral": {
126 + "mistral-7b": {
127 "model": "mistralai/Mistral-7B-Instruct-v0.3",
128 "arguments": [
129 "--enforce-eager",
@@ -145,6 +145,32 @@ CONFIGS: dict[str, ServerConfig] = {
145 "call the tool. Otherwise, answer the user's query directly "
146 "without calling a tool. DO NOT CALL A TOOL THAT IS IRRELEVANT "
147 "to the user's question - just respond to it normally.",
148 + "supports_parallel": True,
149 + },
150 + "mistral-small-3.2": {
151 + "model": "mistralai/Mistral-Small-3.2-24B-Instruct-2506",
152 + "arguments": [
153 + "--enforce-eager",
154 + "--no-enable-prefix-caching",
155 + "--tool-call-parser",
156 + "mistral",
157 + "--tokenizer-mode",
158 + "mistral",
159 + "--config-format",
160 + "mistral",
161 + "--load-format",
162 + "mistral",
163 + "--tensor-parallel-size",
164 + "4",
165 + '--ignore-patterns="consolidated.safetensors"',
166 + ],
167 + "system_prompt": "You are a helpful assistant with access to tools. If a tool"
168 + " that you have would be helpful to answer a user query, "
169 + "call the tool. Otherwise, answer the user's query directly "
170 + "without calling a tool. DO NOT CALL A TOOL THAT IS IRRELEVANT "
171 + "to the user's question - just respond to it normally.",
172 + "supports_parallel": True,
173 + "extended": True,
174 },
175 # FIXME: This test currently fails, need to debug why.
176 # "granite20b": {
vllm/entrypoints/openai/tool_parsers/mistral_tool_parser.py
+397 -201
@@ -3,12 +3,12 @@
3
4 import json
5 from collections.abc import Sequence
6 +from enum import Enum, auto
7 from random import choices
8 from string import ascii_letters, digits
9
9 -import partial_json_parser
10 +import ijson
11 import regex as re
11 -from partial_json_parser.core.options import Allow
12 from pydantic import Field
13
14 from vllm.entrypoints.openai.protocol import (
@@ -23,7 +23,6 @@ from vllm.entrypoints.openai.protocol import (
23 from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
24 ToolParser,
25 )
26 -from vllm.entrypoints.openai.tool_parsers.utils import extract_intermediate_diff
26 from vllm.logger import init_logger
27 from vllm.tokenizers import MistralTokenizer, TokenizerLike
28
@@ -32,6 +31,22 @@ logger = init_logger(__name__)
31 ALPHANUMERIC = ascii_letters + digits
32
33
34 +class StreamingState(Enum):
35 + """Enum for tracking the current streaming parsing state."""
36 +
37 + WAITING_FOR_TOOL_START = auto()
38 + WAITING_FOR_TOOL_KEY = (
39 + auto()
40 + ) # waiting for the "name" or "arguments" key to be complete
41 + PARSING_NAME = auto()
42 + PARSING_NAME_COMPLETED = auto()
43 + WAITING_FOR_ARGUMENTS_START = auto()
44 + PARSING_ARGUMENTS = auto()
45 + PARSING_ARGUMENTS_COMPLETED = auto()
46 + TOOL_COMPLETE = auto()
47 + ALL_TOOLS_COMPLETE = auto()
48 +
49 +
50 class MistralToolCall(ToolCall):
51 id: str = Field(default_factory=lambda: MistralToolCall.generate_random_id())
52
@@ -46,8 +61,8 @@ class MistralToolCall(ToolCall):
61 return id.isalnum() and len(id) == 9
62
63
49 -def _is_fn_name_regex_support(model_tokenizer: TokenizerLike) -> bool:
50 - return (
64 +def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool:
65 + return not (
66 isinstance(model_tokenizer, MistralTokenizer) and model_tokenizer.version >= 11
67 )
68
@@ -69,16 +84,22 @@ class MistralToolParser(ToolParser):
84
85 # initialize properties used for state when parsing tool calls in
86 # streaming mode
72 - self.prev_tool_call_arr: list[dict] = []
87 self.current_tool_id: int = -1
74 - self.current_tool_name_sent: bool = False
75 - self.streamed_args_for_tool: list[
76 - str
77 - ] = [] # map what has been streamed for each tool so far to a list
88 + self.streaming_state: StreamingState = StreamingState.WAITING_FOR_TOOL_START
89 +
90 + # For streaming pre v11 tokenizer tool calls
91 + self.current_tool_name: str | None = None
92 + self.current_tool_mistral_id: str | None = None
93 + self.starting_new_tool = False
94 + if _is_pre_v11_tokeniser(self.model_tokenizer):
95 + self.parse_coro = ijson.parse_coro(
96 + self.update_stream_state_pre_v11_tokenizer()
97 + )
98 +
99 self.bot_token = "[TOOL_CALLS]"
100 self.bot_token_id = self.vocab.get(self.bot_token)
101 self.tool_call_regex = re.compile(r"\[{.*}\]", re.DOTALL)
81 - if _is_fn_name_regex_support(self.model_tokenizer):
102 + if not _is_pre_v11_tokeniser(self.model_tokenizer):
103 self.fn_name_regex = re.compile(
104 r"([a-zA-Z0-9_-]+)(\{[\s\S]*?\}+)", re.DOTALL
105 )
@@ -131,18 +152,19 @@ class MistralToolParser(ToolParser):
152 # jsons is difficult
153 try:
154 if self.fn_name_regex:
134 - matches = self.fn_name_regex.findall(tool_content)
135 -
155 function_call_arr = []
137 - for match in matches:
138 - fn_name = match[0]
139 - args = match[1]
140 -
141 - # fn_name is encoded outside serialized json dump
142 - # only arguments are serialized
143 - function_call_arr.append(
144 - {"name": fn_name, "arguments": json.loads(args)}
145 - )
156 + for single_tool_content in model_output.split(self.bot_token):
157 + matches = self.fn_name_regex.findall(single_tool_content)
158 +
159 + for match in matches:
160 + fn_name = match[0]
161 + args = match[1]
162 +
163 + # fn_name is encoded outside serialized json dump
164 + # only arguments are serialized
165 + function_call_arr.append(
166 + {"name": fn_name, "arguments": json.loads(args)}
167 + )
168 else:
169 function_call_arr = json.loads(tool_content)
170 except json.JSONDecodeError:
@@ -193,198 +215,372 @@ class MistralToolParser(ToolParser):
215 delta_token_ids: Sequence[int],
216 request: ChatCompletionRequest,
217 ) -> DeltaMessage | None:
196 - # if the tool call token is not in the tokens generated so far, append
197 - # output to contents since it's not a tool
198 - if self.bot_token not in current_text:
218 + if self.bot_token_id not in current_token_ids:
219 + # if the tool call token is not in the tokens generated so far,
220 + # append output to contents since it's not a tool
221 return DeltaMessage(content=delta_text)
222
201 - # if the tool call token ID IS in the tokens generated so far, that
223 + # if the tool call token IS in the tokens generated so far, that
224 # means we're parsing as tool calls now
203 -
204 - # handle if we detected the BOT token which means the start of tool
205 - # calling
206 - if self.bot_token_id in delta_token_ids and len(delta_token_ids) == 1:
207 - # if it's the only token, return None, so we don't send a chat
208 - # completion any don't send a control token
209 - return None
210 -
211 - # bit mask flags for partial JSON parsing. If the name hasn't been
212 - # sent yet, don't allow sending
213 - # an incomplete string since OpenAI only ever (as far as I have
214 - # seen) allows sending the entire tool/ function name at once.
215 - flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR
225 try:
217 - # replace BOT token with empty string, and convert single quotes
218 - # to double to allow parsing as JSON since mistral uses single
219 - # quotes instead of double for tool calls
220 - parsable_arr = current_text.split(self.bot_token)[-1]
221 -
222 - # tool calls are generated in an array, so do partial JSON
223 - # parsing on the entire array
224 - try:
225 - tool_call_arr: list[dict] = partial_json_parser.loads(
226 - parsable_arr, flags
226 + if _is_pre_v11_tokeniser(self.model_tokenizer):
227 + return self._extract_tool_calls_streaming_pre_v11_tokenizer(
228 + delta_text=delta_text,
229 + delta_token_ids=delta_token_ids,
230 )
228 - except partial_json_parser.core.exceptions.MalformedJSON:
229 - logger.debug("not enough tokens to parse into JSON yet")
230 - return None
231 -
232 - # select as the current tool call the one we're on the state at
231 + else:
232 + return self._extract_tool_calls_streaming(
233 + delta_text=delta_text, delta_token_ids=delta_token_ids
234 + )
235 + except Exception:
236 + logger.exception("Error trying to handle streaming tool call.")
237 + return None
238
234 - current_tool_call: dict = (
235 - tool_call_arr[self.current_tool_id] if len(tool_call_arr) > 0 else {}
236 - )
239 + def _extract_tool_calls_streaming(
240 + self,
241 + delta_text: str,
242 + delta_token_ids: Sequence[int],
243 + ) -> DeltaMessage | None:
244 + """
245 + Extracts tool calls for Mistral models
246 + doing tool calls of the following format:
247 + `[TOOL_CALLS]add{"a": 3.5, "b": 4}`
248 + """
249 + additional_content: str = ""
250 + if self.streaming_state == StreamingState.WAITING_FOR_TOOL_START:
251 + # this is the first tool call
252 + assert self.bot_token_id in delta_token_ids
253 + if not delta_text.startswith(self.bot_token):
254 + additional_content += delta_text.split(self.bot_token)[0]
255 + delta_text = self.bot_token + "".join(
256 + delta_text.split(self.bot_token)[1:]
257 + )
258
238 - # case -- if no tokens have been streamed for the tool, e.g.
239 - # only the array brackets, stream nothing
240 - if len(tool_call_arr) == 0:
259 + delta_tool_calls = self._generate_delta_tool_call(delta_text)
260 + if not additional_content and len(delta_tool_calls) == 0:
261 + if self.streaming_state in [
262 + StreamingState.PARSING_ARGUMENTS,
263 + StreamingState.PARSING_ARGUMENTS_COMPLETED,
264 + StreamingState.TOOL_COMPLETE,
265 + StreamingState.ALL_TOOLS_COMPLETE,
266 + ]:
267 + # Return an empty DeltaMessage once the tool calls are all done
268 + # so that finish_reason gets set.
269 + return DeltaMessage()
270 + else:
271 + # return None when the tool is not likely to be finished
272 + # This can occur when the name is being parsed for example
273 + # and we wait for the name to be complete
274 + # before sending the function name
275 return None
276
243 - # case: we are starting a new tool in the array
244 - # -> array has > 0 length AND length has moved past cursor
245 - elif (
246 - len(tool_call_arr) > 0 and len(tool_call_arr) > self.current_tool_id + 1
247 - ):
248 - # if we're moving on to a new call, first make sure we
249 - # haven't missed anything in the previous one that was
250 - # auto-generated due to JSON completions, but wasn't
251 - # streamed to the client yet.
252 - if self.current_tool_id >= 0:
253 - diff: str | None = current_tool_call.get("arguments")
254 -
255 - if diff:
256 - diff = json.dumps(diff, ensure_ascii=False).replace(
257 - self.streamed_args_for_tool[self.current_tool_id], ""
258 - )
259 - delta = DeltaMessage(
260 - tool_calls=[
261 - DeltaToolCall(
262 - index=self.current_tool_id,
263 - function=DeltaFunctionCall(
264 - arguments=diff
265 - ).model_dump(exclude_none=True),
266 - )
267 - ]
268 - )
269 - self.streamed_args_for_tool[self.current_tool_id] += diff
270 - else:
271 - delta = None
272 - else:
273 - delta = None
274 - # re-set stuff pertaining to progress in the current tool
275 - self.current_tool_id = len(tool_call_arr) - 1
276 - self.current_tool_name_sent = False
277 - self.streamed_args_for_tool.append("")
278 - logger.debug("starting on new tool %d", self.current_tool_id)
279 - return delta
280 -
281 - # case: update an existing tool - this is handled below
282 -
283 - # if the current tool name hasn't been sent, send if available
284 - # - otherwise send nothing
285 - if not self.current_tool_name_sent:
286 - function_name = current_tool_call.get("name")
287 - if function_name:
288 - delta = DeltaMessage(
289 - tool_calls=[
290 - DeltaToolCall(
291 - index=self.current_tool_id,
292 - type="function",
293 - id=MistralToolCall.generate_random_id(),
294 - function=DeltaFunctionCall(
295 - name=function_name
296 - ).model_dump(exclude_none=True),
297 - )
298 - ]
277 + delta = DeltaMessage()
278 + if additional_content:
279 + delta.content = additional_content
280 + if len(delta_tool_calls) > 0:
281 + delta.tool_calls = delta_tool_calls
282 +
283 + # HACK: serving_chat.py inspects the internal state of tool parsers
284 + # when determining its final streaming delta, automatically
285 + # adding autocompleted JSON.
286 + # These two lines avoid that nonsense while ensuring finish_reason
287 + # is set to tool_calls when at least one tool is called.
288 + if delta_tool_calls and not self.prev_tool_call_arr:
289 + self.prev_tool_call_arr = [{"arguments": {}}]
290 + return delta
291 +
292 + def _generate_delta_tool_call(self, delta_text: str) -> list[DeltaToolCall]:
293 + if delta_text == "" or delta_text is None:
294 + return []
295 + delta_function_name = None
296 + tool_id = None
297 + if self.streaming_state not in [
298 + StreamingState.PARSING_NAME,
299 + StreamingState.PARSING_ARGUMENTS,
300 + ] and delta_text.startswith(self.bot_token):
301 + self.current_tool_id += 1
302 + self.streaming_state = StreamingState.PARSING_NAME
303 + delta_text = delta_text.replace(self.bot_token, "", 1)
304 + if self.streaming_state == StreamingState.PARSING_NAME:
305 + if self.current_tool_name is None:
306 + self.current_tool_name = ""
307 + # The name stops where the arguments start
308 + # And the arguments start with the `{` char
309 + if "{" in delta_text:
310 + tool_id = MistralToolCall.generate_random_id()
311 + delta_function_name = delta_text.split("{")[0]
312 + self.current_tool_name += delta_function_name
313 + delta_text = delta_text[len(delta_function_name) :]
314 + self.streaming_state = StreamingState.PARSING_ARGUMENTS
315 + else:
316 + # we want to send the tool name once it's complete
317 + self.current_tool_name += delta_text
318 + return []
319 + if self.streaming_state == StreamingState.PARSING_ARGUMENTS:
320 + next_function_text = None
321 + if self.bot_token in delta_text:
322 + # current tool call is over
323 + delta_arguments = ""
324 + delta_arguments += delta_text.split(self.bot_token)[0]
325 + next_function_text = delta_text[len(delta_arguments) :]
326 + self.streaming_state = StreamingState.TOOL_COMPLETE
327 + else:
328 + delta_arguments = delta_text
329 + ret = []
330 + if self.current_tool_name or delta_arguments:
331 + ret += [
332 + DeltaToolCall(
333 + index=self.current_tool_id,
334 + type="function",
335 + id=tool_id,
336 + function=DeltaFunctionCall(
337 + name=self.current_tool_name, arguments=delta_arguments
338 + ).model_dump(exclude_none=True),
339 )
300 - self.current_tool_name_sent = True
301 - else:
302 - delta = None
303 -
304 - # now we know we're on the same tool call and we're streaming
305 - # arguments
340 + ]
341 + self.current_tool_name = None
342 + if next_function_text:
343 + ret += self._generate_delta_tool_call(next_function_text)
344 + return ret
345 + # Should not happen
346 + return []
347 +
348 + @ijson.coroutine
349 + def update_stream_state_pre_v11_tokenizer(self):
350 + while True:
351 + (prefix, event, value) = yield
352 +
353 + if prefix == "item" and event == "start_map":
354 + self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
355 + if prefix == "item" and event == "map_key" and value == "name":
356 + self.streaming_state = StreamingState.PARSING_NAME
357 + if prefix == "item.name" and event == "string":
358 + self.current_tool_name = value
359 + self.streaming_state = StreamingState.PARSING_NAME_COMPLETED
360 + if prefix == "item" and event == "map_key" and value == "arguments":
361 + self.streaming_state = StreamingState.WAITING_FOR_ARGUMENTS_START
362 + if prefix == "item.arguments" and event == "start_map":
363 + self.streaming_state = StreamingState.PARSING_ARGUMENTS
364 + if prefix == "item.arguments" and event == "end_map":
365 + self.streaming_state = StreamingState.PARSING_ARGUMENTS_COMPLETED
366 + if prefix == "item" and event == "end_map":
367 + self.streaming_state = StreamingState.TOOL_COMPLETE
368 + if prefix == "" and event == "end_array":
369 + self.streaming_state = StreamingState.ALL_TOOLS_COMPLETE
370 +
371 + def _extract_tool_calls_streaming_pre_v11_tokenizer(
372 + self,
373 + delta_text: str,
374 + delta_token_ids: Sequence[int],
375 + ) -> DeltaMessage | None:
376 + """
377 + Extracts tool calls for Mistral models
378 + doing tool calls of the following format:
379 + `[TOOL_CALLS][{"name": "add", "arguments":{"a": 3.5, "b": 4}}`
380 + """
381 + assert self.parse_coro is not None
382 + content = None
383 + delta_tool_calls: list[DeltaToolCall] = []
384 + current_tool_call: DeltaToolCall = DeltaToolCall(
385 + index=self.current_tool_id, type="function"
386 + )
387 + current_tool_call_modified = False
388 + if self.bot_token_id in delta_token_ids:
389 + # this is the first tool call
390 + if not delta_text.startswith(self.bot_token):
391 + content = delta_text.split(self.bot_token)[0]
392 + delta_text = "".join(delta_text.split(self.bot_token)[1:])
393 +
394 + # Cut smartly the delta text to catch the ijson events
395 + # as ijson does not give us the index in the text at each event.
396 + # We need to cut so that we know
397 + # where in the text the events are emitted from.
398 + while len(delta_text) > 0:
399 + streaming_state_before_parse = self.streaming_state
400 +
401 + if self.streaming_state == StreamingState.WAITING_FOR_TOOL_START:
402 + delta_to_be_parsed, delta_text = self._split_delta(
403 + delta_text=delta_text,
404 + stop_after_opening_curly_braces=1,
405 + )
406 + elif self.streaming_state == StreamingState.WAITING_FOR_TOOL_KEY:
407 + # Wait until another key is sent
408 + # or the current tool is completed
409 + delta_to_be_parsed, delta_text = self._split_delta(
410 + delta_text=delta_text,
411 + stop_after_colon=1,
412 + stop_after_opening_curly_braces=1,
413 + # if the tool ends, we want to separate
414 + # at the start of the next tool
415 + )
416 + elif self.streaming_state == StreamingState.PARSING_NAME:
417 + delta_to_be_parsed, delta_text = self._split_delta(
418 + delta_text=delta_text,
419 + stop_after_comma=1,
420 + stop_after_closing_brackets=1,
421 + )
422 + elif self.streaming_state == StreamingState.WAITING_FOR_ARGUMENTS_START:
423 + delta_to_be_parsed, delta_text = self._split_delta(
424 + delta_text=delta_text,
425 + stop_after_opening_curly_braces=1,
426 + )
427 + elif self.streaming_state == StreamingState.PARSING_ARGUMENTS:
428 + delta_to_be_parsed, delta_text = self._split_delta(
429 + delta_text=delta_text,
430 + stop_after_closing_curly_braces=1,
431 + # we could be more clever
432 + # by listening to item.arguments.* start_map events
433 + # and know how many curly braces we can allow
434 + )
435 + elif self.streaming_state in [
436 + StreamingState.PARSING_ARGUMENTS_COMPLETED,
437 + StreamingState.PARSING_NAME_COMPLETED,
438 + ]:
439 + delta_to_be_parsed, delta_text = self._split_delta(
440 + delta_text=delta_text,
441 + stop_after_closing_curly_braces=1,
442 + stop_after_closing_brackets=1,
443 + )
444 + elif self.streaming_state == StreamingState.TOOL_COMPLETE:
445 + delta_to_be_parsed, delta_text = self._split_delta(
446 + delta_text=delta_text,
447 + stop_after_opening_curly_braces=1,
448 + stop_after_closing_brackets=1,
449 + )
450 + elif self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE:
451 + content = delta_text
452 + delta_text = ""
453 else:
307 - prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
308 - "arguments"
454 + delta_to_be_parsed = delta_text
455 + delta_text = ""
456 +
457 + if self.streaming_state != StreamingState.ALL_TOOLS_COMPLETE:
458 + self.parse_coro.send(delta_to_be_parsed.encode("utf-8"))
459 +
460 + # Given the parsed text and the possible streaming state change,
461 + # let's add to the tool delta
462 + if (
463 + (streaming_state_before_parse != self.streaming_state)
464 + and streaming_state_before_parse
465 + in [StreamingState.WAITING_FOR_TOOL_START, StreamingState.TOOL_COMPLETE]
466 + and self.streaming_state
467 + not in [
468 + StreamingState.ALL_TOOLS_COMPLETE,
469 + StreamingState.TOOL_COMPLETE,
470 + StreamingState.WAITING_FOR_TOOL_START,
471 + ]
472 + ):
473 + # starting a new tool call
474 + if current_tool_call_modified:
475 + if self.current_tool_mistral_id is not None:
476 + current_tool_call.id = self.current_tool_mistral_id
477 + self.current_tool_mistral_id = None
478 + delta_tool_calls.append(current_tool_call)
479 + current_tool_call_modified = False
480 + self.current_tool_id += 1
481 + self.current_tool_mistral_id = MistralToolCall.generate_random_id()
482 + current_tool_call = DeltaToolCall(
483 + index=self.current_tool_id,
484 + type="function",
485 )
310 - cur_arguments = current_tool_call.get("arguments")
311 -
312 - new_text = delta_text.replace("'", '"')
313 - if '"}' in new_text:
314 - new_text = new_text[: new_text.rindex('"}')]
315 -
316 - if not cur_arguments and not prev_arguments:
317 - delta = None
318 - elif not cur_arguments and prev_arguments:
319 - logger.error(
320 - "INVARIANT - impossible to have arguments reset mid-arguments"
321 - )
322 - delta = None
323 - elif cur_arguments and not prev_arguments:
324 - cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False)[
325 - :-2
326 - ]
327 - logger.debug("finding %s in %s", new_text, cur_arguments_json)
328 -
329 - if new_text not in cur_arguments_json:
330 - return None
331 - arguments_delta = cur_arguments_json[
332 - : cur_arguments_json.rindex(new_text) + len(new_text)
333 - ]
334 - logger.debug(
335 - "First tokens in arguments received: %s", arguments_delta
336 - )
337 - delta = DeltaMessage(
338 - tool_calls=[
339 - DeltaToolCall(
340 - index=self.current_tool_id,
341 - function=DeltaFunctionCall(
342 - arguments=arguments_delta
343 - ).model_dump(exclude_none=True),
344 - )
345 - ]
346 - )
347 - self.streamed_args_for_tool[self.current_tool_id] += arguments_delta
348 -
349 - elif cur_arguments and prev_arguments:
350 - cur_args_json = json.dumps(cur_arguments, ensure_ascii=False)
351 - prev_args_json = json.dumps(prev_arguments, ensure_ascii=False)
352 - logger.debug(
353 - "Searching for diff between \n%s\n%s",
354 - cur_args_json,
355 - prev_args_json,
486 + if current_tool_call.function is None:
487 + current_tool_call.function = DeltaFunctionCall()
488 +
489 + if self.current_tool_name is not None:
490 + # we have the complete tool name
491 + current_tool_call_modified = True
492 + current_tool_call.function.name = self.current_tool_name
493 + self.current_tool_name = None
494 + if self.streaming_state == StreamingState.PARSING_NAME_COMPLETED:
495 + self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
496 + if self.streaming_state in [
497 + StreamingState.PARSING_ARGUMENTS,
498 + StreamingState.PARSING_ARGUMENTS_COMPLETED,
499 + ]:
500 + if self.streaming_state == StreamingState.PARSING_ARGUMENTS_COMPLETED:
501 + self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
502 + # the delta_to_be_parsed is part of arguments.
503 + current_tool_call_modified = True
504 + if current_tool_call.function.arguments is None:
505 + current_tool_call.function.arguments = delta_to_be_parsed
506 + else:
507 + current_tool_call.function.arguments += delta_to_be_parsed
508 + if streaming_state_before_parse != StreamingState.PARSING_ARGUMENTS:
509 + # It's the first chunk of arg. let's lstrip it
510 + current_tool_call.function.arguments = (
511 + current_tool_call.function.arguments.lstrip()
512 )
513
358 - argument_diff = extract_intermediate_diff(
359 - cur_args_json, prev_args_json
360 - )
361 - logger.debug("got arguments diff: %s", argument_diff)
362 - delta = DeltaMessage(
363 - tool_calls=[
364 - DeltaToolCall(
365 - index=self.current_tool_id,
366 - function=DeltaFunctionCall(
367 - arguments=argument_diff
368 - ).model_dump(exclude_none=True),
369 - )
370 - ]
371 - )
372 - self.streamed_args_for_tool[self.current_tool_id] += argument_diff
373 - else:
374 - # try parsing it with regular JSON - if it works we're
375 - # at the end, and we need to send the difference between
376 - # tokens streamed so far and the valid JSON
377 - delta = None
514 + if current_tool_call_modified:
515 + if self.current_tool_mistral_id is not None:
516 + current_tool_call.id = self.current_tool_mistral_id
517 + self.current_tool_mistral_id = None
518 + delta_tool_calls.append(current_tool_call)
519 +
520 + # HACK: serving_chat.py inspects the internal state of tool parsers
521 + # when determining it's final streaming delta, automatically
522 + # adding autocompleted JSON.
523 + # These two lines avoid that nonsense while ensuring finish_reason
524 + # is set to tool_calls when at least one tool is called.
525 + if delta_tool_calls and not self.prev_tool_call_arr:
526 + self.prev_tool_call_arr = [{"arguments": {}}]
527 +
528 + if content or len(delta_tool_calls) > 0:
529 + delta_message = DeltaMessage()
530 + if content:
531 + delta_message.content = content
532 + if len(delta_tool_calls) > 0:
533 + delta_message.tool_calls = delta_tool_calls
534 + return delta_message
535 + else:
536 + if self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE:
537 + return DeltaMessage()
538 + else:
539 + return None
540
379 - # check to see if the name is defined and has been sent. if so,
380 - # stream the name - otherwise keep waiting
381 - # finish by setting old and returning None as base case
382 - self.prev_tool_call_arr = tool_call_arr
383 - return delta
541 + def _split_delta(
542 + self,
543 + delta_text: str,
544 + stop_after_quotes: int = -1,
545 + stop_after_opening_curly_braces: int = -1,
546 + stop_after_closing_curly_braces: int = -1,
547 + stop_after_closing_brackets: int = -1,
548 + stop_after_colon: int = -1,
549 + stop_after_comma=-1,
550 + ) -> tuple[str, str]:
551 + delta_to_be_parsed = ""
552 + for i, c in enumerate(delta_text):
553 + if c in ['"', "'"]:
554 + delta_to_be_parsed += c
555 + stop_after_quotes -= 1
556 + if stop_after_quotes == 0:
557 + return (delta_to_be_parsed, delta_text[i + 1 :])
558 + elif c == "{":
559 + delta_to_be_parsed += c
560 + stop_after_opening_curly_braces -= 1
561 + if stop_after_opening_curly_braces == 0:
562 + return (delta_to_be_parsed, delta_text[i + 1 :])
563 + elif c == "}":
564 + delta_to_be_parsed += c
565 + stop_after_closing_curly_braces -= 1
566 + if stop_after_closing_curly_braces == 0:
567 + return (delta_to_be_parsed, delta_text[i + 1 :])
568 + elif c == "]":
569 + delta_to_be_parsed += c
570 + stop_after_closing_brackets -= 1
571 + if stop_after_closing_brackets == 0:
572 + return (delta_to_be_parsed, delta_text[i + 1 :])
573 + elif c == ":":
574 + delta_to_be_parsed += c
575 + stop_after_colon -= 1
576 + if stop_after_colon == 0:
577 + return (delta_to_be_parsed, delta_text[i + 1 :])
578 + elif c == ",":
579 + delta_to_be_parsed += c
580 + stop_after_comma -= 1
581 + if stop_after_comma == 0:
582 + return (delta_to_be_parsed, delta_text[i + 1 :])
583 + else:
584 + delta_to_be_parsed += c
585
385 - except Exception:
386 - logger.exception("Error trying to handle streaming tool call.")
387 - logger.debug(
388 - "Skipping chunk as a result of tool streaming extraction error"
389 - )
390 - return None
586 + return (delta_to_be_parsed, "")