vllm.tool_parsers.hy_v4_tool_parser ¶

class ExtractResult(TypedDict):
    """Return of ``extract_tool_calls`` (non-streaming)."""

    tools_called: bool
    content: str | None
    tool_calls: list[ToolCallDict]

class HYV4ToolExtractor:
    """Pure tool-call parsing logic for HYV4. Returns plain dicts; no vLLM types.

    Holds only parser state: the structural token strings/ids, the compiled
    regexes, and the streaming incremental state. Stateless argument parsing is
    delegated to the module-level utilities above.

    Streaming has two entry paths sharing one incremental argument parser,
    selected by the ``guided`` flag of :meth:`extract_tool_calls_streaming`:
    normal / auto tool choice detects markers on atomic token ids so ordinary
    ``<`` in content is never withheld, while structural-tag guided decoding
    detects them on the string level because the grammar may split a marker
    across sub-word tokens.
    """

    def __init__(self, vocab: dict[str, int], token_suffix: str, strict: bool):
        self._strict = strict

        # Kept so the wrapper can build a matching structural tag for
        # constrained decoding without re-scanning the tokenizer vocab.
        self.token_suffix: str = token_suffix

        # Running state read by the vLLM serving layer (mirrored by the wrapper).
        self.prev_tool_call_arr: list[dict] = []
        self.current_tool_id: int = -1
        self.streamed_args_for_tool: list[str] = []

        # Streaming state: send tool name first, then return arguments at once
        self._streaming_tool_name: str | None = None  # tool name being streamed

        # State fields for incremental argument streaming
        self._completed_args: dict = {}  # closed {key: parsed_value}
        self._current_arg_key: str | None = None  # key being collected
        self._current_arg_is_string: bool = False  # is current arg pure string?
        self._streamed_json_len: int = 0  # bytes of JSON already sent

        self.tool_calls_start_token: str = f"<tool_calls{token_suffix}>"
        self.tool_calls_end_token: str = f"</tool_calls{token_suffix}>"

        self.tool_call_start_token: str = f"<tool_call{token_suffix}>"
        self.tool_call_end_token: str = f"</tool_call{token_suffix}>"

        self.arg_key_start_token: str = f"<arg_key{token_suffix}>"
        self.arg_key_end_token: str = f"</arg_key{token_suffix}>"

        self.arg_value_start_token: str = f"<arg_value{token_suffix}>"
        self.arg_value_end_token: str = f"</arg_value{token_suffix}>"

        logger.debug(
            "HYV4ToolExtractor structural tokens: %s ... %s, suffix=%r, strict=%s",
            self.tool_calls_start_token,
            self.tool_calls_end_token,
            token_suffix,
            strict,
        )

        # The suffix comes from the tokenizer vocab, so it must not be treated
        # as a pattern.
        self.tool_call_regex = re.compile(
            rf"{re.escape(self.tool_call_start_token)}(.*?)"
            rf"{re.escape(self.tool_call_end_token)}",
            re.DOTALL,
        )

        self.func_args_regex = re.compile(
            rf"{re.escape(self.arg_key_start_token)}(.*?)"
            rf"{re.escape(self.arg_key_end_token)}"
            rf"{re.escape(self.arg_value_start_token)}(.*?)"
            rf"{re.escape(self.arg_value_end_token)}",
            re.DOTALL,
        )

        self.tool_calls_start_token_id = vocab.get(self.tool_calls_start_token)
        self.tool_calls_end_token_id = vocab.get(self.tool_calls_end_token)

        self.tool_call_start_token_id = vocab.get(self.tool_call_start_token)
        self.tool_call_end_token_id = vocab.get(self.tool_call_end_token)
        self._buffer = ""

        if (
            self.tool_calls_start_token_id is None
            or self.tool_calls_end_token_id is None
        ):
            raise RuntimeError(
                "HYV4 tool extractor could not locate tool call "
                "start/end tokens in the tokenizer!"
            )

    def _extract_tool_calls(
        self,
        model_output: str,
        tools: list[ToolSchema] | None,
    ) -> list[ToolCallDict]:
        try:
            tool_calls: list[ToolCallDict] = []
            function_calls = self.tool_call_regex.findall(model_output)
            for tool_call_body in function_calls:
                arg_start = tool_call_body.find(self.arg_key_start_token)
                if arg_start == -1:
                    function_name = tool_call_body.strip()
                    function_args = ""
                else:
                    function_name = tool_call_body[:arg_start].strip()
                    function_args = tool_call_body[arg_start:].strip()

                arg_pairs = self.func_args_regex.findall(function_args)
                arg_dict = {}
                for key, value in arg_pairs:
                    arg_dict[key] = _parse_value(value, function_name, key, tools)
                tool_calls.append(
                    ToolCallDict(
                        name=function_name,
                        arguments=json.dumps(arg_dict, ensure_ascii=False),
                    )
                )
            return tool_calls
        except Exception:
            logger.exception("Error in extracting tool call from response.")
            return []

    def _extract_tool_calls_strict(
        self,
        model_output: str,
        tools: list[ToolSchema] | None,
    ) -> list[ToolCallDict]:
        start_idx = model_output.find(self.tool_calls_start_token)
        end_idx = model_output.find(self.tool_calls_end_token)
        if start_idx == -1 or end_idx == -1 or end_idx < start_idx:
            raise ValueError("Malformed tool block wrapper.")

        tool_block = model_output[
            start_idx + len(self.tool_calls_start_token) : end_idx
        ]
        if tool_block.count(self.tool_call_start_token) != tool_block.count(
            self.tool_call_end_token
        ):
            raise ValueError("Malformed tool_call wrapper count.")

        tool_calls: list[ToolCallDict] = []
        function_calls = self.tool_call_regex.findall(tool_block)
        if tool_block.count(self.tool_call_start_token) != len(function_calls):
            raise ValueError("Malformed tool_call body.")

        for tool_call_body in function_calls:
            arg_start = tool_call_body.find(self.arg_key_start_token)
            if arg_start == -1:
                function_name = tool_call_body.strip()
                function_args = ""
            else:
                function_name = tool_call_body[:arg_start].strip()
                function_args = tool_call_body[arg_start:].strip()

            if not function_name:
                raise ValueError("Empty function name in tool call.")

            if function_args.count(self.arg_key_start_token) != function_args.count(
                self.arg_key_end_token
            ) or function_args.count(self.arg_value_start_token) != function_args.count(
                self.arg_value_end_token
            ):
                raise ValueError("Malformed argument tag count.")

            arg_pairs = self.func_args_regex.findall(function_args)
            remainder = self.func_args_regex.sub("", function_args)
            if remainder.strip():
                raise ValueError("Unparsed argument payload remains.")

            arg_dict = {}
            for key, value in arg_pairs:
                arg_dict[key] = _parse_value(value, function_name, key, tools)

            tool_calls.append(
                ToolCallDict(
                    name=function_name,
                    arguments=json.dumps(arg_dict, ensure_ascii=False),
                )
            )

        if not tool_calls:
            raise ValueError("No valid tool calls extracted.")

        return tool_calls

    def extract_tool_calls(
        self,
        model_output: str,
        tools: list[ToolSchema] | None,
    ) -> ExtractResult:
        # sanity check; avoid unnecessary processing
        if self.tool_calls_start_token not in model_output:
            return ExtractResult(
                tools_called=False, content=model_output, tool_calls=[]
            )
        else:
            try:
                if self._strict:
                    tool_calls = self._extract_tool_calls_strict(model_output, tools)
                else:
                    tool_calls = self._extract_tool_calls(model_output, tools)
                if not tool_calls:
                    raise ValueError("No valid tool calls extracted.")

                s_index = model_output.find(self.tool_calls_start_token)
                content = model_output[:s_index] if s_index != -1 else model_output
                return ExtractResult(
                    tools_called=True,
                    content=content if content else None,
                    tool_calls=tool_calls,
                )

            except Exception:
                logger.warning(
                    "HYV4ToolExtractor detected malformed tool output; "
                    "falling back to raw content."
                )
                return ExtractResult(
                    tools_called=False,
                    content=model_output,
                    tool_calls=[],
                )

    def _reset_streaming_tool_state(self):
        """Reset the streaming state for a single tool call."""
        self._streaming_tool_name = None
        self._completed_args = {}
        self._current_arg_key = None
        self._current_arg_is_string = False
        self._streamed_json_len = 0

    def extract_tool_calls_streaming(
        self,
        previous_text: str,
        current_text: str,
        delta_text: str,
        previous_token_ids: Sequence[int],
        current_token_ids: Sequence[int],
        delta_token_ids: Sequence[int],
        tools: list[ToolSchema] | None,
        *,
        guided: bool = False,
    ) -> StreamDelta | None:
        """Consume one engine delta and emit content / tool-call deltas.

        Two entry paths share the same incremental argument parser:
        - ``guided=False`` (normal / auto tool choice): the structural markers
          arrive as atomic tokenizer tokens, so detection is done on token ids
          and ordinary ``<`` in model content is never withheld.
        - ``guided=True`` (structural-tag constrained decoding): the grammar
          only constrains strings, so ``<tool_calls...>`` may be split across
          sub-word tokens and must be detected on the string level.

        Args:
            previous_text: Text decoded before this delta.
            current_text: Full text decoded so far.
            delta_text: Text decoded in this delta.
            previous_token_ids: Token ids before this delta.
            current_token_ids: Token ids decoded so far.
            delta_token_ids: Token ids in this delta.
            tools: Plain tool schemas used for argument typing.
            guided: Whether guided decoding may split markers into sub-word
                tokens.

        Returns:
            A streaming delta carrying content and/or the tool calls drained
            from the buffer, or None when nothing can be emitted yet.
        """
        content_delta: str | None = None
        tool_calls: list[StreamToolCall] = []

        if guided:
            # Guided path: a structural-tag grammar constrains strings, not
            # token ids, so the atomic ``tool_calls_start_token_id`` never
            # appears in the token stream and an id-based check would
            # misclassify every tool delta as plain content.
            marker = self.tool_calls_start_token
            marker_pos = current_text.find(marker)

            if marker_pos == -1:
                # Marker not seen yet. The tail of ``current_text`` might still
                # be a partial prefix of the marker split across deltas; hold
                # that tail back and only stream the safe portion as content.
                safe_end = len(current_text) - partial_tag_overlap(current_text, marker)
                # Emit only the newly-arrived, safe content (delta relative to
                # what was previously streamed).
                content = current_text[len(previous_text) : safe_end]
                return StreamDelta(
                    content=content if content else None,
                    tool_calls=[],
                )

            # Marker present. Stream any content that precedes it (only the
            # part we have not already emitted), then buffer the tool payload.
            #
            # The payload is everything after the marker in the *full* text. We
            # keep the buffer in sync with it by only appending the portion that
            # arrived in this delta, mirroring the incremental ``+=`` semantics
            # while remaining correct when the marker itself was split across
            # sub-word tokens.
            content_end = marker_pos
            content = current_text[len(previous_text) : content_end]

            payload_start = marker_pos + len(marker)
            # How much of the payload had already been buffered before this
            # delta.
            prev_payload_len = max(len(previous_text) - payload_start, 0)
            new_payload = current_text[payload_start + prev_payload_len :]
            self._buffer += new_payload

            if content:
                content_delta = content
        else:
            # Atomic-token path: normal / auto generation emits the HYV4
            # structural markers as single tokenizer tokens. Pass ordinary
            # content through exactly as decoded, including comparison
            # operators such as ``<``.
            if self.tool_calls_start_token_id not in current_token_ids:
                return StreamDelta(content=delta_text, tool_calls=[])

            if self.tool_calls_start_token in delta_text:
                text_parts = delta_text.split(self.tool_calls_start_token)
                self._buffer += text_parts[-1]
                if text_parts[0]:
                    content_delta = text_parts[0]
            else:
                self._buffer += delta_text

        # Encountered finish, extract valid arguments
        if (
            current_text.find(self.tool_call_end_token + self.tool_calls_end_token)
            != -1
            and self._buffer.find(self.tool_call_end_token) == -1
        ):
            self._buffer += self.tool_call_end_token + self.tool_calls_end_token

        # Drain every complete tool call that is already buffered so a single
        # engine delta can carry content + N tool calls (batched decode).
        while True:
            cur_text = self._buffer

            # Haven't encountered tool_call start tag yet; drop the inter-tag
            # padding but keep any tail that may be a partial tool_call start
            # tag (guided decoding can split it across deltas).
            start_idx = cur_text.find(self.tool_call_start_token)
            if start_idx == -1 and self._streaming_tool_name is None:
                keep = partial_tag_overlap(cur_text, self.tool_call_start_token)
                self._buffer = cur_text[len(cur_text) - keep :] if keep else ""
                break

            # === Phase 1: Detect tool name (send when first arg_key is seen) =
            # ``name_new`` mirrors the old ``name_delta is not None``: it marks
            # the chunk that first emits this tool's name, which is the only
            # chunk the wrapper mints an id + type for.
            name_new = False
            pending_name: str | None = None
            if self._streaming_tool_name is None:
                arg_idx = cur_text.find(self.arg_key_start_token, start_idx)
                end_idx = cur_text.find(self.tool_call_end_token, start_idx)
                if arg_idx == -1 and end_idx == -1:
                    # tool name not yet complete; keep buffering from
                    # tool_call_start
                    self._buffer = cur_text[start_idx:]
                    break

                name_start = start_idx + len(self.tool_call_start_token)
                name_end = arg_idx if arg_idx != -1 else end_idx

                tool_name = cur_text[name_start:name_end].strip()
                self._streaming_tool_name = tool_name

                if arg_idx != -1:
                    self._buffer = cur_text[arg_idx:]
                else:
                    self._buffer = cur_text[end_idx:]

                # Increment tool_id and mark that a name chunk should be emitted
                self.current_tool_id += 1
                name_new = True
                pending_name = tool_name

                # Check if buffer already has complete arguments
                # (all-in-one-delta); otherwise only the name is ready now.
                if self.tool_call_end_token not in self._buffer:
                    tool_calls.append(
                        StreamToolCall(
                            index=self.current_tool_id,
                            new=True,
                            name=tool_name,
                            arguments=None,
                        )
                    )
                    break
                # Buffer already has a complete tool call; fall through to
                # phase 2 below.

            # === Phase 2: Incremental argument streaming ===
            result = self._extract_streaming_incremental(name_new, pending_name, tools)
            if result:
                tool_calls.extend(result["tool_calls"])

            # Continue only after the current tool call has completed and the
            # remaining structural payload starts with another tool call.
            # A literal ``<tool_call>`` may legitimately appear inside an open
            # string argument; in that case the current tool name is still set
            # and the buffer has not been consumed, so continuing would spin on
            # the same buffer forever.
            if self._streaming_tool_name is None and self._buffer.startswith(
                self.tool_call_start_token
            ):
                continue
            break

        if content_delta is not None or tool_calls:
            return StreamDelta(content=content_delta, tool_calls=tool_calls)
        return None

    def _make_args_delta(self, argument_diff: str) -> StreamDelta:
        """Build an args-only streaming delta (no id/type -> new=False)."""
        return StreamDelta(
            content=None,
            tool_calls=[
                StreamToolCall(
                    index=self.current_tool_id,
                    new=False,
                    name=None,
                    arguments=argument_diff,
                )
            ],
        )

    def _extract_streaming_incremental(
        self,
        name_new: bool,
        pending_name: str | None,
        tools: list[ToolSchema] | None,
    ) -> StreamDelta | None:
        """Incremental phase-2: scan tags in buffer, emit JSON diffs.

        Strategy:
        - Track completed args and emit each one as a JSON fragment.
        - For string-typed args, stream the value character-by-character.
        - Withhold the closing ``}`` until ``</tool_call>`` is seen.

        We build JSON manually via fragments rather than using json.dumps
        with a cursor, because json.dumps of partial-vs-full string values
        produces incompatible prefixes (e.g. ``""}`` vs ``"Hello"}``).
        """
        buf = self._buffer
        is_complete = self.tool_call_end_token in buf

        if is_complete:
            end_idx = buf.find(self.tool_call_end_token)
            args_text = buf[:end_idx]
            remaining = buf[end_idx + len(self.tool_call_end_token) :]
        else:
            args_text = buf
            remaining = ""

        # --- scan all fully closed kv pairs ---
        arg_pairs = self.func_args_regex.findall(args_text)
        for key, value in arg_pairs:
            key = key.strip()
            if key not in self._completed_args:
                self._completed_args[key] = _parse_value(
                    value, self._streaming_tool_name or "", key, tools
                )

        # --- detect partial (unclosed) kv at the tail ---
        last_closed_end = 0
        for m in self.func_args_regex.finditer(args_text):
            last_closed_end = m.end()
        tail = args_text[last_closed_end:]

        partial_key: str | None = None
        partial_value: str | None = None

        ak_start = tail.find(self.arg_key_start_token)
        if ak_start != -1:
            ak_end = tail.find(
                self.arg_key_end_token,
                ak_start + len(self.arg_key_start_token),
            )
            if ak_end != -1:
                partial_key = tail[
                    ak_start + len(self.arg_key_start_token) : ak_end
                ].strip()
                self._current_arg_key = partial_key
                self._current_arg_is_string = _is_only_string_type(
                    self._streaming_tool_name or "",
                    partial_key,
                    tools,
                )

                av_start = tail.find(self.arg_value_start_token, ak_end)
                if av_start != -1:
                    val_content_start = av_start + len(self.arg_value_start_token)
                    if self._current_arg_is_string:
                        raw_value = tail[val_content_start:]
                        # Guided decoding can split ``</arg_value>`` across
                        # deltas; hold back a partial closing tag so it is not
                        # streamed as part of the value.
                        keep = len(raw_value) - partial_tag_overlap(
                            raw_value, self.arg_value_end_token
                        )
                        partial_value = raw_value[:keep]
            else:
                # key not yet closed
                self._current_arg_key = None
                self._current_arg_is_string = False

        # --- build the current JSON snapshot as a string ---
        # We construct JSON manually so we can precisely control
        # what gets sent incrementally.
        snapshot_parts: list[str] = []
        for k, v in self._completed_args.items():
            k_json = json.dumps(k, ensure_ascii=False)
            v_json = json.dumps(v, ensure_ascii=False)
            snapshot_parts.append(f"{k_json}: {v_json}")

        if partial_key is not None and partial_value is not None:
            k_json = json.dumps(partial_key, ensure_ascii=False)
            # For string partial value, we build the JSON string
            # WITHOUT the closing quote, so the prefix stays stable
            # as the value grows.  The closing `"` and `}` will be
            # sent when the value or tool_call closes.
            escaped_val = json.dumps(partial_value, ensure_ascii=False)[1:-1]
            # Note: no closing " here - it's appended only on close
            snapshot_parts.append(f'{k_json}: "{escaped_val}')

        snapshot = "{" + ", ".join(snapshot_parts) + "}"

        # --- compute diff ---
        argument_diff: str | None = None

        if is_complete:
            # Tool call finished - send everything remaining.
            # Build final snapshot with proper JSON (all values closed).
            final_args = dict(self._completed_args)
            final_json = json.dumps(final_args, ensure_ascii=False)
            if self._streamed_json_len < len(final_json):
                argument_diff = final_json[self._streamed_json_len :]
            self._streamed_json_len = len(final_json)

            # Record into prev_tool_call_arr
            self.prev_tool_call_arr.append(
                {
                    "name": self._streaming_tool_name,
                    "arguments": final_args,
                }
            )
            self.streamed_args_for_tool.append(final_json)

            self._reset_streaming_tool_state()
            self._buffer = remaining
        else:
            # Still in progress - withhold the tail.
            # For open strings: snapshot ends with ...partial_val}
            #   we withhold "}" (1 char) - the missing closing " will
            #   be sent when the value closes.
            # For no open string: snapshot ends with ...value"}
            #   we withhold "}" (1 char).
            end = len(snapshot) - 1  # exclude trailing "}"
            if end > self._streamed_json_len:
                argument_diff = snapshot[self._streamed_json_len : end]
                self._streamed_json_len = end

        # --- construct return dict ---
        if name_new and argument_diff:
            return StreamDelta(
                content=None,
                tool_calls=[
                    StreamToolCall(
                        index=self.current_tool_id,
                        new=True,
                        name=pending_name,
                        arguments=argument_diff,
                    )
                ],
            )
        elif name_new:
            return StreamDelta(
                content=None,
                tool_calls=[
                    StreamToolCall(
                        index=self.current_tool_id,
                        new=True,
                        name=pending_name,
                        arguments=None,
                    )
                ],
            )
        elif argument_diff:
            return self._make_args_delta(argument_diff)
        else:
            return None