import ollama

def callModel(prompt):
    response = ollama.chat(
        model='qwen2.5:7b',
        messages=[{'role': 'user', 'content': prompt}],
        tools=[
            {
                "type": "function",
                "function": {
                    "name": "getTrainSchedule",
                    "description": "根据指定的日期、起点城市名称和终点城市名称，查询列车班次",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "queryDate": {
                                "type": "string",
                                "description": "指定日期",
                            },
                            "start": {
                                "type": "string",
                                "description": "起点城市名称",
                            },
                            "end": {
                                "type": "string",
                                "description": "终点城市名称",
                            },
                        },
                        "required": ["queryDate", "start", "end"],
                    },
                }
            }
        ]
    )
    #print(response['message']['content'])
    #return response['message']['content']
    return response



def checkTools(query):
    toolCall = callModel(query)
    return toolCall

def callTools(toolCall):
    with open('tools.py', 'r', encoding='utf-8') as f:
        content = f.read()

    call = toolCall[0]
    print(call.function)
    print(call.function.arguments)
    argsStr = ",".join(str(key)+'="'+call.function.arguments[key]+'"' for key in call.function.arguments)
    print(argsStr)
    content+="\ntoolRes="+str(call.function.name)+"("+argsStr+")"
    print(content)
    exec(content)
    toolResult = locals()["toolRes"]
    return toolResult

def getAnswer(query):
    toolCall = []
    response = checkTools(query)
    if 'tool_calls' in response['message']:
        toolCall = response['message']['tool_calls']
        print("checkTools results:*********************************")
        print(toolCall)
    if len(toolCall)>0:
        toolRes = callTools(toolCall)
        print("callTools results:*********************************")
        print(toolRes)
        prompt = "你是一位可靠的个人助理，用户的问题是："+query+"。 查询工具得到的结果是："+toolRes+"。 根据这些信息，给出合理的回复。"
        finalResult = callModel(prompt)
    else:
        prompt = "你是一位可靠的个人助理，用户的问题是：" + query + "。 根据这些信息，给出合理的回复。"
        finalResult = callModel(prompt)
    return finalResult['message']['content']

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