| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import ast |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import regex as re |
| 8 | from transformers import PreTrainedTokenizerBase |
| 9 | |
| 10 | import vllm.envs as envs |
| 11 | from vllm.entrypoints.generate.base.protocol import ( |
| 12 | DeltaMessage, |
| 13 | ExtractedToolCallInformation, |
| 14 | ) |
| 15 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 16 | ChatCompletionRequest, |
| 17 | ) |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.tool_parsers.abstract_tool_parser import ( |
| 20 | Tool, |
| 21 | ToolParser, |
| 22 | ) |
| 23 | from vllm.tool_parsers.utils import ( |
| 24 | UnexpectedAstError, |
| 25 | compute_tool_delta, |
| 26 | handle_single_tool, |
| 27 | make_valid_python, |
| 28 | ) |
| 29 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | |
| 33 | class Llama4PythonicToolParser(ToolParser): |
| 34 | """Toolcall parser for Llama4 that produce tool calls in a pythonic style |
| 35 | Use --enable-auto-tool-choice --tool-call-parser llama4_pythonic |
| 36 | """ |
| 37 | |
| 38 | # TODO(mdepinet): Possible future improvements: |
| 39 | # 1. Support text + tools separated by either <|python_tag|> or \n\n |
| 40 | # 2. Support tools outside of a list (or separated by a semicolon). |
| 41 | # This depends on item 1 for consistent streaming. |
| 42 | # Neither of these are necessary for e.g. ToolACE, but both would help make |
| 43 | # Llama3.2 models more reliable. |
| 44 | |
| 45 | TOOL_CALL_REGEX = re.compile( |
| 46 | r"\[([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s)?\),\s*)*([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s*)?\)\s*)+\]", |
| 47 | re.DOTALL, |
| 48 | ) |
| 49 | |
| 50 | def __init__( |
| 51 | self, |
| 52 | tokenizer: PreTrainedTokenizerBase, |
| 53 | tools: list[Tool] | None = None, |
| 54 | ): |
| 55 | super().__init__(tokenizer, tools) |
| 56 | |
| 57 | # Rename for readability. This is NOT a tool id. |
| 58 | @property |
| 59 | def current_tool_index(self) -> int: |
| 60 | return self.current_tool_id |
| 61 | |
| 62 | @current_tool_index.setter |
| 63 | def current_tool_index(self, value: int) -> None: |
| 64 | self.current_tool_id = value |
| 65 | |
| 66 | def extract_tool_calls( |
| 67 | self, model_output: str, request: ChatCompletionRequest |
| 68 | ) -> ExtractedToolCallInformation: |
| 69 | """Extract the tool calls from a complete model response.""" |
| 70 | # remove <|python_start|> and <|python_end|> |
| 71 | # as Llama 4 model sometime will output those tokens |
| 72 | if model_output.startswith("<|python_start|>"): |
| 73 | model_output = model_output[len("<|python_start|>") :] |
| 74 | model_output = model_output.replace("<|python_end|>", "") |
| 75 | |
| 76 | is_tool_call_pattern = False |
| 77 | try: |
| 78 | is_tool_call_pattern = ( |
| 79 | self.TOOL_CALL_REGEX.match( |
| 80 | model_output, timeout=envs.VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS |
| 81 | ) |
| 82 | is not None |
| 83 | ) |
| 84 | except TimeoutError: |
| 85 | logger.warning("Regex timeout occurred when matching tool call pattern.") |
| 86 | logger.debug( |
| 87 | "Regex timeout occurred when matching user input: %s", model_output |
| 88 | ) |
| 89 | |
| 90 | if not is_tool_call_pattern: |
| 91 | return ExtractedToolCallInformation( |
| 92 | tools_called=False, tool_calls=[], content=model_output |
| 93 | ) |
| 94 | |
| 95 | try: |
| 96 | module = ast.parse(model_output) |
| 97 | parsed = getattr(module.body[0], "value", None) |
| 98 | if isinstance(parsed, ast.List) and all( |
| 99 | isinstance(e, ast.Call) for e in parsed.elts |
| 100 | ): |
| 101 | return ExtractedToolCallInformation( |
| 102 | tools_called=True, |
| 103 | tool_calls=[ |
| 104 | handle_single_tool(e) # type: ignore |
| 105 | for e in parsed.elts |
| 106 | ], |
| 107 | content=None, |
| 108 | ) |
| 109 | else: |
| 110 | raise UnexpectedAstError("Tool output must be a list of function calls") |
| 111 | except Exception: |
| 112 | logger.exception("Error in extracting tool call from response.") |
| 113 | # Treat as regular text |
| 114 | return ExtractedToolCallInformation( |
| 115 | tools_called=False, tool_calls=[], content=model_output |
| 116 | ) |
| 117 | |
| 118 | def extract_tool_calls_streaming( |
| 119 | self, |
| 120 | previous_text: str, |
| 121 | current_text: str, |
| 122 | delta_text: str, |
| 123 | previous_token_ids: Sequence[int], |
| 124 | current_token_ids: Sequence[int], |
| 125 | delta_token_ids: Sequence[int], |
| 126 | request: ChatCompletionRequest, |
| 127 | ) -> DeltaMessage | None: |
| 128 | if not current_text.startswith("[") and not current_text.startswith( |
| 129 | "<|python_start|>" |
| 130 | ): |
| 131 | return DeltaMessage(content=delta_text) |
| 132 | |
| 133 | try: |
| 134 | # remove <|python_start|> and <|python_end|> |
| 135 | if current_text.startswith("<|python_start|>"): |
| 136 | current_text = current_text[len("<|python_start|>") :] |
| 137 | if current_text.endswith("<|python_end|>"): |
| 138 | current_text = current_text[: current_text.rfind("<|python_end|>")] |
| 139 | valid_and_added_text = make_valid_python(current_text) |
| 140 | if valid_and_added_text is None: |
| 141 | return None |
| 142 | valid_text, added_text = valid_and_added_text |
| 143 | |
| 144 | module = ast.parse(valid_text) |
| 145 | parsed = getattr(module.body[0], "value", None) |
| 146 | if not isinstance(parsed, ast.List) or not all( |
| 147 | isinstance(e, ast.Call) for e in parsed.elts |
| 148 | ): |
| 149 | raise UnexpectedAstError("Tool output must be a list of function calls") |
| 150 | tool_calls = [ |
| 151 | handle_single_tool(e) # type: ignore |
| 152 | for e in parsed.elts |
| 153 | ] |
| 154 | |
| 155 | tool_deltas = [] |
| 156 | for index, new_call in enumerate(tool_calls): |
| 157 | if index < self.current_tool_index: |
| 158 | continue |
| 159 | |
| 160 | self.current_tool_index = index |
| 161 | if len(self.streamed_args_for_tool) == index: |
| 162 | self.streamed_args_for_tool.append("") |
| 163 | |
| 164 | new_call_complete = ( |
| 165 | index < len(tool_calls) - 1 or ")]" not in added_text |
| 166 | ) |
| 167 | if new_call_complete: |
| 168 | self.current_tool_index += 1 |
| 169 | |
| 170 | withheld_suffix = added_text[:-2] if not new_call_complete else "" |
| 171 | if not new_call_complete and added_text[-2] == ")": |
| 172 | # Function call is incomplete. Withhold the closing bracket. |
| 173 | withheld_suffix = withheld_suffix + "}" |
| 174 | # Strings get single quotes in the model-produced string. |
| 175 | # JSON requires double quotes. |
| 176 | withheld_suffix = withheld_suffix.replace("'", '"') |
| 177 | delta = compute_tool_delta( |
| 178 | self.streamed_args_for_tool[index], new_call, index, withheld_suffix |
| 179 | ) |
| 180 | |
| 181 | if delta is not None: |
| 182 | tool_deltas.append(delta) |
| 183 | if ( |
| 184 | delta.function is not None |
| 185 | and delta.function.arguments is not None |
| 186 | ): |
| 187 | self.streamed_args_for_tool[index] += delta.function.arguments |
| 188 | |
| 189 | # HACK: serving_chat.py inspects the internal state of tool parsers |
| 190 | # when determining its final streaming delta, automatically |
| 191 | # adding autocompleted JSON. |
| 192 | # These two lines avoid that nonsense while ensuring finish_reason |
| 193 | # is set to tool_calls when at least one tool is called. |
| 194 | if tool_deltas and not self.prev_tool_call_arr: |
| 195 | self.prev_tool_call_arr = [{"arguments": {}}] |
| 196 | |
| 197 | if tool_deltas: |
| 198 | return DeltaMessage(tool_calls=tool_deltas) |
| 199 | elif not added_text and self.current_tool_id > 0: |
| 200 | # Return an empty DeltaMessage once the tool calls are all done |
| 201 | # so that finish_reason gets set. |
| 202 | return DeltaMessage(content="") |
| 203 | else: |
| 204 | return None |
| 205 | except Exception: |
| 206 | logger.exception("Error trying to handle streaming tool call.") |
| 207 | logger.debug( |
| 208 | "Skipping chunk as a result of tool streaming extraction error" |
| 209 | ) |
| 210 | return None |