| 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 | from json import JSONDecoder |
| 7 | |
| 8 | import partial_json_parser |
| 9 | import regex as re |
| 10 | from partial_json_parser.core.options import Allow |
| 11 | |
| 12 | from vllm.entrypoints.chat_utils import make_tool_call_id |
| 13 | from vllm.entrypoints.generate.base.protocol import ( |
| 14 | DeltaFunctionCall, |
| 15 | DeltaMessage, |
| 16 | DeltaToolCall, |
| 17 | ExtractedToolCallInformation, |
| 18 | FunctionCall, |
| 19 | ToolCall, |
| 20 | ) |
| 21 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 22 | ChatCompletionRequest, |
| 23 | ) |
| 24 | from vllm.logger import init_logger |
| 25 | from vllm.tokenizers import TokenizerLike |
| 26 | from vllm.tool_parsers.abstract_tool_parser import ( |
| 27 | Tool, |
| 28 | ToolParser, |
| 29 | ) |
| 30 | from vllm.tool_parsers.utils import ( |
| 31 | consume_space, |
| 32 | find_common_prefix, |
| 33 | is_complete_json, |
| 34 | partial_json_loads, |
| 35 | ) |
| 36 | |
| 37 | logger = init_logger(__name__) |
| 38 | |
| 39 | |
| 40 | class Granite20bFCToolParser(ToolParser): |
| 41 | """Tool call parser for the granite-20b-functioncalling model intended |
| 42 | for use with the examples/tool_chat_template_granite20b_fc.jinja |
| 43 | template. |
| 44 | |
| 45 | Used when --enable-auto-tool-choice --tool-call-parser granite-20-fc |
| 46 | are all set |
| 47 | """ |
| 48 | |
| 49 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 50 | super().__init__(tokenizer, tools) |
| 51 | |
| 52 | self.bot_token = "<function_call>" |
| 53 | self.tool_start_token = self.bot_token |
| 54 | self.tool_call_regex = re.compile(r"<function_call>\s*") |
| 55 | |
| 56 | def extract_tool_calls( |
| 57 | self, model_output: str, request: ChatCompletionRequest |
| 58 | ) -> ExtractedToolCallInformation: |
| 59 | if self.tool_start_token not in model_output: |
| 60 | return ExtractedToolCallInformation( |
| 61 | tools_called=False, tool_calls=[], content=model_output |
| 62 | ) |
| 63 | |
| 64 | dec = JSONDecoder() |
| 65 | try: |
| 66 | matches = list(self.tool_call_regex.finditer(model_output)) |
| 67 | logger.debug("Found %d tool call matches", len(matches)) |
| 68 | |
| 69 | raw_function_calls = [] |
| 70 | |
| 71 | for i, match in enumerate(matches): |
| 72 | # position after the <function_call> tag |
| 73 | start_of_json = match.end() |
| 74 | # end_index == the start of the next function call |
| 75 | # (if exists) |
| 76 | next_function_call_start = ( |
| 77 | matches[i + 1].start() if i + 1 < len(matches) else None |
| 78 | ) |
| 79 | |
| 80 | raw_function_calls.append( |
| 81 | dec.raw_decode( |
| 82 | model_output[start_of_json:next_function_call_start] |
| 83 | )[0] |
| 84 | ) |
| 85 | |
| 86 | logger.debug("Extracted %d tool calls", len(raw_function_calls)) |
| 87 | tool_calls = [ |
| 88 | ToolCall( |
| 89 | type="function", |
| 90 | function=FunctionCall( |
| 91 | name=function_call["name"], |
| 92 | # function call args are JSON but as a string |
| 93 | arguments=json.dumps( |
| 94 | function_call["arguments"], ensure_ascii=False |
| 95 | ), |
| 96 | ), |
| 97 | ) |
| 98 | for function_call in raw_function_calls |
| 99 | ] |
| 100 | |
| 101 | content = model_output[: model_output.find(self.bot_token)] |
| 102 | return ExtractedToolCallInformation( |
| 103 | tools_called=True, |
| 104 | tool_calls=tool_calls, |
| 105 | content=content if content else None, |
| 106 | ) |
| 107 | |
| 108 | except Exception as e: |
| 109 | logger.error("Error in extracting tool call from response %s", e) |
| 110 | return ExtractedToolCallInformation( |
| 111 | tools_called=False, tool_calls=[], content=model_output |
| 112 | ) |
| 113 | |
| 114 | def extract_tool_calls_streaming( |
| 115 | self, |
| 116 | previous_text: str, |
| 117 | current_text: str, |
| 118 | delta_text: str, |
| 119 | previous_token_ids: Sequence[int], |
| 120 | current_token_ids: Sequence[int], |
| 121 | delta_token_ids: Sequence[int], |
| 122 | request: ChatCompletionRequest, |
| 123 | ) -> DeltaMessage | None: |
| 124 | if len(current_text) < len(self.bot_token) and self.bot_token.startswith( |
| 125 | current_text |
| 126 | ): |
| 127 | return None |
| 128 | |
| 129 | if not current_text.startswith(self.bot_token): |
| 130 | return DeltaMessage(content=delta_text) |
| 131 | |
| 132 | # bit mask flags for partial JSON parsing. If the name hasn't been |
| 133 | # sent yet, don't allow sending |
| 134 | # an incomplete string since OpenAI only ever (as far as I have |
| 135 | # seen) allows sending the entire tool/ function name at once. |
| 136 | flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR |
| 137 | try: |
| 138 | tool_call_arr = [] |
| 139 | is_complete = [] |
| 140 | try: |
| 141 | start_idx = len(self.bot_token) |
| 142 | start_idx = consume_space(start_idx, current_text) |
| 143 | |
| 144 | while start_idx < len(current_text): |
| 145 | (obj, end_idx) = partial_json_loads(current_text[start_idx:], flags) |
| 146 | is_complete.append( |
| 147 | is_complete_json(current_text[start_idx : start_idx + end_idx]) |
| 148 | ) |
| 149 | start_idx += end_idx |
| 150 | start_idx = consume_space(start_idx, current_text) |
| 151 | start_idx += len(self.bot_token) |
| 152 | start_idx = consume_space(start_idx, current_text) |
| 153 | tool_call_arr.append(obj) |
| 154 | except partial_json_parser.core.exceptions.MalformedJSON: |
| 155 | logger.debug("not enough tokens to parse into JSON yet") |
| 156 | return None |
| 157 | |
| 158 | # select as the current tool call the one we're on the state at |
| 159 | current_tool_call: dict = ( |
| 160 | tool_call_arr[self.current_tool_id] if len(tool_call_arr) > 0 else {} |
| 161 | ) |
| 162 | |
| 163 | # case -- if no tokens have been streamed for the tool, e.g. |
| 164 | # only the array brackets, stream nothing |
| 165 | if len(tool_call_arr) == 0: |
| 166 | return None |
| 167 | |
| 168 | # case: we are starting a new tool in the array |
| 169 | # -> array has > 0 length AND length has moved past cursor |
| 170 | elif ( |
| 171 | len(tool_call_arr) > 0 and len(tool_call_arr) > self.current_tool_id + 1 |
| 172 | ): |
| 173 | # if we're moving on to a new call, first make sure we |
| 174 | # haven't missed anything in the previous one that was |
| 175 | # auto-generated due to JSON completions, but wasn't |
| 176 | # streamed to the client yet. |
| 177 | if self.current_tool_id >= 0: |
| 178 | cur_arguments = current_tool_call.get("arguments") |
| 179 | if cur_arguments: |
| 180 | cur_args_json = json.dumps(cur_arguments, ensure_ascii=False) |
| 181 | sent = len(self.streamed_args_for_tool[self.current_tool_id]) |
| 182 | argument_diff = cur_args_json[sent:] |
| 183 | |
| 184 | logger.debug("got arguments diff: %s", argument_diff) |
| 185 | delta = DeltaMessage( |
| 186 | tool_calls=[ |
| 187 | DeltaToolCall( |
| 188 | index=self.current_tool_id, |
| 189 | function=DeltaFunctionCall( |
| 190 | arguments=argument_diff |
| 191 | ).model_dump(exclude_none=True), |
| 192 | ) |
| 193 | ] |
| 194 | ) |
| 195 | self.streamed_args_for_tool[self.current_tool_id] += ( |
| 196 | argument_diff |
| 197 | ) |
| 198 | else: |
| 199 | delta = None |
| 200 | else: |
| 201 | delta = None |
| 202 | # re-set stuff pertaining to progress in the current tool |
| 203 | self.current_tool_id = len(tool_call_arr) - 1 |
| 204 | self.current_tool_name_sent = False |
| 205 | self.streamed_args_for_tool.append("") |
| 206 | logger.debug("starting on new tool %d", self.current_tool_id) |
| 207 | return delta |
| 208 | |
| 209 | # if the current tool name hasn't been sent, send if available |
| 210 | # - otherwise send nothing |
| 211 | elif not self.current_tool_name_sent: |
| 212 | function_name = current_tool_call.get("name") |
| 213 | if function_name: |
| 214 | delta = DeltaMessage( |
| 215 | tool_calls=[ |
| 216 | DeltaToolCall( |
| 217 | index=self.current_tool_id, |
| 218 | type="function", |
| 219 | id=make_tool_call_id(), |
| 220 | function=DeltaFunctionCall( |
| 221 | name=function_name |
| 222 | ).model_dump(exclude_none=True), |
| 223 | ) |
| 224 | ] |
| 225 | ) |
| 226 | self.current_tool_name_sent = True |
| 227 | else: |
| 228 | delta = None |
| 229 | |
| 230 | # now we know we're on the same tool call and we're streaming |
| 231 | # arguments |
| 232 | else: |
| 233 | cur_arguments = current_tool_call.get("arguments") |
| 234 | delta = None |
| 235 | |
| 236 | if cur_arguments: |
| 237 | sent = len(self.streamed_args_for_tool[self.current_tool_id]) |
| 238 | cur_args_json = json.dumps(cur_arguments, ensure_ascii=False) |
| 239 | prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get( |
| 240 | "arguments" |
| 241 | ) |
| 242 | |
| 243 | argument_diff = None |
| 244 | if is_complete[self.current_tool_id]: |
| 245 | argument_diff = cur_args_json[sent:] |
| 246 | elif prev_arguments: |
| 247 | prev_args_json = json.dumps(prev_arguments, ensure_ascii=False) |
| 248 | if cur_args_json != prev_args_json: |
| 249 | prefix = find_common_prefix(prev_args_json, cur_args_json) |
| 250 | argument_diff = prefix[sent:] |
| 251 | |
| 252 | if argument_diff is not None: |
| 253 | delta = DeltaMessage( |
| 254 | tool_calls=[ |
| 255 | DeltaToolCall( |
| 256 | index=self.current_tool_id, |
| 257 | function=DeltaFunctionCall( |
| 258 | arguments=argument_diff |
| 259 | ).model_dump(exclude_none=True), |
| 260 | ) |
| 261 | ] |
| 262 | ) |
| 263 | self.streamed_args_for_tool[self.current_tool_id] += ( |
| 264 | argument_diff |
| 265 | ) |
| 266 | |
| 267 | self.prev_tool_call_arr = tool_call_arr |
| 268 | return delta |
| 269 | |
| 270 | except Exception as e: |
| 271 | logger.error("Error trying to handle streaming tool call: %s", e) |
| 272 | logger.debug( |
| 273 | "Skipping chunk as a result of tool streaming extraction error" |
| 274 | ) |
| 275 | return None |