tiny_gdn/tools.py

"""Nemotron Cascade-2 tool protocol (XML in ChatML content).

Tercet-R is trained on `nvidia/Nemotron-Cascade-2-SFT-Data`, which inlines
available tools, calls, and results as `<tools>` / `<tool_call>` /
`<tool_response>` in message text. That is not OpenAI `tool_calls` JSON and
not this tokenizer's unused `<|tool_call|>` / `<|tool_response|>` specials.

This module matches NVIDIA's Cascade-2 chat template:
https://huggingface.co/nvidia/Nemotron-Cascade-2-30B-A3B/blob/main/chat_template.jinja
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any


TOOL_CALL_OPEN = "<tool_call>"
TOOL_CALL_CLOSE = "</tool_call>"
TOOL_RESPONSE_OPEN = "<tool_response>"
TOOL_RESPONSE_CLOSE = "</tool_response>"
TOOLS_OPEN = "<tools>"
TOOLS_CLOSE = "</tools>"
THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"

CHATML_START_RE = re.compile(r"^<\|im_start\|>[A-Za-z]+\n")
CHATML_END_RE = re.compile(r"\n?<\|im_end\|>\s*$")
TOOL_CALL_BLOCK_RE = re.compile(
    rf"{re.escape(TOOL_CALL_OPEN)}(.*?){re.escape(TOOL_CALL_CLOSE)}",
    re.DOTALL,
)
FUNCTION_BLOCK_RE = re.compile(
    r"<function=([^>\s]+)>(.*?)</function>",
    re.DOTALL,
)
PARAMETER_BLOCK_RE = re.compile(
    r"<parameter=([^>\s]+)>\n?(.*?)\n?</parameter>",
    re.DOTALL,
)
FUNCTION_NAME_RE = re.compile(r"<function=([^>\s]+)")

INFERENCE_ROLE_MAP = {
    "system": "system",
    "user": "user",
    "human": "user",
    "assistant": "assistant",
    "gpt": "assistant",
    "tool": "tool",
    "function": "tool",
}

TOOLS_PREAMBLE = "# Tools\n\nYou have access to the following functions:\n\n"
TOOL_CALL_INSTRUCTIONS = (
    "\n\nIf you choose to call a function ONLY reply in the following format "
    "with NO suffix:\n\n"
    "<tool_call>\n"
    "<function=example_function_name>\n"
    "<parameter=example_parameter_1>\n"
    "value_1\n"
    "</parameter>\n"
    "<parameter=example_parameter_2>\n"
    "This is the value for the second parameter\n"
    "that can span\n"
    "multiple lines\n"
    "</parameter>\n"
    "</function>\n"
    "</tool_call>\n\n"
    "<IMPORTANT>\n"
    "Reminder:\n"
    "- Function calls MUST follow the specified format: an inner "
    "<function=...></function> block must be nested within "
    "<tool_call></tool_call> XML tags\n"
    "- Required parameters MUST be specified\n"
    "- You may provide optional reasoning for your function call in natural "
    "language BEFORE the function call, but NOT after\n"
    "- If there is no function call available, answer the question like "
    "normal with your current knowledge and do not tell the user about "
    "function calls\n"
    "</IMPORTANT>"
)


@dataclass(frozen=True)
class ParsedToolCall:
    name: str
    arguments: dict[str, Any]
    raw: str


def message_text(raw_content: Any) -> str:
    if raw_content is None:
        return ""
    if isinstance(raw_content, str):
        return raw_content
    if isinstance(raw_content, list):
        parts: list[str] = []
        for item in raw_content:
            if isinstance(item, str):
                parts.append(item)
                continue
            if not isinstance(item, dict):
                continue
            part_type = item.get("type")
            if part_type in {None, "text", "input_text", "output_text"}:
                text = item.get("text")
                if isinstance(text, str):
                    parts.append(text)
        return "".join(parts)
    return str(raw_content)


def strip_leaked_chatml(content: str) -> str:
    text = content.strip()
    while True:
        match = CHATML_START_RE.match(text)
        if match is None:
            break
        text = text[match.end() :]
    text = CHATML_END_RE.sub("", text)
    return text.strip()


def has_tools_block(content: str) -> bool:
    return TOOLS_OPEN in content


def wrap_tool_response(content: str) -> str:
    text = strip_leaked_chatml(content)
    if TOOL_RESPONSE_OPEN in text:
        return text
    return f"{TOOL_RESPONSE_OPEN}\n\n\n{text}\n{TOOL_RESPONSE_CLOSE}"


def _xml_value(value: Any) -> str:
    if isinstance(value, dict) or (
        isinstance(value, (list, tuple)) and not isinstance(value, (str, bytes))
    ):
        return json.dumps(value, ensure_ascii=False)
    if value is True or value is False or value is None:
        return str(value)
    return str(value)


def _render_extra_keys(payload: dict[str, Any], handled: set[str]) -> str:
    chunks: list[str] = []
    for key, value in payload.items():
        if key in handled:
            continue
        chunks.append(f"\n<{key}>{_xml_value(value)}</{key}>")
    return "".join(chunks)


def _unwrap_tool(raw_tool: Any) -> dict[str, Any]:
    if not isinstance(raw_tool, dict):
        raise ValueError("Each tool must be an object")
    if isinstance(raw_tool.get("function"), dict):
        tool = dict(raw_tool["function"])
    else:
        tool = dict(raw_tool)
    name = tool.get("name")
    if not isinstance(name, str) or not name.strip():
        raise ValueError("Tool is missing a function name")
    tool["name"] = name.strip()
    return tool


def coerce_tools(raw_tools: Any) -> list[dict[str, Any]]:
    if raw_tools is None:
        return []
    if isinstance(raw_tools, str):
        text = raw_tools.strip()
        if not text:
            return []
        raw_tools = json.loads(text)
    if isinstance(raw_tools, dict):
        raw_tools = [raw_tools]
    if not isinstance(raw_tools, list):
        raise ValueError("tools must be a list of function specs")
    return [_unwrap_tool(item) for item in raw_tools]


def render_function_schema(tool: dict[str, Any]) -> str:
    chunks = [f"\n<function>\n<name>{tool['name']}</name>"]
    description = tool.get("description")
    if isinstance(description, str) and description.strip():
        chunks.append(f"\n<description>{description.strip()}</description>")
    chunks.append("\n<parameters>")
    parameters = tool.get("parameters")
    properties: dict[str, Any] = {}
    if isinstance(parameters, dict):
        raw_properties = parameters.get("properties")
        if isinstance(raw_properties, dict):
            properties = raw_properties
        for param_name, raw_fields in properties.items():
            fields = raw_fields if isinstance(raw_fields, dict) else {}
            chunks.append("\n<parameter>")
            chunks.append(f"\n<name>{param_name}</name>")
            if "type" in fields:
                chunks.append(f"\n<type>{_xml_value(fields['type'])}</type>")
            if isinstance(fields.get("description"), str) and fields["description"].strip():
                chunks.append(
                    f"\n<description>{fields['description'].strip()}</description>"
                )
            if "enum" in fields:
                chunks.append(f"\n<enum>{_xml_value(fields['enum'])}</enum>")
            chunks.append(
                _render_extra_keys(
                    fields,
                    {"name", "type", "description", "enum"},
                )
            )
            chunks.append("\n</parameter>")
        chunks.append(
            _render_extra_keys(parameters, {"type", "properties", "required"})
        )
        if "required" in parameters:
            chunks.append(f"\n<required>{_xml_value(parameters['required'])}</required>")
    chunks.append("\n</parameters>")
    chunks.append(
        _render_extra_keys(
            tool,
            {"type", "name", "description", "parameters"},
        )
    )
    chunks.append("\n</function>")
    return "".join(chunks)


def render_available_tools(raw_tools: Any) -> str:
    tools = coerce_tools(raw_tools)
    if not tools:
        return ""
    body = "".join(render_function_schema(tool) for tool in tools)
    return (
        f"{TOOLS_PREAMBLE}{TOOLS_OPEN}{body}\n{TOOLS_CLOSE}"
        f"{TOOL_CALL_INSTRUCTIONS}"
    )


def inject_available_tools(system_content: str, raw_tools: Any) -> str:
    block = render_available_tools(raw_tools)
    if not block:
        return system_content
    if has_tools_block(system_content):
        return system_content
    if not system_content.strip():
        return block
    return f"{system_content.rstrip()}\n\n{block}"


def json_ready(value: Any) -> Any:
    """Coerce Python literals (sets, tuples) into JSON-serialisable values.

    Some SFT sources store tool arguments as Python literals. `ast.literal_eval`
    turns `{1, 2}` into a `set`, which later `json.dumps` calls reject.
    """

    if isinstance(value, dict):
        return {str(key): json_ready(item) for key, item in value.items()}
    if isinstance(value, (set, frozenset)):
        items = [json_ready(item) for item in value]
        try:
            return sorted(
                items,
                key=lambda item: json.dumps(item, sort_keys=True, default=str),
            )
        except TypeError:
            return items
    if isinstance(value, tuple):
        return [json_ready(item) for item in value]
    if isinstance(value, list):
        return [json_ready(item) for item in value]
    return value


def parse_argument_value(raw: str) -> Any:
    text = raw.strip()
    if not text:
        return ""
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return text


def parse_arguments_payload(raw: Any) -> dict[str, Any]:
    if raw is None:
        return {}
    ready = json_ready(raw)
    if isinstance(ready, dict):
        return ready
    if isinstance(ready, str):
        text = ready.strip()
        if not text:
            return {}
        try:
            loaded = json.loads(text)
        except json.JSONDecodeError:
            return {"value": ready}
        loaded = json_ready(loaded)
        if isinstance(loaded, dict):
            return loaded
        return {"value": loaded}
    return {"value": ready}


def format_tool_call_xml(name: str, arguments: dict[str, Any]) -> str:
    chunks = [f"{TOOL_CALL_OPEN}\n<function={name}>\n"]
    for key, value in arguments.items():
        chunks.append(f"<parameter={key}>\n{_xml_value(value)}\n</parameter>\n")
    chunks.append(f"</function>\n{TOOL_CALL_CLOSE}\n")
    return "".join(chunks)


def format_tool_calls_xml(raw_tool_calls: Any) -> str:
    if not raw_tool_calls:
        return ""
    if not isinstance(raw_tool_calls, list):
        raise ValueError("tool_calls must be a list")
    chunks: list[str] = []
    for raw_call in raw_tool_calls:
        if not isinstance(raw_call, dict):
            raise ValueError("Each tool_call must be an object")
        payload = raw_call.get("function") if isinstance(raw_call.get("function"), dict) else raw_call
        if not isinstance(payload, dict):
            raise ValueError("tool_call is missing a function object")
        name = payload.get("name")
        if not isinstance(name, str) or not name.strip():
            raise ValueError("tool_call is missing a function name")
        arguments = parse_arguments_payload(payload.get("arguments"))
        chunks.append(format_tool_call_xml(name.strip(), arguments))
    return "".join(chunks)


def parse_json_tool_inner(inner: str) -> ParsedToolCall | None:
    text = inner.strip()
    if not text:
        return None
    payload: Any
    try:
        payload = json.loads(text)
    except json.JSONDecodeError:
        try:
            payload = json.loads(text.replace("'", '"'))
        except json.JSONDecodeError:
            return None
    if not isinstance(payload, dict):
        return None
    nested = payload.get("function")
    source = nested if isinstance(nested, dict) else payload
    name = source.get("name")
    if not isinstance(name, str) or not name.strip():
        name = payload.get("name")
    if not isinstance(name, str) or not name.strip():
        return None
    arguments = parse_arguments_payload(
        source.get("arguments", source.get("parameters", payload.get("arguments")))
    )
    return ParsedToolCall(name=name.strip(), arguments=arguments, raw="")


def parse_tool_calls(text: str) -> list[ParsedToolCall]:
    calls: list[ParsedToolCall] = []
    for block in TOOL_CALL_BLOCK_RE.finditer(text):
        inner = block.group(1)
        raw = block.group(0).strip()
        found_xml = False
        for function in FUNCTION_BLOCK_RE.finditer(inner):
            found_xml = True
            name = function.group(1).strip()
            arguments: dict[str, Any] = {}
            for parameter in PARAMETER_BLOCK_RE.finditer(function.group(2)):
                arguments[parameter.group(1).strip()] = parse_argument_value(
                    parameter.group(2)
                )
            calls.append(
                ParsedToolCall(
                    name=name,
                    arguments=json_ready(arguments),
                    raw=raw,
                )
            )
        if found_xml:
            continue
        parsed = parse_json_tool_inner(inner)
        if parsed is not None:
            calls.append(
                ParsedToolCall(name=parsed.name, arguments=parsed.arguments, raw=raw)
            )
    return calls


def openai_tool_calls_from_text(text: str) -> list[dict[str, Any]]:
    encoded: list[dict[str, Any]] = []
    for index, call in enumerate(parse_tool_calls(text)):
        encoded.append(
            {
                "id": f"call_{index}_{call.name}",
                "type": "function",
                "function": {
                    "name": call.name,
                    "arguments": json.dumps(call.arguments, ensure_ascii=False),
                },
            }
        )
    return encoded


def assistant_message_content(message: dict[str, Any]) -> str:
    reasoning = message.get("reasoning_content")
    content = message_text(message.get("content"))
    if isinstance(reasoning, str) and reasoning.strip():
        content = f"{THINK_OPEN}\n{reasoning.strip()}\n{THINK_CLOSE}\n{content}"
    tool_xml = format_tool_calls_xml(message.get("tool_calls"))
    if tool_xml:
        if content.strip():
            return f"{content.rstrip()}\n{tool_xml}"
        return tool_xml
    return content


def _flush_tool_group(
    group: list[str],
    messages: list[dict[str, str]],
) -> None:
    if not group:
        return
    messages.append({"role": "user", "content": "\n".join(group)})
    group.clear()


def prepare_inference_messages(
    raw_messages: list[dict[str, Any]],
    *,
    tools: Any | None = None,
) -> list[dict[str, str]]:
    if not raw_messages:
        raise ValueError("Chat history cannot be empty")

    prepared: list[dict[str, str]] = []
    pending_tool_results: list[str] = []
    for index, raw_message in enumerate(raw_messages):
        if not isinstance(raw_message, dict):
            raise ValueError(f"Unsupported chat message at index {index}")
        raw_role = raw_message.get("role")
        if not isinstance(raw_role, str):
            raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
        role = INFERENCE_ROLE_MAP.get(raw_role.strip().lower())
        if role is None:
            raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
        if role == "assistant":
            content = assistant_message_content(raw_message)
        else:
            content = message_text(raw_message.get("content"))
        if role == "tool":
            content = wrap_tool_response(content)
            if not content.strip():
                raise ValueError(f"Chat content at index {index} must be non-empty")
            pending_tool_results.append(content)
            continue
        _flush_tool_group(pending_tool_results, prepared)
        if not content.strip():
            if role == "system":
                continue
            raise ValueError(f"Chat content at index {index} must be non-empty")
        if role == "system":
            content = strip_leaked_chatml(content)
        prepared.append({"role": role, "content": content})
    _flush_tool_group(pending_tool_results, prepared)

    tools_block = render_available_tools(tools)
    if tools_block:
        if prepared and prepared[0]["role"] == "system":
            prepared[0] = {
                "role": "system",
                "content": inject_available_tools(prepared[0]["content"], tools),
            }
        else:
            prepared.insert(0, {"role": "system", "content": tools_block})

    if not prepared:
        raise ValueError("Chat history cannot be empty")
    return prepared