# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import functools
import json
from collections.abc import Sequence
from enum import Enum, auto
from random import choices
from string import ascii_letters, digits
from typing import TYPE_CHECKING, Any, Literal
import ijson
import regex as re
from mistral_common.protocol.instruct.tool_calls import (
NamedToolChoice as MistralNamedToolChoice,
)
from mistral_common.protocol.instruct.tool_calls import (
Tool as MistralTool,
)
from mistral_common.protocol.instruct.tool_calls import (
ToolChoice as MistralToolChoice,
)
from mistral_common.protocol.instruct.tool_calls import (
ToolChoiceEnum as MistralToolChoiceEnum,
)
from mistral_common.tokens.tokenizers.base import SpecialTokens
from pydantic import Field
from vllm.entrypoints.generate.base.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionNamedToolChoiceParam,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.logger import init_logger
from vllm.parser.engine.events import EventType
from vllm.parser.engine.parser_engine import ParserEngine
from vllm.parser.engine.parser_engine_config import (
ParserEngineConfig,
ParserState,
Transition,
)
from vllm.sampling_params import StructuredOutputsParams
from vllm.utils.mistral import is_mistral_tokenizer
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
from vllm.parser.engine.events import SemanticEvent
from vllm.parser.engine.parser_engine import ToolCallSlot
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import Tool
logger = init_logger(__name__)
_ALPHANUMERIC = ascii_letters + digits
_DEFAULT_JSON_SCHEMA: dict[str, Any] = {
"anyOf": [{"type": "object"}, {"type": "array"}]
}
# Special token strings as emitted/decoded by MistralTokenizer.
_TOOL_CALLS = SpecialTokens.tool_calls.value
_ARGS = SpecialTokens.args.value
_THINK_START_SPECIAL = SpecialTokens.begin_think.value
_THINK_END_SPECIAL = SpecialTokens.end_think.value
# Plain-text reasoning markers used by v11 tokenizers.
_THINK_START_TEXT = ""
_THINK_END_TEXT = ""
_OPEN_BRACE = "{"
class StreamingState(Enum):
"""Streaming parsing state for pre-v11 tool call extraction."""
WAITING_FOR_TOOL_START = auto()
WAITING_FOR_TOOL_KEY = auto()
PARSING_NAME = auto()
PARSING_NAME_COMPLETED = auto()
WAITING_FOR_ARGUMENTS_START = auto()
PARSING_ARGUMENTS = auto()
PARSING_ARGUMENTS_COMPLETED = auto()
TOOL_COMPLETE = auto()
ALL_TOOLS_COMPLETE = auto()
class MistralToolCall(ToolCall):
"""ToolCall with a Mistral-compatible random alphanumeric id."""
id: str = Field(default_factory=lambda: MistralToolCall.generate_random_id())
@staticmethod
def generate_random_id() -> str:
# Mistral Tool Call Ids must be alphanumeric with a length of 9.
# https://github.com/mistralai/mistral-common/blob/21ee9f6cee3441e9bb1e6ed2d10173f90bd9b94b/src/mistral_common/protocol/instruct/validator.py#L299
return "".join(choices(_ALPHANUMERIC, k=9))
@staticmethod
def is_valid_id(id: str) -> bool:
return id.isalnum() and len(id) == 9
def _is_pre_v11_tokeniser(model_tokenizer: TokenizerLike) -> bool:
if is_mistral_tokenizer(model_tokenizer):
return model_tokenizer.version < 11
vocab: dict[str, int] = getattr(model_tokenizer, "get_vocab", lambda: {})()
return _ARGS not in vocab
@functools.cache
def mistral_config(
*,
reasoning_encoding: Literal["special_token", "text", "none"],
name: str = "mistral",
) -> ParserEngineConfig:
"""Return a :class:`ParserEngineConfig` for the Mistral output format.
Args:
reasoning_encoding: Reasoning token format.
``"special_token"`` – v13+ ``[THINK]``/``[/THINK]`` special
tokens placed in both ``terminals`` and ``token_id_terminals``.
``"text"`` – v11 ````/```` plain text placed in
``terminals`` only; no ``token_id_terminals`` for think tokens.
``"none"`` – no reasoning support.
name: Name embedded in the returned config (used for debugging).
Returns:
A frozen :class:`ParserEngineConfig` with ``initial_state=CONTENT``.
"""
if reasoning_encoding == "special_token":
think_start = _THINK_START_SPECIAL
think_end = _THINK_END_SPECIAL
reasoning_terminals: dict[str, str] = {
"THINK_START": think_start,
"THINK_END": think_end,
}
reasoning_token_id_terminals: dict[str, str] = {
"THINK_START": think_start,
"THINK_END": think_end,
}
elif reasoning_encoding == "text":
think_start = _THINK_START_TEXT
think_end = _THINK_END_TEXT
reasoning_terminals = {
"THINK_START": think_start,
"THINK_END": think_end,
}
# Text think markers have no token-id terminals.
reasoning_token_id_terminals = {}
else:
reasoning_terminals = {}
reasoning_token_id_terminals = {}
if reasoning_encoding != "none":
reasoning_transitions: dict[tuple[ParserState, str], Transition] = {
(ParserState.CONTENT, "THINK_START"): Transition(
ParserState.REASONING,
(),
),
# Absorb a duplicate/re-emitted THINK_START inside reasoning.
(ParserState.REASONING, "THINK_START"): Transition(
ParserState.REASONING,
(),
),
(ParserState.REASONING, "THINK_END"): Transition(
ParserState.CONTENT,
(EventType.REASONING_END,),
),
# Absorb stray THINK_END that arrives after reasoning ended.
(ParserState.CONTENT, "THINK_END"): Transition(
ParserState.CONTENT,
(),
),
# [TOOL_CALLS] directly from reasoning implicitly ends it.
(ParserState.REASONING, "TOOL_CALLS"): Transition(
ParserState.TOOL_NAME,
(EventType.REASONING_END, EventType.TOOL_CALL_START),
),
}
else:
reasoning_transitions = {}
return ParserEngineConfig(
name=name,
initial_state=ParserState.CONTENT,
terminals={
**reasoning_terminals,
"TOOL_CALLS": _TOOL_CALLS,
"ARGS": _ARGS,
"OPEN_BRACE": _OPEN_BRACE,
},
token_id_terminals={
**reasoning_token_id_terminals,
"TOOL_CALLS": _TOOL_CALLS,
"ARGS": _ARGS,
},
transitions={
**reasoning_transitions,
# A tool call from content implicitly ends reasoning when enabled.
(ParserState.CONTENT, "TOOL_CALLS"): Transition(
ParserState.TOOL_NAME,
(EventType.REASONING_END, EventType.TOOL_CALL_START)
if reasoning_encoding != "none"
else (EventType.TOOL_CALL_START,),
),
# NAME→ARGS via explicit [ARGS] separator (v11+): consumed, no events.
(ParserState.TOOL_NAME, "ARGS"): Transition(
ParserState.TOOL_ARGS,
(),
),
# NAME→ARGS via "{": carried as ARG_VALUE_CHUNK so the opening
# brace lands in the JSON argument buffer (fallback for name{args}).
(ParserState.TOOL_NAME, "OPEN_BRACE"): Transition(
ParserState.TOOL_ARGS,
(EventType.ARG_VALUE_CHUNK,),
),
# Parallel tool calls: next [TOOL_CALLS] ends current call.
(ParserState.TOOL_ARGS, "TOOL_CALLS"): Transition(
ParserState.TOOL_NAME,
(EventType.TOOL_CALL_END, EventType.TOOL_CALL_START),
),
},
stream_arg_deltas=True,
tool_args_json=True,
strip_trailing_reasoning_whitespace=True,
drop_whitespace_only_content_before_tools=True,
)
class MistralParser(ParserEngine):
"""Mistral parser: engine-based reasoning + ``[TOOL_CALLS]`` tool calls.
Reasoning encoding is auto-detected from the tokenizer:
- ``"special_token"`` – ``[THINK]`` present in vocab (v13+).
- ``"text"`` – tokenizer supports grammar but has no ``[THINK]`` (v11).
- ``"none"`` – no grammar support; reasoning disabled.
Tool calls use the ``[TOOL_CALLS]func_name{...}`` format. The opening
``{`` doubles as the NAME→ARGS separator and is included in the argument
buffer via an ``ARG_VALUE_CHUNK`` event on the transition (mirrors
kimi_k2, JSON args, `tool_args_json=True`).
When the tokenizer does not support grammar (``_reasoning_encoding ==
"none"``), the legacy ``[TOOL_CALLS]``-based extraction path is used
instead of the declarative engine, handling both pre-v11 JSON-array and
v11+ ``funcname{args}`` formats.
"""
def __init__(
self,
tokenizer: TokenizerLike,
tools: list[Tool] | None = None,
**kwargs,
) -> None:
vocab = tokenizer.get_vocab()
self._reasoning_encoding: Literal["special_token", "text", "none"]
if _THINK_START_SPECIAL in vocab:
self._reasoning_encoding = "special_token"
elif getattr(tokenizer, "supports_grammar", False):
self._reasoning_encoding = "text"
else:
self._reasoning_encoding = "none"
kwargs.setdefault(
"parser_engine_config",
mistral_config(reasoning_encoding=self._reasoning_encoding),
)
super().__init__(tokenizer, tools, **kwargs)
self._tool_calls_token_id: int | None = self.vocab.get(_TOOL_CALLS)
# Tool calls use the legacy parser for all tokenizer versions;
# reasoning is handled by the engine.
self.bot_token: str = _TOOL_CALLS
self.bot_token_id: int | None = self._tool_calls_token_id
if self.bot_token_id is None:
raise RuntimeError(
"Mistral parser could not locate the tool call token in the tokenizer!"
)
# Legacy tool-call streaming state.
self.prev_tool_call_arr: list[dict[str, Any]] = []
self.current_tool_id: int = -1
self.streaming_state: StreamingState = StreamingState.WAITING_FOR_TOOL_START
self.tool_call_started: bool = False
self.current_tool_name: str | None = None
self.current_tool_mistral_id: str | None = None
self.starting_new_tool: bool = False
self.streamed_args_for_tool: list[str] = []
self._is_pre_v11: bool = _is_pre_v11_tokeniser(tokenizer)
self.parse_coro = None
if self._is_pre_v11:
self.parse_coro = ijson.parse_coro(
self.update_stream_state_pre_v11_tokenizer()
)
self.tool_call_regex = re.compile(r"\[{.*}\]", re.DOTALL)
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
if not isinstance(request, ResponsesRequest) and request._grammar_from_parser:
return request
so_non_supported_attributes = [
"regex",
"choice",
"grammar",
# whitespace_pattern is not a constraint type but an option;
# Mistral grammar factory does not support it.
"whitespace_pattern",
"structural_tag",
]
any_so_non_supported_active = request.structured_outputs is not None and any(
getattr(request.structured_outputs, attribute) is not None
for attribute in so_non_supported_attributes
)
response_format_non_supported_active = (
isinstance(request, ResponsesRequest)
or request.response_format is not None
and request.response_format.type == "structural_tag"
)
if (
not is_mistral_tokenizer(self.model_tokenizer)
or isinstance(request, ResponsesRequest)
or not self.model_tokenizer.supports_grammar
or any_so_non_supported_active
or response_format_non_supported_active
):
request = super().adjust_request(request)
if request.tools and request.tool_choice != "none":
# Keep special tokens so the [TOOL_CALLS] marker
# survives for tool detection.
request.skip_special_tokens = False
# Inject a guided JSON schema for pre-v11 required/named tool choice
# so the model emits a well-formed bare JSON array instead of
# rambling. tool_choice forces a tool call, so a json_object /
# json_schema response_format is cleared here (the tool schema is
# the sole structured-output constraint), mirroring the base
# ToolParser; a structural_tag response_format is left untouched.
req_tool_choice = request.tool_choice
is_required_or_named = req_tool_choice == "required" or isinstance(
req_tool_choice, ChatCompletionNamedToolChoiceParam
)
response_format = getattr(request, "response_format", None)
response_format_overridable = response_format is None or (
response_format.type in ("text", "json_object", "json_schema")
)
# This runs only in the non-grammar (legacy) branch, so it covers
# both pre-v11 Mistral tokenizers and non-Mistral (e.g. HF-mode)
# tokenizers driving the Mistral tool parser. The remaining guards
# keep it off when the user supplied their own structured output.
if (
not isinstance(request, ResponsesRequest)
and request.tools
and request.structured_outputs is None
and is_required_or_named
and response_format_overridable
):
schema = self._build_guided_schema_pre_v11(request)
if schema is not None:
request.structured_outputs = StructuredOutputsParams(json=schema)
request.response_format = None
return request
json_schema: dict[str, Any] | None = None
if request.structured_outputs is not None:
if request.structured_outputs.json_object is not None:
json_schema = _DEFAULT_JSON_SCHEMA
elif request.structured_outputs.json is not None:
if isinstance(request.structured_outputs.json, str):
json_schema = json.loads(request.structured_outputs.json)
else:
json_schema = request.structured_outputs.json
else:
raise ValueError(
"Unsupported request.structured_outputs for MistralParser. "
"Only `json` and `json_object` are supported."
)
elif (
request.response_format is not None
and request.response_format.type != "text"
):
if request.response_format.type == "json_object":
json_schema = _DEFAULT_JSON_SCHEMA
elif request.response_format.type == "json_schema":
if request.response_format.json_schema is not None:
json_schema = request.response_format.json_schema.json_schema
else:
json_schema = _DEFAULT_JSON_SCHEMA
else:
raise ValueError(
"MistralParser only accepts `text`, `json_object` or "
f"`json_schema`, got {request.response_format=}"
)
request.response_format = None
grammar_factory = self.model_tokenizer.grammar_factory
# Rendering grammar is cached in mistral-common given tools, template and mode.
template = grammar_factory.select_jinja_template()
mistral_tools = (
[MistralTool.from_openai(tool.model_dump()) for tool in request.tools]
if request.tools is not None
else None
)
tool_choice: MistralToolChoice
match request.tool_choice:
case "none" | "auto" | "required":
tool_choice = MistralToolChoiceEnum(request.tool_choice)
case None:
tool_choice = MistralToolChoiceEnum.auto
# _ == Named tool choice
case _:
tool_choice = MistralNamedToolChoice.model_validate(
{
"type": "function",
"function": {"name": request.tool_choice.function.name},
}
)
match tool_choice, json_schema is not None:
case MistralToolChoiceEnum.none, True:
lark_grammar = grammar_factory.get_lark_for_json_schema(
template=template, json_schema=json_schema
)
case _, _:
lark_grammar = grammar_factory.get_lark_from_jinja(
template=template,
mode=tool_choice,
tools=mistral_tools,
json_schema=json_schema,
parallel_tool_calls=request.parallel_tool_calls,
json_only=False,
)
request.structured_outputs = StructuredOutputsParams(grammar=lark_grammar)
request._grammar_from_parser = True
return request
def _build_guided_schema_pre_v11(
self,
request: ChatCompletionRequest,
) -> dict[str, Any] | None:
"""Build a guided JSON schema for pre-v11 required/named tool choice.
The schema enforces the Mistral-native array format
``[{"name": ..., "arguments": {...}}]`` so the model emits a parseable
bare JSON array instead of free-form text.
Args:
request: The chat completion request carrying `tools` and
`tool_choice`.
Returns:
A JSON Schema dict, or ``None`` if the named tool is not found.
"""
tool_choice = request.tool_choice
tools = request.tools or []
extra: dict[str, Any] = {}
if tool_choice == "required":
applicable_tools = tools
else:
# Named tool choice — restrict to the single requested tool.
assert isinstance(tool_choice, ChatCompletionNamedToolChoiceParam)
chosen_name = tool_choice.function.name
applicable_tools = [t for t in tools if t.function.name == chosen_name]
if not applicable_tools:
logger.warning(
"Named tool %r not found in tools list; "
"skipping guided schema injection.",
chosen_name,
)
return None
extra["maxItems"] = 1
any_of = [
{
"type": "object",
"properties": {
"name": {"type": "string", "enum": [tool.function.name]},
"arguments": tool.function.parameters or {"type": "object"},
},
"required": ["name", "arguments"],
}
for tool in applicable_tools
]
return {
"type": "array",
"minItems": 1,
**extra,
"items": {
"type": "object",
"anyOf": any_of,
},
}
def _ensure_tool_id(self, slot: ToolCallSlot, name: str) -> None:
"""Assign a Mistral-compatible 9-char alphanumeric id to `slot`."""
if not slot.id:
slot.id = MistralToolCall.generate_random_id()
def extract_tool_calls_from_content(
self,
content: str,
request: ChatCompletionRequest,
) -> ExtractedToolCallInformation:
if self._is_pre_v11:
return self._legacy_extract_tool_calls(content, request)
return super().extract_tool_calls_from_content(content, request)
def _legacy_extract_tool_calls(
self,
model_output: str,
request: ChatCompletionRequest | None,
) -> ExtractedToolCallInformation:
"""Pre-v11 non-streaming extraction.
Handles ``[TOOL_CALLS][{...}]`` and guided bare-array formats.
"""
if request is None:
tool_choice = None
tools = None
else:
tool_choice = request.tool_choice
tools = request.tools
# tool_choice="none" with tools: never produce tool calls.
if tool_choice == "none" and tools:
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
content: str | None = None
if self.bot_token in model_output:
content_and_raw_tool_calls = model_output.split(self.bot_token)
content = content_and_raw_tool_calls[0]
raw_tool_calls = content_and_raw_tool_calls[1:]
# pre-v11: content[BOT] [{tool_call1},{tool_call2}]
if len(raw_tool_calls) != 1:
raise ValueError(
"Only one BOT token should have been outputted, "
f"but got {model_output}."
)
stringified_tool_calls = raw_tool_calls[0].strip()
elif tool_choice == "required" or isinstance(
tool_choice, ChatCompletionNamedToolChoiceParam
):
# Guided bare-array output (no [TOOL_CALLS] marker).
stringified_tool_calls = model_output.strip()
else:
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
try:
# Use raw_decode to parse the first valid JSON value,
# ignoring trailing tokens the model may emit after
# the tool call array.
tool_calls, _ = json.JSONDecoder().raw_decode(stringified_tool_calls)
except json.JSONDecodeError:
try:
raw_tool_call = self.tool_call_regex.findall(stringified_tool_calls)[0]
tool_calls = json.loads(raw_tool_call)
tool_calls = [
{
"name": tool_call["name"],
"arguments": json.dumps(
tool_call.get("arguments", {}),
ensure_ascii=False,
),
}
for tool_call in tool_calls
]
except (IndexError, json.JSONDecodeError):
logger.exception("Error in extracting tool call from response.")
return ExtractedToolCallInformation(
tools_called=False,
tool_calls=[],
content=stringified_tool_calls,
)
else:
tool_calls = [
{
"name": tool_call["name"],
"arguments": json.dumps(
tool_call.get("arguments", {}),
ensure_ascii=False,
),
}
for tool_call in tool_calls
]
mistral_tool_calls: list[MistralToolCall] = [
MistralToolCall(
type="function",
function=FunctionCall(
name=tool_call["name"],
arguments=tool_call.get("arguments", "{}"),
),
)
for tool_call in tool_calls
]
return ExtractedToolCallInformation(
tools_called=True,
tool_calls=mistral_tool_calls,
content=content if content and content.strip() else None,
)
def extract_tool_calls_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
request: ChatCompletionRequest | ResponsesRequest,
) -> DeltaMessage | None:
if not self._is_pre_v11:
return super().extract_tool_calls_streaming(
previous_text,
current_text,
delta_text,
previous_token_ids,
current_token_ids,
delta_token_ids,
request,
)
# Pre-v11: latch on [TOOL_CALLS] or on the first content of a guided
# required/named request (bare JSON array, no special token).
if self.bot_token_id in delta_token_ids or self.bot_token in delta_text:
self.tool_call_started = True
elif not self.tool_call_started and delta_text:
is_guided = request.tool_choice == "required" or isinstance(
request.tool_choice, ChatCompletionNamedToolChoiceParam
)
if is_guided:
self.tool_call_started = True
if not self.tool_call_started:
return DeltaMessage(content=delta_text)
try:
return self._extract_tool_calls_streaming_pre_v11_tokenizer(
delta_text=delta_text,
delta_token_ids=delta_token_ids,
)
except Exception:
logger.exception("Error trying to handle streaming tool call.")
return None
@ijson.coroutine
def update_stream_state_pre_v11_tokenizer(self):
while True:
(prefix, event, value) = yield
if prefix == "item" and event == "start_map":
self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
self.starting_new_tool = True
if prefix == "item" and event == "map_key" and value == "name":
self.streaming_state = StreamingState.PARSING_NAME
if prefix == "item.name" and event == "string":
self.current_tool_name = value
self.streaming_state = StreamingState.PARSING_NAME_COMPLETED
if prefix == "item" and event == "map_key" and value == "arguments":
self.streaming_state = StreamingState.WAITING_FOR_ARGUMENTS_START
if prefix == "item.arguments" and event == "start_map":
self.streaming_state = StreamingState.PARSING_ARGUMENTS
if prefix == "item.arguments" and event == "end_map":
self.streaming_state = StreamingState.PARSING_ARGUMENTS_COMPLETED
if prefix == "item" and event == "end_map":
self.streaming_state = StreamingState.TOOL_COMPLETE
if prefix == "" and event == "end_array":
self.streaming_state = StreamingState.ALL_TOOLS_COMPLETE
def _extract_tool_calls_streaming_pre_v11_tokenizer(
self,
delta_text: str,
delta_token_ids: Sequence[int],
) -> DeltaMessage | None:
"""Extract tool calls for pre-v11 Mistral models.
Handles ``[TOOL_CALLS][{"name": "add", "arguments":{"a": 3.5}}]``.
"""
assert self.parse_coro is not None
content = None
delta_tool_calls: list[DeltaToolCall] = []
current_tool_call: DeltaToolCall = DeltaToolCall(
index=self.current_tool_id, type="function"
)
current_tool_call_modified = False
if self.bot_token_id in delta_token_ids or self.bot_token in delta_text:
# this is the first tool call
if not delta_text.startswith(self.bot_token):
content = delta_text.split(self.bot_token)[0]
delta_text = "".join(delta_text.split(self.bot_token)[1:])
# ijson gives no text index per event, so split the delta manually
# to know where each event is emitted from.
while len(delta_text) > 0:
streaming_state_before_parse = self.streaming_state
if self.streaming_state == StreamingState.WAITING_FOR_TOOL_START:
delta_to_be_parsed, delta_text = self._split_delta(
delta_text=delta_text,
stop_after_opening_curly_braces=1,
)
elif self.streaming_state == StreamingState.WAITING_FOR_TOOL_KEY:
delta_to_be_parsed, delta_text = self._split_delta(
delta_text=delta_text,
stop_after_colon=1,
stop_after_opening_curly_braces=1,
)
elif self.streaming_state == StreamingState.PARSING_NAME:
delta_to_be_parsed, delta_text = self._split_delta(
delta_text=delta_text,
stop_after_comma=1,
stop_after_closing_brackets=1,
)
elif self.streaming_state == StreamingState.WAITING_FOR_ARGUMENTS_START:
delta_to_be_parsed, delta_text = self._split_delta(
delta_text=delta_text,
stop_after_opening_curly_braces=1,
)
elif self.streaming_state == StreamingState.PARSING_ARGUMENTS:
delta_to_be_parsed, delta_text = self._split_delta(
delta_text=delta_text,
stop_after_closing_curly_braces=1,
)
elif self.streaming_state in [
StreamingState.PARSING_ARGUMENTS_COMPLETED,
StreamingState.PARSING_NAME_COMPLETED,
]:
delta_to_be_parsed, delta_text = self._split_delta(
delta_text=delta_text,
stop_after_closing_curly_braces=1,
stop_after_closing_brackets=1,
)
elif self.streaming_state == StreamingState.TOOL_COMPLETE:
delta_to_be_parsed, delta_text = self._split_delta(
delta_text=delta_text,
stop_after_opening_curly_braces=1,
stop_after_closing_brackets=1,
)
elif self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE:
content = delta_text
delta_text = ""
else:
delta_to_be_parsed = delta_text
delta_text = ""
if self.streaming_state != StreamingState.ALL_TOOLS_COMPLETE:
self.parse_coro.send(delta_to_be_parsed.encode("utf-8"))
# start_map is the authoritative new-tool signal and survives
# batched deltas, unlike comparing pre/post streaming states.
if self.starting_new_tool:
self.starting_new_tool = False
if current_tool_call_modified:
if self.current_tool_mistral_id is not None:
current_tool_call.id = self.current_tool_mistral_id
self.current_tool_mistral_id = None
self._track_streamed_args_pre_v11(current_tool_call)
delta_tool_calls.append(current_tool_call)
current_tool_call_modified = False
self.current_tool_id += 1
self.streamed_args_for_tool.append("")
self.prev_tool_call_arr.append({})
self.current_tool_mistral_id = MistralToolCall.generate_random_id()
current_tool_call = DeltaToolCall(
index=self.current_tool_id,
type="function",
)
if current_tool_call.function is None:
current_tool_call.function = DeltaFunctionCall()
if self.current_tool_name is not None:
current_tool_call_modified = True
current_tool_call.function.name = self.current_tool_name
self.prev_tool_call_arr[self.current_tool_id]["name"] = (
self.current_tool_name
)
self.current_tool_name = None
if self.streaming_state == StreamingState.PARSING_NAME_COMPLETED:
self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
if self.streaming_state in [
StreamingState.PARSING_ARGUMENTS,
StreamingState.PARSING_ARGUMENTS_COMPLETED,
]:
if self.streaming_state == StreamingState.PARSING_ARGUMENTS_COMPLETED:
self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
current_tool_call_modified = True
if current_tool_call.function.arguments is None:
current_tool_call.function.arguments = delta_to_be_parsed
else:
current_tool_call.function.arguments += delta_to_be_parsed
if streaming_state_before_parse != StreamingState.PARSING_ARGUMENTS:
# It's the first chunk of arg. let's lstrip it
current_tool_call.function.arguments = (
current_tool_call.function.arguments.lstrip()
)
if current_tool_call_modified:
if self.current_tool_mistral_id is not None:
current_tool_call.id = self.current_tool_mistral_id
self.current_tool_mistral_id = None
self._track_streamed_args_pre_v11(current_tool_call)
delta_tool_calls.append(current_tool_call)
if content or len(delta_tool_calls) > 0:
delta_message = DeltaMessage()
if content:
delta_message.content = content
if len(delta_tool_calls) > 0:
delta_message.tool_calls = delta_tool_calls
return delta_message
else:
if self.streaming_state == StreamingState.ALL_TOOLS_COMPLETE:
return DeltaMessage()
else:
return None
def _track_streamed_args_pre_v11(self, tool_call: DeltaToolCall) -> None:
r"""Accumulate `tool_call` arguments into the streaming state."""
if tool_call.function is not None and tool_call.function.arguments is not None:
self.streamed_args_for_tool[self.current_tool_id] += (
tool_call.function.arguments
)
self.prev_tool_call_arr[self.current_tool_id]["arguments"] = (
self.streamed_args_for_tool[self.current_tool_id]
)
def _split_delta(
self,
delta_text: str,
stop_after_quotes: int = -1,
stop_after_opening_curly_braces: int = -1,
stop_after_closing_curly_braces: int = -1,
stop_after_closing_brackets: int = -1,
stop_after_colon: int = -1,
stop_after_comma: int = -1,
) -> tuple[str, str]:
delta_to_be_parsed = ""
for i, c in enumerate(delta_text):
if c in ['"', "'"]:
delta_to_be_parsed += c
stop_after_quotes -= 1
if stop_after_quotes == 0:
return (delta_to_be_parsed, delta_text[i + 1 :])
elif c == "{":
delta_to_be_parsed += c
stop_after_opening_curly_braces -= 1
if stop_after_opening_curly_braces == 0:
return (delta_to_be_parsed, delta_text[i + 1 :])
elif c == "}":
delta_to_be_parsed += c
stop_after_closing_curly_braces -= 1
if stop_after_closing_curly_braces == 0:
return (delta_to_be_parsed, delta_text[i + 1 :])
elif c == "]":
delta_to_be_parsed += c
stop_after_closing_brackets -= 1
if stop_after_closing_brackets == 0:
return (delta_to_be_parsed, delta_text[i + 1 :])
elif c == ":":
delta_to_be_parsed += c
stop_after_colon -= 1
if stop_after_colon == 0:
return (delta_to_be_parsed, delta_text[i + 1 :])
elif c == ",":
delta_to_be_parsed += c
stop_after_comma -= 1
if stop_after_comma == 0:
return (delta_to_be_parsed, delta_text[i + 1 :])
else:
delta_to_be_parsed += c
return (delta_to_be_parsed, "")
def is_reasoning_end(self, input_ids: list[int]) -> bool:
if self._reasoning_encoding == "none":
return True
if super().is_reasoning_end(input_ids):
return True
# [TOOL_CALLS] acts as an implicit reasoning-end marker
if self.bot_token_id is not None:
reasoning_start_id = self._reasoning_start_token_id
for i in range(len(input_ids) - 1, -1, -1):
if (
reasoning_start_id is not None
and input_ids[i] == reasoning_start_id
):
return False
if input_ids[i] == self.bot_token_id:
return True
return False
def extract_reasoning(
self,
model_output: str,
request: ChatCompletionRequest | ResponsesRequest,
) -> tuple[str | None, str | None]:
if self._reasoning_encoding == "none":
return None, model_output
return super().extract_reasoning(model_output, request)
def _accept_tool_name(self, name: str) -> bool:
# Once `[ARGS]` or the opening `{` has moved the slot past its name,
# that name is final, empty included. Emitting the call with ""
# surfaces the malformed generation instead of dropping it.
return self._is_valid_tool_name(name)
def _try_extract_name(self, idx: int) -> str | None:
# No JSON-embedded "name" key in this format; the slot name is final.
return self._tool_slots[idx].name
def _extract_name_and_args(self, raw_body: str) -> tuple[str, str]:
# Never a {"name": ..., "arguments": ...} envelope -- a literal "name"
# argument key must not be mistaken for the tool name.
return "", self._extract_args_json(raw_body, "")
def _handle_arg_chunk(
self,
event: SemanticEvent,
deltas: list[DeltaToolCall],
) -> None:
"""Emit the opening ``{`` as an arg delta when the name is first sent.
When the TOOL_NAME→TOOL_ARGS transition fires, ``{`` arrives as an
ARG_VALUE_CHUNK before ``name_sent`` is True. The parent emits the
name delta but not the ``{`` arg delta. This override re-emits the
current chunk so streaming clients receive a valid JSON prefix.
"""
idx = event.tool_index
name_sent_before = (
0 <= idx < len(self._tool_slots) and self._tool_slots[idx].name_sent
)
super()._handle_arg_chunk(event, deltas)
if (
event.value
and not name_sent_before
and 0 <= idx < len(self._tool_slots)
and self._tool_slots[idx].name_sent
):
deltas.append(
DeltaToolCall(
index=idx,
function=DeltaFunctionCall(arguments=event.value),
)
)
def _extract_args_json(self, raw_args: str, func_name: str) -> str:
"""Return the first complete JSON value in ``raw_args``.
v11+ tool calls are emitted as ``name{args}`` with no terminator, so a
model may append ordinary text after the closing brace. Parse the first
JSON value with ``raw_decode`` and drop any trailing output, mirroring
the pre-v11 path (fixes gh#48975).
"""
stripped = raw_args.strip()
if not stripped:
return "{}"
try:
_, end = json.JSONDecoder().raw_decode(stripped)
except json.JSONDecodeError:
return stripped
return stripped[:end]