main
py 551 lines 20.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tool call parser for Apertus models.
4
5 Extracts tool calls from the format:
6 <|tools_prefix|>[{"function_name": {"arg1": "value1", ...}}, ...]<|tools_suffix|>
7
8 Used when --enable-auto-tool-choice --tool-call-parser apertus are set.
9 """
10
11 import json
12 from collections.abc import Sequence
13
14 import regex as re
15 from partial_json_parser.core.options import Allow
16
17 from vllm.entrypoints.chat_utils import make_tool_call_id
18 from vllm.entrypoints.generate.base.protocol import (
19 DeltaFunctionCall,
20 DeltaMessage,
21 DeltaToolCall,
22 ExtractedToolCallInformation,
23 FunctionCall,
24 ToolCall,
25 )
26 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
27 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
28 from vllm.logger import init_logger
29 from vllm.tokenizers import TokenizerLike
30 from vllm.tool_parsers.abstract_tool_parser import (
31 Tool,
32 ToolParser,
33 )
34 from vllm.tool_parsers.utils import (
35 find_common_prefix,
36 partial_json_loads,
37 )
38
39 logger = init_logger(__name__)
40
41 # Apertus special tokens for tool calls
42 TOOL_CALLS_PREFIX = "<|tools_prefix|>"
43 TOOL_CALLS_SUFFIX = "<|tools_suffix|>"
44
45
46 class ApertusToolParser(ToolParser):
47 """Tool call parser for Apertus models.
48
49 Handles the extraction of tool calls from text in both non-streaming
50 (complete string) and streaming (chunked token) environments.
51
52 The expected Apertus function call format is a JSON array of single-key dictionaries
53 sandwiched between special tokens:
54 `<|tools_prefix|>[{"function_name": {"arg1": "value1"}}, ...]<|tools_suffix|>`
55
56 Examples:
57 >>> tokenizer = ... # Mock tokenizer
58 >>> parser = ApertusToolParser(tokenizer)
59 >>> output = 'I will check. <|tools_prefix|>[{"get_weather": '\
60 '{"city": "Paris"}}]<|tools_suffix|>'
61 >>> request = ChatCompletionRequest(...)
62 >>> info = parser.extract_tool_calls(output, request)
63 >>> info.content
64 "I will check."
65 >>> info.tool_calls[0].function.name
66 "get_weather"
67 >>> info.tool_calls[0].function.arguments
68 '{"city": "Paris"}'
69
70 """
71
72 def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
73 """Initializes the ApertusToolParser.
74
75 Args:
76 tokenizer: The model's tokenizer.
77 Must be provided to interact with special tokens.
78 tools: Optional list of tools available for the current request.
79
80 Raises:
81 ValueError: If the `model_tokenizer`
82 is not successfully passed to the base class.
83
84 """
85 super().__init__(tokenizer, tools)
86
87 if not self.model_tokenizer:
88 raise ValueError(
89 "The model tokenizer must be passed to the ToolParser "
90 "constructor during construction."
91 )
92 # Regex to extract tool calls block (suffix is optional for incomplete outputs)
93 self.tool_call_regex = re.compile(
94 rf"{re.escape(TOOL_CALLS_PREFIX)}"
95 rf"(.*?)"
96 rf"(?:{re.escape(TOOL_CALLS_SUFFIX)}|$)",
97 re.DOTALL,
98 )
99
100 self._reset_streaming_state()
101
102 def _reset_streaming_state(self) -> None:
103 """Resets all streaming state variables for a new completion request.
104
105 This clears the delta text buffer and resets the pointers used to
106 track the currently streaming tool index and arguments. Called implicitly
107 during initialization and should be called between separate streams.
108 """
109 self.buffered_delta_text = ""
110 self.current_tool_id = -1
111 self.current_tool_name_sent = False
112 self.streamed_args_for_tool: list[str] = []
113
114 def adjust_request(
115 self, request: ChatCompletionRequest | ResponsesRequest
116 ) -> ChatCompletionRequest | ResponsesRequest:
117 """Adjusts the generation request to ensure special tool tokens are not skipped.
118
119 Forces `skip_special_tokens=False` if tools are actively being evaluated,
120 ensuring the tools special tokens are surfaced to the engine for parsing.
121
122 Args:
123 request: The incoming OpenAI-compatible chat completion request.
124
125 Returns:
126 The potentially modified chat completion request.
127
128 """
129 request = super().adjust_request(request)
130 if request.tools and request.tool_choice != "none":
131 request.skip_special_tokens = False
132 return request
133
134 def _buffer_delta_text(self, delta_text: str) -> str:
135 """Buffers incoming delta chunks to prevent
136 fragmentation of multi-token special tags.
137
138 If a chunk ends with a partial match of
139 `<|tools_prefix|>` or `<|tools_suffix|>`,
140 it holds that part back until the next chunk clarifies if it's the actual tag
141 or just normal text.
142
143 Args:
144 delta_text: The newly generated text chunk
145
146 Returns:
147 The safe, verified text chunk free of partial tag collisions.
148
149 Examples:
150 >>> parser = ApertusToolParser(...)
151 >>> parser._buffer_delta_text("Let me check <|tool" \
152 "Let me check " # "<|tool" is buffered internally
153 >>> parser._buffer_delta_text("s_prefix|>" \
154 "<|tools_prefix|>" # Buffer released on completion
155
156 """
157 self.buffered_delta_text += delta_text
158 text = self.buffered_delta_text
159
160 for tag in (TOOL_CALLS_PREFIX, TOOL_CALLS_SUFFIX):
161 if text.endswith(tag):
162 self.buffered_delta_text = ""
163 return text
164
165 # Evaluate longest possible partial match first
166 for i in range(len(tag) - 1, 0, -1):
167 if text.endswith(tag[:i]):
168 self.buffered_delta_text = text[-i:]
169 return text[:-i]
170
171 self.buffered_delta_text = ""
172 return text
173
174 def extract_tool_calls(
175 self,
176 model_output: str,
177 request: ChatCompletionRequest,
178 ) -> ExtractedToolCallInformation:
179 """Extract tool calls from a complete model response (non-streaming).
180
181 Args:
182 model_output: The full completion string generated by the model.
183 request: The current chat completion
184 request context containing tool schemas.
185
186 Returns:
187 An `ExtractedToolCallInformation` object containing normal text content
188 and a list of fully formatted `ToolCall` objects.
189
190 Examples:
191 >>> output = 'Let me see. <|tools_prefix|>[{"get_weather":' \
192 '{"loc": "Paris"}}]<|tools_suffix|>'
193 >>> info = parser.extract_tool_calls(output, request)
194 >>> info.tools_called
195 True
196 >>> info.content
197 'Let me see.'
198 >>> info.tool_calls[0].function.name
199 'get_weather'
200
201 """
202 match = self.tool_call_regex.search(model_output)
203 if not match:
204 return ExtractedToolCallInformation(
205 tools_called=False, tool_calls=[], content=model_output
206 )
207
208 try:
209 # group(1) might contain trailing text if the suffix is missing
210 matched_text = match.group(1)
211 stripped_text = matched_text.lstrip()
212
213 try:
214 # Use raw_decode to robustly isolate
215 # the valid JSON array from any trailing garbage
216 parsed_json, idx = json.JSONDecoder().raw_decode(stripped_text)
217 trailing_in_group = stripped_text[idx:]
218 except json.JSONDecodeError:
219 # Fallback sequentially to partial parser for token-truncated requests
220 parsed_json, _ = partial_json_loads(matched_text, Allow.ALL)
221 trailing_in_group = ""
222
223 if not isinstance(parsed_json, list):
224 parsed_json = [parsed_json] if parsed_json else []
225
226 tool_calls: list[ToolCall] = []
227 for obj in parsed_json:
228 if isinstance(obj, dict) and obj:
229 name, args = next(iter(obj.items()))
230 tool_calls.append(
231 ToolCall(
232 type="function",
233 id=make_tool_call_id(),
234 function=FunctionCall(
235 name=name,
236 arguments=json.dumps(args, ensure_ascii=False),
237 ),
238 )
239 )
240
241 # Content combines any generated text
242 # prior to and safely after the tool block
243 content_str = model_output[: match.start()].strip()
244
245 # Surface any hallucinated text inside
246 # the regex group (due to missing suffix)
247 if trailing_in_group.strip():
248 trailing = trailing_in_group.replace(TOOL_CALLS_SUFFIX, "").strip()
249 if trailing:
250 content_str = (content_str + "\n" + trailing).strip()
251
252 # Surface text natively generated after the explicit suffix
253 after_suffix = (
254 model_output[match.end() :].replace(TOOL_CALLS_SUFFIX, "").strip()
255 )
256 if after_suffix:
257 content_str = (content_str + "\n" + after_suffix).strip()
258
259 return ExtractedToolCallInformation(
260 tools_called=True,
261 tool_calls=tool_calls,
262 content=content_str if content_str else None,
263 )
264
265 except Exception:
266 logger.exception("Error extracting tool calls from Apertus response")
267 return ExtractedToolCallInformation(
268 tools_called=False, tool_calls=[], content=model_output
269 )
270
271 def extract_tool_calls_streaming(
272 self,
273 previous_text: str,
274 current_text: str,
275 delta_text: str,
276 previous_token_ids: Sequence[int],
277 current_token_ids: Sequence[int],
278 delta_token_ids: Sequence[int],
279 request: ChatCompletionRequest,
280 ) -> DeltaMessage | None:
281 """Handles streaming chunks.
282
283 Args:
284 previous_text: The complete model text generated prior to this chunk.
285 current_text: The complete model text including this chunk.
286 delta_text: The incremental text addition.
287 previous_token_ids: Tokens generated prior to this chunk.
288 current_token_ids: Total tokens generated.
289 delta_token_ids: Incremental token additions.
290 request: The chat completion request.
291
292 Returns:
293 A `DeltaMessage` with updated content or tool argument diffs, or `None` if
294 the chunk shouldn't emit visible changes yet (e.g. it was purely buffered).
295
296 Examples:
297 >>> prev = '<|tools_prefix|>[{"get_weather": {"loc'
298 >>> cur = '<|tools_prefix|>[{"get_weather": {"location": "Paris"}}'
299 >>> delta = 'ation": "Paris"}}'
300 >>> msg = parser.extract_tool_calls_streaming(
301 ... prev, cur, delta, ..., request
302 ... )
303 >>> msg.tool_calls[0].function.arguments
304 'ation": "Paris"}'
305
306 """
307 delta_text = self._buffer_delta_text(delta_text)
308 if not delta_text:
309 return None
310
311 # Fast path: normal text generation before any tools are invoked
312 if TOOL_CALLS_PREFIX not in current_text:
313 return DeltaMessage(content=delta_text)
314
315 try:
316 return self._extract_streaming(current_text, delta_text)
317 except Exception:
318 logger.exception("Error in Apertus streaming tool call extraction")
319 return None
320
321 def _extract_streaming(
322 self, current_text: str, delta_text: str
323 ) -> DeltaMessage | None:
324 """Core streaming logic.
325 Separates visible chat text from JSON blocks and computes diffs.
326
327 Args:
328 current_text: The full generated output string so far.
329 delta_text: The latest chunk of text added.
330
331 Returns:
332 A `DeltaMessage` containing the `content` delta and/or `tool_calls` delta.
333
334 """
335 prefix_idx = current_text.rfind(TOOL_CALLS_PREFIX)
336 suffix_idx = current_text.rfind(TOOL_CALLS_SUFFIX)
337
338 is_inside_tools = prefix_idx > suffix_idx
339
340 json_completed = False
341 json_end_idx: int | None = None
342
343 # Check if the JSON array successfully closed implicitly
344 if is_inside_tools:
345 json_start = prefix_idx + len(TOOL_CALLS_PREFIX)
346 s = current_text[json_start:].lstrip()
347 try:
348 # If raw_decode succeeds,
349 # the JSON array is fully formed and implicitly closed
350 _, idx = json.JSONDecoder().raw_decode(s)
351 json_end_idx = len(current_text) - len(s) + idx
352 json_completed, is_inside_tools = True, False
353 except Exception:
354 pass
355
356 just_finished = (TOOL_CALLS_SUFFIX in delta_text) or json_completed
357
358 # 1. Fast path: Output normal text immediately
359 # if we are completely outside tool block constraints
360 if not is_inside_tools and not just_finished:
361 text = delta_text.replace(TOOL_CALLS_PREFIX, "").replace(
362 TOOL_CALLS_SUFFIX, ""
363 )
364 return DeltaMessage(content=text) if text else None
365
366 # 2. Extract leading and trailing normal text directly adjacent to tool blocks
367 content_str = ""
368 if TOOL_CALLS_PREFIX in delta_text:
369 content_str += delta_text.split(TOOL_CALLS_PREFIX)[0].replace(
370 TOOL_CALLS_SUFFIX, ""
371 )
372
373 if just_finished:
374 if json_completed and json_end_idx is not None:
375 # The tool block finished in this chunk via implicit JSON completion
376 # Ensure we strictly isolate
377 # and extract only trailing text that is part of `delta_text`
378 delta_start_idx = len(current_text) - len(delta_text)
379 content_start = max(json_end_idx, delta_start_idx)
380 if content_start < len(current_text):
381 content_str += current_text[content_start:].replace(
382 TOOL_CALLS_SUFFIX, ""
383 )
384 else:
385 content_str += delta_text.split(TOOL_CALLS_SUFFIX)[-1]
386
387 # 3. Extract the isolated JSON array string for the active block
388 json_start = prefix_idx + len(TOOL_CALLS_PREFIX)
389 json_end = suffix_idx if suffix_idx > prefix_idx else json_end_idx
390 json_str = current_text[json_start:json_end]
391
392 tool_calls = self._parse_and_diff_json(json_str, is_final=not is_inside_tools)
393
394 if tool_calls or content_str:
395 return DeltaMessage(
396 content=content_str if content_str else None,
397 tool_calls=tool_calls if tool_calls else None,
398 )
399
400 return None
401
402 def _parse_and_diff_json(
403 self, json_str: str, is_final: bool
404 ) -> list[DeltaToolCall]:
405 """Parses an isolated, potentially incomplete streaming JSON array and returns
406 newly accumulated tool call diffs.
407
408 Args:
409 json_str: The extracted JSON array string so far
410 (e.g. `[{"weather": {"city": "Par"}]`).
411 is_final: True if the tool block has received its closing`<|tools_suffix|>`
412
413 Returns:
414 A list of `DeltaToolCall`
415 items representing string diffs in function arguments
416 to stream back to the client.
417
418 """
419 try:
420 parsed, _ = partial_json_loads(json_str, Allow.ALL)
421 if not isinstance(parsed, list):
422 parsed = [parsed] if parsed else []
423 except Exception:
424 return []
425
426 if not parsed:
427 return []
428
429 tool_calls: list[DeltaToolCall] = []
430 latest_index = len(parsed) - 1
431
432 # Catch up and finalize any tools we fully skipped over in one large text delta
433 while self.current_tool_id < latest_index:
434 if self.current_tool_id >= 0:
435 if not self.current_tool_name_sent:
436 self._emit_tool_name(parsed, self.current_tool_id, tool_calls)
437
438 delta = self._get_tool_diff(parsed, self.current_tool_id, is_final=True)
439 if delta:
440 tool_calls.append(delta)
441
442 self.current_tool_id += 1
443 self.current_tool_name_sent = False
444 while len(self.streamed_args_for_tool) <= self.current_tool_id:
445 self.streamed_args_for_tool.append("")
446
447 # Stream the currently active tool
448 if self.current_tool_id >= 0:
449 if not self.current_tool_name_sent:
450 self._emit_tool_name(parsed, self.current_tool_id, tool_calls)
451
452 delta = self._get_tool_diff(parsed, self.current_tool_id, is_final)
453 if delta:
454 tool_calls.append(delta)
455
456 return tool_calls
457
458 def _emit_tool_name(
459 self, parsed: list, index: int, tool_calls: list[DeltaToolCall]
460 ) -> None:
461 """Extracts and emits the function name mapped to a new tool call ID.
462
463 Args:
464 parsed: The partially parsed JSON list containing tool dictionaries.
465 index: The active index within the JSON list.
466 tool_calls: The running list of delta chunks to mutate.
467
468 Examples:
469 Appends `DeltaToolCall(index=0,
470 function=DeltaFunctionCall(name="get_weather", ...))`
471 to the `tool_calls` list and marks the name as sent.
472
473 """
474 obj = parsed[index]
475 if isinstance(obj, dict) and obj:
476 name = next(iter(obj))
477 self.current_tool_name_sent = True
478 tool_calls.append(
479 DeltaToolCall(
480 index=index,
481 type="function",
482 id=make_tool_call_id(),
483 function=DeltaFunctionCall(name=name, arguments="").model_dump(
484 exclude_none=True
485 ),
486 )
487 )
488
489 def _get_tool_diff(
490 self, parsed: list, index: int, is_final: bool
491 ) -> DeltaToolCall | None:
492 """Calculates the exact string difference to safely append new tool parameters.
493
494 This ensures characters like `{`, `}`, and `"` don't jump around unevenly
495 in the UI frontend while streaming incomplete JSON arguments.
496
497 Args:
498 parsed: The latest list of parsed JSON objects.
499 index: The active tool's array index.
500 is_final: Whether to emit
501 trailing structural brackets (True if block is done).
502
503 Returns:
504 A `DeltaToolCall` mapping to the arguments diff,
505 or None if no text was appended.
506
507 Examples:
508 >>> # Previous streamed state: '{"city": "Pari'
509 >>> # Current full parse state: '{"city": "Paris"}'
510 >>> # Returns diff (closing bracket suppressed until final):
511 >>> parser._get_tool_diff(parsed, index=0, is_final=False)
512 DeltaToolCall(index=0, function=DeltaFunctionCall(arguments='s'))
513
514 """
515 obj = parsed[index]
516 if not isinstance(obj, dict) or not obj:
517 return None
518
519 name, args = next(iter(obj.items()))
520 if args is None:
521 return None
522
523 args_json = json.dumps(args, ensure_ascii=False)
524
525 # Suppress trailing structural characters
526 # during stream (looks cleaner in frontends)
527 if not is_final:
528 while args_json and args_json[-1] in ("}", '"', "]", " ", ","):
529 args_json = args_json[:-1]
530
531 prev_sent = self.streamed_args_for_tool[index]
532 if args_json == prev_sent:
533 return None
534
535 prefix = find_common_prefix(prev_sent, args_json)
536 if len(prefix) < len(prev_sent):
537 # Backtrack state if partial parser structurally updates a past assumption
538 self.streamed_args_for_tool[index] = prefix
539 return None
540
541 diff = args_json[len(prev_sent) :]
542 if diff:
543 self.streamed_args_for_tool[index] = args_json
544 return DeltaToolCall(
545 index=index,
546 function=DeltaFunctionCall(arguments=diff).model_dump(
547 exclude_none=True
548 ),
549 )
550
551 return None