import gradio as gr
from openai import OpenAI
import json

client = OpenAI(
    base_url="https://jacky2305-llm-api.hf.space/v1",
    api_key="sk-any"
)

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取某个城市的天气",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"}
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "执行数学计算",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "数学表达式"}
                },
                "required": ["expression"]
            }
        }
    }
]

def run_tool(name, args):
    if name == "get_weather":
        return f"{args['city']}今天晴天，28°C"
    if name == "calculate":
        try:
            return str(eval(args["expression"]))
        except:
            return "计算出错"

def test_tool_calling(user_msg):
    log = []
    messages = [{"role": "user", "content": user_msg}]

    resp = client.chat.completions.create(
        model="qwen2.5-3b-instruct",
        messages=messages,
        tools=tools,
        tool_choice="auto",
        max_tokens=300
    )

    msg = resp.choices[0].message
    finish_reason = resp.choices[0].finish_reason
    log.append(f"finish_reason: {finish_reason}")

    if finish_reason != "tool_calls" or not msg.tool_calls:
        log.append(f"⚠️ 模型没有调用工具，直接回答：\n{msg.content}")
        return "\n".join(log)

    for tool_call in msg.tool_calls:
        name = tool_call.function.name
        args = json.loads(tool_call.function.arguments)
        log.append(f"✅ 调用工具：{name}({args})")

        result = run_tool(name, args)
        log.append(f"工具返回：{result}")

        messages.append(msg)
        messages.append({
            "role": "tool",
            "tool_call_id": tool_call.id,
            "content": result
        })

    resp2 = client.chat.completions.create(
        model="qwen2.5-3b-instruct",
        messages=messages,
        max_tokens=300
    )
    log.append(f"\n最终回答：\n{resp2.choices[0].message.content}")
    return "\n".join(log)

with gr.Blocks() as demo:
    gr.Markdown("## Qwen 3B 工具调用测试")
    with gr.Row():
        inp = gr.Textbox(label="输入问题", placeholder="试试：吉隆坡天气怎么样？")
        btn = gr.Button("测试")
    out = gr.Textbox(label="结果日志", lines=15)
    
    gr.Examples(
        examples=[
            ["吉隆坡今天天气怎么样？"],
            ["帮我算一下 1234 * 5678"],
            ["你好"],  # 不需要工具
        ],
        inputs=inp
    )
    
    btn.click(test_tool_calling, inputs=inp, outputs=out)

demo.launch()