| 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.generate.base.protocol import ( |
| 10 | DeltaFunctionCall, |
| 11 | DeltaMessage, |
| 12 | DeltaToolCall, |
| 13 | ExtractedToolCallInformation, |
| 14 | FunctionCall, |
| 15 | ToolCall, |
| 16 | ) |
| 17 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 18 | ChatCompletionRequest, |
| 19 | ) |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.tokenizers import TokenizerLike |
| 22 | from vllm.tool_parsers.abstract_tool_parser import ( |
| 23 | Tool, |
| 24 | ToolParser, |
| 25 | ) |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | |
| 30 | class Ernie45ToolParser(ToolParser): |
| 31 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 32 | """Ernie thinking model format: |
| 33 | abc\n</think>\n\n\n<tool_call>\ndef\n</tool_call>\n |
| 34 | """ |
| 35 | super().__init__(tokenizer, tools) |
| 36 | self.prev_tool_call_arr: list[dict] = [] |
| 37 | self.current_tool_id = -1 |
| 38 | self.streamed_args_for_tool: list[str] = [] |
| 39 | self.think_end_token = "</think>" |
| 40 | self.response_start_token: str = "<response>" |
| 41 | self.response_end_token: str = "</response>" |
| 42 | self.tool_call_start_token = "<tool_call>" |
| 43 | self.tool_call_end_token = "</tool_call>" |
| 44 | self.tool_calls_start_token = self.tool_call_start_token |
| 45 | self.newline_token: str = "<0x0A>" |
| 46 | |
| 47 | self.tool_call_regex = re.compile( |
| 48 | r"<tool_call>\s*(?P<json>\{.*?\})\s*</tool_call>", re.DOTALL |
| 49 | ) |
| 50 | |
| 51 | if not self.model_tokenizer: |
| 52 | raise ValueError( |
| 53 | "The model tokenizer must be passed to the ToolParser " |
| 54 | "constructor during construction." |
| 55 | ) |
| 56 | |
| 57 | self.think_end_token_id = self.vocab.get(self.think_end_token) |
| 58 | self.response_start_token_id = self.vocab.get(self.response_start_token) |
| 59 | self.response_end_token_id = self.vocab.get(self.response_end_token) |
| 60 | self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token) |
| 61 | self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token) |
| 62 | self.newline_token_id = self.vocab.get(self.newline_token) |
| 63 | self.parser_token_ids = [ |
| 64 | self.think_end_token_id, |
| 65 | self.response_start_token_id, |
| 66 | self.response_end_token_id, |
| 67 | ] |
| 68 | |
| 69 | self._buffer = "" |
| 70 | |
| 71 | def extract_tool_calls( |
| 72 | self, |
| 73 | model_output: str, |
| 74 | request: ChatCompletionRequest, |
| 75 | ) -> ExtractedToolCallInformation: |
| 76 | # sanity check; avoid unnecessary processing |
| 77 | if self.tool_calls_start_token not in model_output: |
| 78 | return ExtractedToolCallInformation( |
| 79 | tools_called=False, tool_calls=[], content=model_output |
| 80 | ) |
| 81 | |
| 82 | else: |
| 83 | try: |
| 84 | tool_call_json_list = self.tool_call_regex.findall(model_output) |
| 85 | |
| 86 | tool_calls = [] |
| 87 | for tool_call_json in tool_call_json_list: |
| 88 | tool_call_dict = json.loads(tool_call_json) |
| 89 | args_str = json.dumps( |
| 90 | tool_call_dict.get("arguments", {}), ensure_ascii=False |
| 91 | ) |
| 92 | tool_calls.append( |
| 93 | ToolCall( |
| 94 | type="function", |
| 95 | function=FunctionCall( |
| 96 | name=tool_call_dict.get("name", ""), |
| 97 | arguments=args_str, |
| 98 | ), |
| 99 | ) |
| 100 | ) |
| 101 | |
| 102 | content = model_output[ |
| 103 | : model_output.find(self.tool_calls_start_token) |
| 104 | ].rstrip("\n") |
| 105 | return ExtractedToolCallInformation( |
| 106 | tools_called=True, |
| 107 | tool_calls=tool_calls, |
| 108 | content=content if content else None, |
| 109 | ) |
| 110 | |
| 111 | except Exception: |
| 112 | logger.exception("Error in extracting tool call from response.") |
| 113 | return ExtractedToolCallInformation( |
| 114 | tools_called=False, tool_calls=[], content=model_output |
| 115 | ) |
| 116 | |
| 117 | def extract_tool_calls_streaming( |
| 118 | self, |
| 119 | previous_text: str, |
| 120 | current_text: str, |
| 121 | delta_text: str, |
| 122 | previous_token_ids: Sequence[int], |
| 123 | current_token_ids: Sequence[int], |
| 124 | delta_token_ids: Sequence[int], |
| 125 | request: ChatCompletionRequest, |
| 126 | ) -> DeltaMessage | None: |
| 127 | self._buffer += delta_text |
| 128 | cur_text = self._buffer |
| 129 | start_idx = cur_text.find(self.tool_call_start_token) |
| 130 | if start_idx == -1: |
| 131 | self._buffer = "" |
| 132 | # At least one toolcall has been completed |
| 133 | if self.current_tool_id > 0: |
| 134 | cur_text = "" |
| 135 | if self.current_tool_id == -1 and all( |
| 136 | token_id == self.newline_token_id for token_id in previous_token_ids |
| 137 | ): |
| 138 | cur_text = cur_text.strip("\n") |
| 139 | |
| 140 | # handle <response> </response> when tool_call is not triggered |
| 141 | # cur_text === delta_text |
| 142 | content = cur_text |
| 143 | if self.response_start_token_id in delta_token_ids: |
| 144 | content = content.lstrip("\n") |
| 145 | response_start_idx = content.find(self.response_start_token) |
| 146 | content = content[response_start_idx + len(self.response_start_token) :] |
| 147 | # if have </response>, remove it |
| 148 | response_end_idx = content.rfind(self.response_end_token) |
| 149 | if response_end_idx != -1: |
| 150 | content = content[:response_end_idx] |
| 151 | elif self.response_end_token_id in delta_token_ids: |
| 152 | response_end_idx = content.rfind(self.response_end_token) |
| 153 | content = content[:response_end_idx] |
| 154 | # remove \n after </think> or <response> or </response> |
| 155 | if ( |
| 156 | len(previous_token_ids) > 0 |
| 157 | and previous_token_ids[-1] in self.parser_token_ids |
| 158 | ) and ( |
| 159 | len(delta_token_ids) > 0 and delta_token_ids[0] == self.newline_token_id |
| 160 | ): |
| 161 | content = content.lstrip("\n") |
| 162 | |
| 163 | return DeltaMessage(content=content if content else None) |
| 164 | logger.debug("cur_text = %s", cur_text) |
| 165 | end_idx = cur_text.find(self.tool_call_end_token) |
| 166 | if end_idx != -1: |
| 167 | if self.current_tool_id == -1: |
| 168 | self.current_tool_id = 0 |
| 169 | self.prev_tool_call_arr = [] |
| 170 | self.streamed_args_for_tool = [] |
| 171 | while len(self.prev_tool_call_arr) <= self.current_tool_id: |
| 172 | self.prev_tool_call_arr.append({}) |
| 173 | while len(self.streamed_args_for_tool) <= self.current_tool_id: |
| 174 | self.streamed_args_for_tool.append("") |
| 175 | |
| 176 | extracted_tool_calls = self.extract_tool_calls( |
| 177 | cur_text[: end_idx + len(self.tool_call_end_token)], request |
| 178 | ) |
| 179 | |
| 180 | if len(extracted_tool_calls.tool_calls) == 0: |
| 181 | logger.warning("Failed to extract any tool calls.") |
| 182 | return None |
| 183 | tool_call = extracted_tool_calls.tool_calls[0] |
| 184 | self.prev_tool_call_arr[self.current_tool_id] = { |
| 185 | "name": tool_call.function.name, |
| 186 | "arguments": json.loads(tool_call.function.arguments), |
| 187 | } |
| 188 | self.streamed_args_for_tool[self.current_tool_id] = ( |
| 189 | tool_call.function.arguments |
| 190 | ) |
| 191 | delta = DeltaMessage( |
| 192 | content=extracted_tool_calls.content, |
| 193 | tool_calls=[ |
| 194 | DeltaToolCall( |
| 195 | index=self.current_tool_id, |
| 196 | id=tool_call.id, |
| 197 | type=tool_call.type, |
| 198 | function=DeltaFunctionCall( |
| 199 | name=tool_call.function.name, |
| 200 | arguments=tool_call.function.arguments, |
| 201 | ), |
| 202 | ) |
| 203 | ], |
| 204 | ) |
| 205 | self.current_tool_id += 1 |
| 206 | self._buffer = cur_text[end_idx + len(self.tool_call_end_token) :] |
| 207 | return delta |
| 208 | |
| 209 | self._buffer = cur_text[start_idx:] |
| 210 | content = cur_text[:start_idx].rstrip("\n") |
| 211 | return DeltaMessage(content=content if content else None) |