| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """GLM-4 Tool Call Parser with incremental string streaming support. |
| 4 | |
| 5 | This parser fixes the streaming issue reported in Issue #32829 where long string |
| 6 | parameters (e.g., file content with 4000+ characters of code) are buffered until |
| 7 | complete, causing multi-second delays before the user sees any content. |
| 8 | |
| 9 | The fix streams string values incrementally as they arrive, providing a true |
| 10 | streaming experience for long content. |
| 11 | """ |
| 12 | |
| 13 | import json |
| 14 | from collections.abc import Sequence |
| 15 | from typing import Any |
| 16 | |
| 17 | import partial_json_parser.core.complete |
| 18 | import regex as re |
| 19 | from openai.types.responses import ToolChoiceFunction |
| 20 | from partial_json_parser.core.options import Allow |
| 21 | |
| 22 | from vllm.entrypoints.chat_utils import make_tool_call_id |
| 23 | from vllm.entrypoints.generate.base.protocol import ( |
| 24 | DeltaFunctionCall, |
| 25 | DeltaMessage, |
| 26 | DeltaToolCall, |
| 27 | ExtractedToolCallInformation, |
| 28 | FunctionCall, |
| 29 | ToolCall, |
| 30 | ) |
| 31 | from vllm.entrypoints.openai.chat_completion.protocol import ( |
| 32 | ChatCompletionNamedToolChoiceParam, |
| 33 | ChatCompletionRequest, |
| 34 | ) |
| 35 | from vllm.entrypoints.openai.responses.protocol import ( |
| 36 | ResponsesRequest, |
| 37 | ) |
| 38 | from vllm.logger import init_logger |
| 39 | from vllm.tokenizers import TokenizerLike |
| 40 | from vllm.tool_parsers.abstract_tool_parser import ( |
| 41 | Tool, |
| 42 | ToolParser, |
| 43 | ) |
| 44 | from vllm.tool_parsers.utils import safe_literal_eval |
| 45 | |
| 46 | logger = init_logger(__name__) |
| 47 | |
| 48 | |
| 49 | class PoolsideV1ToolParser(ToolParser): |
| 50 | """Tool parser for GLM-4 models with incremental string streaming. |
| 51 | |
| 52 | This parser emits tool-call deltas incrementally as arguments arrive. |
| 53 | For string-type parameters, content is streamed character-by-character |
| 54 | rather than waiting for the complete </arg_value> tag. |
| 55 | """ |
| 56 | |
| 57 | supports_required_and_named = False |
| 58 | |
| 59 | def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): |
| 60 | super().__init__(tokenizer, tools) |
| 61 | # Stateful streaming fields |
| 62 | self.current_tool_name_sent: bool = False |
| 63 | self.prev_tool_call_arr: list[dict[str, Any]] = [] |
| 64 | self.current_tool_id: int = -1 |
| 65 | self.streamed_args_for_tool: list[str] = [] |
| 66 | |
| 67 | self.tool_call_start_token: str = "<tool_call>" |
| 68 | self.tool_call_end_token: str = "</tool_call>" |
| 69 | self.arg_key_start: str = "<arg_key>" |
| 70 | self.arg_key_end: str = "</arg_key>" |
| 71 | self.arg_val_start: str = "<arg_value>" |
| 72 | self.arg_val_end: str = "</arg_value>" |
| 73 | |
| 74 | self.tool_calls_start_token = self.tool_call_start_token |
| 75 | |
| 76 | self.func_call_regex = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL) |
| 77 | self.func_detail_regex = re.compile( |
| 78 | r"<tool_call>\s*([^\n<]+?)\s*\n?\s*(<arg_key>.*?)?</tool_call>", re.DOTALL |
| 79 | ) |
| 80 | self.func_arg_regex = re.compile( |
| 81 | r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>", re.DOTALL |
| 82 | ) |
| 83 | |
| 84 | if not self.model_tokenizer: |
| 85 | raise ValueError( |
| 86 | "The model tokenizer must be passed to the ToolParser " |
| 87 | "constructor during construction." |
| 88 | ) |
| 89 | |
| 90 | self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token) |
| 91 | self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token) |
| 92 | self._buffer: str = "" |
| 93 | |
| 94 | # Streaming state for incremental tool-call streaming |
| 95 | self._in_tool_call: bool = False |
| 96 | self._current_tool_name: str | None = None |
| 97 | self._pending_key: str | None = None |
| 98 | self._streaming_string_value: bool = False |
| 99 | self._tool_call_ids: list[str] = [] |
| 100 | self._args_started: list[bool] = [] |
| 101 | self._args_closed: list[bool] = [] |
| 102 | self._seen_keys: list[set[str]] = [] |
| 103 | |
| 104 | @staticmethod |
| 105 | def _deserialize(value: str) -> Any: |
| 106 | try: |
| 107 | return json.loads(value) |
| 108 | except json.JSONDecodeError: |
| 109 | pass |
| 110 | |
| 111 | try: |
| 112 | return safe_literal_eval(value) |
| 113 | except (ValueError, SyntaxError): |
| 114 | pass |
| 115 | |
| 116 | return value |
| 117 | |
| 118 | @staticmethod |
| 119 | def _json_escape_string_content(s: str) -> str: |
| 120 | """JSON-escape string content for incremental streaming. |
| 121 | |
| 122 | This escapes the content that goes INSIDE a JSON string (between quotes), |
| 123 | not including the surrounding quotes themselves. |
| 124 | """ |
| 125 | if not s: |
| 126 | return "" |
| 127 | return json.dumps(s, ensure_ascii=False)[1:-1] |
| 128 | |
| 129 | @staticmethod |
| 130 | def _is_string_type( |
| 131 | tool_name: str, |
| 132 | arg_name: str, |
| 133 | tools: list[Tool] | None, |
| 134 | ) -> bool: |
| 135 | if tools is None: |
| 136 | return False |
| 137 | for tool in tools: |
| 138 | # ChatCompletion tools nest under .function; Responses |
| 139 | # FunctionTool is flat (.name/.parameters at the top level). |
| 140 | fn = getattr(tool, "function", tool) |
| 141 | if getattr(fn, "name", None) != tool_name: |
| 142 | continue |
| 143 | params = getattr(fn, "parameters", None) |
| 144 | if params is None: |
| 145 | return False |
| 146 | arg_type = params.get("properties", {}).get(arg_name, {}).get("type", None) |
| 147 | return arg_type == "string" |
| 148 | logger.debug("No tool named '%s'.", tool_name) |
| 149 | return False |
| 150 | |
| 151 | @staticmethod |
| 152 | def _tools_enabled(request: ChatCompletionRequest) -> bool: |
| 153 | """Return whether tool parsing should be applied for this request.""" |
| 154 | try: |
| 155 | tools = getattr(request, "tools", None) |
| 156 | tool_choice = getattr(request, "tool_choice", None) |
| 157 | return bool(tools) and tool_choice != "none" |
| 158 | except Exception: |
| 159 | logger.exception("Failed to determine if tools are enabled.") |
| 160 | return False |
| 161 | |
| 162 | def adjust_request( |
| 163 | self, request: ChatCompletionRequest | ResponsesRequest |
| 164 | ) -> ChatCompletionRequest | ResponsesRequest: |
| 165 | """Adjust request parameters for tool call token handling. |
| 166 | |
| 167 | For required/named tool_choice, skip super().adjust_request() so it |
| 168 | does not install JSON guided decoding. These models emit XML tool |
| 169 | calls (per the chat template), which JSON guidance would break. |
| 170 | """ |
| 171 | if request.tools: |
| 172 | tc = request.tool_choice |
| 173 | if tc == "required" or isinstance( |
| 174 | tc, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) |
| 175 | ): |
| 176 | request.skip_special_tokens = False |
| 177 | return request |
| 178 | request = super().adjust_request(request) |
| 179 | if request.tools and request.tool_choice != "none": |
| 180 | # Ensure tool call tokens (<tool_call>, </tool_call>) are not skipped |
| 181 | # during decoding. Even though they are not marked as special tokens, |
| 182 | # setting skip_special_tokens=False ensures proper handling in |
| 183 | # transformers 5.x where decoding behavior may have changed. |
| 184 | request.skip_special_tokens = False |
| 185 | return request |
| 186 | |
| 187 | def extract_tool_calls( |
| 188 | self, |
| 189 | model_output: str, |
| 190 | request: ChatCompletionRequest, |
| 191 | ) -> ExtractedToolCallInformation: |
| 192 | matched_tool_calls = self.func_call_regex.findall(model_output) |
| 193 | logger.debug("model_output: %s", model_output) |
| 194 | try: |
| 195 | tool_calls: list[ToolCall] = [] |
| 196 | for match in matched_tool_calls: |
| 197 | tc_detail = self.func_detail_regex.search(match) |
| 198 | if not tc_detail: |
| 199 | logger.warning( |
| 200 | "Failed to parse tool call details from: %s", |
| 201 | match, |
| 202 | ) |
| 203 | continue |
| 204 | tc_name = tc_detail.group(1).strip() |
| 205 | tc_args = tc_detail.group(2) |
| 206 | pairs = self.func_arg_regex.findall(tc_args) if tc_args else [] |
| 207 | arg_dct: dict[str, Any] = {} |
| 208 | for key, value in pairs: |
| 209 | arg_key = key.strip() |
| 210 | # Keep string values verbatim; whitespace is significant |
| 211 | # (e.g. code/file content). Only strip non-string types. |
| 212 | if self._is_string_type(tc_name, arg_key, request.tools): |
| 213 | arg_val = value |
| 214 | else: |
| 215 | arg_val = self._deserialize(value.strip()) |
| 216 | logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val) |
| 217 | arg_dct[arg_key] = arg_val |
| 218 | tool_calls.append( |
| 219 | ToolCall( |
| 220 | type="function", |
| 221 | function=FunctionCall( |
| 222 | name=tc_name, |
| 223 | arguments=json.dumps(arg_dct, ensure_ascii=False), |
| 224 | ), |
| 225 | ) |
| 226 | ) |
| 227 | except Exception: |
| 228 | logger.exception("Failed to extract tool call spec") |
| 229 | return ExtractedToolCallInformation( |
| 230 | tools_called=False, tool_calls=[], content=model_output |
| 231 | ) |
| 232 | else: |
| 233 | if len(tool_calls) > 0: |
| 234 | content: str | None = model_output[ |
| 235 | : model_output.find(self.tool_calls_start_token) |
| 236 | ] |
| 237 | # Normalize empty/whitespace-only content to None |
| 238 | if not content or not content.strip(): |
| 239 | content = None |
| 240 | return ExtractedToolCallInformation( |
| 241 | tools_called=True, tool_calls=tool_calls, content=content |
| 242 | ) |
| 243 | return ExtractedToolCallInformation( |
| 244 | tools_called=False, tool_calls=[], content=model_output |
| 245 | ) |
| 246 | |
| 247 | def extract_tool_calls_streaming( |
| 248 | self, |
| 249 | previous_text: str, |
| 250 | current_text: str, |
| 251 | delta_text: str, |
| 252 | previous_token_ids: Sequence[int], |
| 253 | current_token_ids: Sequence[int], |
| 254 | delta_token_ids: Sequence[int], |
| 255 | request: ChatCompletionRequest, |
| 256 | ) -> DeltaMessage | None: |
| 257 | if not self._tools_enabled(request): |
| 258 | return DeltaMessage(content=delta_text) if delta_text else None |
| 259 | |
| 260 | self._buffer += delta_text |
| 261 | |
| 262 | pending_deltas: dict[int, DeltaToolCall] = {} |
| 263 | content: str | None = None |
| 264 | |
| 265 | while True: |
| 266 | if not self._in_tool_call: |
| 267 | start_idx = self._buffer.find(self.tool_call_start_token) |
| 268 | if start_idx == -1: |
| 269 | # Check for partial start token at end of buffer |
| 270 | for i in range(1, len(self.tool_call_start_token)): |
| 271 | if self._buffer.endswith(self.tool_call_start_token[:i]): |
| 272 | out = self._buffer[:-i] |
| 273 | self._buffer = self._buffer[-i:] |
| 274 | if out: |
| 275 | content = (content or "") + out |
| 276 | break |
| 277 | else: |
| 278 | out = self._buffer |
| 279 | self._buffer = "" |
| 280 | if out: |
| 281 | content = (content or "") + out |
| 282 | break |
| 283 | |
| 284 | if start_idx > 0: |
| 285 | content = (content or "") + self._buffer[:start_idx] |
| 286 | self._buffer = self._buffer[start_idx:] |
| 287 | |
| 288 | self._buffer = self._buffer[len(self.tool_call_start_token) :] |
| 289 | self._begin_tool_call() |
| 290 | continue |
| 291 | |
| 292 | # Parse tool name first |
| 293 | if not self.current_tool_name_sent: |
| 294 | nl = self._buffer.find("\n") |
| 295 | ak = self._buffer.find(self.arg_key_start) |
| 296 | end = self._buffer.find(self.tool_call_end_token) |
| 297 | candidates = [i for i in [nl, ak, end] if i != -1] |
| 298 | if not candidates: |
| 299 | break |
| 300 | cut = min(candidates) |
| 301 | tool_name = self._buffer[:cut].strip() |
| 302 | if tool_name == "" and cut == end: |
| 303 | # Handle empty tool call like `<tool_call></tool_call>`. |
| 304 | # Consume the tokens and reset state to avoid infinite loop. |
| 305 | self._buffer = self._buffer[end + len(self.tool_call_end_token) :] |
| 306 | self._finish_tool_call() |
| 307 | self._revert_last_tool_call_state() |
| 308 | continue |
| 309 | |
| 310 | if cut == nl: |
| 311 | self._buffer = self._buffer[nl + 1 :] |
| 312 | else: |
| 313 | self._buffer = self._buffer[cut:] |
| 314 | |
| 315 | self._current_tool_name = tool_name |
| 316 | self.current_tool_name_sent = True |
| 317 | self._update_tool_name(pending_deltas, tool_name) |
| 318 | continue |
| 319 | |
| 320 | assert self._current_tool_name is not None |
| 321 | |
| 322 | # Handle incremental string value streaming |
| 323 | if self._streaming_string_value: |
| 324 | val_end = self._buffer.find(self.arg_val_end) |
| 325 | if val_end != -1: |
| 326 | raw_content = self._buffer[:val_end] |
| 327 | self._buffer = self._buffer[val_end + len(self.arg_val_end) :] |
| 328 | self._streaming_string_value = False |
| 329 | self._pending_key = None |
| 330 | |
| 331 | escaped = self._json_escape_string_content(raw_content) |
| 332 | frag = escaped + '"' |
| 333 | self.streamed_args_for_tool[self.current_tool_id] += frag |
| 334 | self._update_tool_args(pending_deltas, frag) |
| 335 | continue |
| 336 | |
| 337 | # Check for partial </arg_value> at end |
| 338 | safe_len = len(self._buffer) |
| 339 | for i in range(1, len(self.arg_val_end)): |
| 340 | if self._buffer.endswith(self.arg_val_end[:i]): |
| 341 | safe_len = len(self._buffer) - i |
| 342 | break |
| 343 | |
| 344 | if safe_len > 0: |
| 345 | to_emit = self._buffer[:safe_len] |
| 346 | self._buffer = self._buffer[safe_len:] |
| 347 | escaped = self._json_escape_string_content(to_emit) |
| 348 | if escaped: |
| 349 | self.streamed_args_for_tool[self.current_tool_id] += escaped |
| 350 | self._update_tool_args(pending_deltas, escaped) |
| 351 | break |
| 352 | |
| 353 | # If we have a pending key, parse its value |
| 354 | if self._pending_key is not None: |
| 355 | val_pos = self._buffer.find(self.arg_val_start) |
| 356 | if val_pos == -1: |
| 357 | break |
| 358 | if val_pos > 0: |
| 359 | self._buffer = self._buffer[val_pos:] |
| 360 | |
| 361 | key = (self._pending_key or "").strip() |
| 362 | |
| 363 | is_string = self._is_string_type( |
| 364 | self._current_tool_name, key, request.tools |
| 365 | ) |
| 366 | |
| 367 | if is_string: |
| 368 | # String type: stream incrementally |
| 369 | self._buffer = self._buffer[len(self.arg_val_start) :] |
| 370 | |
| 371 | if key in self._seen_keys[self.current_tool_id]: |
| 372 | self._pending_key = None |
| 373 | continue |
| 374 | |
| 375 | self._seen_keys[self.current_tool_id].add(key) |
| 376 | key_json = json.dumps(key, ensure_ascii=False) |
| 377 | |
| 378 | if not self._args_started[self.current_tool_id]: |
| 379 | frag = "{" + key_json + ': "' |
| 380 | self._args_started[self.current_tool_id] = True |
| 381 | else: |
| 382 | frag = ", " + key_json + ': "' |
| 383 | |
| 384 | self.streamed_args_for_tool[self.current_tool_id] += frag |
| 385 | self._streaming_string_value = True |
| 386 | self._update_tool_args(pending_deltas, frag) |
| 387 | continue |
| 388 | |
| 389 | # Non-string type: wait for complete value |
| 390 | val_end = self._buffer.find(self.arg_val_end) |
| 391 | if val_end == -1: |
| 392 | break |
| 393 | |
| 394 | raw_val = self._buffer[len(self.arg_val_start) : val_end].strip() |
| 395 | self._buffer = self._buffer[val_end + len(self.arg_val_end) :] |
| 396 | self._pending_key = None |
| 397 | |
| 398 | frag_or_none = self._append_arg_fragment(key=key, raw_val=raw_val) |
| 399 | if frag_or_none: |
| 400 | self._update_tool_args(pending_deltas, frag_or_none) |
| 401 | continue |
| 402 | |
| 403 | # Parse next arg or close |
| 404 | end_pos = self._buffer.find(self.tool_call_end_token) |
| 405 | key_pos = self._buffer.find(self.arg_key_start) |
| 406 | if end_pos != -1 and (key_pos == -1 or end_pos < key_pos): |
| 407 | self._buffer = self._buffer[end_pos + len(self.tool_call_end_token) :] |
| 408 | frag_or_none = self._close_args_if_needed() |
| 409 | # Finalize prev_tool_call_arr with complete parsed arguments |
| 410 | if self._current_tool_name: |
| 411 | try: |
| 412 | full_args_str = self.streamed_args_for_tool[ |
| 413 | self.current_tool_id |
| 414 | ] |
| 415 | args_dict = json.loads(full_args_str) |
| 416 | self.prev_tool_call_arr[self.current_tool_id] = { |
| 417 | "name": self._current_tool_name, |
| 418 | "arguments": args_dict, |
| 419 | } |
| 420 | except (json.JSONDecodeError, IndexError) as e: |
| 421 | logger.warning( |
| 422 | "Failed to finalize tool call state for tool %d: %s", |
| 423 | self.current_tool_id, |
| 424 | e, |
| 425 | ) |
| 426 | self._finish_tool_call() |
| 427 | if frag_or_none: |
| 428 | self._update_tool_args(pending_deltas, frag_or_none) |
| 429 | continue |
| 430 | |
| 431 | if key_pos == -1: |
| 432 | break |
| 433 | if key_pos > 0: |
| 434 | self._buffer = self._buffer[key_pos:] |
| 435 | key_end = self._buffer.find(self.arg_key_end) |
| 436 | if key_end == -1: |
| 437 | break |
| 438 | key = self._buffer[len(self.arg_key_start) : key_end] |
| 439 | self._buffer = self._buffer[key_end + len(self.arg_key_end) :] |
| 440 | self._pending_key = key |
| 441 | continue |
| 442 | |
| 443 | tool_calls = list(pending_deltas.values()) |
| 444 | if content is None and len(tool_calls) == 0: |
| 445 | wants_logprobs = getattr(request, "logprobs", None) or ( |
| 446 | isinstance(request, ResponsesRequest) |
| 447 | and request.is_include_output_logprobs() |
| 448 | ) |
| 449 | if wants_logprobs: |
| 450 | return DeltaMessage(content="") |
| 451 | return None |
| 452 | return DeltaMessage(content=content, tool_calls=tool_calls) |
| 453 | |
| 454 | def _ensure_tool_state(self) -> None: |
| 455 | while len(self._tool_call_ids) <= self.current_tool_id: |
| 456 | self._tool_call_ids.append( |
| 457 | make_tool_call_id(id_type="random", func_name=None, idx=None) |
| 458 | ) |
| 459 | while len(self.streamed_args_for_tool) <= self.current_tool_id: |
| 460 | self.streamed_args_for_tool.append("") |
| 461 | while len(self.prev_tool_call_arr) <= self.current_tool_id: |
| 462 | self.prev_tool_call_arr.append({}) |
| 463 | while len(self._args_started) <= self.current_tool_id: |
| 464 | self._args_started.append(False) |
| 465 | while len(self._args_closed) <= self.current_tool_id: |
| 466 | self._args_closed.append(False) |
| 467 | while len(self._seen_keys) <= self.current_tool_id: |
| 468 | self._seen_keys.append(set()) |
| 469 | |
| 470 | def _begin_tool_call(self) -> None: |
| 471 | if self.current_tool_id == -1: |
| 472 | self.current_tool_id = 0 |
| 473 | else: |
| 474 | self.current_tool_id += 1 |
| 475 | self._ensure_tool_state() |
| 476 | self.current_tool_name_sent = False |
| 477 | self._current_tool_name = None |
| 478 | self._pending_key = None |
| 479 | self._streaming_string_value = False |
| 480 | self._in_tool_call = True |
| 481 | |
| 482 | def _finish_tool_call(self) -> None: |
| 483 | self._in_tool_call = False |
| 484 | self._current_tool_name = None |
| 485 | self._pending_key = None |
| 486 | self._streaming_string_value = False |
| 487 | |
| 488 | def _revert_last_tool_call_state(self) -> None: |
| 489 | """Revert the state allocation for the last tool call.""" |
| 490 | if self.current_tool_id < 0: |
| 491 | return |
| 492 | self._tool_call_ids.pop() |
| 493 | self.streamed_args_for_tool.pop() |
| 494 | self.prev_tool_call_arr.pop() |
| 495 | self._args_started.pop() |
| 496 | self._args_closed.pop() |
| 497 | self._seen_keys.pop() |
| 498 | self.current_tool_id -= 1 |
| 499 | |
| 500 | def _get_or_create_delta(self, pending: dict[int, DeltaToolCall]) -> DeltaToolCall: |
| 501 | idx = self.current_tool_id |
| 502 | if idx not in pending: |
| 503 | pending[idx] = DeltaToolCall( |
| 504 | index=idx, |
| 505 | function=DeltaFunctionCall(), |
| 506 | ) |
| 507 | delta = pending[idx] |
| 508 | assert delta.function is not None |
| 509 | return delta |
| 510 | |
| 511 | def _update_tool_name( |
| 512 | self, pending: dict[int, DeltaToolCall], tool_name: str |
| 513 | ) -> None: |
| 514 | self.prev_tool_call_arr[self.current_tool_id] = { |
| 515 | "name": self._current_tool_name, |
| 516 | "arguments": {}, |
| 517 | } |
| 518 | delta = self._get_or_create_delta(pending) |
| 519 | delta.id = self._tool_call_ids[self.current_tool_id] |
| 520 | delta.type = "function" |
| 521 | assert delta.function is not None |
| 522 | delta.function.name = tool_name |
| 523 | if delta.function.arguments is None: |
| 524 | delta.function.arguments = "" |
| 525 | |
| 526 | @staticmethod |
| 527 | def _complete_json_prefix( |
| 528 | json_prefix: str, |
| 529 | allowed_partial_types: Allow, |
| 530 | ) -> dict | None: |
| 531 | """Complete a partial JSON prefix into a valid JSON object. |
| 532 | |
| 533 | Returns (formatted_prefix, parsed_dict) or None on failure. |
| 534 | |
| 535 | Note: ``partial_json_parser`` strips trailing whitespace before |
| 536 | parsing (``complete.py:20``), which means the returned slice is |
| 537 | shorter than ``json_prefix`` when it has trailing whitespace. |
| 538 | Since the parser controls the construction of the json_prefix value, |
| 539 | this code relies on it being a valid prefix and we only use the fix for |
| 540 | the completion of the JSON object. |
| 541 | """ |
| 542 | try: |
| 543 | _, partial_str_completion = partial_json_parser.core.complete.fix( |
| 544 | json_prefix, |
| 545 | allowed_partial_types, |
| 546 | ) |
| 547 | return json.loads(json_prefix + partial_str_completion) |
| 548 | except Exception: |
| 549 | return None |
| 550 | |
| 551 | def _update_tool_args( |
| 552 | self, pending: dict[int, DeltaToolCall], fragment: str |
| 553 | ) -> None: |
| 554 | result = self._complete_json_prefix( |
| 555 | self.streamed_args_for_tool[self.current_tool_id], |
| 556 | Allow.ALL, |
| 557 | ) |
| 558 | if result is not None: |
| 559 | self.prev_tool_call_arr[self.current_tool_id]["arguments"] = result |
| 560 | delta = self._get_or_create_delta(pending) |
| 561 | assert delta.function is not None |
| 562 | if delta.function.arguments is None: |
| 563 | delta.function.arguments = "" |
| 564 | delta.function.arguments += fragment |
| 565 | |
| 566 | def _append_arg_fragment( |
| 567 | self, |
| 568 | *, |
| 569 | key: str, |
| 570 | raw_val: str, |
| 571 | ) -> str | None: |
| 572 | key = key.strip() |
| 573 | if not key: |
| 574 | return None |
| 575 | if key in self._seen_keys[self.current_tool_id]: |
| 576 | return None |
| 577 | |
| 578 | # This function is only called for non-string types (already checked |
| 579 | # by _is_string_type in the caller), so we always deserialize. |
| 580 | val_obj: Any = self._deserialize(raw_val) |
| 581 | |
| 582 | key_json = json.dumps(key, ensure_ascii=False) |
| 583 | val_json = json.dumps(val_obj, ensure_ascii=False) |
| 584 | |
| 585 | if not self._args_started[self.current_tool_id]: |
| 586 | fragment = "{" + key_json + ": " + val_json |
| 587 | self._args_started[self.current_tool_id] = True |
| 588 | else: |
| 589 | fragment = ", " + key_json + ": " + val_json |
| 590 | |
| 591 | self._seen_keys[self.current_tool_id].add(key) |
| 592 | self.streamed_args_for_tool[self.current_tool_id] += fragment |
| 593 | return fragment |
| 594 | |
| 595 | def _close_args_if_needed(self) -> str | None: |
| 596 | if self._args_closed[self.current_tool_id]: |
| 597 | return None |
| 598 | self._args_closed[self.current_tool_id] = True |
| 599 | if not self._args_started[self.current_tool_id]: |
| 600 | fragment = "{}" |
| 601 | self.streamed_args_for_tool[self.current_tool_id] = fragment |
| 602 | else: |
| 603 | fragment = "}" |
| 604 | self.streamed_args_for_tool[self.current_tool_id] += fragment |
| 605 | return fragment |