main
py 274 lines 10.8 KB
Raw
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.chat_utils import make_tool_call_id
10 from vllm.entrypoints.generate.base.protocol import (
11 DeltaFunctionCall,
12 DeltaMessage,
13 DeltaToolCall,
14 ExtractedToolCallInformation,
15 FunctionCall,
16 ToolCall,
17 )
18 from vllm.entrypoints.openai.chat_completion.protocol import (
19 ChatCompletionRequest,
20 )
21 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
22 from vllm.logger import init_logger
23 from vllm.tokenizers import TokenizerLike
24 from vllm.tool_parsers.abstract_tool_parser import (
25 Tool,
26 ToolParser,
27 )
28 from vllm.tool_parsers.utils import is_complete_json, partial_tag_overlap
29 from vllm.utils.mistral import is_mistral_tokenizer
30
31 logger = init_logger(__name__)
32
33
34 class Hermes2ProToolParser(ToolParser):
35 structural_tag_model = "hermes"
36 tool_call_start_token: str = "<tool_call>"
37 tool_call_end_token: str = "</tool_call>"
38 tool_call_regex = re.compile(
39 r"<tool_call>(.*?)</tool_call>|<tool_call>(.*)", re.DOTALL
40 )
41 scratch_pad_regex = re.compile(r"<scratch_pad>(.*?)</scratch_pad>", re.DOTALL)
42
43 def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
44 super().__init__(tokenizer, tools)
45
46 if is_mistral_tokenizer(tokenizer):
47 logger.error("Detected Mistral tokenizer when using a Hermes model")
48 self.model_tokenizer = tokenizer.tokenizer
49
50 if not self.model_tokenizer:
51 raise ValueError(
52 "The model tokenizer must be passed to the ToolParser "
53 "constructor during construction."
54 )
55
56 # Streaming state: what has been sent to the client.
57 self._sent_content_idx: int = 0
58
59 def adjust_request(
60 self, request: ChatCompletionRequest | ResponsesRequest
61 ) -> ChatCompletionRequest | ResponsesRequest:
62 request = super().adjust_request(request)
63 if request.tools and request.tool_choice != "none":
64 # do not skip special tokens because the tool_call tokens are
65 # marked "special" in some models. Since they are skipped
66 # prior to the call to the tool parser, it breaks tool calling.
67 request.skip_special_tokens = False
68 return request
69
70 def extract_tool_calls(
71 self,
72 model_output: str,
73 request: ChatCompletionRequest,
74 ) -> ExtractedToolCallInformation:
75 # sanity check; avoid unnecessary processing
76 if self.tool_call_start_token not in model_output:
77 return ExtractedToolCallInformation(
78 tools_called=False, tool_calls=[], content=model_output
79 )
80
81 else:
82 try:
83 # there are two possible captures - between tags, or between a
84 # tag and end-of-string so the result of
85 # findall is an array of tuples where one is a function call and
86 # the other is None
87 function_call_tuples = self.tool_call_regex.findall(model_output)
88
89 # load the JSON, and then use it to build the Function and
90 # Tool Call
91 raw_function_calls = [
92 json.loads(match[0] if match[0] else match[1])
93 for match in function_call_tuples
94 ]
95 tool_calls = [
96 ToolCall(
97 type="function",
98 function=FunctionCall(
99 name=function_call["name"],
100 # function call args are JSON but as a string
101 arguments=json.dumps(
102 function_call["arguments"], ensure_ascii=False
103 ),
104 ),
105 )
106 for function_call in raw_function_calls
107 ]
108
109 content = model_output[: model_output.find(self.tool_call_start_token)]
110 return ExtractedToolCallInformation(
111 tools_called=True,
112 tool_calls=tool_calls,
113 content=content if content else None,
114 )
115
116 except Exception:
117 logger.exception("Error in extracting tool call from response.")
118 return ExtractedToolCallInformation(
119 tools_called=False, tool_calls=[], content=model_output
120 )
121
122 def _extract_content(self, current_text: str) -> str | None:
123 """Return unsent non-tool-call text, or None.
124
125 Holds back any suffix that could be a partial <tool_call> tag.
126 """
127 if self.tool_call_start_token not in current_text:
128 overlap_length = partial_tag_overlap(
129 current_text, self.tool_call_start_token
130 )
131 sendable_idx = len(current_text) - overlap_length
132 else:
133 sendable_idx = current_text.index(self.tool_call_start_token)
134
135 if sendable_idx > self._sent_content_idx:
136 content = current_text[self._sent_content_idx : sendable_idx]
137 self._sent_content_idx = sendable_idx
138 return content
139 return None
140
141 def _extract_tool_call_jsons(self, text: str) -> list[tuple[str, bool]]:
142 """Extract (json_text, is_complete) for each <tool_call> region."""
143 results: list[tuple[str, bool]] = []
144 pos = 0
145 while True:
146 start = text.find(self.tool_call_start_token, pos)
147 if start == -1:
148 break
149 json_start = start + len(self.tool_call_start_token)
150 json_end = text.find(self.tool_call_end_token, json_start)
151 if json_end != -1:
152 results.append((text[json_start:json_end].strip(), True))
153 pos = json_end + len(self.tool_call_end_token)
154 else:
155 raw = text[json_start:]
156 # Strip partial </tool_call> suffix if present.
157 overlap = partial_tag_overlap(raw, self.tool_call_end_token)
158 if overlap:
159 raw = raw[:-overlap]
160 tc_json = raw.strip()
161 # Valid JSON without closing tag = complete body,
162 # tag tokens just haven't arrived yet.
163 is_complete = is_complete_json(tc_json) if tc_json else False
164 results.append((tc_json, is_complete))
165 break
166 return results
167
168 @staticmethod
169 def _extract_tool_name(tc_json: str) -> str | None:
170 """Extract tool name, or None if the name isn't complete yet."""
171 match = re.search(r'"name"\s*:\s*"([^"]+)"', tc_json)
172 return match.group(1) if match else None
173
174 @staticmethod
175 def _extract_tool_args(tc_json: str, is_complete: bool) -> str | None:
176 """Extract tool arguments from the tool call JSON.
177
178 Given {"name": "f", "arguments": {"x": 1}}, returns '{"x": 1}'.
179 When is_complete, strips the trailing '}' that closes the outer
180 object (not the arguments). For partial JSON, returns as-is.
181 """
182 match = re.search(r'"arguments"\s*:\s*', tc_json)
183 if not match:
184 return None
185 raw = tc_json[match.end() :]
186 if is_complete:
187 raw = raw.rstrip()
188 if raw.endswith("}"):
189 raw = raw[:-1].rstrip()
190 return raw
191
192 def _compute_args_diff(
193 self, index: int, tc_json: str, is_complete: bool
194 ) -> str | None:
195 """Return new argument text not yet sent for tool `index`, or None."""
196 args = self._extract_tool_args(tc_json, is_complete)
197 if args is None or len(args) <= len(self.streamed_args_for_tool[index]):
198 return None
199 diff = args[len(self.streamed_args_for_tool[index]) :]
200 self.streamed_args_for_tool[index] = args
201 self.prev_tool_call_arr[index]["arguments"] = args
202 return diff
203
204 def extract_tool_calls_streaming(
205 self,
206 previous_text: str,
207 current_text: str,
208 delta_text: str,
209 previous_token_ids: Sequence[int],
210 current_token_ids: Sequence[int],
211 delta_token_ids: Sequence[int],
212 request: ChatCompletionRequest,
213 ) -> DeltaMessage | None:
214 """Incrementally stream tool call deltas from accumulated output.
215
216 On each invocation, re-parses the full ``current_text`` to find
217 ``<tool_call>`` regions, then diffs against previously sent state
218 to emit only new content, tool names, or argument fragments.
219
220 Returns a ``DeltaMessage`` containing either plain content (for
221 text preceding any tool call) or one or more ``DeltaToolCall``
222 entries, or ``None`` if there is nothing new to send yet."""
223 try:
224 # Extract any content before tool calls.
225 content = self._extract_content(current_text)
226 tool_call_jsons = self._extract_tool_call_jsons(current_text)
227 tool_call_deltas: list[DeltaToolCall] = []
228
229 for i, (tc_json, is_complete) in enumerate(tool_call_jsons):
230 if i >= len(self.prev_tool_call_arr):
231 self.prev_tool_call_arr.append({})
232 self.streamed_args_for_tool.append("")
233
234 # Stream back tool name.
235 if "name" not in self.prev_tool_call_arr[i]:
236 name = self._extract_tool_name(tc_json)
237 if not name:
238 # Can't skip to tool i+1 if i isn't ready
239 break
240 self.prev_tool_call_arr[i]["name"] = name
241 tool_call_deltas.append(
242 DeltaToolCall(
243 index=i,
244 type="function",
245 id=make_tool_call_id(),
246 function=DeltaFunctionCall(name=name).model_dump(
247 exclude_none=True
248 ),
249 )
250 )
251
252 # Stream back new tool args by diffing against what was sent.
253 args_diff = self._compute_args_diff(i, tc_json, is_complete)
254 if args_diff:
255 tool_call_deltas.append(
256 DeltaToolCall(
257 index=i,
258 function=DeltaFunctionCall(arguments=args_diff).model_dump(
259 exclude_none=True
260 ),
261 )
262 )
263
264 if content or tool_call_deltas:
265 return DeltaMessage(
266 content=content,
267 tool_calls=tool_call_deltas,
268 )
269
270 return None
271
272 except Exception:
273 logger.exception("Error trying to handle streaming tool call.")
274 return None