| 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 | from partial_json_parser.core.options import Allow |
| 9 | |
| 10 | from vllm.entrypoints.chat_utils import make_tool_call_id |
| 11 | from vllm.entrypoints.generate.base.protocol import ( |
| 12 | DeltaFunctionCall, |
| 13 | DeltaMessage, |
| 14 | DeltaToolCall, |
| 15 | ExtractedToolCallInformation, |
| 16 | FunctionCall, |
| 17 | ToolCall, |
| 18 | ) |
| 19 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 20 | ChatCompletionRequest, |
| 21 | ) |
| 22 | from vllm.entrypoints.openai.responses.protocol import ResponsesRequest |
| 23 | from vllm.logger import init_logger |
| 24 | from vllm.tokenizers import TokenizerLike |
| 25 | from vllm.tool_parsers.abstract_tool_parser import ( |
| 26 | Tool, |
| 27 | ToolParser, |
| 28 | ) |
| 29 | from vllm.tool_parsers.utils import extract_intermediate_diff, is_complete_json |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | |
| 34 | class Internlm2ToolParser(ToolParser): |
| 35 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 36 | super().__init__(tokenizer, tools) |
| 37 | self.position = 0 |
| 38 | |
| 39 | def adjust_request( |
| 40 | self, request: ChatCompletionRequest | ResponsesRequest |
| 41 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 42 | request = super().adjust_request(request) |
| 43 | if request.tools and request.tool_choice != "none": |
| 44 | # do not skip special tokens because internlm use the special |
| 45 | # tokens to indicate the start and end of the tool calls |
| 46 | # information. |
| 47 | request.skip_special_tokens = False |
| 48 | return request |
| 49 | |
| 50 | def get_arguments(self, obj): |
| 51 | if "parameters" in obj: |
| 52 | return obj.get("parameters") |
| 53 | elif "arguments" in obj: |
| 54 | return obj.get("arguments") |
| 55 | return None |
| 56 | |
| 57 | def extract_tool_calls_streaming( |
| 58 | self, |
| 59 | previous_text: str, |
| 60 | current_text: str, |
| 61 | delta_text: str, |
| 62 | previous_token_ids: Sequence[int], |
| 63 | current_token_ids: Sequence[int], |
| 64 | delta_token_ids: Sequence[int], |
| 65 | request: ChatCompletionRequest, |
| 66 | ) -> DeltaMessage | None: |
| 67 | if "<|action_start|>" not in current_text: |
| 68 | self.position = len(current_text) |
| 69 | return DeltaMessage(content=delta_text) |
| 70 | # if the tool call is sent, return an empty delta message |
| 71 | # to make sure the finish_reason will be sent correctly. |
| 72 | if self.current_tool_id > 0: |
| 73 | return DeltaMessage(content="") |
| 74 | |
| 75 | last_pos = self.position |
| 76 | if "<|action_start|><|plugin|>" not in current_text[last_pos:]: |
| 77 | return None |
| 78 | |
| 79 | new_delta = current_text[last_pos:] |
| 80 | text, action = new_delta.split("<|action_start|><|plugin|>") |
| 81 | |
| 82 | if len(text) > 0: |
| 83 | self.position = self.position + len(text) |
| 84 | return DeltaMessage(content=text) |
| 85 | |
| 86 | action = action.strip() |
| 87 | action = action.split("<|action_end|>".strip())[0] |
| 88 | |
| 89 | # bit mask flags for partial JSON parsing. If the name hasn't been |
| 90 | # sent yet, don't allow sending |
| 91 | # an incomplete string since OpenAI only ever (as far as I have |
| 92 | # seen) allows sending the entire tool/ function name at once. |
| 93 | flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR |
| 94 | |
| 95 | try: |
| 96 | parsable_arr = action |
| 97 | |
| 98 | # tool calls are generated in an object in internlm2 |
| 99 | # it's not support parallel tool calls |
| 100 | try: |
| 101 | tool_call_arr: dict = partial_json_parser.loads(parsable_arr, flags) |
| 102 | except partial_json_parser.core.exceptions.MalformedJSON: |
| 103 | logger.debug("not enough tokens to parse into JSON yet") |
| 104 | return None |
| 105 | |
| 106 | # if the current tool name hasn't been sent, send if available |
| 107 | # - otherwise send nothing |
| 108 | if not self.current_tool_name_sent: |
| 109 | function_name = tool_call_arr.get("name") |
| 110 | if function_name: |
| 111 | self.current_tool_id = self.current_tool_id + 1 |
| 112 | delta = DeltaMessage( |
| 113 | tool_calls=[ |
| 114 | DeltaToolCall( |
| 115 | index=self.current_tool_id, |
| 116 | type="function", |
| 117 | id=make_tool_call_id(), |
| 118 | function=DeltaFunctionCall( |
| 119 | name=function_name |
| 120 | ).model_dump(exclude_none=True), |
| 121 | ) |
| 122 | ] |
| 123 | ) |
| 124 | self.current_tool_name_sent = True |
| 125 | self.streamed_args_for_tool.append("") |
| 126 | else: |
| 127 | delta = None |
| 128 | # now we know we're on the same tool call and we're streaming |
| 129 | # arguments |
| 130 | else: |
| 131 | prev_arguments = self.get_arguments( |
| 132 | self.prev_tool_call_arr[self.current_tool_id] |
| 133 | ) |
| 134 | cur_arguments = self.get_arguments(tool_call_arr) |
| 135 | |
| 136 | # not arguments generated |
| 137 | if not cur_arguments and not prev_arguments: |
| 138 | delta = None |
| 139 | # will never happen |
| 140 | elif not cur_arguments and prev_arguments: |
| 141 | logger.error( |
| 142 | "INVARIANT - impossible to have arguments reset mid-arguments" |
| 143 | ) |
| 144 | delta = None |
| 145 | # first time to get parameters |
| 146 | elif cur_arguments and not prev_arguments: |
| 147 | cur_arguments_json = json.dumps(cur_arguments, ensure_ascii=False) |
| 148 | |
| 149 | match_start = cur_arguments_json.find(delta_text) |
| 150 | if match_start != -1: |
| 151 | arguments_delta = cur_arguments_json[ |
| 152 | : match_start + len(delta_text) |
| 153 | ] |
| 154 | elif is_complete_json(parsable_arr): |
| 155 | # Complete in this delta: send whole, don't drop. |
| 156 | arguments_delta = cur_arguments_json |
| 157 | else: |
| 158 | # Still partial: wait for more text. |
| 159 | return None |
| 160 | delta = DeltaMessage( |
| 161 | tool_calls=[ |
| 162 | DeltaToolCall( |
| 163 | index=self.current_tool_id, |
| 164 | function=DeltaFunctionCall( |
| 165 | arguments=arguments_delta |
| 166 | ).model_dump(exclude_none=True), |
| 167 | ) |
| 168 | ] |
| 169 | ) |
| 170 | self.streamed_args_for_tool[self.current_tool_id] += arguments_delta |
| 171 | # both prev and cur parameters, send the increase parameters |
| 172 | elif cur_arguments and prev_arguments: |
| 173 | cur_args_json = json.dumps(cur_arguments, ensure_ascii=False) |
| 174 | prev_args_json = json.dumps(prev_arguments, ensure_ascii=False) |
| 175 | |
| 176 | argument_diff = extract_intermediate_diff( |
| 177 | cur_args_json, prev_args_json |
| 178 | ) |
| 179 | |
| 180 | delta = DeltaMessage( |
| 181 | tool_calls=[ |
| 182 | DeltaToolCall( |
| 183 | index=self.current_tool_id, |
| 184 | function=DeltaFunctionCall( |
| 185 | arguments=argument_diff |
| 186 | ).model_dump(exclude_none=True), |
| 187 | ) |
| 188 | ] |
| 189 | ) |
| 190 | self.streamed_args_for_tool[self.current_tool_id] += argument_diff |
| 191 | |
| 192 | # check to see if the name is defined and has been sent. if so, |
| 193 | # stream the name - otherwise keep waiting |
| 194 | # finish by setting old and returning None as base case |
| 195 | tool_call_arr["arguments"] = self.get_arguments(tool_call_arr) |
| 196 | self.prev_tool_call_arr = [tool_call_arr] |
| 197 | return delta |
| 198 | except Exception: |
| 199 | logger.exception("Error trying to handle streaming tool call.") |
| 200 | logger.debug( |
| 201 | "Skipping chunk as a result of tool streaming extraction error" |
| 202 | ) |
| 203 | return None |
| 204 | |
| 205 | def extract_tool_calls( |
| 206 | self, |
| 207 | model_output: str, |
| 208 | request: ChatCompletionRequest, |
| 209 | ) -> ExtractedToolCallInformation: |
| 210 | text = model_output |
| 211 | tools = self.tools |
| 212 | if "<|action_start|><|plugin|>" in text: |
| 213 | text, action = text.split("<|action_start|><|plugin|>") |
| 214 | action = action.split("<|action_end|>".strip())[0] |
| 215 | action = action[action.find("{") :] |
| 216 | action_dict = json.loads(action) |
| 217 | name, parameters = ( |
| 218 | action_dict["name"], |
| 219 | json.dumps( |
| 220 | action_dict.get("parameters", action_dict.get("arguments", {})), |
| 221 | ensure_ascii=False, |
| 222 | ), |
| 223 | ) |
| 224 | |
| 225 | if not tools or name not in [t.function.name for t in tools]: |
| 226 | ExtractedToolCallInformation( |
| 227 | tools_called=False, tool_calls=[], content=text |
| 228 | ) |
| 229 | |
| 230 | tool_calls = [ |
| 231 | ToolCall(function=FunctionCall(name=name, arguments=parameters)) |
| 232 | ] |
| 233 | return ExtractedToolCallInformation( |
| 234 | tools_called=True, |
| 235 | tool_calls=tool_calls, |
| 236 | content=text if len(text) > 0 else None, |
| 237 | ) |
| 238 | |
| 239 | return ExtractedToolCallInformation( |
| 240 | tools_called=False, tool_calls=[], content=text |
| 241 | ) |