| 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 | from typing import Any, Protocol, TypeVar |
| 7 | |
| 8 | import regex as re |
| 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 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | |
| 33 | def dump_args(args: None | dict[str, Any] | str) -> str | None: |
| 34 | if args is None or isinstance(args, str): |
| 35 | return args |
| 36 | else: |
| 37 | return json.dumps(args, ensure_ascii=False) |
| 38 | |
| 39 | |
| 40 | class _FunctionCallCtor(Protocol): |
| 41 | def __init__(self, *, name: str, arguments: str | None): ... |
| 42 | |
| 43 | |
| 44 | FuncT = TypeVar("FuncT", bound=_FunctionCallCtor) |
| 45 | |
| 46 | |
| 47 | class Granite4ToolParser(ToolParser): |
| 48 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 49 | super().__init__(tokenizer, tools) |
| 50 | |
| 51 | self.prev_tool_call_arr: list[dict] = [] |
| 52 | self.current_tool_id: int = -1 |
| 53 | self.streamed_args_for_tool = list[str]() |
| 54 | |
| 55 | self.look_ahead = "" |
| 56 | self.in_tc = False |
| 57 | |
| 58 | self.tc_start = "<tool_call>" |
| 59 | self.tc_end = "</tool_call>" |
| 60 | self.start_regex = re.compile(self.tc_start) |
| 61 | self.end_regex = re.compile(self.tc_end) |
| 62 | |
| 63 | def adjust_request( |
| 64 | self, request: ChatCompletionRequest | ResponsesRequest |
| 65 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 66 | request = super().adjust_request(request) |
| 67 | if request.tools and request.tool_choice != "none": |
| 68 | # do not skip special tokens because the tool_call tokens are |
| 69 | # marked "special" in some models. Since they are skipped |
| 70 | # prior to the call to the tool parser, it breaks tool calling. |
| 71 | request.skip_special_tokens = False |
| 72 | return request |
| 73 | |
| 74 | def _collect_results( |
| 75 | self, text_segments: list[str], tc_segments: list[str], cls: type[FuncT] |
| 76 | ) -> tuple[str, list[FuncT]]: |
| 77 | tool_calls_json: list[dict[str, Any]] = [ |
| 78 | json.loads(tc_text) for tc_text in tc_segments |
| 79 | ] |
| 80 | tool_calls = [] |
| 81 | for tc in tool_calls_json: |
| 82 | assert isinstance(tc, dict) |
| 83 | self.prev_tool_call_arr.append(tc) |
| 84 | tool_calls.append( |
| 85 | cls( |
| 86 | name=tc["name"], |
| 87 | arguments=dump_args(tc["arguments"]), |
| 88 | ) |
| 89 | ) |
| 90 | return "".join(text_segments), tool_calls |
| 91 | |
| 92 | def extract_tool_calls( |
| 93 | self, |
| 94 | model_output: str, |
| 95 | request: ChatCompletionRequest, |
| 96 | ) -> ExtractedToolCallInformation: |
| 97 | msg = ExtractedToolCallInformation( |
| 98 | tools_called=False, tool_calls=[], content=model_output |
| 99 | ) |
| 100 | try: |
| 101 | delimiters = [("TC_START", self.tc_start), ("TC_END", self.tc_end)] |
| 102 | pattern = "|".join(f"(?P<{name}>{pattern})" for name, pattern in delimiters) |
| 103 | regex = re.compile(pattern) |
| 104 | |
| 105 | text_segments = list[str]() |
| 106 | tc_segments = list[str]() |
| 107 | last_cut_loc = 0 |
| 108 | |
| 109 | for match in regex.finditer(model_output): |
| 110 | match_type = match.lastgroup |
| 111 | if match_type == "TC_START": |
| 112 | assert not self.in_tc, "Two tool call start tokens found in a row" |
| 113 | if preceding_text := model_output[last_cut_loc : match.start()]: |
| 114 | text_segments.append(preceding_text) |
| 115 | self.in_tc = True |
| 116 | elif match_type == "TC_END": |
| 117 | assert self.in_tc, ( |
| 118 | "Tool call end token found without corresponding start token" |
| 119 | ) |
| 120 | tool_text = model_output[last_cut_loc : match.start()] |
| 121 | assert tool_text, ( |
| 122 | "Expected the model to generate text between tool call tokens" |
| 123 | ) |
| 124 | tc_segments.append(tool_text) |
| 125 | self.in_tc = False |
| 126 | else: |
| 127 | raise ValueError("Unexpected match") |
| 128 | last_cut_loc = match.end() |
| 129 | assert not self.in_tc, "The model generated an incomplete tool call" |
| 130 | if final_text := model_output[last_cut_loc:]: |
| 131 | text_segments.append(final_text) |
| 132 | |
| 133 | content, tool_call_funcs = self._collect_results( |
| 134 | text_segments, tc_segments, FunctionCall |
| 135 | ) |
| 136 | tool_calls = [ |
| 137 | ToolCall( |
| 138 | type="function", |
| 139 | function=func, |
| 140 | ) |
| 141 | for func in tool_call_funcs |
| 142 | ] |
| 143 | msg.tools_called = bool(tool_calls) |
| 144 | msg.tool_calls = tool_calls |
| 145 | msg.content = content or None |
| 146 | except Exception: |
| 147 | logger.exception("Error in extracting tool call from response.") |
| 148 | return msg |
| 149 | |
| 150 | def _tool_extraction_step( |
| 151 | self, |
| 152 | delta_text: str, |
| 153 | ) -> tuple[bool, str, str]: |
| 154 | start_token_pos = start_token_end = end_token_pos = end_token_end = -1 |
| 155 | |
| 156 | if start_match := self.start_regex.search(delta_text, partial=True): |
| 157 | if not start_match.partial: |
| 158 | start_token_pos, start_token_end = start_match.span() |
| 159 | elif start_match.end() > start_match.start(): |
| 160 | start_token_pos = -2 |
| 161 | |
| 162 | if end_match := self.end_regex.search(delta_text): |
| 163 | end_token_pos, end_token_end = end_match.span() |
| 164 | |
| 165 | # Done means that we've exhausted the current buffer |
| 166 | # and need more output from the model |
| 167 | done = True |
| 168 | content = tc_text = "" |
| 169 | |
| 170 | if start_token_pos < 0: |
| 171 | # just streaming text so far |
| 172 | if start_token_pos == -2: |
| 173 | # There is a partial match |
| 174 | content = delta_text[: start_match.start()] |
| 175 | self.look_ahead = delta_text[start_match.start() :] |
| 176 | else: |
| 177 | content = delta_text |
| 178 | |
| 179 | elif not self.in_tc: |
| 180 | # we're entering a new tool call |
| 181 | self.in_tc = True |
| 182 | |
| 183 | content = delta_text[:start_token_pos] |
| 184 | if end_token_pos > 0: |
| 185 | self.start_in_tc = False |
| 186 | tc_text = delta_text[start_token_end:end_token_pos] |
| 187 | self.look_ahead = delta_text[end_token_end:] |
| 188 | done = False # There could be more content already buffered |
| 189 | else: |
| 190 | self.look_ahead = delta_text[start_token_pos:] |
| 191 | |
| 192 | elif end_token_pos < 0: |
| 193 | # we're in between the start and the end token |
| 194 | assert self.in_tc |
| 195 | self.look_ahead = delta_text |
| 196 | else: |
| 197 | # We have found the end |
| 198 | assert self.in_tc |
| 199 | tc_text = delta_text[start_token_end:end_token_pos] |
| 200 | self.in_tc = False |
| 201 | self.look_ahead = delta_text[end_token_end:] |
| 202 | done = False # There could be more content already buffered |
| 203 | return done, content, tc_text |
| 204 | |
| 205 | def extract_tool_calls_streaming( |
| 206 | self, |
| 207 | previous_text: str, |
| 208 | current_text: str, |
| 209 | delta_text: str, |
| 210 | previous_token_ids: Sequence[int], |
| 211 | current_token_ids: Sequence[int], |
| 212 | delta_token_ids: Sequence[int], |
| 213 | request: ChatCompletionRequest, |
| 214 | ) -> DeltaMessage | None: |
| 215 | try: |
| 216 | done = False |
| 217 | text_segments = list[str]() |
| 218 | tc_segments = list[str]() |
| 219 | |
| 220 | while not done: |
| 221 | delta_text = self.look_ahead + delta_text |
| 222 | self.look_ahead = "" |
| 223 | done, content, tc_text = self._tool_extraction_step(delta_text) |
| 224 | if content: |
| 225 | text_segments.append(content) |
| 226 | if tc_text: |
| 227 | tc_segments.append(tc_text) |
| 228 | delta_text = "" |
| 229 | |
| 230 | content, tool_call_funcs = self._collect_results( |
| 231 | text_segments, tc_segments, DeltaFunctionCall |
| 232 | ) |
| 233 | |
| 234 | delta_tool_calls = list[DeltaToolCall]() |
| 235 | for function in tool_call_funcs: |
| 236 | self.current_tool_id += 1 |
| 237 | delta_tool_calls.append( |
| 238 | DeltaToolCall( |
| 239 | id=make_tool_call_id(), |
| 240 | type="function", |
| 241 | index=self.current_tool_id, |
| 242 | function=function.model_dump(exclude_none=True), |
| 243 | ) |
| 244 | ) |
| 245 | self.streamed_args_for_tool.append(function.arguments or "") |
| 246 | |
| 247 | assert self.current_tool_id + 1 == len(self.prev_tool_call_arr) |
| 248 | assert self.current_tool_id + 1 == len(self.streamed_args_for_tool) |
| 249 | |
| 250 | msg = DeltaMessage(content=content or None, tool_calls=delta_tool_calls) |
| 251 | if msg.content or msg.tool_calls: |
| 252 | return msg |
| 253 | |
| 254 | except Exception: |
| 255 | logger.exception("Error trying to handle streaming tool call.") |
| 256 | return None |