| 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 Olmo3PythonicToolParser(ToolParser): |
| 34 | """Tool call parser for Olmo 3 models that produce tool calls as |
| 35 | newline-separated pythonic strings. |
| 36 | Used when --enable-auto-tool-choice --tool-call-parser pythonic are all set |
| 37 | Code copied from pythonic_tool_parser.py and updated to handle |
| 38 | - newline separated pythonic tool calls. |
| 39 | - argument values being null/true/false instead of Pythonic literals. |
| 40 | """ |
| 41 | |
| 42 | # TODO(mdepinet): Possible future improvements: |
| 43 | # 1. Support text + tools separated by either <|python_tag|> or \n\n |
| 44 | # 2. Support tools outside of a list (or separated by a semicolon). |
| 45 | # This depends on item 1 for consistent streaming. |
| 46 | # Neither of these are necessary for e.g. ToolACE, but both would help make |
| 47 | # Llama3.2 models more reliable. |
| 48 | |
| 49 | TOOL_CALL_REGEX = re.compile( |
| 50 | 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*)+\]", |
| 51 | re.DOTALL, |
| 52 | ) |
| 53 | |
| 54 | def __init__( |
| 55 | self, |
| 56 | tokenizer: PreTrainedTokenizerBase, |
| 57 | tools: list[Tool] | None = None, |
| 58 | ): |
| 59 | super().__init__(tokenizer, tools) |
| 60 | |
| 61 | # Rename for readability. This is NOT a tool id. |
| 62 | @property |
| 63 | def current_tool_index(self) -> int: |
| 64 | return self.current_tool_id |
| 65 | |
| 66 | @current_tool_index.setter |
| 67 | def current_tool_index(self, value: int) -> None: |
| 68 | self.current_tool_id = value |
| 69 | |
| 70 | def extract_tool_calls( |
| 71 | self, model_output: str, request: ChatCompletionRequest |
| 72 | ) -> ExtractedToolCallInformation: |
| 73 | """Extract the tool calls from a complete model response.""" |
| 74 | original_model_output = model_output |
| 75 | # Remove xml tags. |
| 76 | match = re.search( |
| 77 | r"<function_calls>(.*?)</function_calls>", model_output, re.DOTALL |
| 78 | ) |
| 79 | if match: |
| 80 | model_output = match.group(1).strip() |
| 81 | # Make the newline separated function calls into a list. |
| 82 | model_output = ", ".join( |
| 83 | [line.strip() for line in model_output.splitlines() if line.strip()] |
| 84 | ) |
| 85 | model_output = f"[{model_output}]" |
| 86 | |
| 87 | is_tool_call_pattern = False |
| 88 | try: |
| 89 | is_tool_call_pattern = ( |
| 90 | self.TOOL_CALL_REGEX.match( |
| 91 | model_output, timeout=envs.VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS |
| 92 | ) |
| 93 | is not None |
| 94 | ) |
| 95 | except TimeoutError: |
| 96 | logger.warning("Regex timeout occurred when matching tool call pattern.") |
| 97 | logger.debug( |
| 98 | "Regex timeout occurred when matching user input: %s", model_output |
| 99 | ) |
| 100 | |
| 101 | if not is_tool_call_pattern: |
| 102 | return ExtractedToolCallInformation( |
| 103 | tools_called=False, tool_calls=[], content=original_model_output |
| 104 | ) |
| 105 | |
| 106 | try: |
| 107 | module = ast.parse(model_output) |
| 108 | parsed = getattr(module.body[0], "value", None) |
| 109 | if isinstance(parsed, ast.List) and all( |
| 110 | isinstance(e, ast.Call) for e in parsed.elts |
| 111 | ): |
| 112 | return ExtractedToolCallInformation( |
| 113 | tools_called=True, |
| 114 | tool_calls=[ |
| 115 | handle_single_tool(e) # type: ignore |
| 116 | for e in parsed.elts |
| 117 | ], |
| 118 | content=None, |
| 119 | ) |
| 120 | else: |
| 121 | raise UnexpectedAstError("Tool output must be a list of function calls") |
| 122 | except Exception: |
| 123 | logger.exception("Error in extracting tool call from response.") |
| 124 | # Treat as regular text |
| 125 | return ExtractedToolCallInformation( |
| 126 | tools_called=False, tool_calls=[], content=original_model_output |
| 127 | ) |
| 128 | |
| 129 | def extract_tool_calls_streaming( |
| 130 | self, |
| 131 | previous_text: str, |
| 132 | current_text: str, |
| 133 | delta_text: str, |
| 134 | previous_token_ids: Sequence[int], |
| 135 | current_token_ids: Sequence[int], |
| 136 | delta_token_ids: Sequence[int], |
| 137 | request: ChatCompletionRequest, |
| 138 | ) -> DeltaMessage | None: |
| 139 | # All function calls start with the <function_calls> tag. |
| 140 | # But since this is streaming, we may have seen only part of the tag. |
| 141 | if not current_text.startswith("<"): |
| 142 | return DeltaMessage(content=delta_text) |
| 143 | |
| 144 | try: |
| 145 | # Remove xml tags. |
| 146 | if current_text.startswith("<function_calls>"): |
| 147 | current_text = current_text[len("<function_calls>") :] |
| 148 | if current_text.endswith("</function_calls>"): |
| 149 | current_text = current_text[: -len("</function_calls>")] |
| 150 | |
| 151 | valid_and_added_text = make_valid_python(current_text) |
| 152 | if valid_and_added_text is None: |
| 153 | return None |
| 154 | valid_text, added_text = valid_and_added_text |
| 155 | |
| 156 | # Make the newline separated function calls into a list. |
| 157 | valid_text = ", ".join( |
| 158 | [line.strip() for line in valid_text.splitlines() if line.strip()] |
| 159 | ) |
| 160 | valid_text = f"[{valid_text}]" |
| 161 | module = ast.parse(valid_text) |
| 162 | parsed = getattr(module.body[0], "value", None) |
| 163 | if not isinstance(parsed, ast.List) or not all( |
| 164 | isinstance(e, ast.Call) for e in parsed.elts |
| 165 | ): |
| 166 | raise UnexpectedAstError( |
| 167 | "Tool output must be a sequence of newline-separated calls" |
| 168 | ) |
| 169 | tool_calls = [ |
| 170 | handle_single_tool(e) # type: ignore |
| 171 | for e in parsed.elts |
| 172 | ] |
| 173 | |
| 174 | tool_deltas = [] |
| 175 | for index, new_call in enumerate(tool_calls): |
| 176 | if index < self.current_tool_index: |
| 177 | continue |
| 178 | |
| 179 | self.current_tool_index = index |
| 180 | if len(self.streamed_args_for_tool) == index: |
| 181 | self.streamed_args_for_tool.append("") |
| 182 | |
| 183 | new_call_complete = index < len(tool_calls) - 1 or ")" not in added_text |
| 184 | if new_call_complete: |
| 185 | self.current_tool_index += 1 |
| 186 | |
| 187 | withheld_suffix = added_text[:-1] if not new_call_complete else "" |
| 188 | if not new_call_complete and added_text[-1] == ")": |
| 189 | # Function call is incomplete. Withhold the closing bracket. |
| 190 | withheld_suffix = withheld_suffix + "}" |
| 191 | # Strings get single quotes in the model-produced string. |
| 192 | # JSON requires double quotes. |
| 193 | withheld_suffix = withheld_suffix.replace("'", '"') |
| 194 | delta = compute_tool_delta( |
| 195 | self.streamed_args_for_tool[index], new_call, index, withheld_suffix |
| 196 | ) |
| 197 | |
| 198 | if delta is not None: |
| 199 | tool_deltas.append(delta) |
| 200 | if ( |
| 201 | delta.function is not None |
| 202 | and delta.function.arguments is not None |
| 203 | ): |
| 204 | self.streamed_args_for_tool[index] += delta.function.arguments |
| 205 | |
| 206 | # HACK: serving_chat.py inspects the internal state of tool parsers |
| 207 | # when determining its final streaming delta, automatically |
| 208 | # adding autocompleted JSON. |
| 209 | # These two lines avoid that nonsense while ensuring finish_reason |
| 210 | # is set to tool_calls when at least one tool is called. |
| 211 | if tool_deltas and not self.prev_tool_call_arr: |
| 212 | self.prev_tool_call_arr = [{"arguments": {}}] |
| 213 | |
| 214 | if tool_deltas: |
| 215 | return DeltaMessage(tool_calls=tool_deltas) |
| 216 | elif not added_text and self.current_tool_id > 0: |
| 217 | # Return an empty DeltaMessage once the tool calls are all done |
| 218 | # so that finish_reason gets set. |
| 219 | return DeltaMessage(content="") |
| 220 | else: |
| 221 | return None |
| 222 | except Exception: |
| 223 | logger.exception("Error trying to handle streaming tool call.") |
| 224 | logger.debug( |
| 225 | "Skipping chunk as a result of tool streaming extraction error" |
| 226 | ) |
| 227 | return None |