@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 293f036e6d
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 14:30 UTC

Add gigachat 3.1 tool parser + fix gigachat3 tool parser (#36664)

Signed-off-by: Viacheslav Barinov <viacheslav.teh@gmail.com>

Viacheslav committed Mar 17, 2026 at 15:03 UTC 293f036e6d83ba05236d948e9800bc6d4d58a727
2 files changed +274 -88
tests/entrypoints/openai/tool_parsers/test_gigachat3_tool_parser.py
+193 -26
@@ -13,6 +13,13 @@ from vllm.entrypoints.openai.engine.protocol import FunctionCall
13 from vllm.tokenizers import TokenizerLike
14 from vllm.tool_parsers import ToolParser, ToolParserManager
15
16 +MSG_SEP_TOKEN = "<|message_sep|>\n\n"
17 +ROLE_SEP_TOKEN = "<|role_sep|>\n"
18 +EOS_TOKEN = "</s>"
19 +TOOL_HEADER_GIGACHAT3 = f"function call{ROLE_SEP_TOKEN}"
20 +TOOL_HEADER_GIGACHAT31 = "<|function_call|>"
21 +
22 +
23 SIMPLE_ARGS_DICT = {
24 "action": "create",
25 "id": "preferences",
@@ -24,7 +31,10 @@ SIMPLE_FUNCTION_JSON = json.dumps(
31 },
32 ensure_ascii=False,
33 )
27 -SIMPLE_FUNCTION_OUTPUT = "function call" + SIMPLE_FUNCTION_JSON
34 +SIMPLE_FUNCTION_OUTPUT_GIGACHAT3 = (
35 + f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{SIMPLE_FUNCTION_JSON}"
36 +)
37 +SIMPLE_FUNCTION_OUTPUT_GIGACHAT31 = f"{TOOL_HEADER_GIGACHAT31}{SIMPLE_FUNCTION_JSON}"
38 SIMPLE_FUNCTION_CALL = FunctionCall(
39 name="manage_user_memory",
40 arguments=json.dumps(SIMPLE_ARGS_DICT, ensure_ascii=False),
@@ -38,7 +48,12 @@ PARAMETERLESS_FUNCTION_JSON = json.dumps(
48 },
49 ensure_ascii=False,
50 )
41 -PARAMETERLESS_FUNCTION_OUTPUT = "function call" + PARAMETERLESS_FUNCTION_JSON
51 +PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3 = (
52 + f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{PARAMETERLESS_FUNCTION_JSON}"
53 +)
54 +PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31 = (
55 + f"{TOOL_HEADER_GIGACHAT31}{PARAMETERLESS_FUNCTION_JSON}"
56 +)
57 PARAMETERLESS_FUNCTION_CALL = FunctionCall(
58 name="manage_user_memory",
59 arguments=json.dumps({}, ensure_ascii=False),
@@ -62,17 +77,38 @@ COMPLEX_FUNCTION_JSON = json.dumps(
77 },
78 ensure_ascii=False,
79 )
65 -COMPLEX_FUNCTION_OUTPUT = "function call" + COMPLEX_FUNCTION_JSON
80 +COMPLEX_FUNCTION_OUTPUT_GIGACHAT3 = (
81 + f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{COMPLEX_FUNCTION_JSON}"
82 +)
83 +COMPLEX_FUNCTION_OUTPUT_GIGACHAT31 = f"{TOOL_HEADER_GIGACHAT31}{COMPLEX_FUNCTION_JSON}"
84 COMPLEX_FUNCTION_CALL = FunctionCall(
85 name="manage_user_memory",
86 arguments=json.dumps(COMPLEX_ARGS_DICT, ensure_ascii=False),
87 )
88
89
90 +CONTENT_TEXT = "I'll check that for you."
91 +MIXED_OUTPUT_GIGACHAT3 = f"{CONTENT_TEXT}{SIMPLE_FUNCTION_OUTPUT_GIGACHAT3}"
92 +MIXED_OUTPUT_GIGACHAT31 = f"{CONTENT_TEXT}{SIMPLE_FUNCTION_OUTPUT_GIGACHAT31}"
93 +
94 +
95 +@pytest.fixture(name="gigachat_tokenizer")
96 +def fixture_gigachat_tokenizer(default_tokenizer: TokenizerLike):
97 + default_tokenizer.add_tokens(
98 + [
99 + MSG_SEP_TOKEN,
100 + ROLE_SEP_TOKEN,
101 + TOOL_HEADER_GIGACHAT31,
102 + EOS_TOKEN,
103 + ]
104 + )
105 + return default_tokenizer
106 +
107 +
108 @pytest.mark.parametrize("streaming", [True, False])
73 -def test_no_tool_call(streaming: bool, default_tokenizer: TokenizerLike):
109 +def test_no_tool_call(streaming: bool, gigachat_tokenizer: TokenizerLike):
110 tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")(
75 - default_tokenizer
111 + gigachat_tokenizer
112 )
113 model_output = "How can I help you today?"
114 content, tool_calls = run_tool_extraction(
@@ -85,45 +121,143 @@ def test_no_tool_call(streaming: bool, default_tokenizer: TokenizerLike):
121 TEST_CASES = [
122 pytest.param(
123 True,
88 - SIMPLE_FUNCTION_OUTPUT,
124 + SIMPLE_FUNCTION_OUTPUT_GIGACHAT3,
125 + [SIMPLE_FUNCTION_CALL],
126 + None,
127 + id="simple_streaming_gigachat3",
128 + ),
129 + pytest.param(
130 + False,
131 + SIMPLE_FUNCTION_OUTPUT_GIGACHAT3,
132 + [SIMPLE_FUNCTION_CALL],
133 + None,
134 + id="simple_nonstreaming_gigachat3",
135 + ),
136 + pytest.param(
137 + True,
138 + PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3,
139 + [PARAMETERLESS_FUNCTION_CALL],
140 + None,
141 + id="parameterless_streaming_gigachat3",
142 + ),
143 + pytest.param(
144 + False,
145 + PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3,
146 + [PARAMETERLESS_FUNCTION_CALL],
147 + None,
148 + id="parameterless_nonstreaming_gigachat3",
149 + ),
150 + pytest.param(
151 + True,
152 + COMPLEX_FUNCTION_OUTPUT_GIGACHAT3,
153 + [COMPLEX_FUNCTION_CALL],
154 + None,
155 + id="complex_streaming_gigachat3",
156 + ),
157 + pytest.param(
158 + False,
159 + COMPLEX_FUNCTION_OUTPUT_GIGACHAT3,
160 + [COMPLEX_FUNCTION_CALL],
161 + None,
162 + id="complex_nonstreaming_gigachat3",
163 + ),
164 + pytest.param(
165 + True,
166 + MIXED_OUTPUT_GIGACHAT3,
167 + [SIMPLE_FUNCTION_CALL],
168 + CONTENT_TEXT,
169 + id="mixed_content_streaming_gigachat3",
170 + ),
171 + pytest.param(
172 + False,
173 + MIXED_OUTPUT_GIGACHAT3,
174 + [SIMPLE_FUNCTION_CALL],
175 + CONTENT_TEXT,
176 + id="mixed_content_nonstreaming_gigachat3",
177 + ),
178 + pytest.param(
179 + True,
180 + MIXED_OUTPUT_GIGACHAT3 + EOS_TOKEN,
181 + [SIMPLE_FUNCTION_CALL],
182 + CONTENT_TEXT,
183 + id="mixed_content_streaming_with_eos_gigachat3",
184 + ),
185 + pytest.param(
186 + False,
187 + MIXED_OUTPUT_GIGACHAT3 + EOS_TOKEN,
188 + [SIMPLE_FUNCTION_CALL],
189 + CONTENT_TEXT,
190 + id="mixed_content_nonstreaming_with_eos_gigachat3",
191 + ),
192 + pytest.param(
193 + True,
194 + SIMPLE_FUNCTION_OUTPUT_GIGACHAT31,
195 [SIMPLE_FUNCTION_CALL],
196 None,
91 - id="simple_streaming",
197 + id="simple_streaming_gigachat31",
198 ),
199 pytest.param(
200 False,
95 - SIMPLE_FUNCTION_OUTPUT,
201 + SIMPLE_FUNCTION_OUTPUT_GIGACHAT31,
202 [SIMPLE_FUNCTION_CALL],
203 None,
98 - id="simple_nonstreaming",
204 + id="simple_nonstreaming_gigachat31",
205 ),
206 pytest.param(
207 True,
102 - PARAMETERLESS_FUNCTION_OUTPUT,
208 + PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31,
209 [PARAMETERLESS_FUNCTION_CALL],
210 None,
105 - id="parameterless_streaming",
211 + id="parameterless_streaming_gigachat31",
212 ),
213 pytest.param(
214 False,
109 - PARAMETERLESS_FUNCTION_OUTPUT,
215 + PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31,
216 [PARAMETERLESS_FUNCTION_CALL],
217 None,
112 - id="parameterless_nonstreaming",
218 + id="parameterless_nonstreaming_gigachat31",
219 ),
220 pytest.param(
221 True,
116 - COMPLEX_FUNCTION_OUTPUT,
222 + COMPLEX_FUNCTION_OUTPUT_GIGACHAT31,
223 [COMPLEX_FUNCTION_CALL],
224 None,
119 - id="complex_streaming",
225 + id="complex_streaming_gigachat31",
226 ),
227 pytest.param(
228 False,
123 - COMPLEX_FUNCTION_OUTPUT,
229 + COMPLEX_FUNCTION_OUTPUT_GIGACHAT31,
230 [COMPLEX_FUNCTION_CALL],
231 None,
126 - id="complex_nonstreaming",
232 + id="complex_nonstreaming_gigachat31",
233 + ),
234 + pytest.param(
235 + True,
236 + MIXED_OUTPUT_GIGACHAT31,
237 + [SIMPLE_FUNCTION_CALL],
238 + CONTENT_TEXT,
239 + id="mixed_content_streaming_gigachat31",
240 + ),
241 + pytest.param(
242 + False,
243 + MIXED_OUTPUT_GIGACHAT31,
244 + [SIMPLE_FUNCTION_CALL],
245 + CONTENT_TEXT,
246 + id="mixed_content_nonstreaming_gigachat31",
247 + ),
248 + pytest.param(
249 + True,
250 + MIXED_OUTPUT_GIGACHAT31 + EOS_TOKEN,
251 + [SIMPLE_FUNCTION_CALL],
252 + CONTENT_TEXT,
253 + id="mixed_content_streaming_with_eos_gigachat31",
254 + ),
255 + pytest.param(
256 + False,
257 + MIXED_OUTPUT_GIGACHAT31 + EOS_TOKEN,
258 + [SIMPLE_FUNCTION_CALL],
259 + CONTENT_TEXT,
260 + id="mixed_content_nonstreaming_with_eos_gigachat31",
261 ),
262 ]
263
@@ -136,14 +270,16 @@ def test_tool_call(
270 model_output: str,
271 expected_tool_calls: list[FunctionCall],
272 expected_content: str | None,
139 - default_tokenizer: TokenizerLike,
273 + gigachat_tokenizer: TokenizerLike,
274 ):
275 tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")(
142 - default_tokenizer
276 + gigachat_tokenizer
277 )
278 content, tool_calls = run_tool_extraction(
279 tool_parser, model_output, streaming=streaming
280 )
281 + if content == "":
282 + content = None
283 assert content == expected_content
284 assert len(tool_calls) == len(expected_tool_calls)
285 for actual, expected in zip(tool_calls, expected_tool_calls):
@@ -154,15 +290,46 @@ def test_tool_call(
290 assert actual_args == expected_args
291
292
157 -def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike):
293 +@pytest.mark.parametrize(
294 + "model_output_deltas",
295 + [
296 + pytest.param(
297 + [
298 + CONTENT_TEXT[:3],
299 + CONTENT_TEXT[3:5],
300 + CONTENT_TEXT[5:],
301 + MSG_SEP_TOKEN,
302 + TOOL_HEADER_GIGACHAT3,
303 + COMPLEX_FUNCTION_JSON[:40],
304 + COMPLEX_FUNCTION_JSON[40:-1],
305 + COMPLEX_FUNCTION_JSON[-1],
306 + ],
307 + id="gigachat3",
308 + ),
309 + pytest.param(
310 + [
311 + CONTENT_TEXT[:3],
312 + CONTENT_TEXT[3:5],
313 + CONTENT_TEXT[5:],
314 + TOOL_HEADER_GIGACHAT31,
315 + COMPLEX_FUNCTION_JSON[:40],
316 + COMPLEX_FUNCTION_JSON[40:-1],
317 + COMPLEX_FUNCTION_JSON[-1],
318 + ],
319 + id="gigachat31",
320 + ),
321 + ],
322 +)
323 +def test_streaming_tool_call_with_large_steps(
324 + model_output_deltas: list[str],
325 + gigachat_tokenizer: TokenizerLike,
326 +):
327 + """
328 + Test that the closing braces are streamed correctly.
329 + """
330 tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")(
159 - default_tokenizer
331 + gigachat_tokenizer
332 )
161 - model_output_deltas = [
162 - "function call",
163 - COMPLEX_FUNCTION_JSON[:40],
164 - COMPLEX_FUNCTION_JSON[40:],
165 - ]
333 reconstructor = run_tool_extraction_streaming(
334 tool_parser,
335 model_output_deltas,
vllm/tool_parsers/gigachat3_tool_parser.py
+81 -62
@@ -25,7 +25,12 @@ from vllm.tool_parsers.abstract_tool_parser import ToolParser
25 logger = init_logger(__name__)
26
27 REGEX_FUNCTION_CALL = re.compile(
28 - r"function call(?:<\|role_sep\|>\n)?(\{.*)",
28 + r"(?:function call<\|role_sep\|>\n|<\|function_call\|>)(.*)",
29 + re.DOTALL,
30 +)
31 +
32 +REGEX_CONTENT_PATTERN = re.compile(
33 + r"^(.*?)(?:<\|message_sep\|>|<\|function_call\|>)",
34 re.DOTALL,
35 )
36
@@ -47,57 +52,67 @@ class GigaChat3ToolParser(ToolParser):
52 self.tool_name_sent: bool = False
53 self.tool_id: str | None = None
54 self.prev_tool_call_arr: list[dict] = []
50 - self.content_buffer: str = ""
51 - self.trigger_start = "function call{"
55 + self.end_content: bool = False
56 + self.streamed_args_for_tool: list[str] = []
57 +
58 + def adjust_request(self, request: ChatCompletionRequest) -> ChatCompletionRequest:
59 + request = super().adjust_request(request)
60 + if request.tools and request.tool_choice != "none":
61 + request.skip_special_tokens = False
62 + return request
63
64 def extract_tool_calls(
65 self,
66 model_output: str,
67 request: ChatCompletionRequest,
68 ) -> ExtractedToolCallInformation:
58 - match = REGEX_FUNCTION_CALL.search(model_output)
59 - if not match:
60 - return ExtractedToolCallInformation(
61 - tools_called=False,
62 - tool_calls=[],
63 - content=model_output,
64 - )
65 - json_candidate = match.group(1).strip()
66 - try:
67 - data = json.loads(json_candidate)
68 - except json.JSONDecodeError:
69 - return ExtractedToolCallInformation(
70 - tools_called=False,
71 - tool_calls=[],
72 - content=model_output,
73 - )
74 - if not (isinstance(data, dict) and "name" in data and "arguments" in data):
69 + function_call = None
70 + content = None
71 + if model_output.rstrip().endswith("</s>"):
72 + model_output = model_output[: model_output.rfind("</s>")]
73 + m_func = REGEX_FUNCTION_CALL.search(model_output)
74 + if m_func:
75 + try:
76 + function_call = json.loads(m_func.group(1), strict=False)
77 + if (
78 + isinstance(function_call, dict)
79 + and "name" in function_call
80 + and "arguments" in function_call
81 + ):
82 + if not isinstance(function_call["arguments"], dict):
83 + function_call = None
84 + else:
85 + function_call = None
86 + except json.JSONDecodeError:
87 + return ExtractedToolCallInformation(
88 + tools_called=False,
89 + tool_calls=[],
90 + content=model_output,
91 + )
92 + m_content = REGEX_CONTENT_PATTERN.search(model_output)
93 + content = m_content.group(1) if m_content else model_output
94 + if not function_call:
95 return ExtractedToolCallInformation(
96 tools_called=False,
97 tool_calls=[],
78 - content=model_output,
98 + content=content if content else None,
99 )
80 - name = data["name"]
81 - args = data["arguments"]
100 + name = function_call["name"]
101 + args = function_call["arguments"]
102 if not isinstance(args, str):
83 - args = json.dumps(args, ensure_ascii=False)
84 -
85 - tool_calls = [
86 - ToolCall(
87 - type="function",
88 - function=FunctionCall(
89 - name=name,
90 - arguments=args,
91 - ),
92 - )
93 - ]
94 - prefix = model_output[: match.start()]
95 - content = prefix.rstrip() if prefix and prefix.strip() else None
96 -
103 + args = json.dumps(function_call["arguments"], ensure_ascii=False)
104 return ExtractedToolCallInformation(
105 tools_called=True,
99 - tool_calls=tool_calls,
100 - content=content,
106 + tool_calls=[
107 + ToolCall(
108 + type="function",
109 + function=FunctionCall(
110 + name=name,
111 + arguments=args,
112 + ),
113 + )
114 + ],
115 + content=content if content else None,
116 )
117
118 def extract_tool_calls_streaming(
@@ -110,39 +125,37 @@ class GigaChat3ToolParser(ToolParser):
125 delta_token_ids: Sequence[int],
126 request: ChatCompletionRequest,
127 ) -> DeltaMessage | None:
128 + content = None
129 func_name = None
130 cur_args = None
131 + m_func = REGEX_FUNCTION_CALL.search(current_text)
132 if not self.tool_started:
116 - match = REGEX_FUNCTION_CALL.search(current_text)
117 - if match:
118 - self.tool_started = True
119 - self.content_buffer = ""
133 + m_content = REGEX_CONTENT_PATTERN.search(delta_text)
134 + if m_content:
135 + content = m_content.group(1)
136 + self.end_content = True
137 else:
121 - self.content_buffer += delta_text
122 - clean_buffer = self.content_buffer.lstrip()
123 - is_prefix = self.trigger_start.startswith(clean_buffer)
124 - starts_with_trigger = clean_buffer.startswith(self.trigger_start)
125 - if is_prefix or starts_with_trigger:
126 - return None
127 - else:
128 - flush_text = self.content_buffer
129 - self.content_buffer = ""
130 - return DeltaMessage(content=flush_text)
131 -
132 - match = REGEX_FUNCTION_CALL.search(current_text)
133 - if not match:
138 + if not self.end_content:
139 + content = delta_text
140 + if m_func:
141 + self.tool_started = True
142 + if content:
143 + return DeltaMessage(content=content)
144 + if not m_func:
145 return None
135 - json_tail = match.group(1).strip()
146 + json_tail = m_func.group(1).strip()
147 name_match = NAME_REGEX.search(json_tail)
148 if name_match:
149 func_name = name_match.group(1)
150 args_match = ARGS_REGEX.search(json_tail)
151 if args_match:
152 cur_args = args_match.group(1).strip()
153 + if cur_args.endswith("</s>"):
154 + cur_args = cur_args[: -len("</s>")]
155 if cur_args.endswith("}"): # last '}' end of json
156 try:
157 candidate = cur_args[:-1].strip()
145 - json.loads(candidate)
158 + json.loads(candidate, strict=False)
159 cur_args = candidate
160 except json.JSONDecodeError:
161 pass
@@ -165,11 +178,10 @@ class GigaChat3ToolParser(ToolParser):
178 ).model_dump(exclude_none=True),
179 )
180 ],
168 - content=None,
181 )
182 if cur_args is None:
183 return None
172 - prev_args = self.prev_tool_call_arr[0].get("arguments", "")
184 + prev_args = self.prev_tool_call_arr[0].get("arguments_str", "")
185 if not prev_args:
186 delta_args = cur_args
187 elif cur_args.startswith(prev_args):
@@ -178,7 +190,15 @@ class GigaChat3ToolParser(ToolParser):
190 return None
191 if not delta_args:
192 return None
181 - self.prev_tool_call_arr[0]["arguments"] = cur_args
193 + self.prev_tool_call_arr[0]["arguments_str"] = cur_args
194 + try:
195 + args_dict = json.loads(cur_args, strict=False)
196 + self.prev_tool_call_arr[0]["arguments"] = args_dict
197 + except json.JSONDecodeError:
198 + self.prev_tool_call_arr[0]["arguments"] = {}
199 + if len(self.streamed_args_for_tool) <= 0:
200 + self.streamed_args_for_tool.append("")
201 + self.streamed_args_for_tool[0] = cur_args
202 return DeltaMessage(
203 tool_calls=[
204 DeltaToolCall(
@@ -188,5 +208,4 @@ class GigaChat3ToolParser(ToolParser):
208 ).model_dump(exclude_none=True),
209 )
210 ],
191 - content=None,
211 )