| 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 partial_json_parser |
| 8 | import regex as re |
| 9 | from partial_json_parser.core.options import Allow |
| 10 | |
| 11 | from vllm.entrypoints.chat_utils import make_tool_call_id |
| 12 | from vllm.entrypoints.generate.base.protocol import ( |
| 13 | DeltaFunctionCall, |
| 14 | DeltaMessage, |
| 15 | DeltaToolCall, |
| 16 | ExtractedToolCallInformation, |
| 17 | FunctionCall, |
| 18 | ToolCall, |
| 19 | ) |
| 20 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 21 | ChatCompletionRequest, |
| 22 | ) |
| 23 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 24 | from vllm.logger import init_logger |
| 25 | from vllm.tokenizers import TokenizerLike |
| 26 | from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser |
| 27 | from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json |
| 28 | from vllm.utils.mistral import is_mistral_tokenizer |
| 29 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | |
| 33 | class JambaToolParser(ToolParser): |
| 34 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 35 | super().__init__(tokenizer, tools) |
| 36 | |
| 37 | if is_mistral_tokenizer(self.model_tokenizer): |
| 38 | raise ValueError( |
| 39 | "Detected a MistralTokenizer tokenizer when using a Jamba model" |
| 40 | ) |
| 41 | |
| 42 | self.current_tool_name_sent: bool = False |
| 43 | self.prev_tool_call_arr: list[dict] = [] |
| 44 | self.current_tool_id: int = -1 |
| 45 | self.streamed_args_for_tool: list[ |
| 46 | str |
| 47 | ] = [] # map what has been streamed for each tool so far to a list |
| 48 | |
| 49 | self.tool_calls_start_token: str = "<tool_calls>" |
| 50 | self.tool_calls_end_token: str = "</tool_calls>" |
| 51 | |
| 52 | self.tool_calls_regex = re.compile( |
| 53 | rf"{self.tool_calls_start_token}(.*?){self.tool_calls_end_token}", re.DOTALL |
| 54 | ) |
| 55 | |
| 56 | if not self.model_tokenizer: |
| 57 | raise ValueError( |
| 58 | "The model tokenizer must be passed to the ToolParser " |
| 59 | "constructor during construction." |
| 60 | ) |
| 61 | self.tool_calls_start_token_id = self.vocab.get(self.tool_calls_start_token) |
| 62 | self.tool_calls_end_token_id = self.vocab.get(self.tool_calls_end_token) |
| 63 | if ( |
| 64 | self.tool_calls_start_token_id is None |
| 65 | or self.tool_calls_end_token_id is None |
| 66 | ): |
| 67 | raise RuntimeError( |
| 68 | "Jamba Tool parser could not locate tool calls start/end " |
| 69 | "tokens in the tokenizer!" |
| 70 | ) |
| 71 | |
| 72 | def adjust_request( |
| 73 | self, request: ChatCompletionRequest | ResponsesRequest |
| 74 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 75 | request = super().adjust_request(request) |
| 76 | if request.tools and request.tool_choice != "none": |
| 77 | # do not skip special tokens because jamba use the special |
| 78 | # tokens to indicate the start and end of the tool calls |
| 79 | # information. |
| 80 | request.skip_special_tokens = False |
| 81 | return request |
| 82 | |
| 83 | def extract_tool_calls( |
| 84 | self, model_output: str, request: ChatCompletionRequest |
| 85 | ) -> ExtractedToolCallInformation: |
| 86 | # sanity check; avoid unnecessary processing |
| 87 | if self.tool_calls_start_token not in model_output: |
| 88 | return ExtractedToolCallInformation( |
| 89 | tools_called=False, tool_calls=[], content=model_output |
| 90 | ) |
| 91 | |
| 92 | else: |
| 93 | try: |
| 94 | # use a regex to find the tool call between the tags |
| 95 | function_calls = self.tool_calls_regex.findall(model_output)[0] |
| 96 | |
| 97 | # load the JSON, and then use it to build the Function and |
| 98 | # Tool Call |
| 99 | raw_function_calls = json.loads(function_calls) |
| 100 | tool_calls = [ |
| 101 | ToolCall( |
| 102 | type="function", |
| 103 | function=FunctionCall( |
| 104 | name=function_call["name"], |
| 105 | # function call args are JSON but as a string |
| 106 | arguments=json.dumps( |
| 107 | function_call["arguments"], ensure_ascii=False |
| 108 | ), |
| 109 | ), |
| 110 | ) |
| 111 | for function_call in raw_function_calls |
| 112 | ] |
| 113 | |
| 114 | content = model_output[: model_output.find(self.tool_calls_start_token)] |
| 115 | return ExtractedToolCallInformation( |
| 116 | tools_called=True, |
| 117 | tool_calls=tool_calls, |
| 118 | content=content if (len(content) > 0 and content != " ") else None, |
| 119 | ) |
| 120 | |
| 121 | except Exception: |
| 122 | logger.exception("Error in extracting tool call from response.") |
| 123 | return ExtractedToolCallInformation( |
| 124 | tools_called=False, tool_calls=[], content=model_output |
| 125 | ) |
| 126 | |
| 127 | def extract_tool_calls_streaming( |
| 128 | self, |
| 129 | previous_text: str, |
| 130 | current_text: str, |
| 131 | delta_text: str, |
| 132 | previous_token_ids: Sequence[int], |
| 133 | current_token_ids: Sequence[int], |
| 134 | delta_token_ids: Sequence[int], |
| 135 | request: ChatCompletionRequest, |
| 136 | ) -> DeltaMessage | None: |
| 137 | # if the tool call token is not in the tokens generated so far, append |
| 138 | # output to contents since it's not a tool |
| 139 | if self.tool_calls_start_token not in current_text: |
| 140 | return DeltaMessage(content=delta_text) |
| 141 | |
| 142 | # if the tool call token ID IS in the tokens generated so far, that |
| 143 | # means we're parsing as tool calls now |
| 144 | |
| 145 | # handle if we detected the start of tool calls token which means |
| 146 | # the start of tool calling |
| 147 | if ( |
| 148 | self.tool_calls_start_token_id in delta_token_ids |
| 149 | and len(delta_token_ids) == 1 |
| 150 | ): |
| 151 | # if it's the only token, return None, so we don't send a chat |
| 152 | # completion and don't send a control token |
| 153 | return None |
| 154 | |
| 155 | # bit mask flags for partial JSON parsing. If the name hasn't been |
| 156 | # sent yet, don't allow sending |
| 157 | # an incomplete string since OpenAI only ever (as far as I have |
| 158 | # seen) allows sending the entire tool/ function name at once. |
| 159 | flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR |
| 160 | try: |
| 161 | # Extract the tool calls between the special tool call tokens |
| 162 | parsable_arr = current_text.split(self.tool_calls_start_token)[-1].split( |
| 163 | self.tool_calls_end_token |
| 164 | )[0] |
| 165 | |
| 166 | # tool calls are generated in an array, so do partial JSON |
| 167 | # parsing on the entire array |
| 168 | try: |
| 169 | tool_call_arr: list[dict] = partial_json_parser.loads( |
| 170 | parsable_arr, flags |
| 171 | ) |
| 172 | except partial_json_parser.core.exceptions.MalformedJSON: |
| 173 | logger.debug("not enough tokens to parse into JSON yet") |
| 174 | return None |
| 175 | |
| 176 | # select as the current tool call the one we're on the state at |
| 177 | |
| 178 | current_tool_call: dict = ( |
| 179 | tool_call_arr[self.current_tool_id] if len(tool_call_arr) > 0 else {} |
| 180 | ) |
| 181 | |
| 182 | # case -- if no tokens have been streamed for the tool, e.g. |
| 183 | # only the array brackets, stream nothing |
| 184 | if len(tool_call_arr) == 0: |
| 185 | return None |
| 186 | |
| 187 | # case: we are starting a new tool in the array |
| 188 | # -> array has > 0 length AND length has moved past cursor |
| 189 | elif ( |
| 190 | len(tool_call_arr) > 0 and len(tool_call_arr) > self.current_tool_id + 1 |
| 191 | ): |
| 192 | # if we're moving on to a new call, first make sure we |
| 193 | # haven't missed anything in the previous one that was |
| 194 | # auto-generated due to JSON completions, but wasn't |
| 195 | # streamed to the client yet. |
| 196 | if self.current_tool_id >= 0: |
| 197 | diff: str | None = current_tool_call.get("arguments") |
| 198 | |
| 199 | if diff: |
| 200 | diff = json.dumps(diff, ensure_ascii=False).replace( |
| 201 | self.streamed_args_for_tool[self.current_tool_id], "" |
| 202 | ) |
| 203 | delta = DeltaMessage( |
| 204 | tool_calls=[ |
| 205 | DeltaToolCall( |
| 206 | index=self.current_tool_id, |
| 207 | function=DeltaFunctionCall( |
| 208 | arguments=diff |
| 209 | ).model_dump(exclude_none=True), |
| 210 | ) |
| 211 | ] |
| 212 | ) |
| 213 | self.streamed_args_for_tool[self.current_tool_id] += diff |
| 214 | else: |
| 215 | delta = None |
| 216 | else: |
| 217 | delta = None |
| 218 | # re-set stuff pertaining to progress in the current tool |
| 219 | self.current_tool_id = len(tool_call_arr) - 1 |
| 220 | self.current_tool_name_sent = False |
| 221 | self.streamed_args_for_tool.append("") |
| 222 | logger.debug("starting on new tool %d", self.current_tool_id) |
| 223 | return delta |
| 224 | |
| 225 | # case: update an existing tool - this is handled below |
| 226 | |
| 227 | # if the current tool name hasn't been sent, send if available |
| 228 | # - otherwise send nothing |
| 229 | if not self.current_tool_name_sent: |
| 230 | function_name = current_tool_call.get("name") |
| 231 | if function_name: |
| 232 | delta = DeltaMessage( |
| 233 | tool_calls=[ |
| 234 | DeltaToolCall( |
| 235 | index=self.current_tool_id, |
| 236 | type="function", |
| 237 | id=make_tool_call_id(), |
| 238 | function=DeltaFunctionCall( |
| 239 | name=function_name |
| 240 | ).model_dump(exclude_none=True), |
| 241 | ) |
| 242 | ] |
| 243 | ) |
| 244 | self.current_tool_name_sent = True |
| 245 | else: |
| 246 | delta = None |
| 247 | |
| 248 | # now we know we're on the same tool call and we're streaming |
| 249 | # arguments |
| 250 | else: |
| 251 | prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get( |
| 252 | "arguments" |
| 253 | ) |
| 254 | cur_arguments = current_tool_call.get("arguments") |
| 255 | |
| 256 | new_text = delta_text.replace("'", '"') |
| 257 | |
| 258 | if not cur_arguments and not prev_arguments: |
| 259 | delta = None |
| 260 | elif not cur_arguments and prev_arguments: |
| 261 | logger.error( |
| 262 | "INVARIANT - impossible to have arguments reset mid-arguments" |
| 263 | ) |
| 264 | delta = None |
| 265 | elif cur_arguments and not prev_arguments: |
| 266 | cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) |
| 267 | logger.debug("finding %s in %s", new_text, cur_arguments_json) |
| 268 | |
| 269 | # `new_text` may not appear verbatim in the re-serialized JSON. |
| 270 | match_start = cur_arguments_json.find(new_text) |
| 271 | if match_start != -1: |
| 272 | arguments_delta = cur_arguments_json[ |
| 273 | : match_start + len(new_text) |
| 274 | ] |
| 275 | elif is_complete_json(parsable_arr): |
| 276 | # Complete in this delta: send whole, don't drop. |
| 277 | arguments_delta = cur_arguments_json |
| 278 | else: |
| 279 | # Still partial: wait for more text. |
| 280 | return None |
| 281 | logger.debug( |
| 282 | "First tokens in arguments received: %s", arguments_delta |
| 283 | ) |
| 284 | delta = DeltaMessage( |
| 285 | tool_calls=[ |
| 286 | DeltaToolCall( |
| 287 | index=self.current_tool_id, |
| 288 | function=DeltaFunctionCall( |
| 289 | arguments=arguments_delta |
| 290 | ).model_dump(exclude_none=True), |
| 291 | ) |
| 292 | ] |
| 293 | ) |
| 294 | self.streamed_args_for_tool[self.current_tool_id] += arguments_delta |
| 295 | |
| 296 | elif cur_arguments and prev_arguments: |
| 297 | cur_args_json = json.dumps(cur_arguments, ensure_ascii=False) |
| 298 | prev_args_json = json.dumps(prev_arguments, ensure_ascii=False) |
| 299 | logger.debug( |
| 300 | "Searching for diff between \n%s\n%s", |
| 301 | cur_args_json, |
| 302 | prev_args_json, |
| 303 | ) |
| 304 | |
| 305 | argument_diff = extract_intermediate_diff( |
| 306 | cur_args_json, prev_args_json |
| 307 | ) |
| 308 | logger.debug("got arguments diff: %s", argument_diff) |
| 309 | delta = DeltaMessage( |
| 310 | tool_calls=[ |
| 311 | DeltaToolCall( |
| 312 | index=self.current_tool_id, |
| 313 | function=DeltaFunctionCall( |
| 314 | arguments=argument_diff |
| 315 | ).model_dump(exclude_none=True), |
| 316 | ) |
| 317 | ] |
| 318 | ) |
| 319 | self.streamed_args_for_tool[self.current_tool_id] += argument_diff |
| 320 | else: |
| 321 | # try parsing it with regular JSON - if it works we're |
| 322 | # at the end, and we need to send the difference between |
| 323 | # tokens streamed so far and the valid JSON |
| 324 | delta = None |
| 325 | |
| 326 | # check to see if the name is defined and has been sent. if so, |
| 327 | # stream the name - otherwise keep waiting |
| 328 | # finish by setting old and returning None as base case |
| 329 | self.prev_tool_call_arr = tool_call_arr |
| 330 | return delta |
| 331 | |
| 332 | except Exception: |
| 333 | logger.exception("Error trying to handle streaming tool call.") |
| 334 | logger.debug( |
| 335 | "Skipping chunk as a result of tool streaming extraction error" |
| 336 | ) |
| 337 | return None |