xgrammar-gemma4-test/gemma_parser.py

"""Minimal Gemma4 tool-call parser.

Pure-function port of sglang's Gemma4Detector parsing helpers
(sglang/python/sglang/srt/function_call/gemma4_detector.py), so generated
outputs can be validated exactly the way sglang would parse them, without
importing sglang.
"""

TOOL_CALL_START = "<|tool_call>"
TOOL_CALL_END = "<tool_call|>"
STRING_DELIM = '<|"|>'


def _parse_value(value_str: str) -> object:
    value_str = value_str.strip()
    if not value_str:
        return value_str
    if value_str == "true":
        return True
    if value_str == "false":
        return False
    try:
        if "." in value_str:
            return float(value_str)
        return int(value_str)
    except ValueError:
        pass
    return value_str


def _parse_array(arr_str: str) -> list:
    items: list = []
    i = 0
    n = len(arr_str)
    while i < n:
        while i < n and arr_str[i] in (" ", ",", "\n", "\t"):
            i += 1
        if i >= n:
            break
        if arr_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
            i += len(STRING_DELIM)
            end_pos = arr_str.find(STRING_DELIM, i)
            if end_pos == -1:
                items.append(arr_str[i:])
                break
            items.append(arr_str[i:end_pos])
            i = end_pos + len(STRING_DELIM)
        elif arr_str[i] == "{":
            depth = 1
            obj_start = i + 1
            i += 1
            while i < n and depth > 0:
                if arr_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
                    i += len(STRING_DELIM)
                    next_delim = arr_str.find(STRING_DELIM, i)
                    i = next_delim + len(STRING_DELIM) if next_delim != -1 else n
                    continue
                if arr_str[i] == "{":
                    depth += 1
                elif arr_str[i] == "}":
                    depth -= 1
                i += 1
            items.append(parse_args(arr_str[obj_start : i - 1]))
        elif arr_str[i] == "[":
            depth = 1
            sub_start = i + 1
            i += 1
            while i < n and depth > 0:
                if arr_str[i] == "[":
                    depth += 1
                elif arr_str[i] == "]":
                    depth -= 1
                i += 1
            items.append(_parse_array(arr_str[sub_start : i - 1]))
        else:
            val_start = i
            while i < n and arr_str[i] not in (",", "]"):
                i += 1
            items.append(_parse_value(arr_str[val_start:i]))
    return items


def parse_args(args_str: str) -> dict:
    """Parse Gemma4's key:value argument format into a Python dict."""
    if not args_str or not args_str.strip():
        return {}
    result: dict = {}
    i = 0
    n = len(args_str)
    while i < n:
        while i < n and args_str[i] in (" ", ",", "\n", "\t"):
            i += 1
        if i >= n:
            break
        key_start = i
        while i < n and args_str[i] != ":":
            i += 1
        if i >= n:
            break
        key = args_str[key_start:i].strip()
        i += 1
        if i >= n:
            result[key] = ""
            break
        while i < n and args_str[i] in (" ", "\n", "\t"):
            i += 1
        if i >= n:
            result[key] = ""
            break
        if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
            i += len(STRING_DELIM)
            val_start = i
            end_pos = args_str.find(STRING_DELIM, i)
            if end_pos == -1:
                result[key] = args_str[val_start:]
                break
            result[key] = args_str[val_start:end_pos]
            i = end_pos + len(STRING_DELIM)
        elif args_str[i] == "{":
            depth = 1
            obj_start = i + 1
            i += 1
            while i < n and depth > 0:
                if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
                    i += len(STRING_DELIM)
                    next_delim = args_str.find(STRING_DELIM, i)
                    i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
                    continue
                if args_str[i] == "{":
                    depth += 1
                elif args_str[i] == "}":
                    depth -= 1
                i += 1
            result[key] = parse_args(args_str[obj_start : i - 1])
        elif args_str[i] == "[":
            depth = 1
            arr_start = i + 1
            i += 1
            while i < n and depth > 0:
                if args_str[i : i + len(STRING_DELIM)] == STRING_DELIM:
                    i += len(STRING_DELIM)
                    next_delim = args_str.find(STRING_DELIM, i)
                    i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
                    continue
                if args_str[i] == "[":
                    depth += 1
                elif args_str[i] == "]":
                    depth -= 1
                i += 1
            result[key] = _parse_array(args_str[arr_start : i - 1])
        else:
            val_start = i
            while i < n and args_str[i] not in (",", "}", "]"):
                i += 1
            result[key] = _parse_value(args_str[val_start:i])
    return result


def _find_matching_brace(text: str) -> int:
    """Index of the matching '}' in text (which starts just after '{'), or -1."""
    depth = 1
    i = 0
    n = len(text)
    dlen = len(STRING_DELIM)
    while i < n and depth > 0:
        if text[i : i + dlen] == STRING_DELIM:
            i += dlen
            next_delim = text.find(STRING_DELIM, i)
            if next_delim == -1:
                return -1
            i = next_delim + dlen
            continue
        if text[i] == "{":
            depth += 1
        elif text[i] == "}":
            depth -= 1
        i += 1
    return (i - 1) if depth == 0 else -1


def extract_tool_calls(text: str) -> list:
    """Extract [(func_name, args_dict_or_None)] from text.

    args is None when the call block is malformed (unparseable), mirroring
    where sglang's detector would fail.
    """
    results = []
    search_from = 0
    while True:
        start = text.find(TOOL_CALL_START, search_from)
        if start == -1:
            break
        end = text.find(TOOL_CALL_END, start)
        if end == -1:
            # unterminated call block
            results.append((None, None))
            break
        inner = text[start + len(TOOL_CALL_START) : end]
        if inner.startswith("call:"):
            brace = inner.find("{")
            if brace != -1:
                func_name = inner[5:brace]
                args_content = inner[brace + 1 :]
                match_idx = _find_matching_brace(args_content)
                args_str = args_content[:match_idx] if match_idx != -1 else args_content
                try:
                    results.append((func_name, parse_args(args_str)))
                except Exception:
                    results.append((func_name, None))
            else:
                results.append((None, None))
        else:
            results.append((None, None))
        search_from = end + len(TOOL_CALL_END)
    return results