| 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 Sequence |
| 6 | |
| 7 | import regex as re |
| 8 | |
| 9 | from vllm.entrypoints.chat_utils import make_tool_call_id |
| 10 | from vllm.entrypoints.generate.base.protocol import ( |
| 11 | DeltaFunctionCall, |
| 12 | DeltaMessage, |
| 13 | DeltaToolCall, |
| 14 | ExtractedToolCallInformation, |
| 15 | FunctionCall, |
| 16 | ToolCall, |
| 17 | ) |
| 18 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 19 | ChatCompletionRequest, |
| 20 | ) |
| 21 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 22 | from vllm.logger import init_logger |
| 23 | from vllm.tokenizers import TokenizerLike |
| 24 | from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | |
| 29 | class FunctionGemmaToolParser(ToolParser): |
| 30 | """Tool parser for Google's FunctionGemma model (google/functiongemma-270m-it). |
| 31 | |
| 32 | Handles the FunctionGemma function call format: |
| 33 | <start_function_call>call:func_name{param:<escape>value<escape>}<end_function_call> |
| 34 | """ |
| 35 | |
| 36 | # FunctionGemma tokens |
| 37 | tool_call_start_token: str = "<start_function_call>" |
| 38 | tool_call_end_token: str = "<end_function_call>" |
| 39 | |
| 40 | # Regex patterns |
| 41 | tool_call_regex: re.Pattern = re.compile( |
| 42 | r"<start_function_call>call:(\w+)\{(.*?)\}<end_function_call>" |
| 43 | r"|<start_function_call>call:(\w+)\{(.*)", |
| 44 | re.DOTALL, |
| 45 | ) |
| 46 | arg_regex: re.Pattern = re.compile( |
| 47 | r"(\w+):<escape>(.*?)<escape>", |
| 48 | re.DOTALL, |
| 49 | ) |
| 50 | |
| 51 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 52 | super().__init__(tokenizer, tools) |
| 53 | |
| 54 | # Streaming state |
| 55 | self.current_tool_name_sent: bool = False |
| 56 | self.prev_tool_call_arr: list[dict] = [] |
| 57 | self.current_tool_id: int = -1 |
| 58 | self.streamed_args_for_tool: list[str] = [] |
| 59 | self.buffered_delta_text = "" |
| 60 | |
| 61 | def _parse_arguments(self, args_str: str) -> dict: |
| 62 | """Parse FunctionGemma argument string into a dictionary.""" |
| 63 | arguments: dict = {} |
| 64 | if not args_str: |
| 65 | return arguments |
| 66 | |
| 67 | matches = self.arg_regex.findall(args_str) |
| 68 | for key, value in matches: |
| 69 | try: |
| 70 | parsed_value = json.loads(value) |
| 71 | arguments[key] = parsed_value |
| 72 | except json.JSONDecodeError: |
| 73 | arguments[key] = value |
| 74 | |
| 75 | return arguments |
| 76 | |
| 77 | def adjust_request( |
| 78 | self, request: ChatCompletionRequest | ResponsesRequest |
| 79 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 80 | request = super().adjust_request(request) |
| 81 | if request.tools and request.tool_choice != "none": |
| 82 | request.skip_special_tokens = False |
| 83 | return request |
| 84 | |
| 85 | def extract_tool_calls( |
| 86 | self, |
| 87 | model_output: str, |
| 88 | request: ChatCompletionRequest, |
| 89 | ) -> ExtractedToolCallInformation: |
| 90 | if self.tool_call_start_token not in model_output: |
| 91 | return ExtractedToolCallInformation( |
| 92 | tools_called=False, tool_calls=[], content=model_output |
| 93 | ) |
| 94 | |
| 95 | try: |
| 96 | matches = self.tool_call_regex.findall(model_output) |
| 97 | |
| 98 | if not matches: |
| 99 | return ExtractedToolCallInformation( |
| 100 | tools_called=False, tool_calls=[], content=model_output |
| 101 | ) |
| 102 | |
| 103 | tool_calls: list[ToolCall] = [] |
| 104 | |
| 105 | for match in matches: |
| 106 | func_name = match[0] if match[0] else match[2] |
| 107 | args_str = match[1] if match[1] else match[3] |
| 108 | |
| 109 | if not func_name: |
| 110 | continue |
| 111 | |
| 112 | arguments = self._parse_arguments(args_str) |
| 113 | |
| 114 | tool_calls.append( |
| 115 | ToolCall( |
| 116 | type="function", |
| 117 | function=FunctionCall( |
| 118 | name=func_name, |
| 119 | arguments=json.dumps(arguments, ensure_ascii=False), |
| 120 | ), |
| 121 | ) |
| 122 | ) |
| 123 | |
| 124 | if tool_calls: |
| 125 | content_end = model_output.find(self.tool_call_start_token) |
| 126 | content = ( |
| 127 | model_output[:content_end].strip() if content_end > 0 else None |
| 128 | ) |
| 129 | |
| 130 | return ExtractedToolCallInformation( |
| 131 | tools_called=True, |
| 132 | tool_calls=tool_calls, |
| 133 | content=content if content else None, |
| 134 | ) |
| 135 | |
| 136 | return ExtractedToolCallInformation( |
| 137 | tools_called=False, tool_calls=[], content=model_output |
| 138 | ) |
| 139 | |
| 140 | except Exception: |
| 141 | logger.exception("Error extracting tool calls from FunctionGemma response") |
| 142 | return ExtractedToolCallInformation( |
| 143 | tools_called=False, tool_calls=[], content=model_output |
| 144 | ) |
| 145 | |
| 146 | def _buffer_delta_text(self, delta_text: str) -> str: |
| 147 | """Buffer incoming delta text to handle multi-token special sequences.""" |
| 148 | potential_start = "<start_function_call>" |
| 149 | potential_end = "<end_function_call>" |
| 150 | |
| 151 | combined = self.buffered_delta_text + delta_text |
| 152 | |
| 153 | if combined.endswith(potential_start) or combined.endswith(potential_end): |
| 154 | self.buffered_delta_text = "" |
| 155 | return combined |
| 156 | |
| 157 | for tag in [potential_start, potential_end]: |
| 158 | for i in range(1, len(tag)): |
| 159 | if combined.endswith(tag[:i]): |
| 160 | self.buffered_delta_text = combined[-(i):] |
| 161 | return combined[:-i] |
| 162 | |
| 163 | self.buffered_delta_text = "" |
| 164 | return combined |
| 165 | |
| 166 | def extract_tool_calls_streaming( |
| 167 | self, |
| 168 | previous_text: str, |
| 169 | current_text: str, |
| 170 | delta_text: str, |
| 171 | previous_token_ids: Sequence[int], |
| 172 | current_token_ids: Sequence[int], |
| 173 | delta_token_ids: Sequence[int], |
| 174 | request: ChatCompletionRequest, |
| 175 | ) -> DeltaMessage | None: |
| 176 | delta_text = self._buffer_delta_text(delta_text) |
| 177 | current_text = previous_text + delta_text |
| 178 | |
| 179 | if self.tool_call_start_token not in current_text: |
| 180 | if delta_text: |
| 181 | return DeltaMessage(content=delta_text) |
| 182 | return None |
| 183 | |
| 184 | try: |
| 185 | start_count = current_text.count(self.tool_call_start_token) |
| 186 | end_count = current_text.count(self.tool_call_end_token) |
| 187 | prev_start_count = previous_text.count(self.tool_call_start_token) |
| 188 | prev_end_count = previous_text.count(self.tool_call_end_token) |
| 189 | |
| 190 | if self.tool_call_start_token not in current_text: |
| 191 | return DeltaMessage(content=delta_text) |
| 192 | |
| 193 | # Starting a new function call |
| 194 | if start_count > prev_start_count and start_count > end_count: |
| 195 | self.current_tool_id += 1 |
| 196 | self.current_tool_name_sent = False |
| 197 | self.streamed_args_for_tool.append("") |
| 198 | self.prev_tool_call_arr.append({}) |
| 199 | logger.debug("Starting new tool call %d", self.current_tool_id) |
| 200 | return None |
| 201 | |
| 202 | # In the middle of a function call |
| 203 | if start_count > end_count: |
| 204 | last_start = current_text.rfind(self.tool_call_start_token) |
| 205 | partial_call = current_text[ |
| 206 | last_start + len(self.tool_call_start_token) : |
| 207 | ] |
| 208 | |
| 209 | if partial_call.startswith("call:"): |
| 210 | func_part = partial_call[5:] |
| 211 | |
| 212 | if "{" in func_part: |
| 213 | func_name = func_part.split("{")[0] |
| 214 | args_part = ( |
| 215 | func_part.split("{", 1)[1] if "{" in func_part else "" |
| 216 | ) |
| 217 | |
| 218 | if not self.current_tool_name_sent and func_name: |
| 219 | self.current_tool_name_sent = True |
| 220 | self.prev_tool_call_arr[self.current_tool_id] = { |
| 221 | "name": func_name, |
| 222 | "arguments": {}, |
| 223 | } |
| 224 | return DeltaMessage( |
| 225 | tool_calls=[ |
| 226 | DeltaToolCall( |
| 227 | index=self.current_tool_id, |
| 228 | type="function", |
| 229 | id=make_tool_call_id(), |
| 230 | function=DeltaFunctionCall( |
| 231 | name=func_name |
| 232 | ).model_dump(exclude_none=True), |
| 233 | ) |
| 234 | ] |
| 235 | ) |
| 236 | |
| 237 | if self.current_tool_name_sent and args_part: |
| 238 | current_args = self._parse_arguments(args_part) |
| 239 | if current_args: |
| 240 | current_args_json = json.dumps( |
| 241 | current_args, ensure_ascii=False |
| 242 | ) |
| 243 | prev_streamed = self.streamed_args_for_tool[ |
| 244 | self.current_tool_id |
| 245 | ] |
| 246 | |
| 247 | if len(current_args_json) > len(prev_streamed): |
| 248 | diff = current_args_json[len(prev_streamed) :] |
| 249 | self.streamed_args_for_tool[ |
| 250 | self.current_tool_id |
| 251 | ] = current_args_json |
| 252 | self.prev_tool_call_arr[self.current_tool_id][ |
| 253 | "arguments" |
| 254 | ] = current_args |
| 255 | |
| 256 | return DeltaMessage( |
| 257 | tool_calls=[ |
| 258 | DeltaToolCall( |
| 259 | index=self.current_tool_id, |
| 260 | function=DeltaFunctionCall( |
| 261 | arguments=diff |
| 262 | ).model_dump(exclude_none=True), |
| 263 | ) |
| 264 | ] |
| 265 | ) |
| 266 | |
| 267 | return None |
| 268 | |
| 269 | # Function call just ended |
| 270 | if end_count > prev_end_count: |
| 271 | if self.current_tool_id >= 0 and self.current_tool_id < len( |
| 272 | self.prev_tool_call_arr |
| 273 | ): |
| 274 | all_calls = self.tool_call_regex.findall(current_text) |
| 275 | args = {} |
| 276 | if self.current_tool_id < len(all_calls): |
| 277 | match = all_calls[self.current_tool_id] |
| 278 | if match[0]: |
| 279 | args_str = match[1] |
| 280 | args = self._parse_arguments(args_str) |
| 281 | self.prev_tool_call_arr[self.current_tool_id][ |
| 282 | "arguments" |
| 283 | ] = args |
| 284 | |
| 285 | if args: |
| 286 | args_json = json.dumps(args, ensure_ascii=False) |
| 287 | prev_streamed = self.streamed_args_for_tool[ |
| 288 | self.current_tool_id |
| 289 | ] |
| 290 | if len(args_json) > len(prev_streamed): |
| 291 | diff = args_json[len(prev_streamed) :] |
| 292 | self.streamed_args_for_tool[self.current_tool_id] = ( |
| 293 | args_json |
| 294 | ) |
| 295 | return DeltaMessage( |
| 296 | tool_calls=[ |
| 297 | DeltaToolCall( |
| 298 | index=self.current_tool_id, |
| 299 | function=DeltaFunctionCall( |
| 300 | arguments=diff |
| 301 | ).model_dump(exclude_none=True), |
| 302 | ) |
| 303 | ] |
| 304 | ) |
| 305 | return None |
| 306 | |
| 307 | if delta_text: |
| 308 | return DeltaMessage(content=delta_text) |
| 309 | return None |
| 310 | |
| 311 | except Exception: |
| 312 | logger.exception("Error in streaming tool call extraction") |
| 313 | return None |