from transformers import AutoTokenizer
from trl.chat_template_utils import parse_response

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

qwen3_schema = {
    "x-regex": r"^(?:<think>\n?(?P<reasoning_content>.+?)\n?</think>\s*)?(?P<content>.*?)(?=(?:<tool_call>|<\|im_end\|>|$))(?P<tool_calls>(?:<tool_call>.+?</tool_call>\s*)+)?\s*(?:<\|im_end\|>|$)",
    "type": "object",
    "properties": {
        "role": {"const": "assistant"},
        "content": {"type": "string"},
        "reasoning_content": {"type": "string"},
        "tool_calls": {
            "type": "array",
            "x-regex-iterator": r"<tool_call>\s*(.+?)\s*</tool_call>",
            "items": {
                "x-parser": "json",
                "x-parser-args": {"transform": "{type: 'function', function: @}"},
                "type": "object",
                "properties": {
                    "type": {"const": "function"},
                    "function": {
                        "type": "object",
                        "properties": {
                            "name": {"type": "string"},
                            "arguments": {
                                "type": "object",
                                "additionalProperties": {},
                            },
                        },
                    },
                },
            },
        },
    },
}

tokenizer.response_schema = qwen3_schema

text = """<tool_call>
{"name": "multiply", "arguments": {"a": 3, "b": 4}}
</tool_call>
<|im_end|>"""
parse_response(tokenizer, tokenizer(text)["input_ids"])
# {'role': 'assistant',
#  'content': '',
#  'tool_calls': [{'type': 'function',
#    'function': {'name': 'multiply', 'arguments': {'a': 3, 'b': 4}}}]}

text = """<tool_call>
{"name": "multiply", "arguments": {"a": 3, "b": 4}}
</tool_call>
<tool_call>
{"name": "addition", "arguments": {"a": 3, "b": 4}}
</tool_call>
<tool_call>
{"name": "divide","arguments": {"a": 3, "b": 4}}
</tool_call>
<|im_end|>"""
parse_response(tokenizer, tokenizer(text)["input_ids"])
# {
#     "role": "assistant",
#     "content": "",
#     "tool_calls": [
#         {"type": "function", "function": {"name": "multiply", "arguments": {"a": 3, "b": 4}}},
#         {"type": "function", "function": {"name": "addition", "arguments": {"a": 3, "b": 4}}},
#         {"type": "function", "function": {"name": "divide", "arguments": {"a": 3, "b": 4}}},
#     ],
# }