| 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 | from unittest.mock import MagicMock, patch |
| 7 | |
| 8 | import partial_json_parser |
| 9 | import pytest |
| 10 | import regex as re |
| 11 | from mistral_common.protocol.instruct.messages import AssistantMessage |
| 12 | from mistral_common.protocol.instruct.request import InstructRequest |
| 13 | from mistral_common.protocol.instruct.tool_calls import ( |
| 14 | FunctionCall, |
| 15 | ToolCall, |
| 16 | ) |
| 17 | from mistral_common.protocol.instruct.tool_calls import ( |
| 18 | NamedToolChoice as MistralNamedToolChoice, |
| 19 | ) |
| 20 | from mistral_common.protocol.instruct.tool_calls import ( |
| 21 | ToolChoice as MistralToolChoice, |
| 22 | ) |
| 23 | from mistral_common.protocol.instruct.tool_calls import ( |
| 24 | ToolChoiceEnum as MistralToolChoiceEnum, |
| 25 | ) |
| 26 | from partial_json_parser.core.options import Allow |
| 27 | |
| 28 | from vllm.entrypoints.generate.base.protocol import ( |
| 29 | DeltaFunctionCall, |
| 30 | DeltaMessage, |
| 31 | DeltaToolCall, |
| 32 | ExtractedToolCallInformation, |
| 33 | StructuralTagResponseFormat, |
| 34 | ) |
| 35 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 36 | ChatCompletionRequest, |
| 37 | ) |
| 38 | from vllm.parser.engine.events import EventType |
| 39 | from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine |
| 40 | from vllm.parser.mistral import _DEFAULT_JSON_SCHEMA, MistralParser, mistral_config |
| 41 | from vllm.parser.parser_manager import ParserManager |
| 42 | from vllm.sampling_params import StructuredOutputsParams |
| 43 | from vllm.tokenizers import TokenizerLike, get_tokenizer |
| 44 | from vllm.tokenizers.detokenizer_utils import detokenize_incrementally |
| 45 | from vllm.tokenizers.mistral import MistralTokenizer |
| 46 | from vllm.tool_parsers.mistral_tool_parser import MistralToolParser |
| 47 | |
| 48 | _DUMMY_REQUEST = ChatCompletionRequest(messages=[], model="test") |
| 49 | # Minimal request for the engine driver: tool_choice=None falls through to |
| 50 | # auto extraction without triggering "tools must be set" validation. |
| 51 | _ENGINE_REQUEST = ChatCompletionRequest(messages=[], model="test", tool_choice=None) |
| 52 | |
| 53 | |
| 54 | @pytest.fixture(scope="module") |
| 55 | def mistral_pre_v11_tokenizer(): |
| 56 | MODEL = "mistralai/Mistral-7B-Instruct-v0.3" |
| 57 | return get_tokenizer(tokenizer_name=MODEL) |
| 58 | |
| 59 | |
| 60 | @pytest.fixture(scope="module") |
| 61 | def mistral_tokenizer(): |
| 62 | MODEL = "mistralai/Mistral-Small-3.2-24B-Instruct-2506" |
| 63 | return get_tokenizer(tokenizer_name=MODEL, tokenizer_mode="mistral") |
| 64 | |
| 65 | |
| 66 | @pytest.fixture |
| 67 | def mistral_pre_v11_tool_parser(mistral_pre_v11_tokenizer): |
| 68 | return MistralToolParser(mistral_pre_v11_tokenizer) |
| 69 | |
| 70 | |
| 71 | @pytest.fixture |
| 72 | def mistral_tool_parser(mistral_tokenizer): |
| 73 | return MistralToolParser(mistral_tokenizer) |
| 74 | |
| 75 | |
| 76 | @pytest.fixture |
| 77 | def non_mistral_parser() -> MistralToolParser: |
| 78 | mock_tokenizer = MagicMock() |
| 79 | mock_tokenizer.get_vocab.return_value = {"[TOOL_CALLS]": 1} |
| 80 | # Ensure the legacy (non-grammar) path is taken by MistralParser. |
| 81 | mock_tokenizer.supports_grammar = False |
| 82 | return MistralToolParser(mock_tokenizer) |
| 83 | |
| 84 | |
| 85 | def assert_tool_calls( |
| 86 | actual_tool_calls: list[ToolCall] | list[DeltaToolCall], |
| 87 | expected_tool_calls: list[ToolCall], |
| 88 | ): |
| 89 | assert len(actual_tool_calls) == len(expected_tool_calls) |
| 90 | |
| 91 | for actual_tool_call, expected_tool_call in zip( |
| 92 | actual_tool_calls, expected_tool_calls |
| 93 | ): |
| 94 | assert isinstance(actual_tool_call.id, str) |
| 95 | assert len(actual_tool_call.id) == 9 |
| 96 | |
| 97 | if isinstance(actual_tool_call, ToolCall): |
| 98 | assert actual_tool_call.type == "function" |
| 99 | elif isinstance(actual_tool_call, DeltaToolCall): |
| 100 | assert actual_tool_call.function is not None |
| 101 | assert actual_tool_call.function.name is not None |
| 102 | assert actual_tool_call.function.arguments is not None |
| 103 | assert actual_tool_call.function is not None |
| 104 | assert actual_tool_call.function.name == expected_tool_call.function.name, ( |
| 105 | f"got wrong function name:${actual_tool_call.function.name}" |
| 106 | ) |
| 107 | actual_args = actual_tool_call.function.arguments |
| 108 | # Streaming deltas (DeltaToolCall) may have partial JSON because the |
| 109 | # engine holds back the final closing brace until finish() is called. |
| 110 | # Apply partial_json_parser so the comparison works for both the |
| 111 | # legacy (complete JSON) and engine (potentially partial) paths. |
| 112 | if isinstance(actual_tool_call, DeltaToolCall) and actual_args is not None: |
| 113 | actual_args = partial_json_parser.ensure_json( |
| 114 | actual_args, Allow.OBJ | Allow.STR |
| 115 | ) |
| 116 | assert actual_args == expected_tool_call.function.arguments, ( |
| 117 | f"got wrong function argument:${actual_tool_call.function.arguments}" |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | def fix_tool_call_tokenization( |
| 122 | tokens: list[int], |
| 123 | mistral_tool_parser: MistralToolParser, |
| 124 | mistral_tokenizer: TokenizerLike, |
| 125 | ): |
| 126 | """Replace textual token sequences for the control markers ([TOOL_CALLS] |
| 127 | and [ARGS]) with their single special token IDs. |
| 128 | |
| 129 | Real generation emits these as control tokens; encoding the marker *string* |
| 130 | can instead yield literal text tokens. The engine detects markers by id, so |
| 131 | the test stream must carry the special ids. |
| 132 | """ |
| 133 | vocab = mistral_tokenizer.get_vocab() |
| 134 | markers: list[tuple[str, int]] = [ |
| 135 | ( |
| 136 | mistral_tool_parser._parser_engine.bot_token, |
| 137 | mistral_tool_parser._parser_engine.bot_token_id, |
| 138 | ) |
| 139 | ] |
| 140 | args_id = vocab.get("[ARGS]") |
| 141 | if args_id is not None: |
| 142 | markers.append(("[ARGS]", args_id)) |
| 143 | |
| 144 | # (textual token sequence, special id) pairs to replace. |
| 145 | replacements: list[tuple[list[int], int]] = [] |
| 146 | for text, special_id in markers: |
| 147 | if special_id is None: |
| 148 | continue |
| 149 | textual = mistral_tokenizer.encode(text=text, add_special_tokens=False) |
| 150 | if textual: |
| 151 | replacements.append((textual, special_id)) |
| 152 | |
| 153 | result_tokens: list[int] = [] |
| 154 | i = 0 |
| 155 | while i < len(tokens): |
| 156 | for textual, special_id in replacements: |
| 157 | if tokens[i : i + len(textual)] == textual: |
| 158 | result_tokens.append(special_id) |
| 159 | i += len(textual) |
| 160 | break |
| 161 | else: |
| 162 | result_tokens.append(tokens[i]) |
| 163 | i += 1 |
| 164 | |
| 165 | return result_tokens |
| 166 | |
| 167 | |
| 168 | def encode_mistral_output(tokenizer: MistralTokenizer, text: str) -> list[int]: |
| 169 | """Encode generated output with control markers as their special-token ids. |
| 170 | |
| 171 | Real generation emits ``[TOOL_CALLS]``/``[ARGS]`` as single control tokens. |
| 172 | Encoding the marker *strings* can instead glue them to neighbours (e.g. |
| 173 | ``}[``) and yield literal text tokens, which the id-based engine never sees. |
| 174 | Split on the markers and insert their special ids so the token stream |
| 175 | matches real generation. |
| 176 | """ |
| 177 | vocab = tokenizer.get_vocab() |
| 178 | marker_ids = { |
| 179 | "[TOOL_CALLS]": vocab.get("[TOOL_CALLS]"), |
| 180 | "[ARGS]": vocab.get("[ARGS]"), |
| 181 | } |
| 182 | ids: list[int] = [] |
| 183 | for part in re.split(r"(\[TOOL_CALLS\]|\[ARGS\])", text): |
| 184 | if part in marker_ids: |
| 185 | # Fall back to text encoding here and the id-based engine path |
| 186 | # silently stops being exercised while the tests still pass. |
| 187 | assert marker_ids[part] is not None, f"{part} missing from vocab" |
| 188 | ids.append(marker_ids[part]) |
| 189 | elif part: |
| 190 | ids.extend(tokenizer.encode(part, add_special_tokens=False)) |
| 191 | return ids |
| 192 | |
| 193 | |
| 194 | def stream_delta_message_generator( |
| 195 | mistral_tool_parser: MistralToolParser, |
| 196 | mistral_tokenizer: TokenizerLike, |
| 197 | model_output: str | None, |
| 198 | tools: list[tuple[str, str]] | None, |
| 199 | chunk_size: int = 1, |
| 200 | driver: str = "legacy", |
| 201 | reasoning_parser: str | None = None, |
| 202 | ) -> Generator[DeltaMessage, None, None]: |
| 203 | if ( |
| 204 | isinstance(mistral_tokenizer, MistralTokenizer) |
| 205 | and mistral_tokenizer.version >= 11 |
| 206 | ): |
| 207 | # With the newer versions of the tokenizer, |
| 208 | # we cannot tokenize free text |
| 209 | # so we need to create a list of messages to get tokenized |
| 210 | assert tools is not None |
| 211 | assistant_msg = AssistantMessage( |
| 212 | tool_calls=[ |
| 213 | ToolCall( |
| 214 | function=FunctionCall( |
| 215 | name=name, |
| 216 | arguments=arg, |
| 217 | ) |
| 218 | ) |
| 219 | for (name, arg) in tools |
| 220 | ], |
| 221 | ) |
| 222 | request = InstructRequest( |
| 223 | messages=[assistant_msg], |
| 224 | ) |
| 225 | all_token_ids = mistral_tokenizer.instruct.encode_instruct(request).tokens |
| 226 | else: |
| 227 | # Older versions of the tokenizer are |
| 228 | # able to encode directly the model's output (free text) into tokens |
| 229 | assert model_output is not None |
| 230 | all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False) |
| 231 | |
| 232 | all_token_ids = fix_tool_call_tokenization( |
| 233 | all_token_ids, mistral_tool_parser, mistral_tokenizer |
| 234 | ) |
| 235 | |
| 236 | if isinstance(mistral_tokenizer, MistralTokenizer): |
| 237 | # Real serving streams only generated tokens; encode_instruct adds |
| 238 | # framing (BOS/EOS) the parser never receives. Strip them so the |
| 239 | # id-based engine detection sees a realistic stream. |
| 240 | if all_token_ids and all_token_ids[0] == mistral_tokenizer.bos_token_id: |
| 241 | all_token_ids = all_token_ids[1:] |
| 242 | if all_token_ids and all_token_ids[-1] == mistral_tokenizer.eos_token_id: |
| 243 | all_token_ids = all_token_ids[:-1] |
| 244 | |
| 245 | if driver == "engine": |
| 246 | _parser_cls = ParserManager.get_parser( |
| 247 | tool_parser_name="mistral", |
| 248 | reasoning_parser_name=reasoning_parser, |
| 249 | enable_auto_tools=True, |
| 250 | ) |
| 251 | _engine_parser = _parser_cls(mistral_tokenizer, None) |
| 252 | if reasoning_parser is not None: |
| 253 | _req_with_tools = ChatCompletionRequest( |
| 254 | messages=[], |
| 255 | model="test", |
| 256 | tools=SAMPLE_TOOLS_DICTS, |
| 257 | tool_choice="auto", |
| 258 | ) |
| 259 | _engine_req = _engine_parser.adjust_request(_req_with_tools) |
| 260 | else: |
| 261 | _engine_req = _ENGINE_REQUEST |
| 262 | |
| 263 | previous_text = "" |
| 264 | previous_tokens = None |
| 265 | prefix_offset = 0 |
| 266 | read_offset = 0 |
| 267 | pending_text = "" |
| 268 | pending_token_ids: list[int] = [] |
| 269 | for i, delta_token in enumerate(all_token_ids): |
| 270 | (new_tokens, delta_text, new_prefix_offset, new_read_offset) = ( |
| 271 | detokenize_incrementally( |
| 272 | tokenizer=mistral_tokenizer, |
| 273 | all_input_ids=all_token_ids[: i + 1], |
| 274 | prev_tokens=previous_tokens, |
| 275 | prefix_offset=prefix_offset, |
| 276 | read_offset=read_offset, |
| 277 | skip_special_tokens=isinstance(mistral_tokenizer, MistralTokenizer), |
| 278 | spaces_between_special_tokens=True, |
| 279 | ) |
| 280 | ) |
| 281 | previous_tokens = ( |
| 282 | previous_tokens + new_tokens if previous_tokens else new_tokens |
| 283 | ) |
| 284 | prefix_offset = new_prefix_offset |
| 285 | read_offset = new_read_offset |
| 286 | |
| 287 | # Buffer tokens so each streamed delta can carry ``chunk_size`` tokens, |
| 288 | # reproducing the multi-token deltas produced by async scheduling / |
| 289 | # stream_interval > 1. |
| 290 | pending_text += delta_text |
| 291 | pending_token_ids.append(delta_token) |
| 292 | if len(pending_token_ids) < chunk_size and i != len(all_token_ids) - 1: |
| 293 | continue |
| 294 | |
| 295 | previous_token_ids = all_token_ids[: i + 1 - len(pending_token_ids)] |
| 296 | current_token_ids = all_token_ids[: i + 1] |
| 297 | current_text = previous_text + pending_text |
| 298 | |
| 299 | if driver == "legacy": |
| 300 | delta_message = mistral_tool_parser.extract_tool_calls_streaming( |
| 301 | previous_text, |
| 302 | current_text, |
| 303 | pending_text, |
| 304 | previous_token_ids, |
| 305 | current_token_ids, |
| 306 | pending_token_ids, |
| 307 | request=_DUMMY_REQUEST, |
| 308 | ) |
| 309 | else: # engine: per-delta windows, mirroring parse_delta behaviour |
| 310 | delta_message = _engine_parser.parse_delta( |
| 311 | delta_text=pending_text, |
| 312 | delta_token_ids=list(pending_token_ids), |
| 313 | request=_engine_req, |
| 314 | prompt_token_ids=[1], |
| 315 | finished=(i == len(all_token_ids) - 1), |
| 316 | ) |
| 317 | if delta_message: |
| 318 | yield delta_message |
| 319 | |
| 320 | previous_text = current_text |
| 321 | pending_text = "" |
| 322 | pending_token_ids = [] |
| 323 | |
| 324 | # For the legacy driver the engine holds back the final closing brace |
| 325 | # in _args_buffer until finish_streaming() is called, mirroring how |
| 326 | # the real serving layer flushes at end of stream. |
| 327 | if driver == "legacy": |
| 328 | flush_delta = mistral_tool_parser.finish_streaming() |
| 329 | if flush_delta: |
| 330 | yield flush_delta |
| 331 | |
| 332 | |
| 333 | @pytest.mark.parametrize( |
| 334 | "parser_fixture", |
| 335 | ["mistral_pre_v11_tool_parser", "mistral_tool_parser"], |
| 336 | ids=["pre_v11", "v11"], |
| 337 | ) |
| 338 | def test_extract_tool_calls_no_tools(parser_fixture, request): |
| 339 | parser = request.getfixturevalue(parser_fixture) |
| 340 | model_output = "This is a test" |
| 341 | result = parser.extract_tool_calls(model_output, request=_DUMMY_REQUEST) |
| 342 | assert result == ExtractedToolCallInformation( |
| 343 | tools_called=False, tool_calls=[], content=model_output |
| 344 | ) |
| 345 | |
| 346 | |
| 347 | @pytest.mark.parametrize( |
| 348 | ids=[ |
| 349 | "single_tool_add", |
| 350 | "single_tool_weather", |
| 351 | "argument_before_name", |
| 352 | "argument_before_name_and_name_in_argument", |
| 353 | "multiple_tools", |
| 354 | "content_before_tool", |
| 355 | "trailing_data_after_json", |
| 356 | ], |
| 357 | argnames=["model_output", "expected_tool_calls", "expected_content"], |
| 358 | argvalues=[ |
| 359 | ( |
| 360 | """[TOOL_CALLS][{"name": "add", "arguments":{"a": 3.5, "b": 4}}]""", # noqa: E501 |
| 361 | [ |
| 362 | ToolCall( |
| 363 | function=FunctionCall( |
| 364 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 365 | ) |
| 366 | ) |
| 367 | ], |
| 368 | None, |
| 369 | ), |
| 370 | ( |
| 371 | """[TOOL_CALLS] [{"name": "get_current_weather", "arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501 |
| 372 | [ |
| 373 | ToolCall( |
| 374 | function=FunctionCall( |
| 375 | name="get_current_weather", |
| 376 | arguments=json.dumps( |
| 377 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 378 | ), |
| 379 | ) |
| 380 | ) |
| 381 | ], |
| 382 | None, |
| 383 | ), |
| 384 | ( |
| 385 | """[TOOL_CALLS] [{"arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501 |
| 386 | [ |
| 387 | ToolCall( |
| 388 | function=FunctionCall( |
| 389 | name="get_current_weather", |
| 390 | arguments=json.dumps( |
| 391 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 392 | ), |
| 393 | ) |
| 394 | ) |
| 395 | ], |
| 396 | None, |
| 397 | ), |
| 398 | ( |
| 399 | """[TOOL_CALLS] [{"arguments":{"name": "John Doe"}, "name": "get_age"}]""", # noqa: E501 |
| 400 | [ |
| 401 | ToolCall( |
| 402 | function=FunctionCall( |
| 403 | name="get_age", |
| 404 | arguments=json.dumps( |
| 405 | { |
| 406 | "name": "John Doe", |
| 407 | } |
| 408 | ), |
| 409 | ) |
| 410 | ) |
| 411 | ], |
| 412 | None, |
| 413 | ), |
| 414 | ( |
| 415 | """[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, {"name": "get_current_weather", "arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501 |
| 416 | [ |
| 417 | ToolCall( |
| 418 | function=FunctionCall( |
| 419 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 420 | ) |
| 421 | ), |
| 422 | ToolCall( |
| 423 | function=FunctionCall( |
| 424 | name="get_current_weather", |
| 425 | arguments=json.dumps( |
| 426 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 427 | ), |
| 428 | ) |
| 429 | ), |
| 430 | ], |
| 431 | None, |
| 432 | ), |
| 433 | ( |
| 434 | """Hello[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]""", # noqa: E501 |
| 435 | [ |
| 436 | ToolCall( |
| 437 | function=FunctionCall( |
| 438 | name="add", arguments=json.dumps({"a": 1, "b": 2}) |
| 439 | ) |
| 440 | ) |
| 441 | ], |
| 442 | "Hello", |
| 443 | ), |
| 444 | ( |
| 445 | """[TOOL_CALLS] [{"name": "get_current_weather", "arguments":{"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]\nextra trailing data""", # noqa: E501 |
| 446 | [ |
| 447 | ToolCall( |
| 448 | function=FunctionCall( |
| 449 | name="get_current_weather", |
| 450 | arguments=json.dumps( |
| 451 | { |
| 452 | "city": "Dallas", |
| 453 | "state": "TX", |
| 454 | "unit": "fahrenheit", |
| 455 | } |
| 456 | ), |
| 457 | ) |
| 458 | ) |
| 459 | ], |
| 460 | None, |
| 461 | ), |
| 462 | ], |
| 463 | ) |
| 464 | def test_extract_tool_calls_pre_v11_tokenizer( |
| 465 | mistral_pre_v11_tool_parser, model_output, expected_tool_calls, expected_content |
| 466 | ): |
| 467 | extracted_tool_calls = mistral_pre_v11_tool_parser.extract_tool_calls( |
| 468 | model_output, request=_DUMMY_REQUEST |
| 469 | ) |
| 470 | assert extracted_tool_calls.tools_called |
| 471 | |
| 472 | assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) |
| 473 | |
| 474 | assert extracted_tool_calls.content == expected_content |
| 475 | |
| 476 | |
| 477 | def test_extract_tool_calls_pre_v11_multiple_bot_tokens_raises( |
| 478 | mistral_pre_v11_tool_parser, |
| 479 | ): |
| 480 | model_output = ( |
| 481 | '[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1}}]' |
| 482 | '[TOOL_CALLS] [{"name": "sub", "arguments":{"b": 2}}]' |
| 483 | ) |
| 484 | with pytest.raises(ValueError, match="Only one BOT token"): |
| 485 | mistral_pre_v11_tool_parser.extract_tool_calls( |
| 486 | model_output, request=_DUMMY_REQUEST |
| 487 | ) |
| 488 | |
| 489 | |
| 490 | def test_extract_tool_calls_pre_v11_regex_fallback( |
| 491 | mistral_pre_v11_tool_parser, |
| 492 | ): |
| 493 | """The regex fallback path finds valid JSON via regex when the primary |
| 494 | raw_decode fails on leading junk. It should re-serialize arguments |
| 495 | and return a valid tool call.""" |
| 496 | model_output = ( |
| 497 | '[TOOL_CALLS] junk [{"name": "add", "arguments":{"a": 1, "b": 2}}] trail' |
| 498 | ) |
| 499 | result = mistral_pre_v11_tool_parser.extract_tool_calls( |
| 500 | model_output, request=_DUMMY_REQUEST |
| 501 | ) |
| 502 | assert result.tools_called |
| 503 | assert len(result.tool_calls) == 1 |
| 504 | assert result.tool_calls[0].function.name == "add" |
| 505 | assert result.tool_calls[0].function.arguments == json.dumps({"a": 1, "b": 2}) |
| 506 | |
| 507 | |
| 508 | def test_extract_tool_calls_pre_v11_regex_fallback_fails( |
| 509 | mistral_pre_v11_tool_parser, |
| 510 | ): |
| 511 | model_output = "[TOOL_CALLS] not json at all" |
| 512 | result = mistral_pre_v11_tool_parser.extract_tool_calls( |
| 513 | model_output, request=_DUMMY_REQUEST |
| 514 | ) |
| 515 | assert result == ExtractedToolCallInformation( |
| 516 | tools_called=False, tool_calls=[], content="not json at all" |
| 517 | ) |
| 518 | |
| 519 | |
| 520 | @pytest.mark.parametrize( |
| 521 | ids=[ |
| 522 | "single_tool_add", |
| 523 | "single_tool_weather", |
| 524 | "multiple_tool_calls", |
| 525 | "single_tool_add_args", |
| 526 | "single_tool_weather_args", |
| 527 | "multiple_tool_calls_args", |
| 528 | "content_before_tool_args", |
| 529 | "complex", |
| 530 | "wrong_json", |
| 531 | "trailing_text_after_json", |
| 532 | "trailing_text_after_args_json", |
| 533 | ], |
| 534 | argnames=["model_output", "expected_tool_calls", "expected_content"], |
| 535 | argvalues=[ |
| 536 | ( |
| 537 | """[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501 |
| 538 | [ |
| 539 | ToolCall( |
| 540 | function=FunctionCall( |
| 541 | name="add_this_and_that", |
| 542 | arguments=json.dumps({"a": 3.5, "b": 4}), |
| 543 | ) |
| 544 | ) |
| 545 | ], |
| 546 | None, |
| 547 | ), |
| 548 | ( |
| 549 | """[TOOL_CALLS]get_current_weather{"city": "San Francisco", "state": "CA", "unit": "celsius"}""", # noqa: E501 |
| 550 | [ |
| 551 | ToolCall( |
| 552 | function=FunctionCall( |
| 553 | name="get_current_weather", |
| 554 | arguments=json.dumps( |
| 555 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 556 | ), |
| 557 | ) |
| 558 | ) |
| 559 | ], |
| 560 | None, |
| 561 | ), |
| 562 | ( |
| 563 | """[TOOL_CALLS]add{"a": 3.5, "b": 4}[TOOL_CALLS]multiply{"a": 3, "b": 6}""", # noqa: E501 |
| 564 | [ |
| 565 | ToolCall( |
| 566 | function=FunctionCall( |
| 567 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 568 | ) |
| 569 | ), |
| 570 | ToolCall( |
| 571 | function=FunctionCall( |
| 572 | name="multiply", arguments=json.dumps({"a": 3, "b": 6}) |
| 573 | ) |
| 574 | ), |
| 575 | ], |
| 576 | None, |
| 577 | ), |
| 578 | ( |
| 579 | # v11+ emits an explicit [ARGS] separator between name and args. |
| 580 | """[TOOL_CALLS]add_this_and_that[ARGS]{"a": 3.5, "b": 4}""", # noqa: E501 |
| 581 | [ |
| 582 | ToolCall( |
| 583 | function=FunctionCall( |
| 584 | name="add_this_and_that", |
| 585 | arguments=json.dumps({"a": 3.5, "b": 4}), |
| 586 | ) |
| 587 | ) |
| 588 | ], |
| 589 | None, |
| 590 | ), |
| 591 | ( |
| 592 | """[TOOL_CALLS]get_current_weather[ARGS]{"city": "San Francisco", "state": "CA", "unit": "celsius"}""", # noqa: E501 |
| 593 | [ |
| 594 | ToolCall( |
| 595 | function=FunctionCall( |
| 596 | name="get_current_weather", |
| 597 | arguments=json.dumps( |
| 598 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 599 | ), |
| 600 | ) |
| 601 | ) |
| 602 | ], |
| 603 | None, |
| 604 | ), |
| 605 | ( |
| 606 | """[TOOL_CALLS]add[ARGS]{"a": 3.5, "b": 4}[TOOL_CALLS]multiply[ARGS]{"a": 3, "b": 6}""", # noqa: E501 |
| 607 | [ |
| 608 | ToolCall( |
| 609 | function=FunctionCall( |
| 610 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 611 | ) |
| 612 | ), |
| 613 | ToolCall( |
| 614 | function=FunctionCall( |
| 615 | name="multiply", arguments=json.dumps({"a": 3, "b": 6}) |
| 616 | ) |
| 617 | ), |
| 618 | ], |
| 619 | None, |
| 620 | ), |
| 621 | ( |
| 622 | """hi[TOOL_CALLS]add[ARGS]{"a": 1, "b": 2}""", # noqa: E501 |
| 623 | [ |
| 624 | ToolCall( |
| 625 | function=FunctionCall( |
| 626 | name="add", arguments=json.dumps({"a": 1, "b": 2}) |
| 627 | ) |
| 628 | ) |
| 629 | ], |
| 630 | "hi", |
| 631 | ), |
| 632 | ( |
| 633 | # Complex |
| 634 | """hi{hi[TOOL_CALLS]bash{"command": "print(\\"hello world!\\")\\nre.compile(r\'{}\')""", # noqa: E501 |
| 635 | [ |
| 636 | ToolCall( |
| 637 | function=FunctionCall( |
| 638 | name="bash", |
| 639 | arguments=json.dumps( |
| 640 | {"command": "print(\"hello world!\")\nre.compile(r'{}')"} |
| 641 | )[:-2], |
| 642 | ) |
| 643 | ) |
| 644 | ], |
| 645 | "hi{hi", |
| 646 | ), |
| 647 | ( |
| 648 | # Wrong json |
| 649 | """hi{hi[TOOL_CALLS]bash{"command": "print(\\"hello world!\\")\\nre.compile(r\'{}\')"}""", # noqa: E501 |
| 650 | [ |
| 651 | ToolCall( |
| 652 | function=FunctionCall( |
| 653 | name="bash", |
| 654 | arguments=json.dumps( |
| 655 | {"command": "print(\"hello world!\")\nre.compile(r'{}')"} |
| 656 | ), |
| 657 | ) |
| 658 | ) |
| 659 | ], |
| 660 | "hi{hi", |
| 661 | ), |
| 662 | ( |
| 663 | # gh#48975: trailing text after the JSON object must not leak into |
| 664 | # arguments (v11+ name{args} format has no terminator). |
| 665 | """[TOOL_CALLS]get_current_weather{"city": "San Francisco"}Estimating the weather...""", # noqa: E501 |
| 666 | [ |
| 667 | ToolCall( |
| 668 | function=FunctionCall( |
| 669 | name="get_current_weather", |
| 670 | arguments=json.dumps({"city": "San Francisco"}), |
| 671 | ) |
| 672 | ) |
| 673 | ], |
| 674 | None, |
| 675 | ), |
| 676 | ( |
| 677 | # gh#48975 with the explicit [ARGS] separator. |
| 678 | """[TOOL_CALLS]get_current_weather[ARGS]{"city": "San Francisco"}Estimating the weather...""", # noqa: E501 |
| 679 | [ |
| 680 | ToolCall( |
| 681 | function=FunctionCall( |
| 682 | name="get_current_weather", |
| 683 | arguments=json.dumps({"city": "San Francisco"}), |
| 684 | ) |
| 685 | ) |
| 686 | ], |
| 687 | None, |
| 688 | ), |
| 689 | ], |
| 690 | ) |
| 691 | def test_extract_tool_calls( |
| 692 | mistral_tool_parser, model_output, expected_tool_calls, expected_content |
| 693 | ): |
| 694 | extracted_tool_calls = mistral_tool_parser.extract_tool_calls( |
| 695 | model_output, request=_DUMMY_REQUEST |
| 696 | ) |
| 697 | assert extracted_tool_calls.tools_called |
| 698 | |
| 699 | assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) |
| 700 | |
| 701 | assert extracted_tool_calls.content == expected_content |
| 702 | |
| 703 | |
| 704 | def test_extract_tool_calls_v11_without_args_skipped(mistral_tool_parser): |
| 705 | # Before Stage 2: the legacy path skipped tool calls with no arg separator, |
| 706 | # returning tools_called=True with an empty list. After Stage 2 the engine |
| 707 | # path is used for v11; it recognises the name and emits a tool call with |
| 708 | # empty arguments ("{}"). The semantic result is equivalent — a valid zero- |
| 709 | # argument call — and is strictly more correct than silently dropping it. |
| 710 | model_output = "[TOOL_CALLS]toolname_no_args" |
| 711 | result = mistral_tool_parser.extract_tool_calls( |
| 712 | model_output, request=_DUMMY_REQUEST |
| 713 | ) |
| 714 | assert result.tools_called |
| 715 | assert len(result.tool_calls) == 1 |
| 716 | assert result.tool_calls[0].function.name == "toolname_no_args" |
| 717 | assert result.tool_calls[0].function.arguments == "{}" |
| 718 | assert result.content is None |
| 719 | |
| 720 | |
| 721 | def test_extract_tool_calls_malformed_name_before_marker_emits_empty_name( |
| 722 | mistral_tool_parser, |
| 723 | ): |
| 724 | model_output = 'bash[TOOL_CALLS]{"command": "cat file.py"}' |
| 725 | result = mistral_tool_parser.extract_tool_calls( |
| 726 | model_output, request=_DUMMY_REQUEST |
| 727 | ) |
| 728 | assert result.tools_called |
| 729 | assert len(result.tool_calls) == 1 |
| 730 | assert result.tool_calls[0].function.name == "" |
| 731 | assert result.tool_calls[0].function.arguments == json.dumps( |
| 732 | {"command": "cat file.py"} |
| 733 | ) |
| 734 | assert result.content == "bash" |
| 735 | |
| 736 | |
| 737 | def test_extract_tool_calls_well_formed_name_unaffected(mistral_tool_parser): |
| 738 | model_output = '[TOOL_CALLS]bash[ARGS]{"command": "ls -la"}' |
| 739 | result = mistral_tool_parser.extract_tool_calls( |
| 740 | model_output, request=_DUMMY_REQUEST |
| 741 | ) |
| 742 | assert result.tools_called |
| 743 | assert len(result.tool_calls) == 1 |
| 744 | assert result.tool_calls[0].function.name == "bash" |
| 745 | assert result.tool_calls[0].function.arguments == json.dumps({"command": "ls -la"}) |
| 746 | assert result.content is None |
| 747 | |
| 748 | |
| 749 | def test_extract_tool_calls_empty_name_unparsable_args_no_crash( |
| 750 | mistral_tool_parser, |
| 751 | ): |
| 752 | # Empty name AND args that never form valid JSON: still emitted (raw |
| 753 | # text as arguments) instead of raising. |
| 754 | model_output = "bash[TOOL_CALLS]{not valid json at all" |
| 755 | result = mistral_tool_parser.extract_tool_calls( |
| 756 | model_output, request=_DUMMY_REQUEST |
| 757 | ) |
| 758 | assert result.tools_called |
| 759 | assert len(result.tool_calls) == 1 |
| 760 | assert result.tool_calls[0].function.name == "" |
| 761 | assert result.tool_calls[0].function.arguments == "{not valid json at all" |
| 762 | assert result.content == "bash" |
| 763 | |
| 764 | |
| 765 | def test_extract_tool_calls_streaming_malformed_name_before_marker( |
| 766 | mistral_tool_parser, |
| 767 | ): |
| 768 | # Streaming counterpart: once the args start the slot can never go back to |
| 769 | # naming, so the (empty) name is already final and streams immediately |
| 770 | # alongside the args. |
| 771 | model_output = 'bash[TOOL_CALLS]{"command": "cat file.py"}' |
| 772 | mid_delta = mistral_tool_parser.extract_tool_calls_streaming( |
| 773 | previous_text="", |
| 774 | current_text=model_output, |
| 775 | delta_text=model_output, |
| 776 | previous_token_ids=[], |
| 777 | current_token_ids=[], |
| 778 | delta_token_ids=[], |
| 779 | request=_DUMMY_REQUEST, |
| 780 | ) |
| 781 | assert mid_delta is not None |
| 782 | assert mid_delta.content == "bash" |
| 783 | assert len(mid_delta.tool_calls) == 1 |
| 784 | assert mid_delta.tool_calls[0].function.name == "" |
| 785 | assert mid_delta.tool_calls[0].function.arguments == '{"command": "cat file.py"' |
| 786 | |
| 787 | final_delta = mistral_tool_parser.finish_streaming() |
| 788 | assert final_delta is not None |
| 789 | assert len(final_delta.tool_calls) == 1 |
| 790 | assert final_delta.tool_calls[0].function.name is None |
| 791 | assert final_delta.tool_calls[0].function.arguments == "}" |
| 792 | |
| 793 | |
| 794 | def test_extract_tool_calls_parallel_malformed_then_well_formed( |
| 795 | mistral_tool_parser, |
| 796 | ): |
| 797 | # A malformed empty-name call followed by a well-formed one must not |
| 798 | # regress into brace-splitting: nested braces in the second call's |
| 799 | # arguments must stay intact, and both calls must be correctly separated. |
| 800 | model_output = ( |
| 801 | 'bash[TOOL_CALLS]{"command": "ls"}' |
| 802 | '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}' |
| 803 | ) |
| 804 | result = mistral_tool_parser.extract_tool_calls( |
| 805 | model_output, request=_DUMMY_REQUEST |
| 806 | ) |
| 807 | assert result.tools_called |
| 808 | assert len(result.tool_calls) == 2 |
| 809 | assert result.tool_calls[0].function.name == "" |
| 810 | assert result.tool_calls[0].function.arguments == json.dumps({"command": "ls"}) |
| 811 | assert result.tool_calls[1].function.name == "submit" |
| 812 | assert result.tool_calls[1].function.arguments == json.dumps( |
| 813 | {"answer": {"nested": 1}} |
| 814 | ) |
| 815 | assert result.content == "bash" |
| 816 | |
| 817 | |
| 818 | def _test_extract_tool_calls_streaming( |
| 819 | tool_parser, |
| 820 | tokenizer, |
| 821 | model_output, |
| 822 | tools, |
| 823 | expected_tool_calls, |
| 824 | expected_content, |
| 825 | driver: str = "legacy", |
| 826 | reasoning_parser: str | None = None, |
| 827 | ): |
| 828 | other_content: str = "" |
| 829 | function_names: list[str] = [] |
| 830 | function_args_strs: list[str] = [] |
| 831 | tool_call_idx: int = -1 |
| 832 | tool_call_ids: list[str | None] = [] |
| 833 | |
| 834 | for delta_message in stream_delta_message_generator( |
| 835 | tool_parser, |
| 836 | tokenizer, |
| 837 | model_output, |
| 838 | tools, |
| 839 | driver=driver, |
| 840 | reasoning_parser=reasoning_parser, |
| 841 | ): |
| 842 | # role should never be streamed from tool parser |
| 843 | assert not delta_message.role |
| 844 | |
| 845 | if delta_message.content: |
| 846 | other_content += delta_message.content |
| 847 | |
| 848 | streamed_tool_calls = delta_message.tool_calls |
| 849 | |
| 850 | if streamed_tool_calls and len(streamed_tool_calls) > 0: |
| 851 | # On the final finished=True delta the engine's _flush_engine_parsers |
| 852 | # may append the flushed '}' as a separate DeltaToolCall for the |
| 853 | # same index without coalescing. Merge entries sharing an index so |
| 854 | # the "one diff per delta" invariant can be checked after merging. |
| 855 | if len(streamed_tool_calls) > 1: |
| 856 | assert len({tc.index for tc in streamed_tool_calls}) == 1, ( |
| 857 | "only within-chunk finish/flush of one tool is allowed; " |
| 858 | "distinct indices in one delta indicate a regression: " |
| 859 | f"{[tc.index for tc in streamed_tool_calls]}" |
| 860 | ) |
| 861 | by_index: dict[int, DeltaToolCall] = {} |
| 862 | for tc in streamed_tool_calls: |
| 863 | existing = by_index.get(tc.index) |
| 864 | if existing is None: |
| 865 | by_index[tc.index] = tc |
| 866 | else: |
| 867 | if tc.id and not existing.id: |
| 868 | existing.id = tc.id |
| 869 | if tc.type and not existing.type: |
| 870 | existing.type = tc.type |
| 871 | if tc.function and existing.function and tc.function.arguments: |
| 872 | existing.function.arguments = ( |
| 873 | existing.function.arguments or "" |
| 874 | ) + tc.function.arguments |
| 875 | streamed_tool_calls = list(by_index.values()) |
| 876 | |
| 877 | # make sure only one diff is present - correct even for parallel |
| 878 | assert len(streamed_tool_calls) == 1 |
| 879 | tool_call = streamed_tool_calls[0] |
| 880 | |
| 881 | # if a new tool is being called, set up empty arguments |
| 882 | if tool_call.index != tool_call_idx: |
| 883 | tool_call_idx = tool_call.index |
| 884 | function_args_strs.append("") |
| 885 | tool_call_ids.append(None) |
| 886 | |
| 887 | # if a tool call ID is streamed, make sure one hasn't been already |
| 888 | if tool_call.id and not tool_call_ids[tool_call.index]: |
| 889 | tool_call_ids[tool_call.index] = tool_call.id |
| 890 | |
| 891 | # if parts of the function start being streamed |
| 892 | if tool_call.function: |
| 893 | # if the function name is defined, set it. it should be streamed |
| 894 | # IN ENTIRETY, exactly one time. |
| 895 | if tool_call.function.name: |
| 896 | assert isinstance(tool_call.function.name, str) |
| 897 | function_names.append(tool_call.function.name) |
| 898 | |
| 899 | if tool_call.function.arguments: |
| 900 | # make sure they're a string and then add them to the list |
| 901 | assert isinstance(tool_call.function.arguments, str) |
| 902 | |
| 903 | function_args_strs[tool_call.index] += tool_call.function.arguments |
| 904 | |
| 905 | assert other_content == expected_content |
| 906 | |
| 907 | actual_tool_calls = [ |
| 908 | ToolCall( |
| 909 | id=tool_call_id, |
| 910 | function=FunctionCall( |
| 911 | name=function_name, |
| 912 | arguments=partial_json_parser.ensure_json( |
| 913 | function_args_str, Allow.OBJ | Allow.STR |
| 914 | ), |
| 915 | ), |
| 916 | ) |
| 917 | for tool_call_id, function_name, function_args_str in zip( |
| 918 | tool_call_ids, function_names, function_args_strs |
| 919 | ) |
| 920 | ] |
| 921 | assert_tool_calls(actual_tool_calls, expected_tool_calls) |
| 922 | |
| 923 | |
| 924 | @pytest.mark.parametrize("driver", ["legacy", "engine"]) |
| 925 | @pytest.mark.parametrize( |
| 926 | ids=[ |
| 927 | "no_tools", |
| 928 | "single_tool_add", |
| 929 | "single_tool_add_strings", |
| 930 | "single_tool_weather", |
| 931 | "argument_before_name", |
| 932 | "argument_before_name_and_name_in_argument", |
| 933 | "multiple_tools", |
| 934 | "trailing_data_after_json", |
| 935 | ], |
| 936 | argnames=["model_output", "expected_tool_calls", "expected_content"], |
| 937 | argvalues=[ |
| 938 | ("""This is a test""", [], """This is a test"""), |
| 939 | ( |
| 940 | """[TOOL_CALLS] [ {"name":"add" , "arguments" : {"a": 3, "b": 4} } ]""", # noqa: E501 |
| 941 | [ |
| 942 | ToolCall( |
| 943 | function=FunctionCall( |
| 944 | name="add", arguments=json.dumps({"a": 3, "b": 4}) |
| 945 | ) |
| 946 | ) |
| 947 | ], |
| 948 | "", |
| 949 | ), |
| 950 | ( |
| 951 | """[TOOL_CALLS] [{"name": "add", "arguments":{"a": "3", "b": "4"}}]""", # noqa: E501 |
| 952 | [ |
| 953 | ToolCall( |
| 954 | function=FunctionCall( |
| 955 | name="add", arguments=json.dumps({"a": "3", "b": "4"}) |
| 956 | ) |
| 957 | ) |
| 958 | ], |
| 959 | "", |
| 960 | ), |
| 961 | ( |
| 962 | """[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501 |
| 963 | [ |
| 964 | ToolCall( |
| 965 | function=FunctionCall( |
| 966 | name="get_current_weather", |
| 967 | arguments=json.dumps( |
| 968 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 969 | ), |
| 970 | ) |
| 971 | ) |
| 972 | ], |
| 973 | "", |
| 974 | ), |
| 975 | ( |
| 976 | """[TOOL_CALLS] [{"arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501 |
| 977 | [ |
| 978 | ToolCall( |
| 979 | function=FunctionCall( |
| 980 | name="get_current_weather", |
| 981 | arguments=json.dumps( |
| 982 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 983 | ), |
| 984 | ) |
| 985 | ) |
| 986 | ], |
| 987 | "", |
| 988 | ), |
| 989 | ( |
| 990 | """[TOOL_CALLS] [{"arguments": {"name": "John Doe"}, "name": "get_age"}]""", # noqa: E501 |
| 991 | [ |
| 992 | ToolCall( |
| 993 | function=FunctionCall( |
| 994 | name="get_age", |
| 995 | arguments=json.dumps( |
| 996 | { |
| 997 | "name": "John Doe", |
| 998 | } |
| 999 | ), |
| 1000 | ) |
| 1001 | ) |
| 1002 | ], |
| 1003 | "", |
| 1004 | ), |
| 1005 | ( |
| 1006 | """[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, {"name": "get_current_weather", "arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501 |
| 1007 | [ |
| 1008 | ToolCall( |
| 1009 | function=FunctionCall( |
| 1010 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 1011 | ) |
| 1012 | ), |
| 1013 | ToolCall( |
| 1014 | function=FunctionCall( |
| 1015 | name="get_current_weather", |
| 1016 | arguments=json.dumps( |
| 1017 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 1018 | ), |
| 1019 | ) |
| 1020 | ), |
| 1021 | ], |
| 1022 | "", |
| 1023 | ), |
| 1024 | ( |
| 1025 | """[TOOL_CALLS] [{"name": "get_current_weather", "arguments":{"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}]\nextra trailing data""", # noqa: E501 |
| 1026 | [ |
| 1027 | ToolCall( |
| 1028 | function=FunctionCall( |
| 1029 | name="get_current_weather", |
| 1030 | arguments=json.dumps( |
| 1031 | { |
| 1032 | "city": "Dallas", |
| 1033 | "state": "TX", |
| 1034 | "unit": "fahrenheit", |
| 1035 | } |
| 1036 | ), |
| 1037 | ) |
| 1038 | ) |
| 1039 | ], |
| 1040 | "\nextra trailing data", |
| 1041 | ), |
| 1042 | ], |
| 1043 | ) |
| 1044 | def test_extract_tool_calls_streaming_pre_v11_tokenizer( |
| 1045 | mistral_pre_v11_tool_parser, |
| 1046 | mistral_pre_v11_tokenizer, |
| 1047 | model_output, |
| 1048 | expected_tool_calls, |
| 1049 | expected_content, |
| 1050 | driver, |
| 1051 | ): |
| 1052 | _test_extract_tool_calls_streaming( |
| 1053 | mistral_pre_v11_tool_parser, |
| 1054 | mistral_pre_v11_tokenizer, |
| 1055 | model_output, |
| 1056 | None, |
| 1057 | expected_tool_calls, |
| 1058 | expected_content, |
| 1059 | driver=driver, |
| 1060 | ) |
| 1061 | |
| 1062 | |
| 1063 | @pytest.mark.parametrize( |
| 1064 | "driver,reasoning_parser", |
| 1065 | [("legacy", None), ("engine", None), ("engine", "mistral")], |
| 1066 | ids=["legacy", "engine_no_reasoning", "engine_with_reasoning"], |
| 1067 | ) |
| 1068 | @pytest.mark.parametrize( |
| 1069 | ids=[ |
| 1070 | "single_tool_add", |
| 1071 | "single_tool_add_strings", |
| 1072 | "multiple_tools", |
| 1073 | ], |
| 1074 | argnames=["tools", "expected_tool_calls", "expected_content"], |
| 1075 | argvalues=[ |
| 1076 | ( |
| 1077 | [("add", '{"a": 3, "b": 4}')], |
| 1078 | # [TOOL_CALLS]add{"a": 3, "b": 4} |
| 1079 | [ |
| 1080 | ToolCall( |
| 1081 | function=FunctionCall( |
| 1082 | name="add", arguments=json.dumps({"a": 3, "b": 4}) |
| 1083 | ) |
| 1084 | ) |
| 1085 | ], |
| 1086 | "", |
| 1087 | ), |
| 1088 | ( |
| 1089 | [("add_two_strings", '{"a": "3", "b": "4"}')], |
| 1090 | # [TOOL_CALLS]add_two_strings{"a": "3", "b": "4"} |
| 1091 | [ |
| 1092 | ToolCall( |
| 1093 | function=FunctionCall( |
| 1094 | name="add_two_strings", |
| 1095 | arguments=json.dumps({"a": "3", "b": "4"}), |
| 1096 | ) |
| 1097 | ) |
| 1098 | ], |
| 1099 | "", |
| 1100 | ), |
| 1101 | ( |
| 1102 | [ |
| 1103 | ("add", '{"a": 3.5, "b": 4}'), |
| 1104 | ( |
| 1105 | "get_current_weather", |
| 1106 | '{"city": "San Francisco", "state": "CA", "unit": "celsius"}', # noqa: E501 |
| 1107 | ), |
| 1108 | ], |
| 1109 | # [TOOL_CALLS]add{"a": 3.5, "b": 4}[TOOL_CALLS]get_current_weather{"city": "San Francisco", "state": "CA", "unit": "celsius"} # noqa: E501 |
| 1110 | [ |
| 1111 | ToolCall( |
| 1112 | function=FunctionCall( |
| 1113 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 1114 | ) |
| 1115 | ), |
| 1116 | ToolCall( |
| 1117 | function=FunctionCall( |
| 1118 | name="get_current_weather", |
| 1119 | arguments=json.dumps( |
| 1120 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 1121 | ), |
| 1122 | ) |
| 1123 | ), |
| 1124 | ], |
| 1125 | "", |
| 1126 | ), |
| 1127 | ], |
| 1128 | ) |
| 1129 | def test_extract_tool_calls_streaming( |
| 1130 | mistral_tool_parser, |
| 1131 | mistral_tokenizer, |
| 1132 | tools, |
| 1133 | expected_tool_calls, |
| 1134 | expected_content, |
| 1135 | driver, |
| 1136 | reasoning_parser, |
| 1137 | ): |
| 1138 | _test_extract_tool_calls_streaming( |
| 1139 | mistral_tool_parser, |
| 1140 | mistral_tokenizer, |
| 1141 | None, |
| 1142 | tools, |
| 1143 | expected_tool_calls, |
| 1144 | expected_content, |
| 1145 | driver=driver, |
| 1146 | reasoning_parser=reasoning_parser, |
| 1147 | ) |
| 1148 | |
| 1149 | |
| 1150 | # Drives extract_tool_calls_streaming directly (bespoke detokenization loop), |
| 1151 | # not via stream_delta_message_generator — driver parametrization not needed. |
| 1152 | def test_extract_tool_calls_streaming_v11_no_tools( |
| 1153 | mistral_tool_parser, mistral_tokenizer |
| 1154 | ): |
| 1155 | model_output = "This is a test" |
| 1156 | # add_special_tokens=False: real serving streams only generated tokens. |
| 1157 | all_token_ids = mistral_tokenizer.encode(model_output, add_special_tokens=False) |
| 1158 | skip_special = isinstance(mistral_tokenizer, MistralTokenizer) |
| 1159 | collected_content = "" |
| 1160 | previous_text = "" |
| 1161 | previous_tokens = None |
| 1162 | prefix_offset = 0 |
| 1163 | read_offset = 0 |
| 1164 | for i in range(len(all_token_ids)): |
| 1165 | current_token_ids = all_token_ids[: i + 1] |
| 1166 | previous_token_ids = all_token_ids[:i] |
| 1167 | delta_token_ids = [all_token_ids[i]] |
| 1168 | |
| 1169 | new_tokens, delta_text, prefix_offset, read_offset = detokenize_incrementally( |
| 1170 | tokenizer=mistral_tokenizer, |
| 1171 | all_input_ids=current_token_ids, |
| 1172 | prev_tokens=previous_tokens, |
| 1173 | prefix_offset=prefix_offset, |
| 1174 | read_offset=read_offset, |
| 1175 | skip_special_tokens=skip_special, |
| 1176 | spaces_between_special_tokens=True, |
| 1177 | ) |
| 1178 | current_text = previous_text + delta_text |
| 1179 | previous_tokens = ( |
| 1180 | previous_tokens + new_tokens if previous_tokens else new_tokens |
| 1181 | ) |
| 1182 | |
| 1183 | delta_message = mistral_tool_parser.extract_tool_calls_streaming( |
| 1184 | previous_text=previous_text, |
| 1185 | current_text=current_text, |
| 1186 | delta_text=delta_text, |
| 1187 | previous_token_ids=previous_token_ids, |
| 1188 | current_token_ids=current_token_ids, |
| 1189 | delta_token_ids=delta_token_ids, |
| 1190 | request=_DUMMY_REQUEST, |
| 1191 | ) |
| 1192 | if delta_message and delta_message.content: |
| 1193 | collected_content += delta_message.content |
| 1194 | if delta_message: |
| 1195 | assert not delta_message.tool_calls |
| 1196 | |
| 1197 | previous_text = current_text |
| 1198 | |
| 1199 | assert collected_content == model_output |
| 1200 | |
| 1201 | |
| 1202 | def test_mistral_parser_drops_eos_from_output(mistral_tokenizer): |
| 1203 | """EOS token must never surface as content/reasoning; literal EOS string |
| 1204 | must be preserved when the EOS token id is absent. |
| 1205 | |
| 1206 | The engine drops special tokens by id, so: |
| 1207 | Case A: the real EOS token id causes the text to be dropped. |
| 1208 | Case B: the same EOS string with a non-EOS token id is preserved. |
| 1209 | """ |
| 1210 | if not isinstance(mistral_tokenizer, MistralTokenizer): |
| 1211 | pytest.skip("Requires MistralTokenizer") |
| 1212 | |
| 1213 | eos_id: int = mistral_tokenizer.eos_token_id |
| 1214 | eos_text: str = mistral_tokenizer.decode([eos_id]) |
| 1215 | assert eos_text |
| 1216 | |
| 1217 | # Case A: real EOS token — text must be dropped. |
| 1218 | parser_a = MistralParser(mistral_tokenizer) |
| 1219 | parser_a.initialize_streaming() |
| 1220 | events_a = parser_a._feed(eos_text, [eos_id]) |
| 1221 | events_a.extend(parser_a._engine.finish()) |
| 1222 | delta_a = parser_a._events_to_delta(events_a, finished=True) |
| 1223 | content_a = (delta_a.content if delta_a else None) or "" |
| 1224 | reasoning_a = (delta_a.reasoning if delta_a else None) or "" |
| 1225 | assert eos_text not in content_a |
| 1226 | assert eos_text not in reasoning_a |
| 1227 | |
| 1228 | # Case B: EOS text with a non-EOS token id — text must be preserved. |
| 1229 | # This proves content equal to the EOS string is not silently dropped |
| 1230 | # when it did not come from the real EOS token. |
| 1231 | non_eos_ids = mistral_tokenizer.encode(text="hello", add_special_tokens=False) |
| 1232 | non_eos_id = next(tid for tid in non_eos_ids if tid != eos_id) |
| 1233 | parser_b = MistralParser(mistral_tokenizer) |
| 1234 | parser_b.initialize_streaming() |
| 1235 | events_b = parser_b._feed(eos_text, [non_eos_id]) |
| 1236 | events_b.extend(parser_b._engine.finish()) |
| 1237 | delta_b = parser_b._events_to_delta(events_b, finished=True) |
| 1238 | content_b = (delta_b.content if delta_b else None) or "" |
| 1239 | reasoning_b = (delta_b.reasoning if delta_b else None) or "" |
| 1240 | assert eos_text in content_b or eos_text in reasoning_b |
| 1241 | |
| 1242 | |
| 1243 | @pytest.mark.parametrize( |
| 1244 | "parser_fixture, tokenizer_fixture, model_output," |
| 1245 | " expected_tool_calls, expected_content", |
| 1246 | [ |
| 1247 | pytest.param( |
| 1248 | "mistral_tool_parser", |
| 1249 | "mistral_tokenizer", |
| 1250 | """[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501 |
| 1251 | [ |
| 1252 | ToolCall( |
| 1253 | function=FunctionCall( |
| 1254 | name="add_this_and_that", |
| 1255 | arguments=json.dumps({"a": 3.5, "b": 4}), |
| 1256 | ) |
| 1257 | ) |
| 1258 | ], |
| 1259 | "", |
| 1260 | id="v11-single_tool_add", |
| 1261 | ), |
| 1262 | pytest.param( |
| 1263 | "mistral_tool_parser", |
| 1264 | "mistral_tokenizer", |
| 1265 | """[TOOL_CALLS]get_current_weather{"city": "San Francisco", "state": "CA", "unit": "celsius"}""", # noqa: E501 |
| 1266 | [ |
| 1267 | ToolCall( |
| 1268 | function=FunctionCall( |
| 1269 | name="get_current_weather", |
| 1270 | arguments=json.dumps( |
| 1271 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 1272 | ), |
| 1273 | ) |
| 1274 | ) |
| 1275 | ], |
| 1276 | "", |
| 1277 | id="v11-single_tool_weather", |
| 1278 | ), |
| 1279 | pytest.param( |
| 1280 | "mistral_tool_parser", |
| 1281 | "mistral_tokenizer", |
| 1282 | """[TOOL_CALLS]add{"a": 3.5, "b": 4}[TOOL_CALLS]multiply{"a": 3, "b": 6}""", # noqa: E501 |
| 1283 | [ |
| 1284 | ToolCall( |
| 1285 | function=FunctionCall( |
| 1286 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 1287 | ) |
| 1288 | ), |
| 1289 | ToolCall( |
| 1290 | function=FunctionCall( |
| 1291 | name="multiply", arguments=json.dumps({"a": 3, "b": 6}) |
| 1292 | ) |
| 1293 | ), |
| 1294 | ], |
| 1295 | "", |
| 1296 | id="v11-multiple_tool_calls", |
| 1297 | ), |
| 1298 | pytest.param( |
| 1299 | "mistral_tool_parser", |
| 1300 | "mistral_tokenizer", |
| 1301 | """bla[TOOL_CALLS]add_this_and_that{"a": 3.5, "b": 4}""", # noqa: E501 |
| 1302 | [ |
| 1303 | ToolCall( |
| 1304 | function=FunctionCall( |
| 1305 | name="add_this_and_that", |
| 1306 | arguments=json.dumps({"a": 3.5, "b": 4}), |
| 1307 | ) |
| 1308 | ) |
| 1309 | ], |
| 1310 | "bla", |
| 1311 | id="v11-content_before_tool", |
| 1312 | ), |
| 1313 | pytest.param( |
| 1314 | "mistral_tool_parser", |
| 1315 | "mistral_tokenizer", |
| 1316 | """[TOOL_CALLS]add_this_and_that[ARGS]{"a": 3.5, "b": 4}""", # noqa: E501 |
| 1317 | [ |
| 1318 | ToolCall( |
| 1319 | function=FunctionCall( |
| 1320 | name="add_this_and_that", |
| 1321 | arguments=json.dumps({"a": 3.5, "b": 4}), |
| 1322 | ) |
| 1323 | ) |
| 1324 | ], |
| 1325 | "", |
| 1326 | id="v11-single_tool_add_args", |
| 1327 | ), |
| 1328 | pytest.param( |
| 1329 | "mistral_tool_parser", |
| 1330 | "mistral_tokenizer", |
| 1331 | """[TOOL_CALLS]add[ARGS]{"a": 3.5, "b": 4}[TOOL_CALLS]multiply[ARGS]{"a": 3, "b": 6}""", # noqa: E501 |
| 1332 | [ |
| 1333 | ToolCall( |
| 1334 | function=FunctionCall( |
| 1335 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 1336 | ) |
| 1337 | ), |
| 1338 | ToolCall( |
| 1339 | function=FunctionCall( |
| 1340 | name="multiply", arguments=json.dumps({"a": 3, "b": 6}) |
| 1341 | ) |
| 1342 | ), |
| 1343 | ], |
| 1344 | "", |
| 1345 | id="v11-multiple_tool_calls_args", |
| 1346 | ), |
| 1347 | pytest.param( |
| 1348 | "mistral_tool_parser", |
| 1349 | "mistral_tokenizer", |
| 1350 | """bla[TOOL_CALLS]add_this_and_that[ARGS]{"a": 3.5, "b": 4}""", # noqa: E501 |
| 1351 | [ |
| 1352 | ToolCall( |
| 1353 | function=FunctionCall( |
| 1354 | name="add_this_and_that", |
| 1355 | arguments=json.dumps({"a": 3.5, "b": 4}), |
| 1356 | ) |
| 1357 | ) |
| 1358 | ], |
| 1359 | "bla", |
| 1360 | id="v11-content_before_tool_args", |
| 1361 | ), |
| 1362 | pytest.param( |
| 1363 | "mistral_tool_parser", |
| 1364 | "mistral_tokenizer", |
| 1365 | """hi{hi[TOOL_CALLS]bash{"command": "print(\\"hello world!\\")\\nre.compile(r\'{}\')"}""", # noqa: E501 |
| 1366 | [ |
| 1367 | ToolCall( |
| 1368 | function=FunctionCall( |
| 1369 | name="bash", |
| 1370 | arguments=json.dumps( |
| 1371 | {"command": "print(\"hello world!\")\nre.compile(r'{}')"} |
| 1372 | ), |
| 1373 | ) |
| 1374 | ) |
| 1375 | ], |
| 1376 | "hi{hi", |
| 1377 | id="v11-complex", |
| 1378 | ), |
| 1379 | pytest.param( |
| 1380 | "mistral_pre_v11_tool_parser", |
| 1381 | "mistral_pre_v11_tokenizer", |
| 1382 | """This is a test""", |
| 1383 | [], |
| 1384 | """This is a test""", |
| 1385 | id="pre_v11-no_tools", |
| 1386 | ), |
| 1387 | pytest.param( |
| 1388 | "mistral_pre_v11_tool_parser", |
| 1389 | "mistral_pre_v11_tokenizer", |
| 1390 | """[TOOL_CALLS] [ {"name":"add" , "arguments" : {"a": 3, "b": 4} } ]""", # noqa: E501 |
| 1391 | [ |
| 1392 | ToolCall( |
| 1393 | function=FunctionCall( |
| 1394 | name="add", arguments=json.dumps({"a": 3, "b": 4}) |
| 1395 | ) |
| 1396 | ) |
| 1397 | ], |
| 1398 | "", |
| 1399 | id="pre_v11-single_tool_add", |
| 1400 | ), |
| 1401 | pytest.param( |
| 1402 | "mistral_pre_v11_tool_parser", |
| 1403 | "mistral_pre_v11_tokenizer", |
| 1404 | """[TOOL_CALLS] [{"name": "add", "arguments":{"a": "3", "b": "4"}}]""", # noqa: E501 |
| 1405 | [ |
| 1406 | ToolCall( |
| 1407 | function=FunctionCall( |
| 1408 | name="add", arguments=json.dumps({"a": "3", "b": "4"}) |
| 1409 | ) |
| 1410 | ) |
| 1411 | ], |
| 1412 | "", |
| 1413 | id="pre_v11-single_tool_add_strings", |
| 1414 | ), |
| 1415 | pytest.param( |
| 1416 | "mistral_pre_v11_tool_parser", |
| 1417 | "mistral_pre_v11_tokenizer", |
| 1418 | """[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}}]""", # noqa: E501 |
| 1419 | [ |
| 1420 | ToolCall( |
| 1421 | function=FunctionCall( |
| 1422 | name="get_current_weather", |
| 1423 | arguments=json.dumps( |
| 1424 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 1425 | ), |
| 1426 | ) |
| 1427 | ) |
| 1428 | ], |
| 1429 | "", |
| 1430 | id="pre_v11-single_tool_weather", |
| 1431 | ), |
| 1432 | pytest.param( |
| 1433 | "mistral_pre_v11_tool_parser", |
| 1434 | "mistral_pre_v11_tokenizer", |
| 1435 | """[TOOL_CALLS] [{"arguments": {"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501 |
| 1436 | [ |
| 1437 | ToolCall( |
| 1438 | function=FunctionCall( |
| 1439 | name="get_current_weather", |
| 1440 | arguments=json.dumps( |
| 1441 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 1442 | ), |
| 1443 | ) |
| 1444 | ) |
| 1445 | ], |
| 1446 | "", |
| 1447 | id="pre_v11-argument_before_name", |
| 1448 | ), |
| 1449 | pytest.param( |
| 1450 | "mistral_pre_v11_tool_parser", |
| 1451 | "mistral_pre_v11_tokenizer", |
| 1452 | """[TOOL_CALLS] [{"arguments": {"name": "John Doe"}, "name": "get_age"}]""", # noqa: E501 |
| 1453 | [ |
| 1454 | ToolCall( |
| 1455 | function=FunctionCall( |
| 1456 | name="get_age", |
| 1457 | arguments=json.dumps( |
| 1458 | { |
| 1459 | "name": "John Doe", |
| 1460 | } |
| 1461 | ), |
| 1462 | ) |
| 1463 | ) |
| 1464 | ], |
| 1465 | "", |
| 1466 | id="pre_v11-argument_before_name_and_name_in_argument", |
| 1467 | ), |
| 1468 | pytest.param( |
| 1469 | "mistral_pre_v11_tool_parser", |
| 1470 | "mistral_pre_v11_tokenizer", |
| 1471 | """[TOOL_CALLS] [{"arguments": {"a": 3.5, "b": 4}, "name": "add"}, {"arguments":{"city": "San Francisco", "state": "CA", "unit": "celsius"}, "name": "get_current_weather"}]""", # noqa: E501 |
| 1472 | [ |
| 1473 | ToolCall( |
| 1474 | function=FunctionCall( |
| 1475 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 1476 | ) |
| 1477 | ), |
| 1478 | ToolCall( |
| 1479 | function=FunctionCall( |
| 1480 | name="get_current_weather", |
| 1481 | arguments=json.dumps( |
| 1482 | {"city": "San Francisco", "state": "CA", "unit": "celsius"} |
| 1483 | ), |
| 1484 | ) |
| 1485 | ), |
| 1486 | ], |
| 1487 | "", |
| 1488 | id="pre_v11-multiple_tools", |
| 1489 | ), |
| 1490 | pytest.param( |
| 1491 | "mistral_pre_v11_tool_parser", |
| 1492 | "mistral_pre_v11_tokenizer", |
| 1493 | """Some text[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]""", # noqa: E501 |
| 1494 | [ |
| 1495 | ToolCall( |
| 1496 | function=FunctionCall( |
| 1497 | name="add", arguments=json.dumps({"a": 1, "b": 2}) |
| 1498 | ) |
| 1499 | ) |
| 1500 | ], |
| 1501 | "Some text", |
| 1502 | id="pre_v11-content_before_tool", |
| 1503 | ), |
| 1504 | ], |
| 1505 | ) |
| 1506 | def test_extract_tool_calls_streaming_one_chunk( |
| 1507 | parser_fixture, |
| 1508 | tokenizer_fixture, |
| 1509 | model_output, |
| 1510 | expected_tool_calls, |
| 1511 | expected_content, |
| 1512 | request, |
| 1513 | ): |
| 1514 | tool_parser = request.getfixturevalue(parser_fixture) |
| 1515 | tokenizer = request.getfixturevalue(tokenizer_fixture) |
| 1516 | |
| 1517 | # Real serving streams only generated tokens (no BOS/EOS framing) with the |
| 1518 | # control markers as special ids. |
| 1519 | if isinstance(tokenizer, MistralTokenizer) and tokenizer.version >= 11: |
| 1520 | all_token_ids = encode_mistral_output(tokenizer, model_output) |
| 1521 | else: |
| 1522 | all_token_ids = tokenizer.encode(model_output, add_special_tokens=False) |
| 1523 | all_token_ids = fix_tool_call_tokenization( |
| 1524 | all_token_ids, tool_parser, tokenizer |
| 1525 | ) |
| 1526 | |
| 1527 | delta_message = tool_parser.extract_tool_calls_streaming( |
| 1528 | previous_text="", |
| 1529 | current_text=model_output, |
| 1530 | delta_text=model_output, |
| 1531 | previous_token_ids=[], |
| 1532 | current_token_ids=all_token_ids, |
| 1533 | delta_token_ids=all_token_ids, |
| 1534 | request=_DUMMY_REQUEST, |
| 1535 | ) |
| 1536 | # The engine buffers the final closing brace ('}') in _args_buffer until |
| 1537 | # finish_streaming() is called, mirroring real serving. Flush it and |
| 1538 | # merge the result so the asserted arguments are complete. |
| 1539 | flush_delta = tool_parser.finish_streaming() |
| 1540 | if flush_delta is not None and flush_delta.tool_calls: |
| 1541 | if delta_message is not None and delta_message.tool_calls: |
| 1542 | for flush_tc in flush_delta.tool_calls: |
| 1543 | for existing_tc in delta_message.tool_calls: |
| 1544 | if existing_tc.index == flush_tc.index and flush_tc.function: |
| 1545 | existing_tc.function = ( |
| 1546 | existing_tc.function or DeltaFunctionCall() |
| 1547 | ) |
| 1548 | existing_tc.function.arguments = ( |
| 1549 | existing_tc.function.arguments or "" |
| 1550 | ) + (flush_tc.function.arguments or "") |
| 1551 | elif delta_message is not None: |
| 1552 | delta_message.tool_calls = flush_delta.tool_calls |
| 1553 | else: |
| 1554 | delta_message = flush_delta |
| 1555 | |
| 1556 | assert isinstance(delta_message, DeltaMessage) |
| 1557 | assert len(delta_message.tool_calls) == len(expected_tool_calls) |
| 1558 | |
| 1559 | assert_tool_calls(delta_message.tool_calls, expected_tool_calls) |
| 1560 | |
| 1561 | if delta_message.content is None: |
| 1562 | assert expected_content == "" |
| 1563 | else: |
| 1564 | assert delta_message.content == expected_content |
| 1565 | |
| 1566 | |
| 1567 | @pytest.mark.parametrize( |
| 1568 | "parser_fixture, model_output, fake_count, two_phase", |
| 1569 | [ |
| 1570 | pytest.param( |
| 1571 | "mistral_pre_v11_tool_parser", |
| 1572 | '[TOOL_CALLS] [{"name": "add", "arguments":{"a": 1, "b": 2}}]', |
| 1573 | 30, |
| 1574 | False, |
| 1575 | id="pre_v11", |
| 1576 | ), |
| 1577 | ], |
| 1578 | ) |
| 1579 | def test_fast_detokenization_text_detection( |
| 1580 | parser_fixture, model_output, fake_count, two_phase, request |
| 1581 | ): |
| 1582 | """Regression: bot_token in text but not token_ids (PR #37209). |
| 1583 | |
| 1584 | Only the pre-v11 legacy path detects the marker from text; v11+ routes |
| 1585 | through the engine and detects the marker by token id. |
| 1586 | """ |
| 1587 | parser = request.getfixturevalue(parser_fixture) |
| 1588 | # Token IDs that do NOT contain bot_token_id. |
| 1589 | fake_token_ids = list(range(99, 99 + fake_count)) |
| 1590 | |
| 1591 | if two_phase: |
| 1592 | # First delta: pure content, no bot token yet |
| 1593 | delta_message_before = parser.extract_tool_calls_streaming( |
| 1594 | previous_text="", |
| 1595 | current_text="Hello", |
| 1596 | delta_text="Hello", |
| 1597 | previous_token_ids=[], |
| 1598 | current_token_ids=[99], |
| 1599 | delta_token_ids=[99], |
| 1600 | request=_DUMMY_REQUEST, |
| 1601 | ) |
| 1602 | assert delta_message_before is not None |
| 1603 | assert delta_message_before.content == "Hello" |
| 1604 | assert not delta_message_before.tool_calls |
| 1605 | |
| 1606 | previous_text = "Hello" |
| 1607 | current_text = "Hello" + model_output |
| 1608 | previous_token_ids = [99] |
| 1609 | delta_token_ids = fake_token_ids[1:] |
| 1610 | else: |
| 1611 | previous_text = "" |
| 1612 | current_text = model_output |
| 1613 | previous_token_ids = [] |
| 1614 | delta_token_ids = fake_token_ids |
| 1615 | |
| 1616 | delta_message = parser.extract_tool_calls_streaming( |
| 1617 | previous_text=previous_text, |
| 1618 | current_text=current_text, |
| 1619 | delta_text=model_output, |
| 1620 | previous_token_ids=previous_token_ids, |
| 1621 | current_token_ids=fake_token_ids, |
| 1622 | delta_token_ids=delta_token_ids, |
| 1623 | request=_DUMMY_REQUEST, |
| 1624 | ) |
| 1625 | assert delta_message is not None |
| 1626 | assert delta_message.tool_calls is not None |
| 1627 | assert len(delta_message.tool_calls) == 1 |
| 1628 | assert delta_message.tool_calls[0].function is not None |
| 1629 | assert delta_message.tool_calls[0].function.name == "add" |
| 1630 | |
| 1631 | |
| 1632 | def test_extract_tool_calls_streaming_exception_returns_none( |
| 1633 | mistral_pre_v11_tool_parser, |
| 1634 | ): |
| 1635 | # v11 routes through the ParserEngine and does not swallow exceptions; |
| 1636 | # only the pre-v11 legacy state-machine path has explicit exception handling. |
| 1637 | parser = mistral_pre_v11_tool_parser |
| 1638 | patched_method = "_extract_tool_calls_streaming_pre_v11_tokenizer" |
| 1639 | current_text = '[TOOL_CALLS] [{"name":"a","arguments":{}}]' |
| 1640 | with patch.object( |
| 1641 | parser._parser_engine, patched_method, side_effect=RuntimeError("boom") |
| 1642 | ): |
| 1643 | result = parser.extract_tool_calls_streaming( |
| 1644 | previous_text="", |
| 1645 | current_text=current_text, |
| 1646 | delta_text=current_text, |
| 1647 | previous_token_ids=[], |
| 1648 | current_token_ids=[parser._parser_engine.bot_token_id], |
| 1649 | delta_token_ids=[parser._parser_engine.bot_token_id], |
| 1650 | request=_DUMMY_REQUEST, |
| 1651 | ) |
| 1652 | assert result is None |
| 1653 | |
| 1654 | |
| 1655 | SAMPLE_TOOLS_DICTS = [ |
| 1656 | { |
| 1657 | "type": "function", |
| 1658 | "function": { |
| 1659 | "name": "get_weather", |
| 1660 | "description": "Get the weather", |
| 1661 | "parameters": { |
| 1662 | "type": "object", |
| 1663 | "properties": {"city": {"type": "string"}}, |
| 1664 | "required": ["city"], |
| 1665 | }, |
| 1666 | }, |
| 1667 | }, |
| 1668 | { |
| 1669 | "type": "function", |
| 1670 | "function": { |
| 1671 | "name": "add", |
| 1672 | "description": "Add two numbers", |
| 1673 | "parameters": { |
| 1674 | "type": "object", |
| 1675 | "properties": { |
| 1676 | "a": {"type": "number"}, |
| 1677 | "b": {"type": "number"}, |
| 1678 | }, |
| 1679 | "required": ["a", "b"], |
| 1680 | }, |
| 1681 | }, |
| 1682 | }, |
| 1683 | ] |
| 1684 | |
| 1685 | |
| 1686 | def _make_request(**kwargs) -> ChatCompletionRequest: |
| 1687 | defaults: dict = { |
| 1688 | "messages": [], |
| 1689 | "model": "mistralai/Mistral-Small-3.2-24B-Instruct-2506", |
| 1690 | "tools": SAMPLE_TOOLS_DICTS, |
| 1691 | "tool_choice": "auto", |
| 1692 | } |
| 1693 | defaults.update(kwargs) |
| 1694 | return ChatCompletionRequest(**defaults) |
| 1695 | |
| 1696 | |
| 1697 | @pytest.mark.parametrize( |
| 1698 | "request_kwargs,expected_mode,expected_parallel", |
| 1699 | [ |
| 1700 | ({"tool_choice": "auto"}, MistralToolChoiceEnum.auto, True), |
| 1701 | ({"tool_choice": "none"}, MistralToolChoiceEnum.none, True), |
| 1702 | ({"tool_choice": "required"}, MistralToolChoiceEnum.required, True), |
| 1703 | ({"tool_choice": None, "tools": None}, MistralToolChoiceEnum.auto, True), |
| 1704 | ( |
| 1705 | { |
| 1706 | "tool_choice": { |
| 1707 | "type": "function", |
| 1708 | "function": {"name": "get_weather"}, |
| 1709 | } |
| 1710 | }, |
| 1711 | MistralNamedToolChoice.model_validate( |
| 1712 | {"type": "function", "function": {"name": "get_weather"}} |
| 1713 | ), |
| 1714 | True, |
| 1715 | ), |
| 1716 | ( |
| 1717 | {"tool_choice": "auto", "parallel_tool_calls": False}, |
| 1718 | MistralToolChoiceEnum.auto, |
| 1719 | False, |
| 1720 | ), |
| 1721 | ( |
| 1722 | {"tool_choice": "auto", "response_format": {"type": "text"}}, |
| 1723 | MistralToolChoiceEnum.auto, |
| 1724 | True, |
| 1725 | ), |
| 1726 | ], |
| 1727 | ids=[ |
| 1728 | "auto", |
| 1729 | "none", |
| 1730 | "required", |
| 1731 | "null_tool_choice", |
| 1732 | "named_tool_choice", |
| 1733 | "parallel_false", |
| 1734 | "response_format_text", |
| 1735 | ], |
| 1736 | ) |
| 1737 | def test_adjust_request_grammar_factory( |
| 1738 | mistral_tool_parser: MistralToolParser, |
| 1739 | request_kwargs: dict, |
| 1740 | expected_mode: MistralToolChoice, |
| 1741 | expected_parallel: bool, |
| 1742 | ) -> None: |
| 1743 | request = _make_request(**request_kwargs) |
| 1744 | factory = mistral_tool_parser.model_tokenizer.grammar_factory |
| 1745 | |
| 1746 | with patch.object( |
| 1747 | factory, |
| 1748 | "get_lark_from_jinja", |
| 1749 | wraps=factory.get_lark_from_jinja, |
| 1750 | ) as mock_get_lark: |
| 1751 | result = mistral_tool_parser.adjust_request(request) |
| 1752 | |
| 1753 | mock_get_lark.assert_called_once() |
| 1754 | call_kwargs = mock_get_lark.call_args |
| 1755 | |
| 1756 | assert call_kwargs.kwargs["mode"] == expected_mode |
| 1757 | assert call_kwargs.kwargs["json_schema"] is None |
| 1758 | assert call_kwargs.kwargs["parallel_tool_calls"] == expected_parallel |
| 1759 | |
| 1760 | assert result.structured_outputs is not None |
| 1761 | assert isinstance(result.structured_outputs.grammar, str) |
| 1762 | assert len(result.structured_outputs.grammar) > 0 |
| 1763 | |
| 1764 | |
| 1765 | def test_adjust_request_unsupported_grammar_for_tokenizer(mistral_tokenizer) -> None: |
| 1766 | with patch.object( |
| 1767 | type(mistral_tokenizer), |
| 1768 | "supports_grammar", |
| 1769 | new_callable=lambda: property(lambda self: False), |
| 1770 | ): |
| 1771 | parser = MistralToolParser(mistral_tokenizer) |
| 1772 | request = _make_request() |
| 1773 | result = parser.adjust_request(request) |
| 1774 | |
| 1775 | assert result.structured_outputs is None |
| 1776 | |
| 1777 | |
| 1778 | @pytest.mark.parametrize( |
| 1779 | "tool_choice,expected_skip", |
| 1780 | [("auto", False), ("none", False)], |
| 1781 | ids=["auto_skip_false", "none_skip_false"], |
| 1782 | ) |
| 1783 | def test_adjust_request_non_mistral_tokenizer( |
| 1784 | non_mistral_parser: MistralToolParser, |
| 1785 | tool_choice: str, |
| 1786 | expected_skip: bool, |
| 1787 | ) -> None: |
| 1788 | # MistralParser (ParserEngine) always sets skip_special_tokens=False so |
| 1789 | # that special tokens like [TOOL_CALLS] are visible to the parser even |
| 1790 | # when tool_choice="none". |
| 1791 | request = _make_request(tool_choice=tool_choice) |
| 1792 | result = non_mistral_parser.adjust_request(request) |
| 1793 | |
| 1794 | assert result.skip_special_tokens is expected_skip |
| 1795 | |
| 1796 | |
| 1797 | @pytest.mark.parametrize( |
| 1798 | "so_kwargs", |
| 1799 | [ |
| 1800 | {"regex": r"\d+"}, |
| 1801 | {"choice": ["a", "b"]}, |
| 1802 | { |
| 1803 | "structural_tag": json.dumps( |
| 1804 | { |
| 1805 | "structures": [ |
| 1806 | { |
| 1807 | "begin": "<tool>", |
| 1808 | "schema": {"type": "object"}, |
| 1809 | "end": "</tool>", |
| 1810 | } |
| 1811 | ], |
| 1812 | "triggers": ["<tool>"], |
| 1813 | } |
| 1814 | ) |
| 1815 | }, |
| 1816 | {"grammar": "start: 'hello'"}, |
| 1817 | ], |
| 1818 | ids=["regex", "choice", "structural_tag", "grammar"], |
| 1819 | ) |
| 1820 | def test_adjust_request_unsupported_structured_outputs( |
| 1821 | mistral_tool_parser: MistralToolParser, |
| 1822 | so_kwargs: dict, |
| 1823 | ) -> None: |
| 1824 | request = _make_request( |
| 1825 | structured_outputs=StructuredOutputsParams(**so_kwargs), |
| 1826 | ) |
| 1827 | result = mistral_tool_parser.adjust_request(request) |
| 1828 | |
| 1829 | assert result.structured_outputs == request.structured_outputs |
| 1830 | |
| 1831 | |
| 1832 | def test_adjust_request_unsupported_response_format( |
| 1833 | mistral_tool_parser: MistralToolParser, |
| 1834 | ) -> None: |
| 1835 | request = _make_request( |
| 1836 | response_format=StructuralTagResponseFormat( |
| 1837 | type="structural_tag", |
| 1838 | format={ |
| 1839 | "type": "triggered_tags", |
| 1840 | "tags": [ |
| 1841 | { |
| 1842 | "begin": "<tool>", |
| 1843 | "content": {"type": "any_text"}, |
| 1844 | "end": "</tool>", |
| 1845 | } |
| 1846 | ], |
| 1847 | "triggers": ["<tool>"], |
| 1848 | }, |
| 1849 | ), |
| 1850 | ) |
| 1851 | result = mistral_tool_parser.adjust_request(request) |
| 1852 | assert result.structured_outputs is None |
| 1853 | assert result.response_format == request.response_format |
| 1854 | |
| 1855 | |
| 1856 | @pytest.mark.parametrize( |
| 1857 | "so_kwargs,expected_json_schema", |
| 1858 | [ |
| 1859 | ({"json_object": True}, _DEFAULT_JSON_SCHEMA), |
| 1860 | ({"json": '{"type": "object"}'}, {"type": "object"}), |
| 1861 | ( |
| 1862 | {"json": {"type": "object", "properties": {"x": {"type": "integer"}}}}, |
| 1863 | {"type": "object", "properties": {"x": {"type": "integer"}}}, |
| 1864 | ), |
| 1865 | ], |
| 1866 | ids=["json_object", "json_str", "json_dict"], |
| 1867 | ) |
| 1868 | def test_adjust_request_structured_outputs_generates_grammar( |
| 1869 | mistral_tool_parser: MistralToolParser, |
| 1870 | so_kwargs: dict, |
| 1871 | expected_json_schema: str, |
| 1872 | ) -> None: |
| 1873 | request = _make_request( |
| 1874 | structured_outputs=StructuredOutputsParams(**so_kwargs), |
| 1875 | ) |
| 1876 | factory = mistral_tool_parser.model_tokenizer.grammar_factory |
| 1877 | |
| 1878 | with patch.object( |
| 1879 | factory, |
| 1880 | "get_lark_from_jinja", |
| 1881 | wraps=factory.get_lark_from_jinja, |
| 1882 | ) as mock_get_lark: |
| 1883 | result = mistral_tool_parser.adjust_request(request) |
| 1884 | |
| 1885 | mock_get_lark.assert_called_once() |
| 1886 | assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema |
| 1887 | |
| 1888 | assert result.structured_outputs is not None |
| 1889 | assert isinstance(result.structured_outputs.grammar, str) |
| 1890 | assert len(result.structured_outputs.grammar) > 0 |
| 1891 | |
| 1892 | |
| 1893 | @pytest.mark.parametrize( |
| 1894 | "response_format_kwargs,expected_json_schema", |
| 1895 | [ |
| 1896 | ({"type": "json_object"}, _DEFAULT_JSON_SCHEMA), |
| 1897 | ( |
| 1898 | { |
| 1899 | "type": "json_schema", |
| 1900 | "json_schema": { |
| 1901 | "name": "my_schema", |
| 1902 | "schema": { |
| 1903 | "type": "object", |
| 1904 | "properties": {"x": {"type": "integer"}}, |
| 1905 | }, |
| 1906 | }, |
| 1907 | }, |
| 1908 | {"type": "object", "properties": {"x": {"type": "integer"}}}, |
| 1909 | ), |
| 1910 | ], |
| 1911 | ids=["json_object", "json_schema_with_schema"], |
| 1912 | ) |
| 1913 | def test_adjust_request_response_format_generates_grammar( |
| 1914 | mistral_tool_parser: MistralToolParser, |
| 1915 | response_format_kwargs: dict, |
| 1916 | expected_json_schema: str, |
| 1917 | ) -> None: |
| 1918 | request = _make_request(response_format=response_format_kwargs) |
| 1919 | factory = mistral_tool_parser.model_tokenizer.grammar_factory |
| 1920 | |
| 1921 | with patch.object( |
| 1922 | factory, |
| 1923 | "get_lark_from_jinja", |
| 1924 | wraps=factory.get_lark_from_jinja, |
| 1925 | ) as mock_get_lark: |
| 1926 | result = mistral_tool_parser.adjust_request(request) |
| 1927 | |
| 1928 | mock_get_lark.assert_called_once() |
| 1929 | assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema |
| 1930 | |
| 1931 | assert result.structured_outputs is not None |
| 1932 | assert isinstance(result.structured_outputs.grammar, str) |
| 1933 | assert len(result.structured_outputs.grammar) > 0 |
| 1934 | |
| 1935 | |
| 1936 | @pytest.mark.parametrize( |
| 1937 | "tool_choice, expected_method, not_called_method", |
| 1938 | [ |
| 1939 | ("none", "get_lark_for_json_schema", None), |
| 1940 | ("auto", "get_lark_from_jinja", "get_lark_for_json_schema"), |
| 1941 | ], |
| 1942 | ids=["none_uses_json_schema_factory", "auto_uses_jinja_factory"], |
| 1943 | ) |
| 1944 | def test_adjust_request_tool_choice_with_json_schema_factory_routing( |
| 1945 | mistral_tool_parser: MistralToolParser, |
| 1946 | tool_choice: str, |
| 1947 | expected_method: str, |
| 1948 | not_called_method: str | None, |
| 1949 | ) -> None: |
| 1950 | request = _make_request( |
| 1951 | tool_choice=tool_choice, |
| 1952 | structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), |
| 1953 | ) |
| 1954 | factory = mistral_tool_parser.model_tokenizer.grammar_factory |
| 1955 | |
| 1956 | patches = { |
| 1957 | expected_method: patch.object( |
| 1958 | factory, |
| 1959 | expected_method, |
| 1960 | wraps=getattr(factory, expected_method), |
| 1961 | ), |
| 1962 | } |
| 1963 | if not_called_method: |
| 1964 | patches[not_called_method] = patch.object( |
| 1965 | factory, |
| 1966 | not_called_method, |
| 1967 | wraps=getattr(factory, not_called_method), |
| 1968 | ) |
| 1969 | |
| 1970 | with patches[expected_method] as mock_expected: |
| 1971 | ctx = patches[not_called_method] if not_called_method else None |
| 1972 | if ctx: |
| 1973 | with ctx as mock_not_called: |
| 1974 | result = mistral_tool_parser.adjust_request(request) |
| 1975 | mock_not_called.assert_not_called() |
| 1976 | else: |
| 1977 | result = mistral_tool_parser.adjust_request(request) |
| 1978 | |
| 1979 | mock_expected.assert_called_once() |
| 1980 | assert mock_expected.call_args.kwargs["json_schema"] == {"type": "object"} |
| 1981 | |
| 1982 | assert result.structured_outputs is not None |
| 1983 | assert isinstance(result.structured_outputs.grammar, str) |
| 1984 | assert len(result.structured_outputs.grammar) > 0 |
| 1985 | |
| 1986 | |
| 1987 | def test_grammar_from_parser_default_false() -> None: |
| 1988 | request = _make_request() |
| 1989 | assert request._grammar_from_parser is False |
| 1990 | |
| 1991 | |
| 1992 | def test_grammar_from_parser_set_by_adjust_request( |
| 1993 | mistral_tool_parser: MistralToolParser, |
| 1994 | ) -> None: |
| 1995 | request = _make_request() |
| 1996 | result = mistral_tool_parser.adjust_request(request) |
| 1997 | assert result._grammar_from_parser is True |
| 1998 | |
| 1999 | |
| 2000 | @pytest.mark.parametrize("driver", ["legacy", "engine"]) |
| 2001 | @pytest.mark.parametrize("chunk_size", [2, 3, 4, 5]) |
| 2002 | def test_streaming_pre_v11_parallel_calls_batched_deltas( |
| 2003 | mistral_pre_v11_tool_parser, mistral_pre_v11_tokenizer, chunk_size, driver |
| 2004 | ): |
| 2005 | """A batched delta spanning the boundary between two parallel calls must |
| 2006 | keep them on distinct indices (the bug collapsed both onto index 0).""" |
| 2007 | model_output = ( |
| 2008 | '[TOOL_CALLS] [{"name": "add", "arguments": {"a": 3.5, "b": 4}}, ' |
| 2009 | '{"name": "get_current_weather", "arguments": ' |
| 2010 | '{"city": "San Francisco", "state": "CA", "unit": "celsius"}}]' |
| 2011 | ) |
| 2012 | names: list[str] = [] |
| 2013 | args: list[str] = [] |
| 2014 | idx = -1 |
| 2015 | for delta_message in stream_delta_message_generator( |
| 2016 | mistral_pre_v11_tool_parser, |
| 2017 | mistral_pre_v11_tokenizer, |
| 2018 | model_output, |
| 2019 | tools=None, |
| 2020 | chunk_size=chunk_size, |
| 2021 | driver=driver, |
| 2022 | ): |
| 2023 | for tool_call in delta_message.tool_calls or []: |
| 2024 | if tool_call.index != idx: |
| 2025 | idx = tool_call.index |
| 2026 | args.append("") |
| 2027 | if tool_call.function and tool_call.function.name: |
| 2028 | names.append(tool_call.function.name) |
| 2029 | if tool_call.function and tool_call.function.arguments: |
| 2030 | args[tool_call.index] += tool_call.function.arguments |
| 2031 | |
| 2032 | assert names == ["add", "get_current_weather"] |
| 2033 | assert len(args) == 2 |
| 2034 | # trailing args of the final call are flushed by the serving layer |
| 2035 | assert json.loads(args[0]) == {"a": 3.5, "b": 4} |
| 2036 | |
| 2037 | |
| 2038 | @pytest.mark.parametrize( |
| 2039 | "reasoning_encoding", |
| 2040 | ["text", "special_token"], |
| 2041 | ids=["text_encoding", "special_token_encoding"], |
| 2042 | ) |
| 2043 | def test_content_tool_calls_transition_emits_reasoning_end( |
| 2044 | mistral_tokenizer, |
| 2045 | reasoning_encoding, |
| 2046 | ): |
| 2047 | """(CONTENT, TOOL_CALLS) must emit REASONING_END when reasoning is enabled.""" |
| 2048 | cfg = mistral_config(reasoning_encoding=reasoning_encoding) |
| 2049 | engine = StreamingParserEngine( |
| 2050 | config=cfg, tokenizer=mistral_tokenizer, vocab=mistral_tokenizer.get_vocab() |
| 2051 | ) |
| 2052 | engine.skip_tool_parsing = True |
| 2053 | |
| 2054 | bot_token_id = mistral_tokenizer.get_vocab().get("[TOOL_CALLS]") |
| 2055 | assert bot_token_id is not None |
| 2056 | |
| 2057 | events = engine.feed(delta_text="[TOOL_CALLS]", delta_token_ids=[bot_token_id]) |
| 2058 | event_types = [e.type for e in events] |
| 2059 | |
| 2060 | assert EventType.REASONING_END in event_types, ( |
| 2061 | f"REASONING_END missing for reasoning_encoding={reasoning_encoding!r}; " |
| 2062 | f"got events: {event_types}" |
| 2063 | ) |
| 2064 | # Tool call start must NOT appear in skip_tool_parsing mode |
| 2065 | assert EventType.TOOL_CALL_START not in event_types |
| 2066 | |
| 2067 | |
| 2068 | @pytest.mark.parametrize("chunk_size", [1, 3]) |
| 2069 | @pytest.mark.parametrize( |
| 2070 | "tools,expected_tool_calls,expected_content", |
| 2071 | [ |
| 2072 | ( |
| 2073 | [("add", '{"a": 3, "b": 4}')], |
| 2074 | [ |
| 2075 | ToolCall( |
| 2076 | function=FunctionCall( |
| 2077 | name="add", arguments=json.dumps({"a": 3, "b": 4}) |
| 2078 | ) |
| 2079 | ) |
| 2080 | ], |
| 2081 | "", |
| 2082 | ), |
| 2083 | ( |
| 2084 | [ |
| 2085 | ("add", '{"a": 3.5, "b": 4}'), |
| 2086 | ( |
| 2087 | "get_current_weather", |
| 2088 | '{"city": "San Francisco", "state": "CA", "unit": "celsius"}', |
| 2089 | ), |
| 2090 | ], |
| 2091 | [ |
| 2092 | ToolCall( |
| 2093 | function=FunctionCall( |
| 2094 | name="add", arguments=json.dumps({"a": 3.5, "b": 4}) |
| 2095 | ) |
| 2096 | ), |
| 2097 | ToolCall( |
| 2098 | function=FunctionCall( |
| 2099 | name="get_current_weather", |
| 2100 | arguments=json.dumps( |
| 2101 | { |
| 2102 | "city": "San Francisco", |
| 2103 | "state": "CA", |
| 2104 | "unit": "celsius", |
| 2105 | } |
| 2106 | ), |
| 2107 | ) |
| 2108 | ), |
| 2109 | ], |
| 2110 | "", |
| 2111 | ), |
| 2112 | ], |
| 2113 | ids=["single_tool_add", "parallel_tools"], |
| 2114 | ) |
| 2115 | def test_reasoning_active_no_think_block_no_leak( |
| 2116 | mistral_tool_parser, |
| 2117 | mistral_tokenizer, |
| 2118 | tools, |
| 2119 | expected_tool_calls, |
| 2120 | expected_content, |
| 2121 | chunk_size, |
| 2122 | ): |
| 2123 | """Tool call without a preceding think block must not leak as content |
| 2124 | when the reasoning parser is active.""" |
| 2125 | accumulated_content = "" |
| 2126 | function_names: list[str] = [] |
| 2127 | function_args_strs: list[str] = [] |
| 2128 | tool_call_idx = -1 |
| 2129 | tool_call_ids: list[str | None] = [] |
| 2130 | |
| 2131 | for delta_message in stream_delta_message_generator( |
| 2132 | mistral_tool_parser, |
| 2133 | mistral_tokenizer, |
| 2134 | model_output=None, |
| 2135 | tools=tools, |
| 2136 | chunk_size=chunk_size, |
| 2137 | driver="engine", |
| 2138 | reasoning_parser="mistral", |
| 2139 | ): |
| 2140 | if delta_message.content: |
| 2141 | accumulated_content += delta_message.content |
| 2142 | |
| 2143 | for tool_call in delta_message.tool_calls or []: |
| 2144 | if tool_call.index != tool_call_idx: |
| 2145 | tool_call_idx = tool_call.index |
| 2146 | function_args_strs.append("") |
| 2147 | tool_call_ids.append(None) |
| 2148 | if tool_call.id and not tool_call_ids[tool_call.index]: |
| 2149 | tool_call_ids[tool_call.index] = tool_call.id |
| 2150 | if tool_call.function: |
| 2151 | if tool_call.function.name: |
| 2152 | function_names.append(tool_call.function.name) |
| 2153 | if tool_call.function.arguments: |
| 2154 | function_args_strs[tool_call.index] += tool_call.function.arguments |
| 2155 | |
| 2156 | assert "[TOOL_CALLS]" not in accumulated_content |
| 2157 | assert "[ARGS]" not in accumulated_content |
| 2158 | assert accumulated_content == expected_content |
| 2159 | |
| 2160 | actual_tool_calls = [ |
| 2161 | ToolCall( |
| 2162 | id=tc_id, |
| 2163 | function=FunctionCall( |
| 2164 | name=name, |
| 2165 | arguments=partial_json_parser.ensure_json(args, Allow.OBJ | Allow.STR), |
| 2166 | ), |
| 2167 | ) |
| 2168 | for tc_id, name, args in zip(tool_call_ids, function_names, function_args_strs) |
| 2169 | ] |
| 2170 | assert_tool_calls(actual_tool_calls, expected_tool_calls) |
| 2171 | |
| 2172 | |
| 2173 | # --------------------------------------------------------------------------- |
| 2174 | # Pre-v11 guided schema injection and bare-array extraction tests |
| 2175 | # --------------------------------------------------------------------------- |
| 2176 | |
| 2177 | |
| 2178 | @pytest.mark.parametrize( |
| 2179 | "tool_choice,expected_items_count", |
| 2180 | [ |
| 2181 | ("required", 2), # all tools |
| 2182 | ( |
| 2183 | {"type": "function", "function": {"name": "get_weather"}}, |
| 2184 | 1, |
| 2185 | ), # named: single tool |
| 2186 | ], |
| 2187 | ids=["required", "named"], |
| 2188 | ) |
| 2189 | def test_adjust_request_pre_v11_guided_schema_injected( |
| 2190 | mistral_pre_v11_tool_parser: MistralToolParser, |
| 2191 | tool_choice: object, |
| 2192 | expected_items_count: int, |
| 2193 | ) -> None: |
| 2194 | request = _make_request(tool_choice=tool_choice) |
| 2195 | result = mistral_pre_v11_tool_parser.adjust_request(request) |
| 2196 | |
| 2197 | assert result.structured_outputs is not None |
| 2198 | schema = result.structured_outputs.json |
| 2199 | if isinstance(schema, str): |
| 2200 | schema = json.loads(schema) |
| 2201 | assert schema["type"] == "array" |
| 2202 | assert schema["minItems"] == 1 |
| 2203 | items = schema["items"] |
| 2204 | assert "anyOf" in items |
| 2205 | assert len(items["anyOf"]) == expected_items_count |
| 2206 | for entry in items["anyOf"]: |
| 2207 | props = entry["properties"] |
| 2208 | assert "name" in props |
| 2209 | assert "arguments" in props |
| 2210 | |
| 2211 | |
| 2212 | @pytest.mark.parametrize( |
| 2213 | "response_format", |
| 2214 | [{"type": "json_object"}, {"type": "json_schema", "json_schema": {"name": "s"}}], |
| 2215 | ids=["json_object", "json_schema"], |
| 2216 | ) |
| 2217 | def test_adjust_request_pre_v11_required_clears_response_format( |
| 2218 | mistral_pre_v11_tool_parser: MistralToolParser, |
| 2219 | response_format: dict, |
| 2220 | ) -> None: |
| 2221 | """Required + response_format must inject the tool schema and clear |
| 2222 | response_format so the tool schema is the sole structured-output |
| 2223 | constraint, matching the base ToolParser (otherwise the request either |
| 2224 | hits the "multiple constraints" engine error or, with no constraint, |
| 2225 | rambles to finish_reason='length').""" |
| 2226 | request = _make_request(tool_choice="required", response_format=response_format) |
| 2227 | result = mistral_pre_v11_tool_parser.adjust_request(request) |
| 2228 | |
| 2229 | assert result.response_format is None |
| 2230 | assert result.structured_outputs is not None |
| 2231 | schema = result.structured_outputs.json |
| 2232 | if isinstance(schema, str): |
| 2233 | schema = json.loads(schema) |
| 2234 | assert schema["type"] == "array" |
| 2235 | |
| 2236 | |
| 2237 | def test_adjust_request_non_mistral_tokenizer_required_injects_schema( |
| 2238 | non_mistral_parser: MistralToolParser, |
| 2239 | ) -> None: |
| 2240 | """The guided-schema injection must also fire for non-Mistral (e.g. HF-mode) |
| 2241 | tokenizers driving the Mistral tool parser, mirroring the base ToolParser so |
| 2242 | required + response_format does not fall back to an unconstrained ramble.""" |
| 2243 | request = _make_request( |
| 2244 | tool_choice="required", response_format={"type": "json_object"} |
| 2245 | ) |
| 2246 | result = non_mistral_parser.adjust_request(request) |
| 2247 | |
| 2248 | assert result.response_format is None |
| 2249 | assert result.structured_outputs is not None |
| 2250 | assert result.structured_outputs.json is not None |
| 2251 | |
| 2252 | |
| 2253 | @pytest.mark.parametrize( |
| 2254 | "tool_choice", |
| 2255 | ["auto", "none"], |
| 2256 | ids=["auto", "none"], |
| 2257 | ) |
| 2258 | def test_adjust_request_pre_v11_no_injection_for_auto_none( |
| 2259 | mistral_pre_v11_tool_parser: MistralToolParser, |
| 2260 | tool_choice: str, |
| 2261 | ) -> None: |
| 2262 | request = _make_request(tool_choice=tool_choice) |
| 2263 | result = mistral_pre_v11_tool_parser.adjust_request(request) |
| 2264 | |
| 2265 | assert result.structured_outputs is None |
| 2266 | |
| 2267 | |
| 2268 | def test_legacy_extract_tool_calls_guided_bare_array_required( |
| 2269 | mistral_pre_v11_tool_parser: MistralToolParser, |
| 2270 | ) -> None: |
| 2271 | model_output = '[{"name": "get_current_weather", "arguments": {"city": "Dallas"}}]' |
| 2272 | request = _make_request(tool_choice="required") |
| 2273 | |
| 2274 | result = mistral_pre_v11_tool_parser.extract_tool_calls( |
| 2275 | model_output, request=request |
| 2276 | ) |
| 2277 | |
| 2278 | assert result.tools_called |
| 2279 | assert len(result.tool_calls) == 1 |
| 2280 | tc = result.tool_calls[0] |
| 2281 | assert tc.function.name == "get_current_weather" |
| 2282 | assert json.loads(tc.function.arguments) == {"city": "Dallas"} |
| 2283 | assert isinstance(tc.id, str) |
| 2284 | assert len(tc.id) == 9 |
| 2285 | assert result.content is None |
| 2286 | |
| 2287 | |
| 2288 | def test_legacy_extract_tool_calls_none_with_tools( |
| 2289 | mistral_pre_v11_tool_parser: MistralToolParser, |
| 2290 | ) -> None: |
| 2291 | model_output = ( |
| 2292 | '[TOOL_CALLS] [{"name": "get_current_weather",' |
| 2293 | ' "arguments": {"city": "Dallas"}}]' |
| 2294 | ) |
| 2295 | request = _make_request(tool_choice="none") |
| 2296 | |
| 2297 | result = mistral_pre_v11_tool_parser.extract_tool_calls( |
| 2298 | model_output, request=request |
| 2299 | ) |
| 2300 | |
| 2301 | assert not result.tools_called |
| 2302 | assert result.tool_calls == [] |
| 2303 | assert result.content == model_output |
| 2304 | |
| 2305 | |
| 2306 | def test_guided_streaming_required_pre_v11( |
| 2307 | mistral_pre_v11_tool_parser: MistralToolParser, |
| 2308 | mistral_pre_v11_tokenizer, |
| 2309 | ) -> None: |
| 2310 | model_output = '[{"name": "get_current_weather", "arguments": {"city": "Dallas"}}]' |
| 2311 | request = _make_request(tool_choice="required") |
| 2312 | all_token_ids = mistral_pre_v11_tokenizer.encode( |
| 2313 | model_output, add_special_tokens=False |
| 2314 | ) |
| 2315 | |
| 2316 | function_name: str | None = None |
| 2317 | function_args = "" |
| 2318 | tool_call_id: str | None = None |
| 2319 | previous_text = "" |
| 2320 | previous_tokens = None |
| 2321 | prefix_offset = 0 |
| 2322 | read_offset = 0 |
| 2323 | |
| 2324 | for i, token_id in enumerate(all_token_ids): |
| 2325 | (new_tokens, delta_text, prefix_offset, read_offset) = detokenize_incrementally( |
| 2326 | tokenizer=mistral_pre_v11_tokenizer, |
| 2327 | all_input_ids=all_token_ids[: i + 1], |
| 2328 | prev_tokens=previous_tokens, |
| 2329 | prefix_offset=prefix_offset, |
| 2330 | read_offset=read_offset, |
| 2331 | skip_special_tokens=False, |
| 2332 | spaces_between_special_tokens=True, |
| 2333 | ) |
| 2334 | previous_tokens = ( |
| 2335 | (previous_tokens + new_tokens) if previous_tokens else new_tokens |
| 2336 | ) |
| 2337 | current_text = previous_text + delta_text |
| 2338 | |
| 2339 | delta_message = mistral_pre_v11_tool_parser.extract_tool_calls_streaming( |
| 2340 | previous_text=previous_text, |
| 2341 | current_text=current_text, |
| 2342 | delta_text=delta_text, |
| 2343 | previous_token_ids=all_token_ids[:i], |
| 2344 | current_token_ids=all_token_ids[: i + 1], |
| 2345 | delta_token_ids=[token_id], |
| 2346 | request=request, |
| 2347 | ) |
| 2348 | previous_text = current_text |
| 2349 | |
| 2350 | if delta_message and delta_message.tool_calls: |
| 2351 | for tc in delta_message.tool_calls: |
| 2352 | if tc.id and not tool_call_id: |
| 2353 | tool_call_id = tc.id |
| 2354 | if tc.function: |
| 2355 | if tc.function.name: |
| 2356 | function_name = tc.function.name |
| 2357 | if tc.function.arguments: |
| 2358 | function_args += tc.function.arguments |
| 2359 | |
| 2360 | assert function_name == "get_current_weather" |
| 2361 | assert tool_call_id is not None |
| 2362 | assert len(tool_call_id) == 9 |
| 2363 | assert json.loads(function_args) == {"city": "Dallas"} |
| 2364 | |
| 2365 | |
| 2366 | # --------------------------------------------------------------------------- |
| 2367 | # Malformed-name-before-marker matrix: {malformed, well-formed, mixed |
| 2368 | # parallel} x {streaming, non-streaming}. |
| 2369 | # |
| 2370 | # These tests drive a real MistralParser built from the module's |
| 2371 | # `mistral_tokenizer` fixture (grammar-capable v11+ tokenizer, so the |
| 2372 | # declarative ParserEngine path is exercised, not the pre-v11 legacy state |
| 2373 | # machine) with actual bash/search_replace/submit tools, mirroring the tools |
| 2374 | # present in the real captured traces this bug was found in. |
| 2375 | # --------------------------------------------------------------------------- |
| 2376 | |
| 2377 | _BASH_SEARCH_SUBMIT_TOOLS_DICTS = [ |
| 2378 | { |
| 2379 | "type": "function", |
| 2380 | "function": { |
| 2381 | "name": "bash", |
| 2382 | "description": "Run a bash command", |
| 2383 | "parameters": { |
| 2384 | "type": "object", |
| 2385 | "properties": {"command": {"type": "string"}}, |
| 2386 | "required": ["command"], |
| 2387 | }, |
| 2388 | }, |
| 2389 | }, |
| 2390 | { |
| 2391 | "type": "function", |
| 2392 | "function": { |
| 2393 | "name": "search_replace", |
| 2394 | "description": "Search and replace text in a file", |
| 2395 | "parameters": { |
| 2396 | "type": "object", |
| 2397 | "properties": { |
| 2398 | "path": {"type": "string"}, |
| 2399 | "search": {"type": "string"}, |
| 2400 | "replace": {"type": "string"}, |
| 2401 | }, |
| 2402 | "required": ["path", "search", "replace"], |
| 2403 | }, |
| 2404 | }, |
| 2405 | }, |
| 2406 | { |
| 2407 | "type": "function", |
| 2408 | "function": { |
| 2409 | "name": "submit", |
| 2410 | "description": "Submit the final answer", |
| 2411 | "parameters": { |
| 2412 | "type": "object", |
| 2413 | "properties": {"answer": {"type": "object"}}, |
| 2414 | "required": ["answer"], |
| 2415 | }, |
| 2416 | }, |
| 2417 | }, |
| 2418 | ] |
| 2419 | |
| 2420 | |
| 2421 | def _malformed_matrix_request() -> ChatCompletionRequest: |
| 2422 | return ChatCompletionRequest( |
| 2423 | messages=[], |
| 2424 | model="test", |
| 2425 | tools=_BASH_SEARCH_SUBMIT_TOOLS_DICTS, |
| 2426 | tool_choice="auto", |
| 2427 | ) |
| 2428 | |
| 2429 | |
| 2430 | def _stream_via_parse_delta( |
| 2431 | parser: MistralParser, |
| 2432 | tokenizer: MistralTokenizer, |
| 2433 | model_output: str, |
| 2434 | request: ChatCompletionRequest, |
| 2435 | ) -> tuple[str, list[str | None], list[str]]: |
| 2436 | """Feed `model_output` through `parser.parse_delta` token-by-token. |
| 2437 | |
| 2438 | This is the exact API production serving drives on every generated |
| 2439 | token (see `ParserEngine.parse_delta`), unlike |
| 2440 | `extract_tool_calls_streaming` which some legacy call sites drive with |
| 2441 | pre-computed previous/current text windows. |
| 2442 | """ |
| 2443 | parser.initialize_streaming() |
| 2444 | all_token_ids = encode_mistral_output(tokenizer, model_output) |
| 2445 | |
| 2446 | content = "" |
| 2447 | tool_call_idx = -1 |
| 2448 | names: list[str | None] = [] |
| 2449 | args: list[str] = [] |
| 2450 | |
| 2451 | previous_tokens = None |
| 2452 | prefix_offset = 0 |
| 2453 | read_offset = 0 |
| 2454 | for i, token_id in enumerate(all_token_ids): |
| 2455 | new_tokens, delta_text, prefix_offset, read_offset = detokenize_incrementally( |
| 2456 | tokenizer=tokenizer, |
| 2457 | all_input_ids=all_token_ids[: i + 1], |
| 2458 | prev_tokens=previous_tokens, |
| 2459 | prefix_offset=prefix_offset, |
| 2460 | read_offset=read_offset, |
| 2461 | skip_special_tokens=False, |
| 2462 | spaces_between_special_tokens=True, |
| 2463 | ) |
| 2464 | previous_tokens = ( |
| 2465 | previous_tokens + new_tokens if previous_tokens else new_tokens |
| 2466 | ) |
| 2467 | finished = i == len(all_token_ids) - 1 |
| 2468 | |
| 2469 | delta_message = parser.parse_delta( |
| 2470 | delta_text=delta_text, |
| 2471 | delta_token_ids=[token_id], |
| 2472 | request=request, |
| 2473 | prompt_token_ids=[1], |
| 2474 | finished=finished, |
| 2475 | ) |
| 2476 | if delta_message is None: |
| 2477 | continue |
| 2478 | if delta_message.content: |
| 2479 | content += delta_message.content |
| 2480 | for tool_call in delta_message.tool_calls or []: |
| 2481 | if tool_call.index != tool_call_idx: |
| 2482 | tool_call_idx = tool_call.index |
| 2483 | names.append(None) |
| 2484 | args.append("") |
| 2485 | if tool_call.function: |
| 2486 | if tool_call.function.name is not None: |
| 2487 | names[tool_call.index] = tool_call.function.name |
| 2488 | if tool_call.function.arguments: |
| 2489 | args[tool_call.index] += tool_call.function.arguments |
| 2490 | |
| 2491 | return content, names, args |
| 2492 | |
| 2493 | |
| 2494 | _MALFORMED_MATRIX_CASES = [ |
| 2495 | pytest.param( |
| 2496 | # 2x2 corner: name present, `[ARGS]` token present. |
| 2497 | '[TOOL_CALLS]bash[ARGS]{"command": "ls -la"}', |
| 2498 | ["bash"], |
| 2499 | [json.dumps({"command": "ls -la"})], |
| 2500 | None, |
| 2501 | id="wellformed_single", |
| 2502 | ), |
| 2503 | pytest.param( |
| 2504 | # 2x2 corner: name present, no `[ARGS]` token. |
| 2505 | '[TOOL_CALLS]bash{"command": "ls -la"}', |
| 2506 | ["bash"], |
| 2507 | [json.dumps({"command": "ls -la"})], |
| 2508 | None, |
| 2509 | id="wellformed_single_no_args_token", |
| 2510 | ), |
| 2511 | pytest.param( |
| 2512 | # 2x2 corner: name skipped, no `[ARGS]` token. "bash" lands in |
| 2513 | # content before the marker, then generation jumps straight to |
| 2514 | # args (no NAME slot, no [ARGS]). |
| 2515 | 'bash[TOOL_CALLS]{"command": "cat file.py"}', |
| 2516 | [""], |
| 2517 | [json.dumps({"command": "cat file.py"})], |
| 2518 | "bash", |
| 2519 | id="malformed_single", |
| 2520 | ), |
| 2521 | pytest.param( |
| 2522 | # 2x2 corner: name skipped, `[ARGS]` token present. |
| 2523 | 'bash[TOOL_CALLS][ARGS]{"command": "cat file.py"}', |
| 2524 | [""], |
| 2525 | [json.dumps({"command": "cat file.py"})], |
| 2526 | "bash", |
| 2527 | id="malformed_single_args_token", |
| 2528 | ), |
| 2529 | pytest.param( |
| 2530 | # parallel well-formed x2; the second call's nested braces in its |
| 2531 | # arguments must not confuse call splitting. |
| 2532 | '[TOOL_CALLS]bash[ARGS]{"command": "ls"}' |
| 2533 | '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}', |
| 2534 | ["bash", "submit"], |
| 2535 | [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], |
| 2536 | None, |
| 2537 | id="parallel_wellformed", |
| 2538 | ), |
| 2539 | pytest.param( |
| 2540 | # parallel malformed x2: both calls skip the NAME slot entirely |
| 2541 | # (marker immediately followed by `{`), each ending with an empty |
| 2542 | # name and correctly separated (nested-brace) arguments. |
| 2543 | 'bash[TOOL_CALLS]{"command": "ls"}[TOOL_CALLS]{"answer": {"nested": 1}}', |
| 2544 | ["", ""], |
| 2545 | [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], |
| 2546 | "bash", |
| 2547 | id="parallel_malformed", |
| 2548 | ), |
| 2549 | pytest.param( |
| 2550 | # parallel malformed x2, `[ARGS]` token variant: both calls skip |
| 2551 | # the NAME slot but enter TOOL_ARGS via `[ARGS]` instead of `{`. |
| 2552 | 'bash[TOOL_CALLS][ARGS]{"command": "ls"}' |
| 2553 | '[TOOL_CALLS][ARGS]{"answer": {"nested": 1}}', |
| 2554 | ["", ""], |
| 2555 | [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], |
| 2556 | "bash", |
| 2557 | id="parallel_malformed_args_token", |
| 2558 | ), |
| 2559 | pytest.param( |
| 2560 | # mixed: malformed first call, well-formed second call. The |
| 2561 | # important case: proves calls are split on the `[TOOL_CALLS]` |
| 2562 | # marker, not on braces — the first call's simple args and the |
| 2563 | # second call's nested-brace args must not bleed into each other. |
| 2564 | 'bash[TOOL_CALLS]{"command": "ls"}' |
| 2565 | '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}', |
| 2566 | ["", "submit"], |
| 2567 | [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], |
| 2568 | "bash", |
| 2569 | id="mixed_malformed_then_wellformed", |
| 2570 | ), |
| 2571 | pytest.param( |
| 2572 | # mixed, `[ARGS]` token variant: malformed first call enters |
| 2573 | # TOOL_ARGS via `[ARGS]` instead of `{`, second call well-formed. |
| 2574 | 'bash[TOOL_CALLS][ARGS]{"command": "ls"}' |
| 2575 | '[TOOL_CALLS]submit[ARGS]{"answer": {"nested": 1}}', |
| 2576 | ["", "submit"], |
| 2577 | [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], |
| 2578 | "bash", |
| 2579 | id="mixed_malformed_args_token_then_wellformed", |
| 2580 | ), |
| 2581 | pytest.param( |
| 2582 | # mixed: well-formed first call, malformed second call (NAME slot |
| 2583 | # skipped entirely for the second call). |
| 2584 | '[TOOL_CALLS]bash[ARGS]{"command": "ls"}[TOOL_CALLS]{"answer": {"nested": 1}}', |
| 2585 | ["bash", ""], |
| 2586 | [json.dumps({"command": "ls"}), json.dumps({"answer": {"nested": 1}})], |
| 2587 | None, |
| 2588 | id="mixed_wellformed_then_malformed", |
| 2589 | ), |
| 2590 | ] |
| 2591 | |
| 2592 | |
| 2593 | @pytest.mark.parametrize( |
| 2594 | "model_output,expected_names,expected_args,expected_content", |
| 2595 | _MALFORMED_MATRIX_CASES, |
| 2596 | ) |
| 2597 | def test_malformed_tool_call_matrix_non_streaming( |
| 2598 | mistral_tokenizer, |
| 2599 | model_output, |
| 2600 | expected_names, |
| 2601 | expected_args, |
| 2602 | expected_content, |
| 2603 | ): |
| 2604 | parser = MistralParser(mistral_tokenizer) |
| 2605 | request = _malformed_matrix_request() |
| 2606 | |
| 2607 | result = parser.extract_tool_calls_from_content(model_output, request) |
| 2608 | |
| 2609 | assert result.tools_called |
| 2610 | assert [tc.function.name for tc in result.tool_calls] == expected_names |
| 2611 | assert [tc.function.arguments for tc in result.tool_calls] == expected_args |
| 2612 | assert result.content == expected_content |
| 2613 | |
| 2614 | |
| 2615 | @pytest.mark.parametrize( |
| 2616 | "model_output,expected_names,expected_args,expected_content", |
| 2617 | _MALFORMED_MATRIX_CASES, |
| 2618 | ) |
| 2619 | def test_malformed_tool_call_matrix_streaming( |
| 2620 | mistral_tokenizer, |
| 2621 | model_output, |
| 2622 | expected_names, |
| 2623 | expected_args, |
| 2624 | expected_content, |
| 2625 | ): |
| 2626 | parser = MistralParser(mistral_tokenizer) |
| 2627 | request = _malformed_matrix_request() |
| 2628 | |
| 2629 | content, names, args = _stream_via_parse_delta( |
| 2630 | parser, mistral_tokenizer, model_output, request |
| 2631 | ) |
| 2632 | |
| 2633 | assert content == (expected_content or "") |
| 2634 | assert names == expected_names |
| 2635 | assert args == expected_args |
| 2636 | |
| 2637 | |
| 2638 | # Real payload captured from a production SWE-bench-agent. |
| 2639 | _REAL_LEAK_CONTENT = 'bash[TOOL_CALLS]{"command": "cd /testbed && git diff xarray/core/weighted.py"}ometracediff --git a/xarray/core/weighted.py b/xarray/core/weighted.py\\nindex 5f2d138e..9c38b0b7 100644\\n--- a/xarray/core/weighted.py\\n+++ b/xarray/core/weighted.py\\n@@ -1,4 +1,6 @@\\n from typing import TYPE_CHECKING, Hashable, Iterable, Optional, Union, overload\\n+\\n+import numpy as np\\n \\n from .computation import dot\\n from .options import _get_keep_attrs\\n@@ -105,6 +107,10 @@ class Weighted:\\n )\\n \\n self.obj = obj\\n- self.weights = weights\\n+ # Ensure weights are numeric (float64) to avoid boolean arithmetic issues\\n+ # in dot products and weighted operations\\n+ if not np.issubdtype(weights.dtype, np.number):\\n+ weights = weights.astype(np.float64)\\n+ self.weights = weights\\n "}现在 hareview the changes we made:bash[ARGS]{"command": "cat /testbed/xarray/core/weighted.py | head -120 | tail -20"}' # noqa: E501 |
| 2640 | |
| 2641 | |
| 2642 | def test_extract_tool_calls_real_captured_malformed_output(mistral_tokenizer): |
| 2643 | parser = MistralParser(mistral_tokenizer) |
| 2644 | request = _malformed_matrix_request() |
| 2645 | |
| 2646 | result = parser.extract_tool_calls_from_content(_REAL_LEAK_CONTENT, request) |
| 2647 | |
| 2648 | assert result.tools_called |
| 2649 | assert len(result.tool_calls) == 1 |
| 2650 | assert result.tool_calls[0].function.name == "" |
| 2651 | assert result.tool_calls[0].function.arguments == json.dumps( |
| 2652 | {"command": "cd /testbed && git diff xarray/core/weighted.py"} |
| 2653 | ) |
| 2654 | assert result.content == "bash" |