93269bb43e · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,13 +1,15 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -# code modified from deepseekv3_tool_parser.py

3

4 +import ast

5 +import json

6 from collections.abc import Sequence

6 -from typing import Union

7 +from typing import Any, Optional, Union

8

9 import regex as re

10

11 from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

12 + ChatCompletionToolsParam,

13 DeltaFunctionCall, DeltaMessage,

14 DeltaToolCall,

15 ExtractedToolCallInformation,

@@ -34,36 +36,13 @@ class Glm4MoeModelToolParser(ToolParser):

36

37 self.tool_calls_start_token = self.tool_call_start_token

38

37 - # Updated regex for the XML-based format

38 - self.tool_call_regex = re.compile(

39 - r"<tool_call>\s*"

40 - r"(?P<function_name>[^\n<]+)\s*" # 函数名（到换行或 <）

41 - r"(?P<arguments>(?:\s*<arg_key>[^<]+</arg_key>\s*"

42 - r"<arg_value>[^<]*</arg_value>\s*)*)\s*"

43 - r"</tool_call>",

44 - re.DOTALL,

45 - )

46 -

47 - # Regex for parsing individual arguments

48 - self.arg_regex = re.compile(

49 - r"<arg_key>(?P<key>[^<]+)</arg_key>\s*<arg_value>(?P<value>[^<]*)</arg_value>",

50 - re.DOTALL,

51 - )

52 -

53 - # Streaming regex

54 - self.stream_tool_call_portion_regex = re.compile(

55 - r"(?P<function_name>[^\n<]+)\s*"

56 - r"(?P<arguments>(?:\s*<arg_key>[^<]+</arg_key>\s*"

57 - r"<arg_value>[^<]*</arg_value>\s*)*)",

58 - re.DOTALL,

59 - )

60 -

61 - # For streaming, we also need a regex to match just the function name

62 - self.stream_tool_call_name_regex = re.compile(

63 - r"(?P<function_name>[^\n<]+)",

64 - re.DOTALL,

65 - )

66 -

39 + self.func_call_regex = re.compile(r"<tool_call>.*?</tool_call>",

40 + re.DOTALL)

41 + self.func_detail_regex = re.compile(

42 + r"<tool_call>([^\n]*)\n(.*)</tool_call>", re.DOTALL)

43 + self.func_arg_regex = re.compile(

44 + r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>",

45 + re.DOTALL)

46 if not self.model_tokenizer:

47 raise ValueError(

48 "The model tokenizer must be passed to the ToolParser "

@@ -72,20 +51,7 @@ class Glm4MoeModelToolParser(ToolParser):

51 self.tool_call_start_token_id = self.vocab.get(

52 self.tool_call_start_token)

53 self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)

75 -

76 - def _parse_arguments(self, args_text: str) -> str:

77 - """Parse XML-based arguments into JSON format."""

78 - if not args_text or not args_text.strip():

79 - return "{}"

80 -

81 - args_dict = {}

82 - matches = self.arg_regex.findall(args_text)

83 -

84 - for key, value in matches:

85 - args_dict[key.strip()] = value.strip()

86 -

87 - import json

88 - return json.dumps(args_dict, ensure_ascii=False)

54 + self._buffer = ""

55

56 def extract_tool_calls(

57 self,

@@ -93,52 +59,67 @@ class Glm4MoeModelToolParser(ToolParser):

59 request: ChatCompletionRequest,

60 ) -> ExtractedToolCallInformation:

61

96 - # sanity check; avoid unnecessary processing

97 - if self.tool_calls_start_token not in model_output:

98 - return ExtractedToolCallInformation(tools_called=False,

99 - tool_calls=[],

100 - content=model_output)

101 -

62 + def _is_string_type(

63 + tool_name: str, arg_name: str,

64 + tools: Optional[list[ChatCompletionToolsParam]]) -> bool:

65 + if tools is None:

66 + return False

67 + for tool in tools:

68 + if tool.function.name == tool_name:

69 + if tool.function.parameters is None:

70 + return False

71 + arg_type = tool.function.parameters.get(

72 + "properties", {}).get(arg_name, {}).get("type", None)

73 + return arg_type == "string"

74 + logger.warning("No tool named '%s'.", tool_name)

75 + return False

76 +

77 + def _deserialize(value: str) -> Any:

78 + try:

79 + return json.loads(value)

80 + except Exception:

81 + pass

82 +

83 + try:

84 + return ast.literal_eval(value)

85 + except Exception:

86 + pass

87 + return value

88 +

89 + matched_tool_calls = self.func_call_regex.findall(model_output)

90 + logger.debug("model_output: %s", model_output)

91 try:

103 - # Find all tool calls in the output

104 - function_call_matches = self.tool_call_regex.findall(model_output)

105 -

106 - logger.debug("function_call_matches: %s", function_call_matches)

107 -

108 - if not function_call_matches:

109 - return ExtractedToolCallInformation(

110 - tools_called=False,

111 - tool_calls=[],

112 - content=model_output,

113 - )

114 -

92 tool_calls = []

116 - for i, match in enumerate(function_call_matches):

117 - function_name, function_args_xml = match

118 - function_name = function_name.strip()

119 -

120 - # Parse XML arguments to JSON

121 - function_args_json = self._parse_arguments(function_args_xml)

122 -

93 + for match in matched_tool_calls:

94 + tc_detail = self.func_detail_regex.search(match)

95 + tc_name = tc_detail.group(1)

96 + tc_args = tc_detail.group(2)

97 + pairs = self.func_arg_regex.findall(tc_args)

98 + arg_dct = {}

99 + for key, value in pairs:

100 + arg_key = key.strip()

101 + arg_val = value.strip()

102 + if not _is_string_type(tc_name, arg_key, request.tools):

103 + arg_val = _deserialize(arg_val)

104 + logger.debug("arg_key = %s, arg_val = %s", arg_key,

105 + arg_val)

106 + arg_dct[arg_key] = arg_val

107 tool_calls.append(

124 - ToolCall(

125 - id=f"call_{i}",

126 - type='function',

127 - function=FunctionCall(name=function_name,

128 - arguments=function_args_json),

129 - ))

130 -

131 - # Extract content before the first tool call

132 - content = model_output[:model_output.find(self.

133 - tool_calls_start_token)]

134 - return ExtractedToolCallInformation(

135 - tools_called=bool(tool_calls),

136 - tool_calls=tool_calls,

137 - content=content.strip() if content.strip() else None,

138 - )

139 -

108 + ToolCall(type="function",

109 + function=FunctionCall(

110 + name=tc_name, arguments=json.dumps(arg_dct))))

111 except Exception:

141 - logger.exception("Error in extracting tool call from response.")

112 + logger.exception("Failed to extract tool call spec")

113 + return ExtractedToolCallInformation(tools_called=False,

114 + tool_calls=[],

115 + content=model_output)

116 + else:

117 + if len(tool_calls) > 0:

118 + content = model_output[:model_output.

119 + find(self.tool_calls_start_token)]

120 + return ExtractedToolCallInformation(tools_called=True,

121 + tool_calls=tool_calls,

122 + content=content)

123 return ExtractedToolCallInformation(tools_called=False,

124 tool_calls=[],

125 content=model_output)

@@ -153,250 +134,52 @@ class Glm4MoeModelToolParser(ToolParser):

134 delta_token_ids: Sequence[int],

135 request: ChatCompletionRequest,

136 ) -> Union[DeltaMessage, None]:

156 -

157 - logger.debug("delta_text: %s", delta_text)

158 - logger.debug("delta_token_ids: %s", delta_token_ids)

159 - # check to see if we should be streaming a tool call - is there a

160 - if self.tool_call_start_token_id not in current_token_ids:

161 - logger.debug("No tool call tokens found!")

162 - return DeltaMessage(content=delta_text)

163 - delta_text = delta_text.replace(self.tool_calls_start_token,

164 - "").replace(self.tool_call_end_token,

165 - "")

166 - try:

167 -

168 - # figure out where we are in the parsing by counting tool call

169 - # start & end tags

170 - prev_tool_start_count = previous_token_ids.count(

171 - self.tool_call_start_token_id)

172 - prev_tool_end_count = previous_token_ids.count(

173 - self.tool_call_end_token_id)

174 - cur_tool_start_count = current_token_ids.count(

175 - self.tool_call_start_token_id)

176 - cur_tool_end_count = current_token_ids.count(

177 - self.tool_call_end_token_id)

178 - tool_call_portion = None

179 - text_portion = None

180 -

181 - # case: if we're generating text, OR rounding out a tool call

182 - if (cur_tool_start_count == cur_tool_end_count

183 - and prev_tool_end_count == cur_tool_end_count

184 - and self.tool_call_end_token not in delta_text):

185 - logger.debug("Generating text content! skipping tool parsing.")

186 - return DeltaMessage(content=delta_text)

187 -

188 - if self.tool_call_end_token in delta_text:

189 - logger.debug("tool_call_end_token in delta_text")

190 - full_text = current_text + delta_text

191 - tool_call_portion = full_text.split(

192 - self.tool_call_start_token)[-1].split(

193 - self.tool_call_end_token)[0].rstrip()

194 - delta_text = delta_text.split(

195 - self.tool_call_end_token)[0].rstrip()

196 - text_portion = delta_text.split(

197 - self.tool_call_end_token)[-1].lstrip()

198 -

199 - # case -- we're starting a new tool call

200 - if (cur_tool_start_count > cur_tool_end_count

201 - and cur_tool_start_count > prev_tool_start_count):

202 - if len(delta_token_ids) > 1:

203 - tool_call_portion = current_text.split(

204 - self.tool_call_start_token)[-1]

205 - else:

206 - tool_call_portion = None

207 - delta = None

208 -

209 - text_portion = None

210 -

211 - # set cursors and state appropriately

212 - self.current_tool_id += 1

213 - self.current_tool_name_sent = False

214 - self.streamed_args_for_tool.append("")

215 - logger.debug("Starting on a new tool %s", self.current_tool_id)

216 -

217 - # case -- we're updating an existing tool call

218 - elif (cur_tool_start_count > cur_tool_end_count

219 - and cur_tool_start_count == prev_tool_start_count):

220 -

221 - # get the portion of the text that's the tool call

222 - tool_call_portion = current_text.split(

223 - self.tool_call_start_token)[-1]

224 - text_portion = None

225 -

226 - # case -- the current tool call is being closed.

227 - elif (cur_tool_start_count == cur_tool_end_count

228 - and cur_tool_end_count >= prev_tool_end_count):

229 - if self.prev_tool_call_arr is None or len(

230 - self.prev_tool_call_arr) == 0:

231 - logger.debug(

232 - "attempting to close tool call, but no tool call")

233 - return None

234 - diff = self.prev_tool_call_arr[self.current_tool_id].get(

235 - "arguments")

236 - if diff:

237 - diff = (diff.encode("utf-8").decode("unicode_escape")

238 - if diff is str else diff)

239 - if '"}' not in delta_text:

240 - return None

241 - end_loc = delta_text.rindex('"}')

242 - diff = delta_text[:end_loc] + '"}'

243 - logger.debug(

244 - "Finishing tool and found diff that had not "

245 - "been streamed yet: %s",

246 - diff,

247 - )

248 - self.streamed_args_for_tool[self.current_tool_id] += diff

249 - return DeltaMessage(tool_calls=[

250 - DeltaToolCall(

251 - index=self.current_tool_id,

252 - function=DeltaFunctionCall(

253 - arguments=diff).model_dump(exclude_none=True),

254 - )

255 - ])

256 -

257 - # case -- otherwise we're just generating text

258 - else:

259 - text = delta_text.replace(self.tool_call_start_token, "")

260 - text = text.replace(self.tool_call_end_token, "")

261 - delta = DeltaMessage(tool_calls=[], content=text)

262 - return delta

263 -

264 - current_tool_call = dict()

265 - if tool_call_portion:

266 - current_tool_call_matches = (

267 - self.stream_tool_call_portion_regex.match(

268 - tool_call_portion))

269 - if current_tool_call_matches:

270 - tool_id, tool_args = (current_tool_call_matches.groups())

271 - tool_name = tool_id.split('.')[1].split(':')[0]

272 - current_tool_call['id'] = tool_id

273 - current_tool_call["name"] = tool_name

274 - current_tool_call["arguments"] = tool_args

275 - else:

276 - current_tool_call_name_matches = (

277 - self.stream_tool_call_name_regex.match(

278 - tool_call_portion))

279 - if current_tool_call_name_matches:

280 - tool_id_str, = current_tool_call_name_matches.groups()

281 - tool_name = tool_id_str.split('.')[1].split(':')[0]

282 - current_tool_call['id'] = tool_id_str

283 - current_tool_call["name"] = tool_name

284 - current_tool_call["arguments"] = ""

285 - else:

286 - logger.debug("Not enough token")

287 - return None

288 -

289 - # case - we haven't sent the tool name yet. If it's available, send

290 - # it. otherwise, wait until it's available.

291 - if not self.current_tool_name_sent:

292 - if current_tool_call is None:

293 - return None

294 - function_name: Union[str, None] = current_tool_call.get("name")

295 - tool_id = current_tool_call.get("id")

296 - if function_name:

297 - self.current_tool_name_sent = True

298 - return DeltaMessage(tool_calls=[

299 - DeltaToolCall(

300 - index=self.current_tool_id,

301 - type="function",

302 - id=tool_id,

303 - function=DeltaFunctionCall(

304 - name=function_name).model_dump(

305 - exclude_none=True),

306 - )

307 - ])

308 - else:

309 - return None

310 -

311 - # case -- otherwise, send the tool call delta

312 -

313 - # if the tool call portion is None, send the delta as text

314 - if tool_call_portion is None:

315 - # if there's text but not tool calls, send that -

316 - # otherwise None to skip chunk

317 - delta = (DeltaMessage(

318 - content=delta_text) if text_portion is not None else None)

319 - return delta

320 -

321 - # now, the nitty-gritty of tool calls

322 - # now we have the portion to parse as tool call.

323 -

324 - logger.debug("Trying to parse current tool call with ID %s",

325 - self.current_tool_id)

326 -

327 - # if we're starting a new tool call, push an empty object in as

328 - # a placeholder for the arguments

329 - if len(self.prev_tool_call_arr) <= self.current_tool_id:

137 + self._buffer += delta_text

138 + cur_text = self._buffer

139 + start_idx = cur_text.find(self.tool_call_start_token)

140 + if start_idx == -1:

141 + self._buffer = ""

142 + if self.current_tool_id > 0:

143 + cur_text = ""

144 + return DeltaMessage(content=cur_text)

145 + logger.debug("cur_text = %s", cur_text)

146 + end_idx = cur_text.find(self.tool_call_end_token)

147 + if end_idx != -1:

148 + if self.current_tool_id == -1:

149 + self.current_tool_id = 0

150 + self.prev_tool_call_arr = []

151 + self.streamed_args_for_tool = []

152 + while len(self.prev_tool_call_arr) <= self.current_tool_id:

153 self.prev_tool_call_arr.append({})

154 + while len(self.streamed_args_for_tool) <= self.current_tool_id:

155 + self.streamed_args_for_tool.append("")

156

332 - # main logic for tool parsing here - compare prev. partially-parsed

333 - # JSON to the current partially-parsed JSON

334 - prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(

335 - "arguments")

336 - cur_arguments = current_tool_call.get("arguments")

337 -

338 - logger.debug("diffing old arguments: %s", prev_arguments)

339 - logger.debug("against new ones: %s", cur_arguments)

340 -

341 - # case -- no arguments have been created yet. skip sending a delta.

342 - if not cur_arguments and not prev_arguments:

343 - logger.debug("Skipping text %s - no arguments", delta_text)

344 - delta = None

345 -

346 - # case -- prev arguments are defined, but non are now.

347 - # probably impossible, but not a fatal error - just keep going

348 - elif not cur_arguments and prev_arguments:

349 - logger.error("should be impossible to have arguments reset "

350 - "mid-call. skipping streaming anything.")

351 - delta = None

352 -

353 - # case -- we now have the first info about arguments available from

354 - # autocompleting the JSON

355 - elif cur_arguments and not prev_arguments:

356 -

357 - delta = DeltaMessage(tool_calls=[

358 - DeltaToolCall(

359 - index=self.current_tool_id,

360 - function=DeltaFunctionCall(

361 - arguments=cur_arguments).model_dump(

362 - exclude_none=True),

363 - )

157 + extracted_tool_calls = self.extract_tool_calls(

158 + cur_text[:end_idx + len(self.tool_call_end_token)], request)

159 +

160 + if len(extracted_tool_calls.tool_calls) == 0:

161 + logger.warning("Failed to extract any tool calls.")

162 + return None

163 + tool_call = extracted_tool_calls.tool_calls[0]

164 + self.prev_tool_call_arr[self.current_tool_id] = {

165 + "name": tool_call.function.name,

166 + "arguments": json.loads(tool_call.function.arguments)

167 + }

168 + self.streamed_args_for_tool[

169 + self.current_tool_id] = tool_call.function.arguments

170 + delta = DeltaMessage(

171 + content=extracted_tool_calls.content,

172 + tool_calls=[

173 + DeltaToolCall(index=self.current_tool_id,

174 + id=tool_call.id,

175 + type=tool_call.type,

176 + function=DeltaFunctionCall(

177 + name=tool_call.function.name,

178 + arguments=tool_call.function.arguments))

179 ])

365 - self.streamed_args_for_tool[

366 - self.current_tool_id] = cur_arguments

367 -

368 - # last case -- we have an update to existing arguments.

369 - elif cur_arguments and prev_arguments:

370 - if (isinstance(delta_text, str)

371 - and cur_arguments != prev_arguments

372 - and len(cur_arguments) > len(prev_arguments)

373 - and cur_arguments.startswith(prev_arguments)):

374 - delta_arguments = cur_arguments[len(prev_arguments):]

375 - logger.debug("got diff %s", delta_text)

376 -

377 - delta = DeltaMessage(tool_calls=[

378 - DeltaToolCall(

379 - index=self.current_tool_id,

380 - function=DeltaFunctionCall(

381 - arguments=delta_arguments).model_dump(

382 - exclude_none=True),

383 - )

384 - ])

385 - self.streamed_args_for_tool[

386 - self.current_tool_id] = cur_arguments

387 - else:

388 - delta = None

389 -

390 - # handle saving the state for the current tool into

391 - # the "prev" list for use in diffing for the next iteration

392 - if self.current_tool_id == len(self.prev_tool_call_arr) - 1:

393 - self.prev_tool_call_arr[

394 - self.current_tool_id] = current_tool_call

395 - else:

396 - self.prev_tool_call_arr.append(current_tool_call)

397 -

180 + self.current_tool_id += 1

181 + self._buffer = cur_text[end_idx + len(self.tool_call_end_token):]

182 return delta

183

400 - except Exception:

401 - logger.exception("Error trying to handle streaming tool call.")

402 - return None # do not stream a delta. skip this token ID.

184 + self._buffer = cur_text[start_idx:]

185 + return DeltaMessage(content=cur_text[:start_idx])

Fix GLM tool parser (#21668)