import json
import ollama

def callModel(prompt):
    response = ollama.chat(
        model='qwen2.5:7b',
        messages=[{'role': 'user', 'content': prompt}]
    )
    print(response['message']['content'])
    return response['message']['content']

def checkTools(query):

    prompt = ("""你是一位可靠的个人助理，现在为你提供【工具集】，根据【用户问题】判断是否需要使用其中的工具，并按照【指定格式】输出。

        【工具集】
        [getTrainSchedule]
        getTrainSchedule的作用是根据指定的日期、起点城市名称和终点城市名称，查询列车班次

        【指定格式】
        {"是否需要调用【工具集】中的工具":"是",
        "需要调用的工具名称":"getTrainSchedule",
        "调用工具需要的参数":{"queryDate":"10月12日","start":"上海","end":"北京"}}

        【用户问题】
        """ + query)

    toolCall = callModel(prompt)
    return toolCall

def callTools(toolCall):
    with open('tools.py', 'r', encoding='utf-8') as f:
        content = f.read()
    argsStr = ",".join(str(key)+'="'+toolCall["调用工具需要的参数"][key]+'"' for key in toolCall["调用工具需要的参数"])
    print(argsStr)
    content+="\ntoolRes="+str(toolCall["需要调用的工具名称"])+"("+argsStr+")"
    #toolRes=getTrainSchedule(queryDate="10月12日",start="上海",end="北京")
    print(content)
    exec(content)
    toolResult = locals()["toolRes"]
    return toolResult

def getAnswer(query):
    toolCall = json.loads(checkTools(query))
    print("checkTools results:*********************************")
    print(toolCall)
    if toolCall["是否需要调用【工具集】中的工具"] == "是":
        toolRes = callTools(toolCall)
        print("callTools results:*********************************")
        print(toolRes)
        prompt = "你是一位可靠的个人助理，用户的问题是："+query+"。 查询工具得到的结果是："+toolRes+"。 根据这些信息，给出合理的回复。"
        finalResult = callModel(prompt)
    else:
        prompt = "你是一位可靠的个人助理，用户的问题是：" + query + "。 根据这些信息，给出合理的回复。"
        finalResult = callModel(prompt)
    return finalResult

if __name__ == '__main__':
    query = "我要在本周六到北京旅行，帮我规划一下行程，要求尽可能快捷舒适。"
    result = getAnswer(query)
    print(result)