import os

os.environ['SWIFT_DEBUG'] = '1'
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1,2,3'
system = 'You are a helpful assistant.'

tools = [{
    'type': 'function',
    'function': {
        'name': 'get_current_weather',
        'description': 'Get the current weather in a given location',
        'parameters': {
            'type': 'object',
            'properties': {
                'location': {
                    'type': 'string',
                    'description': 'The city and state, e.g. San Francisco, CA'
                },
                'unit': {
                    'type': 'string',
                    'enum': ['celsius', 'fahrenheit']
                }
            },
            'required': ['location']
        }
    }
}, {
    'name_for_model': 'tool2',
    'name_for_human': '工具2',
    'description': 'Tool2的描述',
}]

glm4_tools = [{
    'type': 'function',
    'function': {
        'name': 'realtime_aqi',
        'description': '天气预报。获取实时空气质量。当前空气质量，PM2.5，PM10信息',
        'parameters': {
            'type': 'object',
            'properties': {
                'city': {
                    'description': '城市名'
                }
            },
            'required': ['city']
        }
    }
}]
glm4_tool_messasges = [
    {
        'role': 'tool',
        'content': '{"city": "北京", "aqi": "10", "unit": "celsius"}'
    },
    {
        'role': 'tool',
        'content': '{"city": "上海", "aqi": "72", "unit": "fahrenheit"}'
    },
]
glm4_query = '北京和上海今天的天气情况'


def _infer(engine, num_tools: int = 1, agent_tools=None, tool_messages=None, query=None):
    if agent_tools is None:
        agent_tools = tools
    if tool_messages is None:
        tool_messages = []
        for _ in range(num_tools):
            tool_messages.append({
                'role': 'tool',
                'content': '{"temperature": 32, "condition": "Sunny", "humidity": 50}'
            })
    stop = [engine.template.agent_template.keyword.observation]
    query = query or "How's the weather in Beijing today?"
    infer_request = InferRequest([{'role': 'user', 'content': query}], tools=agent_tools)
    request_config = RequestConfig(max_tokens=512, stop=stop, temperature=0)
    resp_list = engine.infer([infer_request], request_config=request_config)
    response = resp_list[0].choices[0].message.content
    toolcall = resp_list[0].choices[0].message.tool_calls[0].function
    print(f'response: {response}')
    print(f'toolcall: {toolcall}')
    assert toolcall is not None
    infer_request.messages.append({'role': 'assistant', 'content': response})
    infer_request.messages += tool_messages
    resp_list = engine.infer([infer_request], request_config=request_config)
    response2 = resp_list[0].choices[0].message.content
    print(f'response2: {response2}')
    infer_request.messages.append({'role': 'assistant', 'content': response2})
    return infer_request.messages


def test_react_en():
    agent_template = agent_template_map['react_en']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 1144
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'react_en'
    messages = _infer(engine)
    assert messages[-1]['content'] == (
        'Thought: The current temperature in Beijing is 32 degrees Celsius, and the condition is sunny '
        'with a humidity of 50%.\nFinal Answer: The current temperature in Beijing is 32 degrees Celsius,'
        ' and the condition is sunny with a humidity of 50%.')
    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')


def test_react_zh():
    agent_template = agent_template_map['react_zh']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 712
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'react_zh'
    _infer(engine)


def test_qwen_en():
    agent_template = agent_template_map['qwen_en']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 879
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'qwen_en'
    messages = _infer(engine)
    assert messages[-1]['content'] == (
        '✿RETURN✿: Today in Beijing, the temperature is 32°C with sunny conditions and the humidity '
        'is at 50%. Enjoy the nice weather!')
    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')


def test_qwen_zh():
    agent_template = agent_template_map['qwen_zh']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 577
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'qwen_zh'
    _infer(engine)


def test_qwen_en_parallel():
    agent_template = agent_template_map['qwen_en_parallel']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 1012
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'qwen_en_parallel'
    messages = _infer(engine, num_tools=2)
    assert messages[-1]['content'] == (
        '✿RETURN✿: Today in Beijing, the temperature is 32 degrees Celsius with sunny conditions '
        'and the humidity is at 50%. Enjoy the nice weather!')
    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')


def test_qwen_zh_parallel():
    agent_template = agent_template_map['qwen_zh_parallel']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 688
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'qwen_zh_parallel'
    _infer(engine, num_tools=2)


def test_hermes():
    agent_template = agent_template_map['hermes']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 875
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'hermes'
    messages = _infer(engine, num_tools=2)
    template.template_backend = 'jinja'
    messages2 = _infer(engine, num_tools=2)
    assert messages[-1]['content'] == messages2[-1]['content'] == (
        'Today in Beijing, the temperature is 32 degrees Celsius with sunny conditions '
        'and the humidity is at 50%. Enjoy the nice weather!')
    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert encoded['input_ids'] == encoded2['input_ids']


def test_toolbench():
    agent_template = agent_template_map['toolbench']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 1833
    engine = TransformersEngine('Qwen/Qwen2.5-7B-Instruct')
    template = engine.template
    template._agent_template = 'toolbench'
    _infer(engine)


def test_chatglm4():
    agent_template = agent_template_map['chatglm4']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 846
    engine = TransformersEngine('ZhipuAI/glm-4-9b-chat')
    template = engine.template
    template._agent_template = 'chatglm4'
    _infer(engine, agent_tools=glm4_tools, tool_messages=glm4_tool_messasges, query=glm4_query)


def test_glm4():
    agent_template = agent_template_map['glm4']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 769
    engine = TransformersEngine('ZhipuAI/GLM-4-9B-0414')
    template = engine.template
    template._agent_template = 'glm4'
    messages = _infer(engine, agent_tools=glm4_tools, tool_messages=glm4_tool_messasges, query=glm4_query)
    assert messages[-1]['content'] == '根据天气预报工具，北京今天的空气质量指数为10，属于良好水平；上海今天的空气质量指数为72，属于轻度污染水平。'
    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')


def test_llama3():
    engine = TransformersEngine('LLM-Research/Llama-3.2-3B-Instruct')
    template = engine.template
    template._agent_template = 'llama3'
    messages = _infer(engine)

    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')


def test_llama4():
    engine = TransformersEngine('LLM-Research/Llama-4-Scout-17B-16E-Instruct')
    template = engine.template
    messages = _infer(engine)
    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')


def test_hunyuan():
    engine = TransformersEngine('Tencent-Hunyuan/Hunyuan-1.8B-Instruct')
    template = engine.template
    template.template_backend = 'jinja'
    _infer(engine, num_tools=2)

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert encoded['input_ids'][:-1] == encoded2['input_ids']


def test_glm4_5():
    engine = TransformersEngine('ZhipuAI/GLM-4.5-Air')
    template = engine.template
    template.template_backend = 'jinja'
    _infer(engine, num_tools=2)

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert encoded['input_ids'][:-1] == encoded2['input_ids']


def test_glm4_7():
    engine = TransformersEngine('ZhipuAI/GLM-4.7-FP8', load_model=False)
    template = engine.template

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert encoded['input_ids'][:-1] == encoded2['input_ids']


def test_qwen3_coder():
    engine = TransformersEngine('Qwen/Qwen3-Coder-30B-A3B-Instruct')
    template = engine.template
    template.template_backend = 'jinja'
    _infer(engine, num_tools=2)

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert encoded['input_ids'] == encoded2['input_ids']


def test_qwen3_5():
    engine = TransformersEngine('Qwen/Qwen3.5-35B-A3B')
    template = engine.template
    template.template_backend = 'jinja'
    _infer(engine, num_tools=2)

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    data['messages'].insert(0, {'role': 'system', 'content': 'You are a helpful assistant.'})
    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert encoded['input_ids'] == encoded2['input_ids']


def test_deepseek_v3_1():
    engine = TransformersEngine('deepseek-ai/DeepSeek-V3.1', load_model=False)
    template = engine.template

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    # To test multiple tool calls and responses, we duplicate some messages.
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')

    expected_input_ids = (
        '<｜begin▁of▁sentence｜>\n\n## Tools\n'
        'You have access to the following tools:\n\n'
        '### convert_temperature\n'
        'Description: Convert temperature from one unit to another\n\n'
        "Parameters: {\"type\": \"object\", \"properties\": {\"temperature\": {\"type\": \"number\", "
        "\"description\": \"The temperature value\"}, \"from_unit\": {\"type\": \"string\", \"description\": "
        "\"The unit to convert from\"}, \"to_unit\": {\"type\": \"string\", \"description\": \"The unit "
        "to convert to\"}}, \"required\": [\"temperature\", \"from_unit\", \"to_unit\"]}\n\n"
        '### get_current_date\n'
        'Description: Get the current date\n\n'
        'Parameters: {}\n\n'
        'IMPORTANT: ALWAYS adhere to this exact format for tool use:\n'
        '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>tool_call_name<｜tool▁sep｜>tool_call_arguments<｜tool▁call▁end｜>'
        '{additional_tool_calls}<｜tool▁calls▁end｜>\n\n'
        'Where:\n'
        '- `tool_call_name` must be an exact match to one of the available tools\n'
        "- `tool_call_arguments` must be valid JSON that strictly follows the tool's Parameters Schema\n"
        '- For multiple tool calls, chain them directly without separators or spaces<｜User｜>'
        'Hi, I need to convert a temperature from Celsius to Fahrenheit. The temperature is 30 degrees Celsius.'
        '<｜Assistant｜></think><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>convert_temperature<｜tool▁sep｜>'
        "{\"temperature\": 30, \"from_unit\": \"Celsius\", \"to_unit\": \"Fahrenheit\"}<｜tool▁call▁end｜>"
        '<｜tool▁call▁begin｜>convert_temperature<｜tool▁sep｜>'
        "{\"temperature\": 30, \"from_unit\": \"Celsius\", \"to_unit\": \"Fahrenheit\"}<｜tool▁call▁end｜>"
        '<｜tool▁calls▁end｜><｜end▁of▁sentence｜>'
        "<｜tool▁output▁begin｜>{\"converted_temperature\": 86}<｜tool▁output▁end｜>"
        "<｜tool▁output▁begin｜>{\"converted_temperature\": 86}<｜tool▁output▁end｜>"
        'The converted temperature from 30 degrees Celsius to Fahrenheit is 86 degrees Fahrenheit.<｜end▁of▁sentence｜>')

    # Expected labels string
    expected_labels = (
        '[-100 * 239]</think><｜tool▁calls▁begin｜><｜tool▁call▁begin｜>convert_temperature<｜tool▁sep｜>'
        "{\"temperature\": 30, \"from_unit\": \"Celsius\", \"to_unit\": \"Fahrenheit\"}<｜tool▁call▁end｜>"
        '<｜tool▁call▁begin｜>convert_temperature<｜tool▁sep｜>'
        "{\"temperature\": 30, \"from_unit\": \"Celsius\", \"to_unit\": \"Fahrenheit\"}<｜tool▁call▁end｜>"
        '<｜tool▁calls▁end｜><｜end▁of▁sentence｜>[-100 * 22]'
        'The converted temperature from 30 degrees Celsius to Fahrenheit is 86 degrees Fahrenheit.<｜end▁of▁sentence｜>')

    assert template.safe_decode(encoded['input_ids']) == expected_input_ids
    assert template.safe_decode(encoded['labels']) == expected_labels
    assert encoded['input_ids'][-122:] == encoded2['input_ids'][1:]


def test_youtu():
    agent_template = agent_template_map['youtu']()
    new_system = agent_template._format_tools(tools, system)
    assert len(new_system) == 883
    engine = TransformersEngine('Tencent-YouTu-Research/Youtu-LLM-2B')
    template = engine.template
    template._agent_template = 'youtu'

    stop = [template.agent_template.keyword.observation]
    query = "How's the weather in Beijing today?"
    tool_messages = [{'role': 'tool', 'content': '{"temperature": 32, "condition": "Sunny", "humidity": 50}'}]
    infer_request = InferRequest([{'role': 'user', 'content': query}], tools=tools)
    request_config = RequestConfig(max_tokens=2048, stop=stop, temperature=0)

    # First inference: get tool call
    resp_list = engine.infer([infer_request], request_config=request_config)
    response = resp_list[0].choices[0].message.content
    toolcall = resp_list[0].choices[0].message.tool_calls
    print(f'response: {response}')
    print(f'toolcall: {toolcall}')
    assert toolcall is not None, 'No tool_call generated'
    infer_request.messages.append({'role': 'assistant', 'content': response})
    infer_request.messages += tool_messages

    # Second inference: get final response
    resp_list = engine.infer([infer_request], request_config=request_config)
    response2 = resp_list[0].choices[0].message.content
    print(f'response2: {response2}')
    infer_request.messages.append({'role': 'assistant', 'content': response2})
    messages = infer_request.messages

    template.set_mode('train')
    encoded = template.encode({'messages': messages})
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])
    template.template_backend = 'swift'
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert encoded['input_ids'] == encoded2['input_ids']


def test_deepseek_v4():
    engine = TransformersEngine('deepseek-ai/DeepSeek-V4-Flash', load_model=False)
    template = engine.template

    tools = [{
        'type': 'function',
        'function': {
            'name': 'get_weather',
            'description': 'Get the weather for a specific location',
            'parameters': {
                'type': 'object',
                'properties': {
                    'location': {
                        'type': 'string',
                        'description': 'The city name'
                    },
                    'unit': {
                        'type': 'string',
                        'enum': ['celsius', 'fahrenheit'],
                        'description': 'Temperature unit'
                    }
                },
                'required': ['location']
            }
        }
    }, {
        'type': 'function',
        'function': {
            'name': 'search',
            'description': 'Search the web for information',
            'parameters': {
                'type': 'object',
                'properties': {
                    'query': {
                        'type': 'string',
                        'description': 'Search query'
                    },
                    'num_results': {
                        'type': 'integer',
                        'description': 'Number of results to return'
                    }
                },
                'required': ['query']
            }
        }
    }]
    data = {
        'tools':
        tools,
        'messages': [{
            'role': 'system',
            'content': 'You are a helpful assistant.'
        }, {
            'role': 'user',
            'content': "What's the weather in Beijing?"
        }, {
            'role':
            'assistant',
            'content':
            '<think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>\n\n'
        }, {
            'role':
            'tool_call',
            'content':
            '{"name": "get_weather", "arguments": "{\\"location\\": \\"Beijing\\", \\"unit\\": \\"celsius\\"}"}'
        }, {
            'role': 'tool_response',
            'content': '{"temperature": 22, "condition": "sunny", "humidity": 45}'
        }, {
            'role':
            'assistant',
            'content': ('<think>Got the weather data. Let me format a nice response.</think>'
                        'The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.')
        }]
    }

    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')

    expected_input_ids = (
        '<｜begin▁of▁sentence｜>You are a helpful assistant.\n\n## Tools\n\n'
        'You have access to a set of tools to help answer the user\'s question. '
        'You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:\n\n'
        '<｜DSML｜tool_calls>\n'
        '<｜DSML｜invoke name="$TOOL_NAME">\n'
        '<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</｜DSML｜parameter>\n'
        '...\n'
        '</｜DSML｜invoke>\n'
        '<｜DSML｜invoke name="$TOOL_NAME2">\n'
        '...\n'
        '</｜DSML｜invoke>\n'
        '</｜DSML｜tool_calls>\n\n'
        'String parameters should be specified as is and set `string="true"`. '
        'For all other types (numbers, booleans, arrays, objects), '
        'pass the value in JSON format and set `string="false"`.\n\n'
        'If thinking_mode is enabled (triggered by <think>), '
        'you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.'
        '\n\nOtherwise, output directly after </think> with tool calls or final response.\n\n'
        '### Available Tool Schemas\n\n'
        '{"name": "get_weather", "description": "Get the weather for a specific location", '
        '"parameters": {"type": "object", "properties": {"location": {"type": "string", '
        '"description": "The city name"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], '
        '"description": "Temperature unit"}}, "required": ["location"]}}\n'
        '{"name": "search", "description": "Search the web for information", '
        '"parameters": {"type": "object", "properties": {"query": {"type": "string", '
        '"description": "Search query"}, "num_results": {"type": "integer", '
        '"description": "Number of results to return"}}, "required": ["query"]}}\n\n'
        'You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n'
        '<｜User｜>What\'s the weather in Beijing?<｜Assistant｜>'
        '<think>The user wants to know the weather in Beijing. I should use the get_weather tool.</think>\n\n'
        '<｜DSML｜tool_calls>\n'
        '<｜DSML｜invoke name="get_weather">\n'
        '<｜DSML｜parameter name="location" string="true">Beijing</｜DSML｜parameter>\n'
        '<｜DSML｜parameter name="unit" string="true">celsius</｜DSML｜parameter>\n'
        '</｜DSML｜invoke>\n'
        '</｜DSML｜tool_calls>'
        '<｜end▁of▁sentence｜>'
        '<｜User｜><tool_result>{"temperature": 22, "condition": "sunny", "humidity": 45}</tool_result>'
        '<｜Assistant｜>'
        '<think>Got the weather data. Let me format a nice response.</think>'
        'The weather in Beijing is currently sunny with a temperature of 22°C and 45% humidity.'
        '<｜end▁of▁sentence｜>')

    assert template.safe_decode(encoded['input_ids']) == expected_input_ids


def test_seed_oss():
    engine = TransformersEngine('ByteDance-Seed/Seed-OSS-36B-Instruct', load_model=False)

    template = engine.template
    dataset = load_dataset('AI-ModelScope/function-calling-chatml')[0]
    data = dataset[6]
    # To test multiple tool calls and responses, we duplicate some messages.
    data['messages'].insert(1, data['messages'][1])
    data['messages'].insert(3, data['messages'][3])

    # Incomplete tool function will cause seed template to throw an error.
    data['tools'] = [('{\n'
                      '    "name": "convert_temperature",\n'
                      '    "description": "Convert temperature from one unit to another",\n'
                      '    "parameters": {\n'
                      '        "type": "object",\n'
                      '        "properties": {\n'
                      '            "temperature": {\n'
                      '                "type": "number",\n'
                      '                "description": "The temperature value"\n'
                      '            },\n'
                      '            "from_unit": {\n'
                      '                "type": "string",\n'
                      '                "description": "The unit to convert from"\n'
                      '            },\n'
                      '            "to_unit": {\n'
                      '                "type": "string",\n'
                      '                "description": "The unit to convert to"\n'
                      '            }\n'
                      '        },\n'
                      '        "required": [\n'
                      '            "temperature",\n'
                      '            "from_unit",\n'
                      '            "to_unit"\n'
                      '        ]\n'
                      '    }\n'
                      '}'),
                     ('{\n'
                      '    "name": "get_current_date",\n'
                      '    "description": "Get the current date",\n'
                      '    "parameters":  {\n'
                      '        "type": "object",\n'
                      '        "properties": {\n'
                      '         "date": {\n'
                      '                "type": "number",\n'
                      '                "description": "The date value"}}}\n'
                      '}')]

    data['thinking_budget'] = 0

    template.template_backend = 'swift'
    template.set_mode('train')
    encoded = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded["labels"])}')
    import re
    expected_input_ids = re.sub(
        r'<seed:think>.*?</seed:think>', '', template.safe_decode(encoded['input_ids']), flags=re.DOTALL)
    template.template_backend = 'jinja'
    encoded2 = template.encode(data)
    print(f'input_ids: {template.safe_decode(encoded2["input_ids"])}')
    print(f'labels: {template.safe_decode(encoded2["labels"])}')
    assert template.safe_decode(encoded2['input_ids']) == expected_input_ids


if __name__ == '__main__':
    from swift import InferRequest, RequestConfig, TransformersEngine, agent_template_map, load_dataset

    # test_react_en()
    # test_react_zh()
    # test_qwen_en()
    # test_qwen_zh()
    # test_qwen_en_parallel()
    # test_qwen_zh_parallel()
    # test_hermes()
    # test_toolbench()
    # test_chatglm4()
    # test_glm4()
    # test_llama3()
    # test_llama4()
    # test_hunyuan()
    # test_glm4_5()
    # test_glm4_7()
    # test_qwen3_coder()
    # test_qwen3_5()
    # test_deepseek_v3_1()
    test_deepseek_v4()
    # test_seed_oss()
    # test_youtu()
