File: easytool/data_process.py

import os
import gdown
import shutil
import json
from zipfile import ZipFile

urls = {
    "funcqa": "https://drive.google.com/uc?id=13Sj7uIsyqWXoTh1ejWUviTzeQSES2Omd",
    "restbench": "https://raw.githubusercontent.com/Yifan-Song793/RestGPT/main/datasets/tmdb.json",
    "toolbench": "https://drive.google.com/uc?id=1XFjDxVZdUY7TXYF2yvzx3pJlS2fy78jk",
}


def read_jsonline(address):
    not_mark = []
    with open(address, 'r', encoding="utf-8") as f:
        for jsonstr in f.readlines():
            jsonstr = json.loads(jsonstr)
            not_mark.append(jsonstr)
    return not_mark


def read_json(address):
    with open(address, 'r', encoding='utf-8') as json_file:
        json_data = json.load(json_file)
    return json_data


def toolbench_process(data_file, dataset):
    ls = read_json(data_file)
    all_data = read_json(f"{dataset}/tool_instruction/toolbench_tool_instruction.json")
    all_dic = {}
    for ID in all_data.keys():
        all_dic[all_data[ID]["tool_name"]] = all_data[ID]

    not_in = []
    for data in ls:
        Tool_dic = []
        data_dic = {}
        already = []
        for tool in data['api_list']:
            if tool['tool_name'] in all_dic:
                if all_dic[tool['tool_name']]["ID"] not in already:
                    already.append(all_dic[tool['tool_name']]["ID"])
                    Tool_dic.append({"ID": all_dic[tool['tool_name']]["ID"],
                                     "Description": all_dic[tool['tool_name']]["tool_description"], })
        data["Tool_dic"] = Tool_dic

    json_str = json.dumps(ls, indent=4)
    with open(data_file, 'w', encoding='utf-8') as json_file:
        json.dump(ls, json_file, ensure_ascii=False, indent=4)


def main():
    curr_dir = os.path.dirname(__file__)

    for dataset in [
        "funcqa",
        "restbench",
        "toolbench"
    ]:
        dataset_path = os.path.join(curr_dir, "data_{}".format(dataset), "test_data")

        if not os.path.exists(dataset_path):
            os.makedirs(dataset_path)

        if dataset == "funcqa":
            print("Processing FuncQA dataset ...\n")
            # Processing FuncQA dataset ...
            temp_file = os.path.join(dataset_path, "data_toolkengpt_0918.zip")
            gdown.download(urls[dataset], temp_file, quiet=False)
            zf = ZipFile(temp_file, 'r')
            zf.extract("data/funcqa/funcqa_oh.json", ".")
            zf.extract("data/funcqa/funcqa_mh.json", ".")
            os.rename("data/funcqa/funcqa_oh.json", "{}/funcqa_oh.json".format(dataset_path))
            os.rename("data/funcqa/funcqa_mh.json", "{}/funcqa_mh.json".format(dataset_path))
            os.remove(temp_file)
            shutil.rmtree("data")

            print("FuncQA dataset Done!\n")

        if dataset == "restbench":
            print("Processing RestBench dataset ... \n")
            # Processing RestBench Dataset
            os.system("wget -P {} -c {}".format(dataset_path, urls[dataset]))

            print("RestBench dataset Done!\n")

        if dataset == "toolbench":
            print("Processing ToolBench dataset ... \n")
            # Processing ToolBench Dataset
            temp_file = os.path.join(dataset_path, "data.zip")
            gdown.download(urls[dataset], temp_file, quiet=False)
            zf = ZipFile(temp_file, 'r')
            zf.extract("data/test_instruction/G2_category.json", ".")
            zf.extract("data/test_instruction/G3_instruction.json", ".")

            os.rename("data/test_instruction/G2_category.json", "{}/G2_category.json".format(dataset_path))
            os.rename("data/test_instruction/G3_instruction.json", "{}/G3_instruction.json".format(dataset_path))
            toolbench_process("{}/G2_category.json".format(dataset_path), "data_{}".format(dataset))
            toolbench_process("{}/G3_instruction.json".format(dataset_path), "data_{}".format(dataset))
            os.remove(temp_file)
            shutil.rmtree("data")

            print("Toolbench dataset Done!\n")


if __name__ == '__main__':
    main()



File: easytool/main.py

# — coding: utf-8 –
import openai
import json
import argparse
import os
from tqdm import tqdm
from easytool import funcQA, restbench, toolbench_retrieve, toolbench
from easytool.util import *
openai.api_key = os.environ["OPENAI_API_KEY"]
        
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_name', type=str, default='gpt-3.5-turbo')
    parser.add_argument('--task', type=str, default='funcqa_mh', help='funcqa, toolbench_retrieve, toolbench, restbench')
    parser.add_argument('--data_type', type=str, default='G3', help='G2 or G3 or funcqa_mh or funcqa_oh')
    parser.add_argument('--tool_root_dir', type=str, default='.toolenv/tools/')
    parser.add_argument('--retrieval_num', type=int, default=5)
    
    args = parser.parse_args()
    
    if args.task == 'funcqa':
        dataset = read_json('data_funcqa/tool_instruction/functions_data.json')
        Tool_dic = read_jsonline('data_funcqa/tool_instruction/tool_dic.jsonl')
        test_data = read_json(f"data_funcqa/test_data/{args.data_type}.json")
        progress_file = f"FuncQA_{args.data_type}_{args.model_name}_Easytool.txt"
    
        
    elif 'toolbench' in args.task:
        base_path = args.tool_root_dir
        index = build_index(base_path)
        dataset = read_json('data_toolbench/tool_instruction/toolbench_tool_instruction.json')
        if args.data_type == 'G2':
            test_data = read_json(f'''data_toolbench/test_data/{args.data_type}_category.json''')
        elif args.data_type == 'G3':
            test_data = read_json(f'''data_toolbench/test_data/{args.data_type}_instruction.json''')
        progress_file = f'''{args.data_type}_{args.model_name}_Easytool.txt'''
        
    
    elif args.task == 'restbench':
        Tool_dic = read_json('data_restbench/tool_instruction/tmdb_tool.json')
        dic_tool = {}
        for data in Tool_dic:
            dic_tool[data['ID']] = data
        test_data = read_json('data_restbench/test_data/tmdb.json')
        progress_file = f"restbench_{args.model_name}_Easytool.txt"

    else:
        print("Wrong task name")
        exit()  
        
    start_index = get_last_processed_index(progress_file)
    total_files = len(test_data)
    retrieval_num = args.retrieval_num
    ind = start_index
    model_name = args.model_name
    
    print("-------Start Execution-------")
    if args.data_type == 'funcqa_mh':
        funcQA.task_execution_mh(args.data_type, start_index, total_files, 
                                        retrieval_num, ind, model_name, dataset, 
                                        Tool_dic, test_data, progress_file)
    elif args.data_type == 'funcqa_oh':
        funcQA.task_execution_oh(args.data_type, start_index, total_files, 
                                        retrieval_num, ind, model_name, dataset, 
                                        Tool_dic, test_data, progress_file)
        
        
    elif args.task == 'toolbench_retrieve':
        toolbench_retrieve.task_execution(args.data_type,
            base_path, index, dataset, test_data, progress_file, 
            start_index, total_files, retrieval_num, ind, model_name)

        
    
    elif args.task == 'toolbench':
        toolbench.task_execution(args.data_type,
            base_path, index, dataset, test_data, progress_file, 
            start_index, total_files, retrieval_num, ind, model_name)

        
    
    elif args.task == 'restbench':
        restbench.task_execution(
            Tool_dic, dic_tool, test_data, progress_file, 
            start_index, total_files, retrieval_num, ind, model_name)

    
    else:
        print("Wrong task name")
        exit()



File: easytool/easytool/util.py

# — coding: utf-8 –
import json
import re
import os


def read_jsonline(address):
    not_mark = []
    with open(address, 'r', encoding="utf-8") as f:
        for jsonstr in f.readlines():
            jsonstr = json.loads(jsonstr)
            not_mark.append(jsonstr)
    return not_mark


def save_json(ls, address):
    json_str = json.dumps(ls, indent=4)
    with open(address, 'w', encoding='utf-8') as json_file:
        json.dump(ls, json_file, ensure_ascii=False, indent=4)


def read_json(address):
    with open(address, 'r', encoding='utf-8') as json_file:
        json_data = json.load(json_file)
    return json_data


def remove_key(item, key_to_remove):
    if isinstance(item, dict):
        if key_to_remove in item:
            del item[key_to_remove]
        for key, value in list(item.items()):  # 使用list包裹，防止字典大小改变时引发错误
            item[key] = remove_key(value, key_to_remove)
    elif isinstance(item, list):
        for index, value in enumerate(item):
            item[index] = remove_key(value, key_to_remove)
    return item


def data_clean(dic, key):
    dic = remove_key(dic, key)
    return dic


def lowercase_parameter_keys(input_dict):
    if "parameters" in input_dict and isinstance(input_dict["parameters"], dict):
        # Convert all keys in the "parameters" dictionary to uppercase
        input_dict["parameters"] = {change_name(k.lower()): v for k, v in input_dict["parameters"].items()}
    return input_dict


def build_index(base_path):
    index = {}
    for root, dirs, files in os.walk(base_path):
        for dir_name in dirs:
            if dir_name not in index:
                index[dir_name] = []
            index[dir_name].append(root)
    return index


def change_name(name):
    change_list = ["from", "class", "return", "false", "true", "id", "and", "", "ID"]
    if name in change_list:
        name = "is_" + name.lower()
    return name


def standardize(string):
    res = re.compile("[^\\u4e00-\\u9fa5^a-z^A-Z^0-9^_]")
    string = res.sub("_", string)
    string = re.sub(r"(_)\1+", "_", string).lower()
    while True:
        if len(string) == 0:
            return string
        if string[0] == "_":
            string = string[1:]
        else:
            break
    while True:
        if len(string) == 0:
            return string
        if string[-1] == "_":
            string = string[:-1]
        else:
            break
    if string[0].isdigit():
        string = "get_" + string
    return string


def get_last_processed_index(progress_file):
    """Retrieve the last processed index from the progress file."""
    if os.path.exists(progress_file):
        with open(progress_file, 'r', encoding='utf-8') as f:
            last_index = f.read().strip()
            return int(last_index) if last_index else 0
    else:
        return 0


def update_progress(progress_file, index):
    """Update the last processed index in the progress file."""
    with open(progress_file, 'w', encoding='utf-8') as f:
        f.write(str(index))


if __name__ == '__main__':
    print("util.py")



File: easytool/easytool/__init__.py

# __init__.py


File: easytool/easytool/funcQA.py

# — coding: utf-8 –
import openai
import json
import logging
import sys
import argparse
from langchain.chat_models import ChatOpenAI
from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate
)
from langchain import LLMChain
import numpy as np
import requests
import os
import subprocess
import re
import importlib.util
from sklearn.metrics.pairwise import cosine_similarity
import pickle
from util import *
from tqdm import tqdm

openai.api_key = os.environ["OPENAI_API_KEY"]


def get_last_processed_index(progress_file):
    """Retrieve the last processed index from the progress file."""
    if os.path.exists(progress_file):
        with open(progress_file, 'r', encoding='utf-8') as f:
            last_index = f.read().strip()
            return int(last_index) if last_index else 0
    else:
        return 0


def update_progress(progress_file, index):
    """Update the last processed index in the progress file."""
    with open(progress_file, 'w', encoding='utf-8') as f:
        f.write(str(index))


def choose_tool(question, Tool_dic, tool_used, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "This is the user's question: {question}\n"
        "These are the tools you can select to solve the question:\n"
        "Tool List:\n"
        "{Too_list}\n\n"
        "Please note that: \n"
        "1. You should only chooce one tool the Tool List to solve this question.\n"
        "2. You must ONLY output the ID of the tool you chose in a parsible JSON format. Two example outputs look like:\n"
        "'''\n"
        "Example 1: {{\"ID\": 1}}\n"
        "Example 2: {{\"ID\": 2}}\n"
        "'''\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    Tool_list = []
    for ele in Tool_dic:
        for key in ele.keys():
            if str(key) not in tool_used:
                Tool_list.append(f'''ID: {key}\n{ele[key]}''')
    while True:
        try:
            result = chain.run(question=question,
                               Too_list=Tool_dic)
            clean_answer = eval(result.split("(")[0].strip())
            # clean_answer = lowercase_parameter_keys(clean_answer)
            # print(clean_answer)
            break
        except Exception as e:
            print(f"choose tool fails: {e}")
            print(result)
            if ind > 10:
                return -1
            ind += 1
            continue
    return clean_answer


def task_decompose(question, Tool_dic, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You need to decompose a complex user's question into some simple subtasks and let the model execute it step by step.\n"
        "This is the user's question: {question}\n"
        "This is tool list:\n"
        "{Tool_list}\n"
        "Please note that: \n"
        "1. You should only decompose this complex user's question into some simple subtasks which can be executed easily by using one single tool in the tool list.\n"
        "2. If one subtask need the results from other subtask, you can should write clearly. For example:"
        "{{\"Tasks\": [\"Convert 23 km/h to X km/min by 'divide_'\", \"Multiply X km/min by 45 min to get Y by 'multiply_'\"]}}\n"
        "3. You must ONLY output in a parsible JSON format. An example output looks like:\n"
        "'''\n"
        "{{\"Tasks\": [\"Task 1\", \"Task 2\", ...]}}\n"
        "'''\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    Tool_list = []
    for ele in Tool_dic:
        Tool_list.append(str(ele))
    ind = 0
    while True:
        try:
            result = chain.run(question=question, Tool_list=Tool_list)
            result = eval(result.split('\n\n')[0])
            a = result["Tasks"]
            break
        except Exception as e:
            print(f"task decompose fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return result


def task_topology(question, task_ls, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Given a complex user's question, I have decompose this question into some simple subtasks"
        "I think there exists a logical connections and order amontg the tasks. "
        "Thus you need to help me output this logical connections and order.\n"
        "You must ONLY output in a parsible JSON format with the following format:\n"
        "'''\n"
        "[{{\"task\": task, \"id\", task_id, \"dep\": [dependency_task_id1, dependency_task_id2, ...]}}]\n"
        "'''\n"
        "The \"dep\" field denotes the id of the previous task which generates a new resource upon which the current task depends. If there are no dependencies, set \"dep\" to -1.\n\n"
        "This is user's question: {question}\n"
        "These are subtasks of this question:\n"
        "{task_ls}\n"
        "Output: "
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question, task_ls=task_ls)
            result = eval(result)
            for i in range(len(result)):
                if isinstance(result[i]['dep'], str):
                    temp = []
                    for ele in result[i]['dep'].split(','):
                        temp.append(int(ele))
                    result[i]['dep'] = temp
                elif isinstance(result[i]['dep'], int):
                    result[i]['dep'] = [result[i]['dep']]
                elif isinstance(result[i]['dep'], list):
                    temp = []
                    for ele in result[i]['dep']:
                        temp.append(int(ele))
                    result[i]['dep'] = temp
                elif result[i]['dep'] == -1:
                    result[i]['dep'] = [-1]
            a = result[i]['dep'][0]
            return result
        except Exception as e:
            print(f"task topology fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return result


def answer_generation_direct(task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You need to answer the user's question.\n"
        "This is the user's question: {task}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(task=task)
    return result


def choose_parameter(API_instruction, api, api_dic, question, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "This is an API tool documentation. Given a user's question, you need to output parameters according to the API tool documentation to successfully call the API to solve the user's question.\n"
        "This is API tool documentation: {api_dic}\n"
        "Please note that: \n"
        "1. The Example in the API tool documentation can help you better understand the use of the API.\n"
        "2. Ensure the parameters you output are correct. The output must contain the required parameters, and can contain the optional parameters based on the question. If no paremters in the required parameters and optional parameters, just leave it as {{\"Parameters\":{{}}}}\n"
        "3. If the user's question mentions other APIs, you should ONLY consider the API tool documentation I give and do not consider other APIs.\n"
        "4. If you need to use this API multiple times, please set \"Parameters\" to a list.\n"
        "5. You must ONLY output in a parsible JSON format. Two examples output looks like:\n"
        "'''\n"
        "Example 1: {{\"Parameters\":{{\"input\": [1,2,3]}}}}\n"
        "Example 2: {{\"Parameters\":[{{\"input\": [1,2,3]}}, {{\"input\": [2,3,4]}}]}}\n"
        "'''\n"
        "This is user's question: {question}\n"
        "Output:\n"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(api_dic=api_dic,
                               question=question, )
            clean_answer = eval(
                result.replace(": true", ": True").replace(":true", ": True").replace(":false", ": False").replace(
                    ": false", ": False").replace("```", "").strip())
            a = clean_answer["Parameters"]

            return a
        except Exception as e:
            print(f"Choose Parameter fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return a


def choose_parameter_depend(API_instruction, api, api_dic, question, model_name, previous_log):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Given a user's question and a API tool documentation, you need to output parameters according to the API tool documentation to successfully call the API to solve the user's question.\n"
        "Please note that: \n"
        "1. The Example in the API tool documentation can help you better understand the use of the API.\n"
        "2. Ensure the parameters you output are correct. The output must contain the required parameters, and can contain the optional parameters based on the question. If no paremters in the required parameters and optional parameters, just leave it as {{\"Parameters\":{{}}}}\n"
        "3. If the user's question mentions other APIs, you should ONLY consider the API tool documentation I give and do not consider other APIs.\n"
        "4. The question may have dependencies on answers of other questions, so we will provide logs of previous questions and answers for your reference.\n"
        "5. If you need to use this API multiple times,, please set \"Parameters\" to a list.\n"
        "6. You must ONLY output in a parsible JSON format. Two examples output looks like:\n"
        "'''\n"
        "Example 1: {{\"Parameters\":{{\"input\": [1,2,3]}}}}\n"
        "Example 2: {{\"Parameters\":[{{\"input\": [1,2,3]}}, {{\"input\": [2,3,4]}}]}}\n"
        "'''\n"
        "There are logs of previous questions and answers: \n {previous_log}\n"
        "This is the current user's question: {question}\n"
        "This is API tool documentation: {api_dic}\n"
        "Output:\n"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(api_dic=api_dic,
                               question=question,
                               previous_log=previous_log)
            clean_answer = eval(
                result.replace(": true", ": True").replace(": false", ": False").replace("```", "").strip())
            a = clean_answer["Parameters"]

            return a
        except Exception as e:
            print(f"choose parameter depend fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return a


def Call_function(B, arg, id):
    app_path = 'data_funcqa/funchub/math.py'
    spec = importlib.util.spec_from_file_location('math', app_path)
    app_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(app_module)
    if hasattr(app_module, B):
        function_B = getattr(app_module, B)
        try:
            call_result = function_B(arg['input'])
            return call_result
        except Exception as e:
            try:
                arg = {change_name(k.lower()): v for k, v in arg.items()}
                call_result = function_B(arg['input'])
                return call_result
            except Exception as e:
                try:
                    arg = {change_name(k.lower()): v for k, v in arg.items()}
                    arg = {change_name(k.replace("-", "_")): v for k, v in arg.items()}
                    call_result = function_B(arg['input'])
                    return call_result
                except Exception as e:
                    print(f"fails: {e}")
                    with open('wrong_log.json', 'a+', encoding='utf-8') as f:
                        line = json.dumps({
                            "id": id,
                            "parameters": arg,
                            "wrong": str(e)
                        }, ensure_ascii=False)
                        f.write(line + '\n')
                    return -1
    else:
        with open('wrong_log.json', 'a+', encoding='utf-8') as f:
            line = json.dumps({
                "id": id,
                "parameters": arg,
                "wrong": f"No function named {B} in {app_path}"
            }, ensure_ascii=False)
            f.write(line + '\n')
        return (f"No function named {B} in {app_path}")


def retrieval(question, Tool_dic, dataset, tool_used, ind, model_name, previous_log=None):
    tool_id = choose_tool(question, Tool_dic, tool_used, model_name)
    if tool_id == -1:
        return tool_id, "", "", "", ""
    tool_instruction = dataset[str(tool_id["ID"])]
    API_instruction = tool_instruction["API_description"]
    API_tool = tool_instruction["standardized_name"]

    api_selection = [API_tool]
    api_result = []
    for api in api_selection:
        if previous_log is None:
            parameter = choose_parameter(API_instruction, api,
                                         tool_instruction["Usage"], question, model_name)
        else:
            parameter = choose_parameter_depend(API_instruction, api,
                                                tool_instruction["Usage"],
                                                question, model_name, previous_log)
        if parameter == -1:
            continue
        api_result.append({"api_name": api, "parameters": parameter})
    if len(api_result) == 0:
        call_result = ""
        return tool_id, api_result, call_result, tool_instruction, API_instruction
    if isinstance(api_result, set) or isinstance(api_result, list):
        call_results = []
        for api in api_result:
            if isinstance(api["parameters"], dict):
                parameters = {}
                for key in api["parameters"]:
                    value = api["parameters"][key]
                    key = change_name(key)
                    parameters[key] = value
                call_result = Call_function(API_tool, parameters, ind)
                if call_result == -1:
                    continue
                call_results.append(str(call_result))
            elif isinstance(api["parameters"], list):
                for para_ls in api["parameters"]:
                    parameters = {}
                    for key in para_ls:
                        value = para_ls[key]
                        key = change_name(key)
                        parameters[key] = value
                    call_result = Call_function(API_tool, parameters, ind)
                    if call_result == -1:
                        continue
                    call_results.append(str(call_result))
        call_result = '\n\n'.join(call_results)
    elif isinstance(api_result, dict):
        api = api_result
        if isinstance(api["parameters"], dict):
            parameters = {}
            for key in api["parameters"]:
                value = api["parameters"][key]
                key = change_name(key)
                parameters[key] = value
            call_result = Call_function(API_tool, parameters, ind)
        elif isinstance(api["parameters"], list):
            call_results = []
            for para_ls in api["parameters"]:
                parameters = {}
                for key in para_ls:
                    value = para_ls[key]
                    key = change_name(key)
                    parameters[key] = value
                call_result = Call_function(API_tool, parameters, ind)
                if call_result == -1:
                    continue
                call_results.append(str(call_result))
            call_result = '\n\n'.join(call_results)

    return tool_id, api_result, call_result, tool_instruction, API_instruction


def answer_generation(question, API_instruction, call_result, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You should answer the question based on the response output by the API tool."
        "Please note that:\n"
        "1. Answer the question in natural language based on the API response reasonably and effectively.\n"
        "2. The user cannot directly get API response, "
        "so you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "3. If the API tool does not provide useful information in the response, "
        "please answer with your knowledge.\n"
        "This is the user's question:\n {question}\n"
        "This is the API response:\n {call_result}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question,
                               API_instruction=API_instruction,
                               call_result=call_result, )
            break
        except Exception as e:
            print(f"answer generation fails: {e}")
            if ind > 2:
                return -1
            ind += 1
            continue
    return result


def answer_generation_depend(question, API_instruction, call_result, previous_log, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You should answer the question based on the response output by the API tool."
        "Please note that:\n"
        "1. Try to organize the response into a natural language answer.\n"
        "2. We will not show the API response to the user, "
        "thus you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "3. If the API tool does not provide useful information in the response, "
        "please answer with your knowledge.\n"
        "4. The question may have dependencies on answers of other questions, so we will provide logs of previous questions and answers.\n"
        "There are logs of previous questions and answers: \n {previous_log}\n"
        "This is the user's question: {question}\n"
        "This is the response output by the API tool: \n{call_result}\n"
        "We will not show the API response to the user, "
        "thus you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question,
                               API_instruction=API_instruction,
                               call_result=call_result,
                               previous_log=previous_log)
            break
        except Exception as e:
            print(f"answer generation depend fails: {e}")
            if ind > 2:
                return -1
            ind += 1
            continue
    return result


def answer_summarize(question, answer_task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "We break down a complex user's problems into simple subtasks and provide answers to each simple subtask. "
        "You need to organize these answers to each subtask and form a self-consistent final answer to the user's question\n"
        "This is the user's question: {question}\n"
        "These are subtasks and their answers: {answer_task}\n"
        "Final answer:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(question=question, answer_task=answer_task)
    return result


def answer_check(question, answer, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Please check whether the response can reasonably and accurately answer the question."
        "If can, please output 'YES'; If not, please output 'NO'\n"
        "You need to give reasons first and then decide whether the response can reasonably and accurately answer the question. You must only output in a parsible JSON format. Two example outputs look like:\n"
        "Example 1: {{\"Reason\": \"The reason why you think the response can reasonably and accurately answer the question\", \"Choice\": \"Yes\"}}\n"
        "Example 2: {{\"Reason\": \"The reason why you think the response cannot reasonably and accurately answer the question\", \"Choice\": \"No\"}}\n"
        "This is the user's question: {question}\n"
        "This is the response: {answer}\n"
        "Output: "
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(question=question, answer=answer)
    if 'yes'.lower() in eval(result)["Choice"].lower():
        return 1
    else:
        return -1


def task_execution_mh(data_type, start_index, total_files,
                      retrieval_num, ind, model_name, dataset,
                      Tool_dic, test_data, progress_file):
    with tqdm(total=total_files, desc="Processing files", initial=start_index) as pbar:
        for i, data in enumerate(test_data[start_index:], start=start_index):
            answer_ls = []
            question = data["question"]
            print(question)
            temp = task_decompose(question, Tool_dic, model_name)['Tasks']
            task_ls = []
            for t in range(len(temp)):
                task_ls.append({"task": temp[t], "id": t + 1})
            task_ls = task_topology(question, task_ls, model_name)
            task_depend = {'Original Question': question}
            for task_dic in task_ls:
                task_depend[task_dic['id']] = {'task': task_dic['task'], 'answer': ''}
            answer_task = []
            tool_instruction_ls = []
            api_result_ls = []
            call_result_ls = []
            tool_check_reason_ls = []
            for task_dic in task_ls:
                task = task_dic['task']
                print("Do need tool.")
                tool_used = []
                depend_id = [1]
                for r in range(retrieval_num):
                    if depend_id[0] == -1:
                        tool_id, api_result, call_result, tool_instruction, API_instruction = retrieval(task, Tool_dic,
                                                                                                        dataset,
                                                                                                        tool_used, ind,
                                                                                                        model_name)
                        if len(str(call_result)) > 5000:
                            call_result = str(call_result)[:5000]
                        answer = answer_generation(task, API_instruction, call_result, model_name)
                    else:
                        previous_log = task_depend
                        tool_id, api_result, call_result, tool_instruction, API_instruction = retrieval(task, Tool_dic,
                                                                                                        dataset,
                                                                                                        tool_used, ind,
                                                                                                        model_name,
                                                                                                        previous_log=previous_log)
                        if len(str(call_result)) > 5000:
                            call_result = str(call_result)[:5000]
                        answer = answer_generation_depend(task, API_instruction, call_result, previous_log, model_name)

                    check_index = 1
                    if str(call_result).strip() == '-1' or str(call_result).strip() == '':
                        check_index = -1
                    if check_index == 1:
                        answer_task.append({'task': task, 'answer': answer})
                        tool_instruction_ls.append(tool_instruction)
                        api_result_ls.append(api_result)
                        call_result_ls.append(call_result)
                        break
                    else:
                        answer_ls.append({'task': task, 'answer': answer})
                        try:
                            tool_used.append(str(tool_id["ID"]))
                        except:
                            continue
                        print('****Try Again****')

                task_depend[task_dic['id']]['answer'] = answer
            final_answer = answer_summarize(question, answer_task, model_name)
            check_index = answer_check(question, final_answer, model_name)
            ind = ind + 1
            with open(f"FuncQA_{data_type}_{model_name}_easytool.jsonl", 'a+', encoding='utf-8') as f:
                line = json.dumps({
                    "ID": ind,
                    "question": question,
                    "final_answer": final_answer,
                    "subtask": task_ls,
                    "answer_subtask": answer_task,
                    "answer_wrong": answer_ls,
                    "check_index": check_index,
                    "execute_log": {
                        "api_result_ls": api_result_ls,
                        "call_result_ls": call_result_ls,
                        "tool_check_reason_ls": tool_check_reason_ls,
                        "tool_instruction_ls": tool_instruction_ls,
                    },
                    "check": 0
                }, ensure_ascii=False)
                f.write(line + '\n')

            print(final_answer)
            update_progress(progress_file, i + 1)
            pbar.update(1)


def task_execution_oh(data_type, start_index, total_files,
                      retrieval_num, ind, model_name, dataset,
                      Tool_dic, test_data, progress_file):
    with tqdm(total=total_files, desc="Processing files", initial=start_index) as pbar:
        for i, data in enumerate(test_data[start_index:], start=start_index):
            answer_ls = []
            question = data["question"]
            print(question)
            task_ls = [{"task": question}]
            answer_task = []
            tool_instruction_ls = []
            api_result_ls = []
            call_result_ls = []
            tool_check_reason_ls = []
            for task_dic in task_ls:
                task = task_dic['task']
                print("Do need tool.")
                tool_used = []
                depend_id = [1]
                for r in range(retrieval_num):
                    tool_id, api_result, call_result, tool_instruction, API_instruction = retrieval(task, Tool_dic,
                                                                                                    dataset,
                                                                                                    tool_used, ind,
                                                                                                    model_name)
                    if len(str(call_result)) > 5000:
                        call_result = str(call_result)[:5000]
                    answer = answer_generation(task, API_instruction, call_result, model_name)

                    check_index = 1
                    if str(call_result).strip() == '-1' or str(call_result).strip() == '':
                        check_index = -1
                    if check_index == 1:
                        answer_task.append({'task': task, 'answer': answer})
                        tool_instruction_ls.append(tool_instruction)
                        api_result_ls.append(api_result)
                        call_result_ls.append(call_result)
                        break
                    else:
                        answer_ls.append({'task': task, 'answer': answer})
                        try:
                            tool_used.append(str(tool_id["ID"]))
                        except:
                            continue
                        print('****Try Again****')

            final_answer = answer_summarize(question, answer_task, model_name)
            check_index = answer_check(question, final_answer, model_name)
            ind = ind + 1
            with open(f"FuncQA_{data_type}_{model_name}_easytool.jsonl", 'a+', encoding='utf-8') as f:
                line = json.dumps({
                    "ID": ind,
                    "question": question,
                    "final_answer": final_answer,
                    "subtask": task_ls,
                    "answer_subtask": answer_task,
                    "answer_wrong": answer_ls,
                    "check_index": check_index,
                    "execute_log": {
                        "api_result_ls": api_result_ls,
                        "call_result_ls": call_result_ls,
                        "tool_check_reason_ls": tool_check_reason_ls,
                        "tool_instruction_ls": tool_instruction_ls,
                    },
                    "check": 0
                }, ensure_ascii=False)
                f.write(line + '\n')

            print(final_answer)
            update_progress(progress_file, i + 1)
            pbar.update(1)




File: easytool/easytool/toolbench_retrieve.py

# — coding: utf-8 –
import openai
import json
import logging
import sys
import argparse
from langchain.chat_models import ChatOpenAI
from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate
)
from langchain import LLMChain
import numpy as np
import requests
import os
import subprocess
import re
import importlib.util
from sklearn.metrics.pairwise import cosine_similarity
import pickle
from util import *
from tqdm import tqdm

openai.api_key = os.environ["OPENAI_API_KEY"]


def get_last_processed_index(progress_file):
    """Retrieve the last processed index from the progress file."""
    if os.path.exists(progress_file):
        with open(progress_file, 'r', encoding='utf-8') as f:
            last_index = f.read().strip()
            return int(last_index) if last_index else 0
    else:
        return 0


def update_progress(progress_file, index):
    """Update the last processed index in the progress file."""
    with open(progress_file, 'w', encoding='utf-8') as f:
        f.write(str(index))


def get_embedding(text):
    a = openai.Embedding.create(
        engine="text-embedding-ada-002",
        input=text
    )
    return a['data'][0]["embedding"]


def retrieve_reference(embedded_texts, filenames, question, k):
    input_text = question
    input_embedding = get_embedding(input_text)
    similarities = [cosine_similarity([input_embedding], [emb])[0][0] for emb in embedded_texts]
    top_k_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)[:k]
    return [filenames[i] for i in top_k_indices]


def choose_tool(question, Tool_dic, tool_used, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "This is the user's question: {question}\n"
        "These are the tools you can select to solve the question:\n"
        "Tool List:\n"
        "{Too_list}\n\n"
        "Please note that: \n"
        "1. You should only chooce one tool the Tool List to solve this question.\n"
        "2. You must ONLY output the ID of the tool you chose in a parsible JSON format. An example output looks like:\n"
        "'''\n"
        "Example: {{\"ID\": XX}}\n"
        "'''\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    Tool_list = []
    for ele in Tool_dic:
        for key in ele.keys():
            if str(key) not in tool_used:
                Tool_list.append(f'''ID: {key}\n{ele[key]}''')
    while True:
        try:
            result = chain.run(question=question,
                               Too_list='\n'.join(Tool_list))
            clean_answer = eval(result.split("\n\n")[-1].strip())
            break
        except Exception as e:
            print(f"choose tool fails:{e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return clean_answer


def choose_API(API_instruction, API_list, question, model_name):
    input_execute_rapidapi_api_note = '''
This is an API Tool instruction. Given a question, you should choose APIs from the API list you want to use for this question in this instruction.
you must only output in a parsible Python List Format. An example output looks like:
```
["api1", "api2", ...]
```
'''.strip()
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "{API_instruction}\n"
        "{input_execute_rapidapi_api_note}\n"
        "This is the API list: {API_list}\n"
        "Please note that: \n"
        "1. The APIs you choose must in the API list.\n"
        "2. You must ONLY output in the following parsible Python List Format.\n"
        "```\n"
        "Output_Example: [\"api1\", \"api2\", ...]\n"
        "```\n"
        "Question: {question}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(API_instruction=API_instruction,
                               API_list=API_list,
                               question=question,
                               input_execute_rapidapi_api_note=input_execute_rapidapi_api_note)
            clean_answer = eval(result.replace("```", "").strip().split("\n\n")[0].strip())
            if isinstance(clean_answer, str):
                ls = [clean_answer]
            elif isinstance(clean_answer, list):
                ls = clean_answer
            temp = []
            for ele in ls:
                if ele in API_list:
                    temp.append(ele)
            ls = temp
            return ls
        except Exception as e:
            print(f"Choose API fails:{e}")
            print(result)
            if ind > 10:
                return []
            ind += 1
            continue
    return ls


def choose_parameter(API_instruction, api, api_dic, question, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "This is an API tool documentation. Given a user's question, you need to output parameters according to the API tool documentation to successfully call the API to solve the user's question.\n"
        "This is API tool documentation: {api_dic}\n"
        "Please note that: \n"
        "1. The Example in the API tool documentation can help you better understand the use of the API.\n"
        "2. Ensure the parameters you output are correct. The output must contain the required parameters, and can contain the optional parameters based on the question. If no paremters in the required parameters and optional parameters, just leave it as {{\"Parameters\":{{}}}}\n"
        "3. If the user's question mentions other APIs, you should ONLY consider the API tool documentation I give and do not consider other APIs.\n"
        "4. If you need to use this API multiple times, please set \"Parameters\" to a list.\n"
        "5. You must ONLY output in a parsible JSON format. Two examples output looks like:\n"
        "'''\n"
        "Example 1: {{\"Parameters\":{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}}}\n"
        "Example 2: {{\"Parameters\":[{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}, {{\"keyword\": \"Machine Learning\", \"language\": \"English\"}}]}}\n"
        "'''\n"
        "This is user's question: {question}\n"
        "Output:\n"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(api_dic=api_dic,
                               question=question, )
            clean_answer = eval(
                result.replace(": true", ": True").replace(":true", ": True").replace(":false", ": False").replace(
                    ": false", ": False").replace("```", "").strip())
            a = clean_answer["Parameters"]

            return a
        except Exception as e:
            print(f"Choose Parameter fails:{e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return a


def choose_parameter_depend(API_instruction, api, api_dic, question, previous_log, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Given a user's question and a API tool documentation, you need to output parameters according to the API tool documentation to successfully call the API to solve the user's question.\n"
        "Please note that: \n"
        "1. The Example in the API tool documentation can help you better understand the use of the API.\n"
        "2. Ensure the parameters you output are correct. The output must contain the required parameters, and can contain the optional parameters based on the question. If no paremters in the required parameters and optional parameters, just leave it as {{\"Parameters\":{{}}}}\n"
        "3. If the user's question mentions other APIs, you should ONLY consider the API tool documentation I give and do not consider other APIs.\n"
        "4. The question may have dependencies on answers of other questions, so we will provide logs of previous questions and answers for your reference.\n"
        "5. If you need to use this API multiple times,, please set \"Parameters\" to a list.\n"
        "6. You must ONLY output in a parsible JSON format. Two examples output looks like:\n"
        "'''\n"
        "Example 1: {{\"Parameters\":{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}}}\n"
        "Example 2: {{\"Parameters\":[{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}, {{\"keyword\": \"Machine Learning\", \"language\": \"English\"}}]}}\n"
        "'''\n"
        "There are logs of previous questions and answers: \n {previous_log}\n"
        "This is the current user's question: {question}\n"
        "This is API tool documentation: {api_dic}\n"
        "Output:\n"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(api_dic=api_dic,
                               question=question,
                               previous_log=previous_log)
            clean_answer = eval(
                result.replace(": true", ": True").replace(": false", ": False").replace("```", "").strip())
            a = clean_answer["Parameters"]

            return a
        except Exception as e:
            print(f"choose parameter depend fails:{e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return a


def answer_generation(question, API_instruction, call_result, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You should answer the question based on the response output by the API tool."
        "Please note that:\n"
        "1. Answer the question in natural language based on the API response reasonably and effectively.\n"
        "2. The user cannot directly get API response, "
        "so you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "This is the user's question:\n {question}\n"
        "This is the API response:\n {call_result}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question,
                               call_result=call_result)
            break
        except Exception as e:
            print(f"answer generation fails:{e}")
            if ind > 2:
                return -1
            ind += 1
            continue
    return result


def answer_generation_depend(question, API_instruction, call_result, model_name, previous_log):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You should answer the question based on the response output by the API tool."
        "Please note that:\n"
        "1. Try to organize the response into a natural language answer.\n"
        "2. We will not show the API response to the user, "
        "thus you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "3. The question may have dependencies on answers of other questions, so we will provide logs of previous questions and answers.\n"
        "There are logs of previous questions and answers: \n {previous_log}\n"
        "This is the user's question: {question}\n"
        "This is the response output by the API tool: \n{call_result}\n"
        "We will not show the API response to the user, "
        "thus you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question,
                               call_result=call_result,
                               previous_log=previous_log)
            break
        except Exception as e:
            print(f"answer generation depend fails:{e}")
            if ind > 2:
                return -1
            ind += 1
            continue
    return result


def answer_check(question, answer, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Please check whether the response can reasonably and accurately answer the question."
        "If can, please output 'YES'; If not, please output 'NO'\n"
        "You need to give reasons first and then decide whether the response can reasonably and accurately answer the question. You must only output in a parsible JSON format. Two example outputs look like:\n"
        "Example 1: {{\"Reason\": \"The reason why you think the response can reasonably and accurately answer the question\", \"Choice\": \"Yes\"}}\n"
        "Example 2: {{\"Reason\": \"The reason why you think the response cannot reasonably and accurately answer the question\", \"Choice\": \"No\"}}\n"
        "This is the user's question: {question}\n"
        "This is the response: {answer}\n"
        "Output: "
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(question=question, answer=answer)
    if 'yes'.lower() in str(result).lower():
        return 1
    else:
        return -1


def Call_function(A, B, arg, index, id):
    if A in index:
        for path in index[A]:
            app_path = os.path.join(path, A, 'api.py')
            if os.path.isfile(app_path):
                spec = importlib.util.spec_from_file_location('api', app_path)
                app_module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(app_module)
                arg['toolbench_rapidapi_key'] = os.environ['RAPIDAPI_KEY']
                # Check if B is a function in app
                if hasattr(app_module, B):
                    function_B = getattr(app_module, B)
                    try:
                        call_result = function_B(**arg)
                        return call_result
                    except Exception as e:
                        try:
                            arg = {change_name(k.lower()): v for k, v in arg.items()}
                            call_result = function_B(**arg)
                            return call_result
                        except Exception as e:
                            try:
                                arg = {change_name(k.replace("-", "_")): v for k, v in arg.items()}
                                call_result = function_B(**arg)
                                return call_result
                            except Exception as e:
                                try:
                                    arg = {change_name(k.replace("\\", "")): v for k, v in arg.items()}
                                    call_result = function_B(**arg)
                                    return call_result
                                except Exception as e:
                                    print(f"Call function fails:{e}")
                                    with open('wrong_log.json', 'a+', encoding='utf-8') as f:
                                        line = json.dumps({
                                            "id": id,
                                            "parameters": arg,
                                            "wrong": str(e)
                                        }, ensure_ascii=False)
                                        f.write(line + '\n')
                                    return -1
                else:
                    with open('wrong_log.json', 'a+', encoding='utf-8') as f:
                        line = json.dumps({
                            "id": id,
                            "parameters": arg,
                            "wrong": f"No function named {B} in {app_path}"
                        }, ensure_ascii=False)
                        f.write(line + '\n')
                    return (f"No function named {B} in {app_path}")


def retrieval(question, Tool_dic, dataset, tool_used, ind, model_name, index, previous_log=None):
    tool_id = choose_tool(question, Tool_dic, tool_used, model_name)
    if tool_id == -1:
        return tool_id, "", "", "", ""
    if str(tool_id["ID"]) not in dataset:
        return tool_id, "", "", "", ""
    tool_instruction = dataset[str(tool_id["ID"])]
    API_instruction = tool_instruction["tool_description"]
    API_tool = tool_instruction["standardized_name"]
    API_list = []
    for ele in tool_instruction["tool_guidelines"].keys():
        API_list.append(ele)

    api_selection = choose_API(API_instruction, API_list, question, model_name)
    api_result = []
    if len(api_selection) == 0:
        call_result = ""
        print("No Calling")
        return tool_id, api_result, call_result, tool_instruction, API_instruction
    for api in api_selection:
        if previous_log is None:
            parameter = choose_parameter(API_instruction, api,
                                         tool_instruction["tool_guidelines"][api], question,
                                         model_name)
        else:
            parameter = choose_parameter_depend(API_instruction, api,
                                                tool_instruction["tool_guidelines"][api],
                                                question, previous_log,
                                                model_name)
        if parameter == -1:
            continue
        api_result.append({"api_name": api, "parameters": parameter})
    if len(api_result) == 0:
        call_result = ""
        return tool_id, api_result, call_result, tool_instruction, API_instruction
    if isinstance(api_result, set) or isinstance(api_result, list):
        call_results = []
        for api in api_result:
            api_name = change_name(standardize(api["api_name"]))

            if isinstance(api["parameters"], dict):
                parameters = {}
                for key in api["parameters"]:
                    value = api["parameters"][key]
                    key = change_name(key)
                    parameters[key] = value
                call_result = Call_function(API_tool, api_name, parameters, index, ind)
                if call_result == -1:
                    continue
                call_results.append(str(call_result))
            elif isinstance(api["parameters"], list):
                for para_ls in api["parameters"]:
                    parameters = {}
                    for key in para_ls:
                        value = para_ls[key]
                        key = change_name(key)
                        parameters[key] = value
                    call_result = Call_function(API_tool, api_name, parameters, index, ind)
                    if call_result == -1:
                        continue
                    call_results.append(str(call_result))
        call_result = '\n\n'.join(call_results)
    elif isinstance(api_result, dict):
        api_name = change_name(standardize(api_result["api_name"]))
        api = api_result
        if isinstance(api["parameters"], dict):
            parameters = {}
            for key in api["parameters"]:
                value = api["parameters"][key]
                key = change_name(key)
                parameters[key] = value
            call_result = Call_function(API_tool, api_name, parameters, index, ind)
        elif isinstance(api["parameters"], list):
            call_results = []
            for para_ls in api["parameters"]:
                parameters = {}
                for key in para_ls:
                    value = para_ls[key]
                    key = change_name(key)
                    parameters[key] = value
                call_result = Call_function(API_tool, api_name, parameters, index, ind)
                if call_result == -1:
                    continue
                call_results.append(str(call_result))
            call_result = '\n\n'.join(call_results)

    return tool_id, api_result, call_result, tool_instruction, API_instruction


def task_decompose(question, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You need to decompose a complex user's question into some simple subtasks and let the model execute it step by step.\n"
        "This is the user's question: {question}\n"
        "Please note that: \n"
        "1. You should only decompose this complex user's question into some simple subtasks which can be executed easily by using a single tool.\n"
        "2. Each simple subtask should be expressed into natural language.\n"
        "3. Each subtask should contain the necessary information from the original question and should be complete, explicit and self-consistent.\n"
        "4. You must ONLY output the ID of the tool you chose in a parsible JSON format. An example output looks like:\n"
        "'''\n"
        "{{\"Tasks\": [\"Task 1\", \"Task 2\", ...]}}\n"
        "'''\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question)
            result = eval(result.split('\n\n')[0])
            a = result["Tasks"]
            break
        except Exception as e:
            print(f"task decompose fails:{e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return result


def task_topology(question, task_ls, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Given a complex user's question, I have decompose this question into some simple subtasks"
        "I think there exists a logical connections and order amontg the tasks. "
        "Thus you need to help me output this logical connections and order.\n"
        "You must ONLY output in a parsible JSON format with the following format:\n"
        "'''\n"
        "[{{\"task\": task, \"id\", task_id, \"dep\": [dependency_task_id1, dependency_task_id2, ...]}}]\n"
        "'''\n"
        "The \"dep\" field denotes the id of the previous task which generates a new resource upon which the current task depends. If there are no dependencies, set \"dep\" to -1.\n\n"
        "This is user's question: {question}\n"
        "These are subtasks of this question:\n"
        "{task_ls}\n"
        "Output: "
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question, task_ls=task_ls)
            result = eval(result)
            for i in range(len(result)):
                if isinstance(result[i]['dep'], str):
                    temp = []
                    for ele in result[i]['dep'].split(','):
                        temp.append(int(ele))
                    result[i]['dep'] = temp
                elif isinstance(result[i]['dep'], int):
                    result[i]['dep'] = [result[i]['dep']]
                elif isinstance(result[i]['dep'], list):
                    temp = []
                    for ele in result[i]['dep']:
                        temp.append(int(ele))
                    result[i]['dep'] = temp
                elif result[i]['dep'] == -1:
                    result[i]['dep'] = [-1]
            a = result[i]['dep'][0]
            return result
        except Exception as e:
            print(f"task topology fails:{e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return result


def answer_summarize(question, answer_task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "We break down a complex user's problems into simple subtasks and provide answers to each simple subtask. "
        "You need to organize these answers to each subtask and form a self-consistent final answer to the user's question\n"
        "This is the user's question: {question}\n"
        "These are subtasks and their answers: {answer_task}\n"
        "Final answer:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(question=question, answer_task=answer_task)
    return result


def answer_generation_direct(task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You need to answer the user's question.\n"
        "This is the user's question: {task}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(task=task)
    return result


def tool_check(task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful language model which can use external APIs to solve user's question."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "As a powerful language model, you're equipped to answer user's question with accumulated knowledge. "
        "However, in some cases, you need to use external APIs to answer accurately."
        "Thus, you need to check whether the user's question requires you to call an external API to solve it.\n"
        "Here are some tips to help you check: \n"
        "1. If the user's question requires real-time information, since your knowledge base isn't updated in real-time, any such question will demand an API call.\n"
        "2. If you need to obtain information (e.g., ID, name, phone number, geographical location, rank, etc.), you need to call the database APIs if you are not sure.\n"
        "3. If the question demand a database search or internet research to generate an answer, this is another situation where an API call is necessary.\n"
        "If need, please output 'YES'; If not, please output 'NO'\n"
        "You need to give reasons first and then decide whether to keep it or not. You must only output in a parsible JSON format. Two example outputs look like:\n"
        "Example 1: {{\"Reason\": \"The reason why you think you do not need to call an external API to solve the user's question\", \"Choice\": \"No\"}}\n"
        "Example 2: {{\"Reason\": \"The reason why you think you need to call an external API to solve the user's question\", \"Choice\": \"Yes\"}}\n"
        "This is the user's question: {task}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(task=task)
            result = eval(result)
            a = result["Reason"]
            b = result["Choice"]
            if 'yes' in b.lower():
                return result, -1
            else:
                return result, 1
        except Exception as e:
            print(f"tool check fails:{e}")
            if ind > 10:
                return "", -1
            ind += 1
            continue
    return result, -1


def task_execution(data_type,
                   base_path, index, dataset, test_data, progress_file,
                   start_index, total_files, retrieval_num, ind, model_name):
    with open("data_toolbench/tool_instruction/API_description_embeddings.pkl", "rb") as file:
        filenames, embedded_texts = pickle.load(file)
    with tqdm(total=total_files, desc="Processing files", initial=start_index) as pbar:
        for i, data in enumerate(test_data[start_index:], start=start_index):
            answer_ls = []
            question = data["query"]
            print(question)
            temp = task_decompose(question, model_name)['Tasks']
            task_ls = []
            for t in range(len(temp)):
                task_ls.append({"task": temp[t], "id": t + 1})
            task_ls = task_topology(question, task_ls, model_name)
            task_depend = {}
            for task_dic in task_ls:
                task_depend[task_dic['id']] = {'task': task_dic['task'], 'answer': ''}
            answer_task = []
            api_result_ls = []
            call_result_ls = []
            tool_check_reason_ls = []
            parameter_ls = []
            for task_dic in task_ls:
                task = task_dic['task']
                tool_check_reason, tool_check_result = tool_check(task, model_name)
                tool_check_reason_ls.append(tool_check_reason)
                if tool_check_result == 1:
                    print("Do not need tool.")
                    answer = answer_generation_direct(task)
                    answer_task.append({'task': task, 'answer': answer})
                else:
                    print("Do need tool.")
                    depend_id = task_dic['dep']
                    tool_used = []
                    Tool_dic = [{tool: dataset[str(tool)]["tool_description"]} for tool in
                                retrieve_reference(embedded_texts, filenames, task, k=5)]
                    for r in range(retrieval_num):
                        if depend_id[0] == -1:
                            tool_id, api_result, call_result, tool_instruction, API_instruction = retrieval(task,
                                                                                                            Tool_dic,
                                                                                                            dataset,
                                                                                                            tool_used,
                                                                                                            ind,
                                                                                                            model_name,
                                                                                                            index)
                            call_result = str(call_result)[:1000]
                            answer = answer_generation(task, API_instruction,
                                                       call_result, model_name)
                        else:
                            previous_log = []
                            for ids in depend_id:
                                previous_log.append(task_depend[ids])
                            tool_id, api_result, call_result, tool_instruction, API_instruction = retrieval(task,
                                                                                                            Tool_dic,
                                                                                                            dataset,
                                                                                                            tool_used,
                                                                                                            ind,
                                                                                                            model_name,
                                                                                                            index,
                                                                                                            previous_log=previous_log)
                            call_result = str(call_result)[:1000]
                            answer = answer_generation_depend(task, API_instruction, call_result, model_name,
                                                              previous_log=previous_log)

                        check_index = answer_check(task, answer, model_name)
                        if check_index == 1:
                            answer_task.append({'task': task, 'answer': answer})
                            api_result_ls.append(api_result)
                            call_result_ls.append(call_result)
                            break
                        else:
                            answer_ls.append({'task': task, 'answer': answer})
                            try:
                                tool_used.append(str(tool_id["ID"]))
                            except:
                                continue
                            print('****Try Again****')
                task_depend[task_dic['id']]['answer'] = answer
            final_answer = answer_summarize(question, answer_task, model_name)
            check_index = answer_check(question, final_answer, model_name)

            ind = ind + 1
            with open(f'''{data_type}_{model_name}_retrieve_Easytool.jsonl''', 'a+', encoding='utf-8') as f:
                line = json.dumps({
                    "ID": ind,
                    "question": question,
                    "final_answer": final_answer,
                    "subtask": task_ls,
                    "answer_subtask": answer_task,
                    "answer_wrong": answer_ls,
                    "check_index": check_index,
                    "execute_log": {
                        "api_result_ls": api_result_ls,
                        "parameter_ls": parameter_ls,
                        "call_result_ls": call_result_ls,
                        "tool_check_reason_ls": tool_check_reason_ls,
                    }
                }, ensure_ascii=False)
                f.write(line + '\n')

            print(final_answer)
            update_progress(progress_file, i + 1)
            pbar.update(1)



File: easytool/easytool/toolbench.py

# — coding: utf-8 –
import openai
import json
import logging
import sys
import argparse
from langchain.chat_models import ChatOpenAI
from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate
)
from langchain import LLMChain
import numpy as np
import requests
import os
import subprocess
import re
import importlib.util
from sklearn.metrics.pairwise import cosine_similarity
import pickle
from util import *
from tqdm import tqdm

openai.api_key = os.environ["OPENAI_API_KEY"]


def get_last_processed_index(progress_file):
    """Retrieve the last processed index from the progress file."""
    if os.path.exists(progress_file):
        with open(progress_file, 'r', encoding='utf-8') as f:
            last_index = f.read().strip()
            return int(last_index) if last_index else 0
    else:
        return 0


def update_progress(progress_file, index):
    """Update the last processed index in the progress file."""
    with open(progress_file, 'w', encoding='utf-8') as f:
        f.write(str(index))


def choose_tool(question, Tool_dic, tool_used, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "This is the user's question: {question}\n"
        "These are the tools you can select to solve the question:\n"
        "Tool List:\n"
        "{Too_list}\n\n"
        "Please note that: \n"
        "1. You should only chooce one tool the Tool List to solve this question.\n"
        "2. You must ONLY output the ID of the tool you chose in a parsible JSON format. An example output looks like:\n"
        "'''\n"
        "Example: {{\"ID\": XX}}\n"
        "'''\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    Tool_list = []
    for ele in Tool_dic:
        for key in ele.keys():
            if str(key) not in tool_used:
                Tool_list.append(f'''ID: {key}\n{ele[key]}''')
    while True:
        try:
            result = chain.run(question=question,
                               Too_list='\n'.join(Tool_list))
            clean_answer = eval(result.split("\n\n")[-1].strip())
            break
        except Exception as e:
            print(f"choose tool fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return clean_answer


def choose_API(API_instruction, API_list, question, model_name):
    input_execute_rapidapi_api_note = '''
This is an API Tool instruction. Given a question, you should choose APIs from the API list you want to use for this question in this instruction.
you must only output in a parsible Python List Format. An example output looks like:
```
["api1", "api2", ...]
```
'''.strip()

    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "{API_instruction}\n"
        "{input_execute_rapidapi_api_note}\n"
        "This is the API list: {API_list}\n"
        "Please note that: \n"
        "1. The APIs you choose must in the API list.\n"
        "2. You must ONLY output in the following parsible Python List Format.\n"
        "```\n"
        "Output_Example: [\"api1\", \"api2\", ...]\n"
        "```\n"
        "Question: {question}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(API_instruction=API_instruction,
                               API_list=API_list,
                               question=question,
                               input_execute_rapidapi_api_note=input_execute_rapidapi_api_note)
            clean_answer = eval(result.replace("```", "").strip().split("\n\n")[0].strip())
            if isinstance(clean_answer, str):
                ls = [clean_answer]
            elif isinstance(clean_answer, list):
                ls = clean_answer
            temp = []
            for ele in ls:
                if ele in API_list:
                    temp.append(ele)
            ls = temp
            return ls
        except Exception as e:
            print(f"Choose API fails: {e}")
            print(result)
            if ind > 10:
                return []
            ind += 1
            continue
    return ls


def choose_parameter(API_instruction, api, api_dic, question, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "This is an API tool documentation. Given a user's question, you need to output parameters according to the API tool documentation to successfully call the API to solve the user's question.\n"
        "This is API tool documentation: {api_dic}\n"
        "Please note that: \n"
        "1. The Example in the API tool documentation can help you better understand the use of the API.\n"
        "2. Ensure the parameters you output are correct. The output must contain the required parameters, and can contain the optional parameters based on the question. If no paremters in the required parameters and optional parameters, just leave it as {{\"Parameters\":{{}}}}\n"
        "3. If the user's question mentions other APIs, you should ONLY consider the API tool documentation I give and do not consider other APIs.\n"
        "4. If you need to use this API multiple times, please set \"Parameters\" to a list.\n"
        "5. You must ONLY output in a parsible JSON format. Two examples output looks like:\n"
        "'''\n"
        "Example 1: {{\"Parameters\":{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}}}\n"
        "Example 2: {{\"Parameters\":[{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}, {{\"keyword\": \"Machine Learning\", \"language\": \"English\"}}]}}\n"
        "'''\n"
        "This is user's question: {question}\n"
        "Output:\n"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(api_dic=api_dic,
                               question=question, )
            clean_answer = eval(
                result.replace(": true", ": True").replace(":true", ": True").replace(":false", ": False").replace(
                    ": false", ": False").replace("```", "").strip())
            a = clean_answer["Parameters"]

            return a
        except Exception as e:
            print(f"Choose Parameter fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return a


def choose_parameter_depend(API_instruction, api, api_dic, question, previous_log, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Given a user's question and a API tool documentation, you need to output parameters according to the API tool documentation to successfully call the API to solve the user's question.\n"
        "Please note that: \n"
        "1. The Example in the API tool documentation can help you better understand the use of the API.\n"
        "2. Ensure the parameters you output are correct. The output must contain the required parameters, and can contain the optional parameters based on the question. If no paremters in the required parameters and optional parameters, just leave it as {{\"Parameters\":{{}}}}\n"
        "3. If the user's question mentions other APIs, you should ONLY consider the API tool documentation I give and do not consider other APIs.\n"
        "4. The question may have dependencies on answers of other questions, so we will provide logs of previous questions and answers for your reference.\n"
        "5. If you need to use this API multiple times,, please set \"Parameters\" to a list.\n"
        "6. You must ONLY output in a parsible JSON format. Two examples output looks like:\n"
        "'''\n"
        "Example 1: {{\"Parameters\":{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}}}\n"
        "Example 2: {{\"Parameters\":[{{\"keyword\": \"Artificial Intelligence\", \"language\": \"English\"}}, {{\"keyword\": \"Machine Learning\", \"language\": \"English\"}}]}}\n"
        "'''\n"
        "There are logs of previous questions and answers: \n {previous_log}\n"
        "This is the current user's question: {question}\n"
        "This is API tool documentation: {api_dic}\n"
        "Output:\n"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(api_dic=api_dic,
                               question=question,
                               previous_log=previous_log)
            clean_answer = eval(
                result.replace(": true", ": True").replace(": false", ": False").replace("```", "").strip())
            a = clean_answer["Parameters"]

            return a
        except Exception as e:
            print(f"choose parameter depend fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return a


def answer_generation(question, API_instruction, call_result, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You should answer the question based on the response output by the API tool."
        "Please note that:\n"
        "1. Answer the question in natural language based on the API response reasonably and effectively.\n"
        "2. The user cannot directly get API response, "
        "so you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "This is the user's question:\n {question}\n"
        "This is the API response:\n {call_result}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question,
                               call_result=call_result)
            break
        except Exception as e:
            print(f"answer generation fails: {e}")
            if ind > 2:
                return -1
            ind += 1
            continue
    return result


def answer_generation_depend(question, API_instruction, call_result, model_name, previous_log):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You should answer the question based on the response output by the API tool."
        "Please note that:\n"
        "1. Try to organize the response into a natural language answer.\n"
        "2. We will not show the API response to the user, "
        "thus you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "3. The question may have dependencies on answers of other questions, so we will provide logs of previous questions and answers.\n"
        "There are logs of previous questions and answers: \n {previous_log}\n"
        "This is the user's question: {question}\n"
        "This is the response output by the API tool: \n{call_result}\n"
        "We will not show the API response to the user, "
        "thus you need to make full use of the response and give the information "
        "in the response that can satisfy the user's question in as much detail as possible.\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question,
                               call_result=call_result,
                               previous_log=previous_log)
            break
        except Exception as e:
            print(f"answer generation depend fails: {e}")
            if ind > 2:
                return -1
            ind += 1
            continue
    return result


def answer_check(question, answer, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Please check whether the response can reasonably and accurately answer the question."
        "If can, please output 'YES'; If not, please output 'NO'\n"
        "You need to give reasons first and then decide whether the response can reasonably and accurately answer the question. You must only output in a parsible JSON format. Two example outputs look like:\n"
        "Example 1: {{\"Reason\": \"The reason why you think the response can reasonably and accurately answer the question\", \"Choice\": \"Yes\"}}\n"
        "Example 2: {{\"Reason\": \"The reason why you think the response cannot reasonably and accurately answer the question\", \"Choice\": \"No\"}}\n"
        "This is the user's question: {question}\n"
        "This is the response: {answer}\n"
        "Output: "
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(question=question, answer=answer)
    if 'yes'.lower() in str(result).lower():
        return 1
    else:
        return -1


def Call_function(A, B, arg, index, id):
    if A in index:
        for path in index[A]:
            app_path = os.path.join(path, A, 'api.py')
            if os.path.isfile(app_path):
                spec = importlib.util.spec_from_file_location('api', app_path)
                app_module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(app_module)
                arg['toolbench_rapidapi_key'] = os.environ['RAPIDAPI_KEY']
                # Check if B is a function in app
                if hasattr(app_module, B):
                    function_B = getattr(app_module, B)
                    try:
                        call_result = function_B(**arg)
                        return call_result
                    except Exception as e:
                        try:
                            arg = {change_name(k.lower()): v for k, v in arg.items()}
                            call_result = function_B(**arg)
                            return call_result
                        except Exception as e:
                            try:
                                arg = {change_name(k.replace("-", "_")): v for k, v in arg.items()}
                                call_result = function_B(**arg)
                                return call_result
                            except Exception as e:
                                try:
                                    arg = {change_name(k.replace("\\", "")): v for k, v in arg.items()}
                                    call_result = function_B(**arg)
                                    return call_result
                                except Exception as e:
                                    print(f"Call function fails: {e}")
                                    with open('wrong_log.json', 'a+', encoding='utf-8') as f:
                                        line = json.dumps({
                                            "id": id,
                                            "parameters": arg,
                                            "wrong": str(e)
                                        }, ensure_ascii=False)
                                        f.write(line + '\n')
                                    return -1
                else:
                    with open('wrong_log.json', 'a+', encoding='utf-8') as f:
                        line = json.dumps({
                            "id": id,
                            "parameters": arg,
                            "wrong": f"No function named {B} in {app_path}"
                        }, ensure_ascii=False)
                        f.write(line + '\n')
                    return (f"No function named {B} in {app_path}")


def retrieval(question, Tool_dic, dataset, tool_used, ind, model_name, index, previous_log=None):
    tool_id = choose_tool(question, Tool_dic, tool_used, model_name)
    if tool_id == -1:
        return tool_id, "", "", "", ""
    if str(tool_id["ID"]) not in dataset:
        return tool_id, "", "", "", ""
    tool_instruction = dataset[str(tool_id["ID"])]
    API_instruction = tool_instruction["tool_description"]
    API_tool = tool_instruction["standardized_name"]
    API_list = []
    for ele in tool_instruction["tool_guidelines"].keys():
        API_list.append(ele)

    api_selection = choose_API(API_instruction, API_list, question, model_name)
    api_result = []
    if len(api_selection) == 0:
        call_result = ""
        print("No Calling")
        return tool_id, api_result, call_result, tool_instruction, API_instruction
    for api in api_selection:
        if previous_log is None:
            parameter = choose_parameter(API_instruction, api,
                                         tool_instruction["tool_guidelines"][api], question,
                                         model_name)
        else:
            parameter = choose_parameter_depend(API_instruction, api,
                                                tool_instruction["tool_guidelines"][api],
                                                question, previous_log,
                                                model_name)
        if parameter == -1:
            continue
        api_result.append({"api_name": api, "parameters": parameter})
    if len(api_result) == 0:
        call_result = ""
        return tool_id, api_result, call_result, tool_instruction, API_instruction
    if isinstance(api_result, set) or isinstance(api_result, list):
        call_results = []
        for api in api_result:
            api_name = change_name(standardize(api["api_name"]))

            if isinstance(api["parameters"], dict):
                parameters = {}
                for key in api["parameters"]:
                    value = api["parameters"][key]
                    key = change_name(key)
                    parameters[key] = value
                call_result = Call_function(API_tool, api_name, parameters, index, ind)
                if call_result == -1:
                    continue
                call_results.append(str(call_result))
            elif isinstance(api["parameters"], list):
                for para_ls in api["parameters"]:
                    parameters = {}
                    for key in para_ls:
                        value = para_ls[key]
                        key = change_name(key)
                        parameters[key] = value
                    call_result = Call_function(API_tool, api_name, parameters, index, ind)
                    if call_result == -1:
                        continue
                    call_results.append(str(call_result))
        call_result = '\n\n'.join(call_results)
    elif isinstance(api_result, dict):
        api_name = change_name(standardize(api_result["api_name"]))
        api = api_result
        if isinstance(api["parameters"], dict):
            parameters = {}
            for key in api["parameters"]:
                value = api["parameters"][key]
                key = change_name(key)
                parameters[key] = value
            call_result = Call_function(API_tool, api_name, parameters, index, ind)
        elif isinstance(api["parameters"], list):
            call_results = []
            for para_ls in api["parameters"]:
                parameters = {}
                for key in para_ls:
                    value = para_ls[key]
                    key = change_name(key)
                    parameters[key] = value
                call_result = Call_function(API_tool, api_name, parameters, index, ind)
                if call_result == -1:
                    continue
                call_results.append(str(call_result))
            call_result = '\n\n'.join(call_results)

    return tool_id, api_result, call_result, tool_instruction, API_instruction


def task_decompose(question, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You need to decompose a complex user's question into some simple subtasks and let the model execute it step by step.\n"
        "This is the user's question: {question}\n"
        "Please note that: \n"
        "1. You should only decompose this complex user's question into some simple subtasks which can be executed easily by using a single tool.\n"
        "2. Each simple subtask should be expressed into natural language.\n"
        "3. Each subtask should contain the necessary information from the original question and should be complete, explicit and self-consistent.\n"
        "4. You must ONLY output the ID of the tool you chose in a parsible JSON format. An example output looks like:\n"
        "'''\n"
        "{{\"Tasks\": [\"Task 1\", \"Task 2\", ...]}}\n"
        "'''\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question)
            result = eval(result.split('\n\n')[0])
            a = result["Tasks"]
            break
        except Exception as e:
            print(f"task decompose fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return result


def task_topology(question, task_ls, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "Given a complex user's question, I have decompose this question into some simple subtasks"
        "I think there exists a logical connections and order amontg the tasks. "
        "Thus you need to help me output this logical connections and order.\n"
        "You must ONLY output in a parsible JSON format with the following format:\n"
        "'''\n"
        "[{{\"task\": task, \"id\", task_id, \"dep\": [dependency_task_id1, dependency_task_id2, ...]}}]\n"
        "'''\n"
        "The \"dep\" field denotes the id of the previous task which generates a new resource upon which the current task depends. If there are no dependencies, set \"dep\" to -1.\n\n"
        "This is user's question: {question}\n"
        "These are subtasks of this question:\n"
        "{task_ls}\n"
        "Output: "
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question, task_ls=task_ls)
            result = eval(result)
            for i in range(len(result)):
                if isinstance(result[i]['dep'], str):
                    temp = []
                    for ele in result[i]['dep'].split(','):
                        temp.append(int(ele))
                    result[i]['dep'] = temp
                elif isinstance(result[i]['dep'], int):
                    result[i]['dep'] = [result[i]['dep']]
                elif isinstance(result[i]['dep'], list):
                    temp = []
                    for ele in result[i]['dep']:
                        temp.append(int(ele))
                    result[i]['dep'] = temp
                elif result[i]['dep'] == -1:
                    result[i]['dep'] = [-1]
            a = result[i]['dep'][0]
            return result
        except Exception as e:
            print(f"task topology fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return result


def answer_summarize(question, answer_task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "We break down a complex user's problems into simple subtasks and provide answers to each simple subtask. "
        "You need to organize these answers to each subtask and form a self-consistent final answer to the user's question\n"
        "This is the user's question: {question}\n"
        "These are subtasks and their answers: {answer_task}\n"
        "Final answer:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(question=question, answer_task=answer_task)
    return result


def answer_generation_direct(task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "You need to answer the user's question.\n"
        "This is the user's question: {task}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    result = chain.run(task=task)
    return result


def tool_check(task, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful language model which can use external APIs to solve user's question."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "As a powerful language model, you're equipped to answer user's question with accumulated knowledge. "
        "However, in some cases, you need to use external APIs to answer accurately."
        "Thus, you need to check whether the user's question requires you to call an external API to solve it.\n"
        "Here are some tips to help you check: \n"
        "1. If the user's question requires real-time information, since your knowledge base isn't updated in real-time, any such question will demand an API call.\n"
        "2. If you need to obtain information (e.g., ID, name, phone number, geographical location, rank, etc.), you need to call the database APIs if you are not sure.\n"
        "3. If the question demand a database search or internet research to generate an answer, this is another situation where an API call is necessary.\n"
        "If need, please output 'YES'; If not, please output 'NO'\n"
        "You need to give reasons first and then decide whether to keep it or not. You must only output in a parsible JSON format. Two example outputs look like:\n"
        "Example 1: {{\"Reason\": \"The reason why you think you do not need to call an external API to solve the user's question\", \"Choice\": \"No\"}}\n"
        "Example 2: {{\"Reason\": \"The reason why you think you need to call an external API to solve the user's question\", \"Choice\": \"Yes\"}}\n"
        "This is the user's question: {task}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(task=task)
            result = eval(result)
            a = result["Reason"]
            b = result["Choice"]
            if 'yes' in b.lower():
                return result, -1
            else:
                return result, 1
        except Exception as e:
            print(f"tool check fails: {e}")
            if ind > 10:
                return "", -1
            ind += 1
            continue
    return result, -1


def task_execution(data_type,
                   base_path, index, dataset, test_data, progress_file,
                   start_index, total_files, retrieval_num, ind, model_name):
    with tqdm(total=total_files, desc="Processing files", initial=start_index) as pbar:
        for i, data in enumerate(test_data[start_index:], start=start_index):
            answer_ls = []
            question = data["query"]
            print(question)
            temp = task_decompose(question, model_name)['Tasks']
            task_ls = []
            for t in range(len(temp)):
                task_ls.append({"task": temp[t], "id": t + 1})
            task_ls = task_topology(question, task_ls, model_name)
            task_depend = {}
            for task_dic in task_ls:
                task_depend[task_dic['id']] = {'task': task_dic['task'], 'answer': ''}
            answer_task = []
            api_result_ls = []
            call_result_ls = []
            tool_check_reason_ls = []
            parameter_ls = []
            for task_dic in task_ls:
                task = task_dic['task']
                tool_check_reason, tool_check_result = tool_check(task, model_name)
                tool_check_reason_ls.append(tool_check_reason)
                if tool_check_result == 1:
                    print("Do not need tool.")
                    answer = answer_generation_direct(task)
                    answer_task.append({'task': task, 'answer': answer})
                else:
                    print("Do need tool.")
                    depend_id = task_dic['dep']
                    tool_used = []
                    for r in range(retrieval_num):
                        Tool_dic = data["Tool_dic"]
                        if depend_id[0] == -1:
                            tool_id, api_result, call_result, tool_instruction, API_instruction = retrieval(task,
                                                                                                            Tool_dic,
                                                                                                            dataset,
                                                                                                            tool_used,
                                                                                                            ind,
                                                                                                            model_name,
                                                                                                            index)
                            call_result = str(call_result)[:1000]
                            answer = answer_generation(task, API_instruction,
                                                       call_result, model_name)
                        else:
                            previous_log = []
                            for ids in depend_id:
                                previous_log.append(task_depend[ids])
                            tool_id, api_result, call_result, tool_instruction, API_instruction = retrieval(task,
                                                                                                            Tool_dic,
                                                                                                            dataset,
                                                                                                            tool_used,
                                                                                                            ind,
                                                                                                            model_name,
                                                                                                            index,
                                                                                                            previous_log=previous_log)
                            call_result = str(call_result)[:1000]
                            answer = answer_generation_depend(task, API_instruction, call_result, model_name,
                                                              previous_log=previous_log)

                        check_index = answer_check(task, answer, model_name)
                        if check_index == 1:
                            answer_task.append({'task': task, 'answer': answer})
                            api_result_ls.append(api_result)
                            call_result_ls.append(call_result)
                            break
                        else:
                            answer_ls.append({'task': task, 'answer': answer})
                            try:
                                tool_used.append(str(tool_id["ID"]))
                            except:
                                continue
                            print('****Try Again****')
                task_depend[task_dic['id']]['answer'] = answer
            final_answer = answer_summarize(question, answer_task, model_name)
            check_index = answer_check(question, final_answer, model_name)

            ind = ind + 1
            with open(f'''{data_type}_{model_name}_Easytool.jsonl''', 'a+', encoding='utf-8') as f:
                line = json.dumps({
                    "ID": ind,
                    "question": question,
                    "final_answer": final_answer,
                    "subtask": task_ls,
                    "answer_subtask": answer_task,
                    "answer_wrong": answer_ls,
                    "check_index": check_index,
                    "execute_log": {
                        "api_result_ls": api_result_ls,
                        "parameter_ls": parameter_ls,
                        "call_result_ls": call_result_ls,
                        "tool_check_reason_ls": tool_check_reason_ls,
                    }
                }, ensure_ascii=False)
                f.write(line + '\n')

            print(final_answer)
            update_progress(progress_file, i + 1)
            pbar.update(1)



File: easytool/easytool/restbench.py

# — coding: utf-8 –
import openai
import json
import logging
import sys
import argparse
from langchain.chat_models import ChatOpenAI
from langchain.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate
)
from langchain import LLMChain
import numpy as np
import requests
import os
import subprocess
import re
import importlib.util
from sklearn.metrics.pairwise import cosine_similarity
import pickle
from util import *

from tqdm import tqdm

openai.api_key = os.environ["OPENAI_API_KEY"]


def get_last_processed_index(progress_file):
    """Retrieve the last processed index from the progress file."""
    if os.path.exists(progress_file):
        with open(progress_file, 'r', encoding='utf-8') as f:
            last_index = f.read().strip()
            return int(last_index) if last_index else 0
    else:
        return 0


def update_progress(progress_file, index):
    """Update the last processed index in the progress file."""
    with open(progress_file, 'w', encoding='utf-8') as f:
        f.write(str(index))


def task_decompose(question, Tool_dic, model_name):
    chat = ChatOpenAI(model_name=model_name)
    template = "You are a helpful assistant."
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)
    human_message_prompt = HumanMessagePromptTemplate.from_template(
        "We have spotify database and the following tools:\n"
        "{Tool_dic}"
        "You need to decompose a complex user's question into some simple subtasks and let the model execute it step by step with these tools.\n"
        "Please note that: \n"
        "1. you should break down tasks into appropriate subtasks to use the tools mentioned above.\n"
        "2. You should not only list the subtask, but also list the ID of the tool used to solve this subtask.\n"
        "3. If you think you do not need to use the tool to solve the subtask, just leave it as {{\"ID\": -1}}\n"
        "4. You must consider the logical connections, order and constraints among the tools to achieve a correct tool path."
        "5. You must ONLY output the ID of the tool you chose in a parsible JSON format. Two examples output look like:\n"
        "'''\n"
        "Question: Pause the player"
        "Example 1: [{{\"Task\":\"Get information about the user’s current playback state\", \"ID\":15}}, {{\"Task\":\"Pause playback on the user's account\", \"ID\":19}}]\n"
        "'''\n"
        "This is the user's question: {question}\n"
        "Output:"
    )
    chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt)
    ind = 0
    while True:
        try:
            result = chain.run(question=question, Tool_dic=Tool_dic)
            result = eval(result.split('\n\n')[0])
            break
        except Exception as e:
            print(f"task decompose fails: {e}")
            if ind > 10:
                return -1
            ind += 1
            continue
    return result


def task_execution(
        Tool_dic, dic_tool, test_data, progress_file,
        start_index, total_files, retrieval_num, ind, model_name):
    with tqdm(total=total_files, desc="Processing files", initial=start_index) as pbar:
        for i, data in enumerate(test_data[start_index:], start=start_index):
            question = data["query"]
            print(question)
            task_path = task_decompose(question, Tool_dic, model_name)
            tool_choice_ls = []
            for task in task_path:
                if isinstance(task["ID"], list):
                    for ele in task["ID"]:
                        tool_choice_ls.append(dic_tool[ele]['tool_usage'])
                elif int(task["ID"]) in dic_tool.keys():
                    tool_choice_ls.append(dic_tool[task["ID"]]['tool_usage'])
            ind = ind + 1
            with open(f"restbench_{model_name}_Easytool.jsonl", 'a+', encoding='utf-8') as f:
                line = json.dumps({
                    "ID": ind,
                    "question": question,
                    "task_path": task_path,
                    "tool_choice_ls": tool_choice_ls
                }, ensure_ascii=False)
                f.write(line + '\n')
            print(tool_choice_ls)
            update_progress(progress_file, i + 1)
            pbar.update(1)



File: easytool/data_funcqa/funchub/math.py

import math

# this function is used to round the result to 2 decimal places
# e.g. 52.3523 -> 52.35, 52.0011 -> 52, 0.00000233 -> 0.0000023
def custom_round(x, decimal_places=2):
    str_x = f"{x:.10f}"
    before_decimal = str_x.split('.')[0]
    after_decimal = str_x.split('.')[1]
    leading_zeros = len(after_decimal) - len(after_decimal.lstrip('0'))
    
    if leading_zeros >= 1 and before_decimal == "0":
        return round(x, leading_zeros + 2)
    else:
        return round(x, decimal_places)

# this function converts a number in scientific notation to decimal notation
def scito_decimal(sci_str):
    def split_exponent(number_str):
        parts = number_str.split("e")
        coefficient = parts[0]
        exponent = int(parts[1]) if len(parts) == 2 else 0
        return coefficient, exponent

    def multiplyby_10(number_str, exponent):
        if exponent == 0:
            return number_str

        if exponent > 0:
            index = number_str.index(".") if "." in number_str else len(number_str)
            number_str = number_str.replace(".", "")
            new_index = index + exponent
            number_str += "0" * (new_index - len(number_str))
            if new_index < len(number_str):
                number_str = number_str[:new_index] + "." + number_str[new_index:]
            return number_str

        if exponent < 0:
            index = number_str.index(".") if "." in number_str else len(number_str)
            number_str = number_str.replace(".", "")
            new_index = index + exponent
            number_str = "0" * (-new_index) + number_str
            number_str = "0." + number_str
            return number_str

    coefficient, exponent = split_exponent(sci_str)
    decimal_str = multiplyby_10(coefficient, exponent)

    # remove trailing zeros
    if "." in decimal_str:
        decimal_str = decimal_str.rstrip("0")

    return decimal_str

# normalize the result to 2 decimal places and remove trailing zeros
def normalize(res, round_to=2):
        # we round the result to 2 decimal places
        res = custom_round(res, round_to)
        res = str(res)
        if "." in res:
            while res[-1] == "0":
                res = res[:-1]
            res = res.strip(".")
        
        # scientific notation
        if "e" in res:
            res = scito_decimal(res)

        return res

# 1. add
def add_(args):

    return normalize(sum(args))

# 2. subtract
def subtract_(args):

    res = args[0]
    for arg in args[1:]:
        res -= arg
    return normalize(res)

# 3. multiply
def multiply_(args):

    res = args[0]
    for arg in args[1:]:
        res *= arg
    return normalize(res)

# 4. divide
def divide_(args):

    res = args[0]
    for arg in args[1:]:
        res /= arg
    return normalize(res)

# 5. power
def power_(args):
        
    res = args[0]
    for arg in args[1:]:
        res **= arg
    return normalize(res)

# 6. square root
def sqrt_(args):
    res = args[0]
    return normalize(math.sqrt(res))

# 7. 10th log
def log_(args):
    # if only one argument is passed, it is 10th log
    if len(args) == 1:
        res = args[0]
        return normalize(math.log10(res))
    # if two arguments are passed, it is log with base as the second argument   
    elif len(args) == 2:
        res = args[0]
        base = args[1]
        return normalize(math.log(res, base))
    else:
        raise Exception("Invalid number of arguments passed to log function")

# 8. natural log
def ln_(args):
    res = args[0]
    return normalize(math.log(res))


# 9. choose
def choose_(args):
    n = args[0]
    r = args[1]
    return normalize(math.comb(n, r))

# 10. permutation
def permutate_(args):
    n = args[0]
    r = args[1]
    return normalize(math.perm(n, r))

# 11. greatest common divisor
def gcd_(args):
    res = args[0]
    for arg in args[1:]:
        res = math.gcd(res, arg)
    return normalize(res)

# 12. least common multiple
def lcm_(args):
    res = args[0]
    for arg in args[1:]:
        res = res * arg // math.gcd(res, arg)
    return normalize(res)

# 13. remainder
def remainder_(args):
    dividend = args[0]
    divisor = args[1]
    return normalize(dividend % divisor)


File: hugginggpt/server/awesome_chat.py

import base64
import copy
from io import BytesIO
import io
import os
import random
import time
import traceback
import uuid
import requests
import re
import json
import logging
import argparse
import yaml
from PIL import Image, ImageDraw
from diffusers.utils import load_image
from pydub import AudioSegment
import threading
from queue import Queue
import flask
from flask import request, jsonify
import waitress
from flask_cors import CORS, cross_origin
from get_token_ids import get_token_ids_for_task_parsing, get_token_ids_for_choose_model, count_tokens, get_max_context_length
from huggingface_hub.inference_api import InferenceApi
from huggingface_hub.inference_api import ALL_TASKS

parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="configs/config.default.yaml")
parser.add_argument("--mode", type=str, default="cli")
args = parser.parse_args()

if __name__ != "__main__":
    args.config = "configs/config.gradio.yaml"
    args.mode = "gradio"

config = yaml.load(open(args.config, "r"), Loader=yaml.FullLoader)

os.makedirs("logs", exist_ok=True)
os.makedirs("public/images", exist_ok=True)
os.makedirs("public/audios", exist_ok=True)
os.makedirs("public/videos", exist_ok=True)


logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
if not config["debug"]:
    handler.setLevel(logging.CRITICAL)
logger.addHandler(handler)

log_file = config["log_file"]
if log_file:
    filehandler = logging.FileHandler(log_file)
    filehandler.setLevel(logging.DEBUG)
    filehandler.setFormatter(formatter)
    logger.addHandler(filehandler)

LLM = config["model"]
use_completion = config["use_completion"]

# consistent: wrong msra model name 
LLM_encoding = LLM
if config["dev"] and LLM == "gpt-3.5-turbo":
    LLM_encoding = "text-davinci-003"
task_parsing_highlight_ids = get_token_ids_for_task_parsing(LLM_encoding)
choose_model_highlight_ids = get_token_ids_for_choose_model(LLM_encoding)

# ENDPOINT	MODEL NAME	
# /v1/chat/completions	gpt-4, gpt-4-0314, gpt-4-32k, gpt-4-32k-0314, gpt-3.5-turbo, gpt-3.5-turbo-0301	
# /v1/completions	text-davinci-003, text-davinci-002, text-curie-001, text-babbage-001, text-ada-001, davinci, curie, babbage, ada

if use_completion:
    api_name = "completions"
else:
    api_name = "chat/completions"

API_TYPE = None
# priority: local > azure > openai
if "dev" in config and config["dev"]:
    API_TYPE = "local"
elif "azure" in config:
    API_TYPE = "azure"
elif "openai" in config:
    API_TYPE = "openai"
else:
    logger.warning(f"No endpoint specified in {args.config}. The endpoint will be set dynamically according to the client.")

if args.mode in ["test", "cli"]:
    assert API_TYPE, "Only server mode supports dynamic endpoint."

API_KEY = None
API_ENDPOINT = None
if API_TYPE == "local":
    API_ENDPOINT = f"{config['local']['endpoint']}/v1/{api_name}"
elif API_TYPE == "azure":
    API_ENDPOINT = f"{config['azure']['base_url']}/openai/deployments/{config['azure']['deployment_name']}/{api_name}?api-version={config['azure']['api_version']}"
    API_KEY = config["azure"]["api_key"]
elif API_TYPE == "openai":
    API_ENDPOINT = f"https://api.openai.com/v1/{api_name}"
    if config["openai"]["api_key"].startswith("sk-"):  # Check for valid OpenAI key in config file
        API_KEY = config["openai"]["api_key"]
    elif "OPENAI_API_KEY" in os.environ and os.getenv("OPENAI_API_KEY").startswith("sk-"):  # Check for environment variable OPENAI_API_KEY
        API_KEY = os.getenv("OPENAI_API_KEY")
    else:
        raise ValueError(f"Incorrect OpenAI key. Please check your {args.config} file.")

PROXY = None
if config["proxy"]:
    PROXY = {
        "https": config["proxy"],
    }

inference_mode = config["inference_mode"]

# check the local_inference_endpoint
Model_Server = None
if inference_mode!="huggingface":
    Model_Server = "http://" + config["local_inference_endpoint"]["host"] + ":" + str(config["local_inference_endpoint"]["port"])
    message = f"The server of local inference endpoints is not running, please start it first. (or using `inference_mode: huggingface` in {args.config} for a feature-limited experience)"
    try:
        r = requests.get(Model_Server + "/running")
        if r.status_code != 200:
            raise ValueError(message)
    except:
        raise ValueError(message)


parse_task_demos_or_presteps = open(config["demos_or_presteps"]["parse_task"], "r").read()
choose_model_demos_or_presteps = open(config["demos_or_presteps"]["choose_model"], "r").read()
response_results_demos_or_presteps = open(config["demos_or_presteps"]["response_results"], "r").read()

parse_task_prompt = config["prompt"]["parse_task"]
choose_model_prompt = config["prompt"]["choose_model"]
response_results_prompt = config["prompt"]["response_results"]

parse_task_tprompt = config["tprompt"]["parse_task"]
choose_model_tprompt = config["tprompt"]["choose_model"]
response_results_tprompt = config["tprompt"]["response_results"]

MODELS = [json.loads(line) for line in open("data/p0_models.jsonl", "r").readlines()]
MODELS_MAP = {}
for model in MODELS:
    tag = model["task"]
    if tag not in MODELS_MAP:
        MODELS_MAP[tag] = []
    MODELS_MAP[tag].append(model)
METADATAS = {}
for model in MODELS:
    METADATAS[model["id"]] = model

HUGGINGFACE_HEADERS = {}
if config["huggingface"]["token"] and config["huggingface"]["token"].startswith("hf_"):  # Check for valid huggingface token in config file
    HUGGINGFACE_HEADERS = {
        "Authorization": f"Bearer {config['huggingface']['token']}",
    }
elif "HUGGINGFACE_ACCESS_TOKEN" in os.environ and os.getenv("HUGGINGFACE_ACCESS_TOKEN").startswith("hf_"):  # Check for environment variable HUGGINGFACE_ACCESS_TOKEN
    HUGGINGFACE_HEADERS = {
        "Authorization": f"Bearer {os.getenv('HUGGINGFACE_ACCESS_TOKEN')}",
    }
else:
    raise ValueError(f"Incorrect HuggingFace token. Please check your {args.config} file.")

def convert_chat_to_completion(data):
    messages = data.pop('messages', [])
    tprompt = ""
    if messages[0]['role'] == "system":
        tprompt = messages[0]['content']
        messages = messages[1:]
    final_prompt = ""
    for message in messages:
        if message['role'] == "user":
            final_prompt += ("<im_start>"+ "user" + "\n" + message['content'] + "<im_end>\n")
        elif message['role'] == "assistant":
            final_prompt += ("<im_start>"+ "assistant" + "\n" + message['content'] + "<im_end>\n")
        else:
            final_prompt += ("<im_start>"+ "system" + "\n" + message['content'] + "<im_end>\n")
    final_prompt = tprompt + final_prompt
    final_prompt = final_prompt + "<im_start>assistant"
    data["prompt"] = final_prompt
    data['stop'] = data.get('stop', ["<im_end>"])
    data['max_tokens'] = data.get('max_tokens', max(get_max_context_length(LLM) - count_tokens(LLM_encoding, final_prompt), 1))
    return data

def send_request(data):
    api_key = data.pop("api_key")
    api_type = data.pop("api_type")
    api_endpoint = data.pop("api_endpoint")
    if use_completion:
        data = convert_chat_to_completion(data)
    if api_type == "openai":
        HEADER = {
            "Authorization": f"Bearer {api_key}"
        }
    elif api_type == "azure":
        HEADER = {
            "api-key": api_key,
            "Content-Type": "application/json"
        }
    else:
        HEADER = None
    response = requests.post(api_endpoint, json=data, headers=HEADER, proxies=PROXY)
    if "error" in response.json():
        return response.json()
    logger.debug(response.text.strip())
    if use_completion:
        return response.json()["choices"][0]["text"].strip()
    else:
        return response.json()["choices"][0]["message"]["content"].strip()

def replace_slot(text, entries):
    for key, value in entries.items():
        if not isinstance(value, str):
            value = str(value)
        text = text.replace("{{" + key +"}}", value.replace('"', "'").replace('\n', ""))
    return text

def find_json(s):
    s = s.replace("\'", "\"")
    start = s.find("{")
    end = s.rfind("}")
    res = s[start:end+1]
    res = res.replace("\n", "")
    return res

def field_extract(s, field):
    try:
        field_rep = re.compile(f'{field}.*?:.*?"(.*?)"', re.IGNORECASE)
        extracted = field_rep.search(s).group(1).replace("\"", "\'")
    except:
        field_rep = re.compile(f'{field}:\ *"(.*?)"', re.IGNORECASE)
        extracted = field_rep.search(s).group(1).replace("\"", "\'")
    return extracted

def get_id_reason(choose_str):
    reason = field_extract(choose_str, "reason")
    id = field_extract(choose_str, "id")
    choose = {"id": id, "reason": reason}
    return id.strip(), reason.strip(), choose

def record_case(success, **args):
    if success:
        f = open("logs/log_success.jsonl", "a")
    else:
        f = open("logs/log_fail.jsonl", "a")
    log = args
    f.write(json.dumps(log) + "\n")
    f.close()

def image_to_bytes(img_url):
    img_byte = io.BytesIO()
    type = img_url.split(".")[-1]
    load_image(img_url).save(img_byte, format="png")
    img_data = img_byte.getvalue()
    return img_data

def resource_has_dep(command):
    args = command["args"]
    for _, v in args.items():
        if "<GENERATED>" in v:
            return True
    return False

def fix_dep(tasks):
    for task in tasks:
        args = task["args"]
        task["dep"] = []
        for k, v in args.items():
            if "<GENERATED>" in v:
                dep_task_id = int(v.split("-")[1])
                if dep_task_id not in task["dep"]:
                    task["dep"].append(dep_task_id)
        if len(task["dep"]) == 0:
            task["dep"] = [-1]
    return tasks

def unfold(tasks):
    flag_unfold_task = False
    try:
        for task in tasks:
            for key, value in task["args"].items():
                if "<GENERATED>" in value:
                    generated_items = value.split(",")
                    if len(generated_items) > 1:
                        flag_unfold_task = True
                        for item in generated_items:
                            new_task = copy.deepcopy(task)
                            dep_task_id = int(item.split("-")[1])
                            new_task["dep"] = [dep_task_id]
                            new_task["args"][key] = item
                            tasks.append(new_task)
                        tasks.remove(task)
    except Exception as e:
        print(e)
        traceback.print_exc()
        logger.debug("unfold task failed.")

    if flag_unfold_task:
        logger.debug(f"unfold tasks: {tasks}")
        
    return tasks

def chitchat(messages, api_key, api_type, api_endpoint):
    data = {
        "model": LLM,
        "messages": messages,
        "api_key": api_key,
        "api_type": api_type,
        "api_endpoint": api_endpoint
    }
    return send_request(data)

def parse_task(context, input, api_key, api_type, api_endpoint):
    demos_or_presteps = parse_task_demos_or_presteps
    messages = json.loads(demos_or_presteps)
    messages.insert(0, {"role": "system", "content": parse_task_tprompt})

    # cut chat logs
    start = 0
    while start <= len(context):
        history = context[start:]
        prompt = replace_slot(parse_task_prompt, {
            "input": input,
            "context": history 
        })
        messages.append({"role": "user", "content": prompt})
        history_text = "<im_end>\nuser<im_start>".join([m["content"] for m in messages])
        num = count_tokens(LLM_encoding, history_text)
        if get_max_context_length(LLM) - num > 800:
            break
        messages.pop()
        start += 2
    
    logger.debug(messages)
    data = {
        "model": LLM,
        "messages": messages,
        "temperature": 0,
        "logit_bias": {item: config["logit_bias"]["parse_task"] for item in task_parsing_highlight_ids},
        "api_key": api_key,
        "api_type": api_type,
        "api_endpoint": api_endpoint
    }
    return send_request(data)

def choose_model(input, task, metas, api_key, api_type, api_endpoint):
    prompt = replace_slot(choose_model_prompt, {
        "input": input,
        "task": task,
        "metas": metas,
    })
    demos_or_presteps = replace_slot(choose_model_demos_or_presteps, {
        "input": input,
        "task": task,
        "metas": metas
    })
    messages = json.loads(demos_or_presteps)
    messages.insert(0, {"role": "system", "content": choose_model_tprompt})
    messages.append({"role": "user", "content": prompt})
    logger.debug(messages)
    data = {
        "model": LLM,
        "messages": messages,
        "temperature": 0,
        "logit_bias": {item: config["logit_bias"]["choose_model"] for item in choose_model_highlight_ids}, # 5
        "api_key": api_key,
        "api_type": api_type,
        "api_endpoint": api_endpoint
    }
    return send_request(data)


def response_results(input, results, api_key, api_type, api_endpoint):
    results = [v for k, v in sorted(results.items(), key=lambda item: item[0])]
    prompt = replace_slot(response_results_prompt, {
        "input": input,
    })
    demos_or_presteps = replace_slot(response_results_demos_or_presteps, {
        "input": input,
        "processes": results
    })
    messages = json.loads(demos_or_presteps)
    messages.insert(0, {"role": "system", "content": response_results_tprompt})
    messages.append({"role": "user", "content": prompt})
    logger.debug(messages)
    data = {
        "model": LLM,
        "messages": messages,
        "temperature": 0,
        "api_key": api_key,
        "api_type": api_type,
        "api_endpoint": api_endpoint
    }
    return send_request(data)

def huggingface_model_inference(model_id, data, task):
    task_url = f"https://api-inference.huggingface.co/models/{model_id}" # InferenceApi does not yet support some tasks
    inference = InferenceApi(repo_id=model_id, token=config["huggingface"]["token"])
    
    # NLP tasks
    if task == "question-answering":
        inputs = {"question": data["text"], "context": (data["context"] if "context" in data else "" )}
        result = inference(inputs)
    if task == "sentence-similarity":
        inputs = {"source_sentence": data["text1"], "target_sentence": data["text2"]}
        result = inference(inputs)
    if task in ["text-classification",  "token-classification", "text2text-generation", "summarization", "translation", "conversational", "text-generation"]:
        inputs = data["text"]
        result = inference(inputs)
    
    # CV tasks
    if task == "visual-question-answering" or task == "document-question-answering":
        img_url = data["image"]
        text = data["text"]
        img_data = image_to_bytes(img_url)
        img_base64 = base64.b64encode(img_data).decode("utf-8")
        json_data = {}
        json_data["inputs"] = {}
        json_data["inputs"]["question"] = text
        json_data["inputs"]["image"] = img_base64
        result = requests.post(task_url, headers=HUGGINGFACE_HEADERS, json=json_data).json()
        # result = inference(inputs) # not support

    if task == "image-to-image":
        img_url = data["image"]
        img_data = image_to_bytes(img_url)
        # result = inference(data=img_data) # not support
        HUGGINGFACE_HEADERS["Content-Length"] = str(len(img_data))
        r = requests.post(task_url, headers=HUGGINGFACE_HEADERS, data=img_data)
        result = r.json()
        if "path" in result:
            result["generated image"] = result.pop("path")
    
    if task == "text-to-image":
        inputs = data["text"]
        img = inference(inputs)
        name = str(uuid.uuid4())[:4]
        img.save(f"public/images/{name}.png")
        result = {}
        result["generated image"] = f"/images/{name}.png"

    if task == "image-segmentation":
        img_url = data["image"]
        img_data = image_to_bytes(img_url)
        image = Image.open(BytesIO(img_data))
        predicted = inference(data=img_data)
        colors = []
        for i in range(len(predicted)):
            colors.append((random.randint(100, 255), random.randint(100, 255), random.randint(100, 255), 155))
        for i, pred in enumerate(predicted):
            label = pred["label"]
            mask = pred.pop("mask").encode("utf-8")
            mask = base64.b64decode(mask)
            mask = Image.open(BytesIO(mask), mode='r')
            mask = mask.convert('L')

            layer = Image.new('RGBA', mask.size, colors[i])
            image.paste(layer, (0, 0), mask)
        name = str(uuid.uuid4())[:4]
        image.save(f"public/images/{name}.jpg")
        result = {}
        result["generated image"] = f"/images/{name}.jpg"
        result["predicted"] = predicted

    if task == "object-detection":
        img_url = data["image"]
        img_data = image_to_bytes(img_url)
        predicted = inference(data=img_data)
        image = Image.open(BytesIO(img_data))
        draw = ImageDraw.Draw(image)
        labels = list(item['label'] for item in predicted)
        color_map = {}
        for label in labels:
            if label not in color_map:
                color_map[label] = (random.randint(0, 255), random.randint(0, 100), random.randint(0, 255))
        for label in predicted:
            box = label["box"]
            draw.rectangle(((box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])), outline=color_map[label["label"]], width=2)
            draw.text((box["xmin"]+5, box["ymin"]-15), label["label"], fill=color_map[label["label"]])
        name = str(uuid.uuid4())[:4]
        image.save(f"public/images/{name}.jpg")
        result = {}
        result["generated image"] = f"/images/{name}.jpg"
        result["predicted"] = predicted

    if task in ["image-classification"]:
        img_url = data["image"]
        img_data = image_to_bytes(img_url)
        result = inference(data=img_data)
 
    if task == "image-to-text":
        img_url = data["image"]
        img_data = image_to_bytes(img_url)
        HUGGINGFACE_HEADERS["Content-Length"] = str(len(img_data))
        r = requests.post(task_url, headers=HUGGINGFACE_HEADERS, data=img_data, proxies=PROXY)
        result = {}
        if "generated_text" in r.json()[0]:
            result["generated text"] = r.json()[0].pop("generated_text")
    
    # AUDIO tasks
    if task == "text-to-speech":
        inputs = data["text"]
        response = inference(inputs, raw_response=True)
        # response = requests.post(task_url, headers=HUGGINGFACE_HEADERS, json={"inputs": text})
        name = str(uuid.uuid4())[:4]
        with open(f"public/audios/{name}.flac", "wb") as f:
            f.write(response.content)
        result = {"generated audio": f"/audios/{name}.flac"}
    if task in ["automatic-speech-recognition", "audio-to-audio", "audio-classification"]:
        audio_url = data["audio"]
        audio_data = requests.get(audio_url, timeout=10).content
        response = inference(data=audio_data, raw_response=True)
        result = response.json()
        if task == "audio-to-audio":
            content = None
            type = None
            for k, v in result[0].items():
                if k == "blob":
                    content = base64.b64decode(v.encode("utf-8"))
                if k == "content-type":
                    type = "audio/flac".split("/")[-1]
            audio = AudioSegment.from_file(BytesIO(content))
            name = str(uuid.uuid4())[:4]
            audio.export(f"public/audios/{name}.{type}", format=type)
            result = {"generated audio": f"/audios/{name}.{type}"}
    return result

def local_model_inference(model_id, data, task):
    task_url = f"{Model_Server}/models/{model_id}"
    
    # contronlet
    if model_id.startswith("lllyasviel/sd-controlnet-"):
        img_url = data["image"]
        text = data["text"]
        response = requests.post(task_url, json={"img_url": img_url, "text": text})
        results = response.json()
        if "path" in results:
            results["generated image"] = results.pop("path")
        return results
    if model_id.endswith("-control"):
        img_url = data["image"]
        response = requests.post(task_url, json={"img_url": img_url})
        results = response.json()
        if "path" in results:
            results["generated image"] = results.pop("path")
        return results
        
    if task == "text-to-video":
        response = requests.post(task_url, json=data)
        results = response.json()
        if "path" in results:
            results["generated video"] = results.pop("path")
        return results

    # NLP tasks
    if task == "question-answering" or task == "sentence-similarity":
        response = requests.post(task_url, json=data)
        return response.json()
    if task in ["text-classification",  "token-classification", "text2text-generation", "summarization", "translation", "conversational", "text-generation"]:
        response = requests.post(task_url, json=data)
        return response.json()

    # CV tasks
    if task == "depth-estimation":
        img_url = data["image"]
        response = requests.post(task_url, json={"img_url": img_url})
        results = response.json()
        if "path" in results:
            results["generated image"] = results.pop("path")
        return results
    if task == "image-segmentation":
        img_url = data["image"]
        response = requests.post(task_url, json={"img_url": img_url})
        results = response.json()
        results["generated image"] = results.pop("path")
        return results
    if task == "image-to-image":
        img_url = data["image"]
        response = requests.post(task_url, json={"img_url": img_url})
        results = response.json()
        if "path" in results:
            results["generated image"] = results.pop("path")
        return results
    if task == "text-to-image":
        response = requests.post(task_url, json=data)
        results = response.json()
        if "path" in results:
            results["generated image"] = results.pop("path")
        return results
    if task == "object-detection":
        img_url = data["image"]
        response = requests.post(task_url, json={"img_url": img_url})
        predicted = response.json()
        if "error" in predicted:
            return predicted
        image = load_image(img_url)
        draw = ImageDraw.Draw(image)
        labels = list(item['label'] for item in predicted)
        color_map = {}
        for label in labels:
            if label not in color_map:
                color_map[label] = (random.randint(0, 255), random.randint(0, 100), random.randint(0, 255))
        for label in predicted:
            box = label["box"]
            draw.rectangle(((box["xmin"], box["ymin"]), (box["xmax"], box["ymax"])), outline=color_map[label["label"]], width=2)
            draw.text((box["xmin"]+5, box["ymin"]-15), label["label"], fill=color_map[label["label"]])
        name = str(uuid.uuid4())[:4]
        image.save(f"public/images/{name}.jpg")
        results = {}
        results["generated image"] = f"/images/{name}.jpg"
        results["predicted"] = predicted
        return results
    if task in ["image-classification", "image-to-text", "document-question-answering", "visual-question-answering"]:
        img_url = data["image"]
        text = None
        if "text" in data:
            text = data["text"]
        response = requests.post(task_url, json={"img_url": img_url, "text": text})
        results = response.json()
        return results
    # AUDIO tasks
    if task == "text-to-speech":
        response = requests.post(task_url, json=data)
        results = response.json()
        if "path" in results:
            results["generated audio"] = results.pop("path")
        return results
    if task in ["automatic-speech-recognition", "audio-to-audio", "audio-classification"]:
        audio_url = data["audio"]
        response = requests.post(task_url, json={"audio_url": audio_url})
        return response.json()


def model_inference(model_id, data, hosted_on, task):
    if hosted_on == "unknown":
        localStatusUrl = f"{Model_Server}/status/{model_id}"
        r = requests.get(localStatusUrl)
        logger.debug("Local Server Status: " + str(r.json()))
        if r.status_code == 200 and "loaded" in r.json() and r.json()["loaded"]:
            hosted_on = "local"
        else:
            huggingfaceStatusUrl = f"https://api-inference.huggingface.co/status/{model_id}"
            r = requests.get(huggingfaceStatusUrl, headers=HUGGINGFACE_HEADERS, proxies=PROXY)
            logger.debug("Huggingface Status: " + str(r.json()))
            if r.status_code == 200 and "loaded" in r.json() and r.json()["loaded"]:
                hosted_on = "huggingface"
    try:
        if hosted_on == "local":
            inference_result = local_model_inference(model_id, data, task)
        elif hosted_on == "huggingface":
            inference_result = huggingface_model_inference(model_id, data, task)
    except Exception as e:
        print(e)
        traceback.print_exc()
        inference_result = {"error":{"message": str(e)}}
    return inference_result


def get_model_status(model_id, url, headers, queue = None):
    endpoint_type = "huggingface" if "huggingface" in url else "local"
    if "huggingface" in url:
        r = requests.get(url, headers=headers, proxies=PROXY)
    else:
        r = requests.get(url)
    if r.status_code == 200 and "loaded" in r.json() and r.json()["loaded"]:
        if queue:
            queue.put((model_id, True, endpoint_type))
        return True
    else:
        if queue:
            queue.put((model_id, False, None))
        return False

def get_avaliable_models(candidates, topk=5):
    all_available_models = {"local": [], "huggingface": []}
    threads = []
    result_queue = Queue()

    for candidate in candidates:
        model_id = candidate["id"]

        if inference_mode != "local":
            huggingfaceStatusUrl = f"https://api-inference.huggingface.co/status/{model_id}"
            thread = threading.Thread(target=get_model_status, args=(model_id, huggingfaceStatusUrl, HUGGINGFACE_HEADERS, result_queue))
            threads.append(thread)
            thread.start()
        
        if inference_mode != "huggingface" and config["local_deployment"] != "minimal":
            localStatusUrl = f"{Model_Server}/status/{model_id}"
            thread = threading.Thread(target=get_model_status, args=(model_id, localStatusUrl, {}, result_queue))
            threads.append(thread)
            thread.start()
        
    result_count = len(threads)
    while result_count:
        model_id, status, endpoint_type = result_queue.get()
        if status and model_id not in all_available_models:
            all_available_models[endpoint_type].append(model_id)
        if len(all_available_models["local"] + all_available_models["huggingface"]) >= topk:
            break
        result_count -= 1

    for thread in threads:
        thread.join()

    return all_available_models

def collect_result(command, choose, inference_result):
    result = {"task": command}
    result["inference result"] = inference_result
    result["choose model result"] = choose
    logger.debug(f"inference result: {inference_result}")
    return result


def run_task(input, command, results, api_key, api_type, api_endpoint):
    id = command["id"]
    args = command["args"]
    task = command["task"]
    deps = command["dep"]
    if deps[0] != -1:
        dep_tasks = [results[dep] for dep in deps]
    else:
        dep_tasks = []
    
    logger.debug(f"Run task: {id} - {task}")
    logger.debug("Deps: " + json.dumps(dep_tasks))

    if deps[0] != -1:
        if "image" in args and "<GENERATED>-" in args["image"]:
            resource_id = int(args["image"].split("-")[1])
            if "generated image" in results[resource_id]["inference result"]:
                args["image"] = results[resource_id]["inference result"]["generated image"]
        if "audio" in args and "<GENERATED>-" in args["audio"]:
            resource_id = int(args["audio"].split("-")[1])
            if "generated audio" in results[resource_id]["inference result"]:
                args["audio"] = results[resource_id]["inference result"]["generated audio"]
        if "text" in args and "<GENERATED>-" in args["text"]:
            resource_id = int(args["text"].split("-")[1])
            if "generated text" in results[resource_id]["inference result"]:
                args["text"] = results[resource_id]["inference result"]["generated text"]

    text = image = audio = None
    for dep_task in dep_tasks:
        if "generated text" in dep_task["inference result"]:
            text = dep_task["inference result"]["generated text"]
            logger.debug("Detect the generated text of dependency task (from results):" + text)
        elif "text" in dep_task["task"]["args"]:
            text = dep_task["task"]["args"]["text"]
            logger.debug("Detect the text of dependency task (from args): " + text)
        if "generated image" in dep_task["inference result"]:
            image = dep_task["inference result"]["generated image"]
            logger.debug("Detect the generated image of dependency task (from results): " + image)
        elif "image" in dep_task["task"]["args"]:
            image = dep_task["task"]["args"]["image"]
            logger.debug("Detect the image of dependency task (from args): " + image)
        if "generated audio" in dep_task["inference result"]:
            audio = dep_task["inference result"]["generated audio"]
            logger.debug("Detect the generated audio of dependency task (from results): " + audio)
        elif "audio" in dep_task["task"]["args"]:
            audio = dep_task["task"]["args"]["audio"]
            logger.debug("Detect the audio of dependency task (from args): " + audio)

    if "image" in args and "<GENERATED>" in args["image"]:
        if image:
            args["image"] = image
    if "audio" in args and "<GENERATED>" in args["audio"]:
        if audio:
            args["audio"] = audio
    if "text" in args and "<GENERATED>" in args["text"]:
        if text:
            args["text"] = text

    for resource in ["image", "audio"]:
        if resource in args and not args[resource].startswith("public/") and len(args[resource]) > 0 and not args[resource].startswith("http"):
            args[resource] = f"public/{args[resource]}"
    
    if "-text-to-image" in command['task'] and "text" not in args:
        logger.debug("control-text-to-image task, but text is empty, so we use control-generation instead.")
        control = task.split("-")[0]
        
        if control == "seg":
            task = "image-segmentation"
            command['task'] = task
        elif control == "depth":
            task = "depth-estimation"
            command['task'] = task
        else:
            task = f"{control}-control"

    command["args"] = args
    logger.debug(f"parsed task: {command}")

    if task.endswith("-text-to-image") or task.endswith("-control"):
        if inference_mode != "huggingface":
            if task.endswith("-text-to-image"):
                control = task.split("-")[0]
                best_model_id = f"lllyasviel/sd-controlnet-{control}"
            else:
                best_model_id = task
            hosted_on = "local"
            reason = "ControlNet is the best model for this task."
            choose = {"id": best_model_id, "reason": reason}
            logger.debug(f"chosen model: {choose}")
        else:
            logger.warning(f"Task {command['task']} is not available. ControlNet need to be deployed locally.")
            record_case(success=False, **{"input": input, "task": command, "reason": f"Task {command['task']} is not available. ControlNet need to be deployed locally.", "op":"message"})
            inference_result = {"error": f"service related to ControlNet is not available."}
            results[id] = collect_result(command, "", inference_result)
            return False
    elif task in ["summarization", "translation", "conversational", "text-generation", "text2text-generation"]: # ChatGPT Can do
        best_model_id = "ChatGPT"
        reason = "ChatGPT performs well on some NLP tasks as well."
        choose = {"id": best_model_id, "reason": reason}
        messages = [{
            "role": "user",
            "content": f"[ {input} ] contains a task in JSON format {command}. Now you are a {command['task']} system, the arguments are {command['args']}. Just help me do {command['task']} and give me the result. The result must be in text form without any urls."
        }]
        response = chitchat(messages, api_key, api_type, api_endpoint)
        results[id] = collect_result(command, choose, {"response": response})
        return True
    else:
        if task not in MODELS_MAP:
            logger.warning(f"no available models on {task} task.")
            record_case(success=False, **{"input": input, "task": command, "reason": f"task not support: {command['task']}", "op":"message"})
            inference_result = {"error": f"{command['task']} not found in available tasks."}
            results[id] = collect_result(command, "", inference_result)
            return False

        candidates = MODELS_MAP[task][:10]
        all_avaliable_models = get_avaliable_models(candidates, config["num_candidate_models"])
        all_avaliable_model_ids = all_avaliable_models["local"] + all_avaliable_models["huggingface"]
        logger.debug(f"avaliable models on {command['task']}: {all_avaliable_models}")

        if len(all_avaliable_model_ids) == 0:
            logger.warning(f"no available models on {command['task']}")
            record_case(success=False, **{"input": input, "task": command, "reason": f"no available models: {command['task']}", "op":"message"})
            inference_result = {"error": f"no available models on {command['task']} task."}
            results[id] = collect_result(command, "", inference_result)
            return False
            
        if len(all_avaliable_model_ids) == 1:
            best_model_id = all_avaliable_model_ids[0]
            hosted_on = "local" if best_model_id in all_avaliable_models["local"] else "huggingface"
            reason = "Only one model available."
            choose = {"id": best_model_id, "reason": reason}
            logger.debug(f"chosen model: {choose}")
        else:
            cand_models_info = [
                {
                    "id": model["id"],
                    "inference endpoint": all_avaliable_models.get(
                        "local" if model["id"] in all_avaliable_models["local"] else "huggingface"
                    ),
                    "likes": model.get("likes"),
                    "description": model.get("description", "")[:config["max_description_length"]],
                    # "language": model.get("meta").get("language") if model.get("meta") else None,
                    "tags": model.get("meta").get("tags") if model.get("meta") else None,
                }
                for model in candidates
                if model["id"] in all_avaliable_model_ids
            ]

            choose_str = choose_model(input, command, cand_models_info, api_key, api_type, api_endpoint)
            logger.debug(f"chosen model: {choose_str}")
            try:
                choose = json.loads(choose_str)
                reason = choose["reason"]
                best_model_id = choose["id"]
                hosted_on = "local" if best_model_id in all_avaliable_models["local"] else "huggingface"
            except Exception as e:
                logger.warning(f"the response [ {choose_str} ] is not a valid JSON, try to find the model id and reason in the response.")
                choose_str = find_json(choose_str)
                best_model_id, reason, choose  = get_id_reason(choose_str)
                hosted_on = "local" if best_model_id in all_avaliable_models["local"] else "huggingface"
    inference_result = model_inference(best_model_id, args, hosted_on, command['task'])

    if "error" in inference_result:
        logger.warning(f"Inference error: {inference_result['error']}")
        record_case(success=False, **{"input": input, "task": command, "reason": f"inference error: {inference_result['error']}", "op":"message"})
        results[id] = collect_result(command, choose, inference_result)
        return False
    
    results[id] = collect_result(command, choose, inference_result)
    return True

def chat_huggingface(messages, api_key, api_type, api_endpoint, return_planning = False, return_results = False):
    start = time.time()
    context = messages[:-1]
    input = messages[-1]["content"]
    logger.info("*"*80)
    logger.info(f"input: {input}")

    task_str = parse_task(context, input, api_key, api_type, api_endpoint)

    if "error" in task_str:
        record_case(success=False, **{"input": input, "task": task_str, "reason": f"task parsing error: {task_str['error']['message']}", "op":"report message"})
        return {"message": task_str["error"]["message"]}

    task_str = task_str.strip()
    logger.info(task_str)

    try:
        tasks = json.loads(task_str)
    except Exception as e:
        logger.debug(e)
        response = chitchat(messages, api_key, api_type, api_endpoint)
        record_case(success=False, **{"input": input, "task": task_str, "reason": "task parsing fail", "op":"chitchat"})
        return {"message": response}
    
    if task_str == "[]":  # using LLM response for empty task
        record_case(success=False, **{"input": input, "task": [], "reason": "task parsing fail: empty", "op": "chitchat"})
        response = chitchat(messages, api_key, api_type, api_endpoint)
        return {"message": response}

    if len(tasks) == 1 and tasks[0]["task"] in ["summarization", "translation", "conversational", "text-generation", "text2text-generation"]:
        record_case(success=True, **{"input": input, "task": tasks, "reason": "chitchat tasks", "op": "chitchat"})
        response = chitchat(messages, api_key, api_type, api_endpoint)
        return {"message": response}

    tasks = unfold(tasks)
    tasks = fix_dep(tasks)
    logger.debug(tasks)
    
    if return_planning:
        return tasks

    results = {}
    threads = []
    tasks = tasks[:]
    d = dict()
    retry = 0
    while True:
        num_thread = len(threads)
        for task in tasks:
            # logger.debug(f"d.keys(): {d.keys()}, dep: {dep}")
            for dep_id in task["dep"]:
                if dep_id >= task["id"]:
                    task["dep"] = [-1]
                    break
            dep = task["dep"]
            if dep[0] == -1 or len(list(set(dep).intersection(d.keys()))) == len(dep):
                tasks.remove(task)
                thread = threading.Thread(target=run_task, args=(input, task, d, api_key, api_type, api_endpoint))
                thread.start()
                threads.append(thread)
        if num_thread == len(threads):
            time.sleep(0.5)
            retry += 1
        if retry > 160:
            logger.debug("User has waited too long, Loop break.")
            break
        if len(tasks) == 0:
            break
    for thread in threads:
        thread.join()
    
    results = d.copy()

    logger.debug(results)
    if return_results:
        return results
    
    response = response_results(input, results, api_key, api_type, api_endpoint).strip()

    end = time.time()
    during = end - start

    answer = {"message": response}
    record_case(success=True, **{"input": input, "task": task_str, "results": results, "response": response, "during": during, "op":"response"})
    logger.info(f"response: {response}")
    return answer

def test():
    # single round examples
    inputs = [
        "Given a collection of image A: /examples/a.jpg, B: /examples/b.jpg, C: /examples/c.jpg, please tell me how many zebras in these picture?"
        "Can you give me a picture of a small bird flying in the sky with trees and clouds. Generate a high definition image if possible.",
        "Please answer all the named entities in the sentence: Iron Man is a superhero appearing in American comic books published by Marvel Comics. The character was co-created by writer and editor Stan Lee, developed by scripter Larry Lieber, and designed by artists Don Heck and Jack Kirby.",
        "please dub for me: 'Iron Man is a superhero appearing in American comic books published by Marvel Comics. The character was co-created by writer and editor Stan Lee, developed by scripter Larry Lieber, and designed by artists Don Heck and Jack Kirby.'"
        "Given an image: https://huggingface.co/datasets/mishig/sample_images/resolve/main/palace.jpg, please answer the question: What is on top of the building?",
        "Please generate a canny image based on /examples/f.jpg"
        ]
        
    for input in inputs:
        messages = [{"role": "user", "content": input}]
        chat_huggingface(messages, API_KEY, API_TYPE, API_ENDPOINT, return_planning = False, return_results = False)
    
    # multi rounds example
    messages = [
        {"role": "user", "content": "Please generate a canny image based on /examples/f.jpg"},
        {"role": "assistant", "content": """Sure. I understand your request. Based on the inference results of the models, I have generated a canny image for you. The workflow I used is as follows: First, I used the image-to-text model (nlpconnect/vit-gpt2-image-captioning) to convert the image /examples/f.jpg to text. The generated text is "a herd of giraffes and zebras grazing in a field". Second, I used the canny-control model (canny-control) to generate a canny image from the text. Unfortunately, the model failed to generate the canny image. Finally, I used the canny-text-to-image model (lllyasviel/sd-controlnet-canny) to generate a canny image from the text. The generated image is located at /images/f16d.png. I hope this answers your request. Is there anything else I can help you with?"""},
        {"role": "user", "content": """then based on the above canny image and a prompt "a photo of a zoo", generate a new image."""},
    ]
    chat_huggingface(messages, API_KEY, API_TYPE, API_ENDPOINT, return_planning = False, return_results = False)

def cli():
    messages = []
    print("Welcome to Jarvis! A collaborative system that consists of an LLM as the controller and numerous expert models as collaborative executors. Jarvis can plan tasks, schedule Hugging Face models, generate friendly responses based on your requests, and help you with many things. Please enter your request (`exit` to exit).")
    while True:
        message = input("[ User ]: ")
        if message == "exit":
            break
        messages.append({"role": "user", "content": message})
        answer = chat_huggingface(messages, API_KEY, API_TYPE, API_ENDPOINT, return_planning=False, return_results=False)
        print("[ Jarvis ]: ", answer["message"])
        messages.append({"role": "assistant", "content": answer["message"]})


def server():
    http_listen = config["http_listen"]
    host = http_listen["host"]
    port = http_listen["port"]

    app = flask.Flask(__name__, static_folder="public", static_url_path="/")
    app.config['DEBUG'] = False
    CORS(app)
    
    @cross_origin()
    @app.route('/tasks', methods=['POST'])
    def tasks():
        data = request.get_json()
        messages = data["messages"]
        api_key = data.get("api_key", API_KEY)
        api_endpoint = data.get("api_endpoint", API_ENDPOINT)
        api_type = data.get("api_type", API_TYPE)
        if api_key is None or api_type is None or api_endpoint is None:
            return jsonify({"error": "Please provide api_key, api_type and api_endpoint"}) 
        response = chat_huggingface(messages, api_key, api_type, api_endpoint, return_planning=True)
        return jsonify(response)

    @cross_origin()
    @app.route('/results', methods=['POST'])
    def results():
        data = request.get_json()
        messages = data["messages"]
        api_key = data.get("api_key", API_KEY)
        api_endpoint = data.get("api_endpoint", API_ENDPOINT)
        api_type = data.get("api_type", API_TYPE)
        if api_key is None or api_type is None or api_endpoint is None:
            return jsonify({"error": "Please provide api_key, api_type and api_endpoint"}) 
        response = chat_huggingface(messages, api_key, api_type, api_endpoint, return_results=True)
        return jsonify(response)

    @cross_origin()
    @app.route('/hugginggpt', methods=['POST'])
    def chat():
        data = request.get_json()
        messages = data["messages"]
        api_key = data.get("api_key", API_KEY)
        api_endpoint = data.get("api_endpoint", API_ENDPOINT)
        api_type = data.get("api_type", API_TYPE)
        if api_key is None or api_type is None or api_endpoint is None:
            return jsonify({"error": "Please provide api_key, api_type and api_endpoint"}) 
        response = chat_huggingface(messages, api_key, api_type, api_endpoint)
        return jsonify(response)
    print("server running...")
    waitress.serve(app, host=host, port=port)

if __name__ == "__main__":
    if args.mode == "test":
        test()
    elif args.mode == "server":
        server()
    elif args.mode == "cli":
        cli()



File: hugginggpt/server/get_token_ids.py

import tiktoken

encodings = {
    "gpt-4": tiktoken.get_encoding("cl100k_base"),
    "gpt-4-32k": tiktoken.get_encoding("cl100k_base"),
    "gpt-3.5-turbo": tiktoken.get_encoding("cl100k_base"),
    "gpt-3.5-turbo-0301": tiktoken.get_encoding("cl100k_base"),
    "text-davinci-003": tiktoken.get_encoding("p50k_base"),
    "text-davinci-002": tiktoken.get_encoding("p50k_base"),
    "text-davinci-001": tiktoken.get_encoding("r50k_base"),
    "text-curie-001": tiktoken.get_encoding("r50k_base"),
    "text-babbage-001": tiktoken.get_encoding("r50k_base"),
    "text-ada-001": tiktoken.get_encoding("r50k_base"),
    "davinci": tiktoken.get_encoding("r50k_base"),
    "curie": tiktoken.get_encoding("r50k_base"),
    "babbage": tiktoken.get_encoding("r50k_base"),
    "ada": tiktoken.get_encoding("r50k_base"),
}

max_length = {
    "gpt-4": 8192,
    "gpt-4-32k": 32768,
    "gpt-3.5-turbo": 4096,
    "gpt-3.5-turbo-0301": 4096,
    "text-davinci-003": 4096,
    "text-davinci-002": 4096,
    "text-davinci-001": 2049,
    "text-curie-001": 2049,
    "text-babbage-001": 2049,
    "text-ada-001": 2049,
    "davinci": 2049,
    "curie": 2049,
    "babbage": 2049,
    "ada": 2049
}

def count_tokens(model_name, text):
    return len(encodings[model_name].encode(text))

def get_max_context_length(model_name):
    return max_length[model_name]

def get_token_ids_for_task_parsing(model_name):
    text = '''{"task": "text-classification",  "token-classification", "text2text-generation", "summarization", "translation",  "question-answering", "conversational", "text-generation", "sentence-similarity", "tabular-classification", "object-detection", "image-classification", "image-to-image", "image-to-text", "text-to-image", "visual-question-answering", "document-question-answering", "image-segmentation", "text-to-speech", "text-to-video", "automatic-speech-recognition", "audio-to-audio", "audio-classification", "canny-control", "hed-control", "mlsd-control", "normal-control", "openpose-control", "canny-text-to-image", "depth-text-to-image", "hed-text-to-image", "mlsd-text-to-image", "normal-text-to-image", "openpose-text-to-image", "seg-text-to-image", "args", "text", "path", "dep", "id", "<GENERATED>-"}'''
    res = encodings[model_name].encode(text)
    res = list(set(res))
    return res

def get_token_ids_for_choose_model(model_name):
    text = '''{"id": "reason"}'''
    res = encodings[model_name].encode(text)
    res = list(set(res))
    return res


File: hugginggpt/server/models_server.py

import argparse
import logging
import random
import uuid
import numpy as np
from transformers import pipeline
from diffusers import DiffusionPipeline, StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
from diffusers.utils import load_image
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5ForSpeechToSpeech
from transformers import BlipProcessor, BlipForConditionalGeneration
from transformers import TrOCRProcessor, VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from datasets import load_dataset
from PIL import Image
import flask
from flask import request, jsonify
import waitress
from flask_cors import CORS
import io
from torchvision import transforms
import torch
import torchaudio
from speechbrain.pretrained import WaveformEnhancement
import joblib
from huggingface_hub import hf_hub_url, cached_download
from transformers import AutoImageProcessor, TimesformerForVideoClassification
from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation, AutoFeatureExtractor
from controlnet_aux import OpenposeDetector, MLSDdetector, HEDdetector, CannyDetector, MidasDetector
from controlnet_aux.open_pose.body import Body
from controlnet_aux.mlsd.models.mbv2_mlsd_large import MobileV2_MLSD_Large
from controlnet_aux.hed import Network
from transformers import DPTForDepthEstimation, DPTFeatureExtractor
import warnings
import time
from espnet2.bin.tts_inference import Text2Speech
import soundfile as sf
from asteroid.models import BaseModel
import traceback
import os
import yaml

warnings.filterwarnings("ignore")

parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="configs/config.default.yaml")
args = parser.parse_args()

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

config = yaml.load(open(args.config, "r"), Loader=yaml.FullLoader)

# host = config["local_inference_endpoint"]["host"]
port = config["local_inference_endpoint"]["port"]

local_deployment = config["local_deployment"]
device = config.get("device", "cuda:0") 

PROXY = None
if config["proxy"]:
    PROXY = {
        "https": config["proxy"],
    }

app = flask.Flask(__name__)
CORS(app)

start = time.time()

local_fold = "models"
# if args.config.endswith(".dev"):
#     local_fold = "models_dev"


def load_pipes(local_deployment):
    other_pipes = {}
    standard_pipes = {}
    controlnet_sd_pipes = {}
    if local_deployment in ["full"]:
        other_pipes = {
            "nlpconnect/vit-gpt2-image-captioning":{
                "model": VisionEncoderDecoderModel.from_pretrained(f"{local_fold}/nlpconnect/vit-gpt2-image-captioning"),
                "feature_extractor": ViTImageProcessor.from_pretrained(f"{local_fold}/nlpconnect/vit-gpt2-image-captioning"),
                "tokenizer": AutoTokenizer.from_pretrained(f"{local_fold}/nlpconnect/vit-gpt2-image-captioning"),
                "device": device
            },
            # "Salesforce/blip-image-captioning-large": {
            #     "model": BlipForConditionalGeneration.from_pretrained(f"{local_fold}/Salesforce/blip-image-captioning-large"),
            #     "processor": BlipProcessor.from_pretrained(f"{local_fold}/Salesforce/blip-image-captioning-large"),
            #     "device": device
            # },
            "damo-vilab/text-to-video-ms-1.7b": {
                "model": DiffusionPipeline.from_pretrained(f"{local_fold}/damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16"),
                "device": device
            },
            # "facebook/maskformer-swin-large-ade": {
            #     "model": MaskFormerForInstanceSegmentation.from_pretrained(f"{local_fold}/facebook/maskformer-swin-large-ade"),
            #     "feature_extractor" : AutoFeatureExtractor.from_pretrained("facebook/maskformer-swin-large-ade"),
            #     "device": device
            # },
            # "microsoft/trocr-base-printed": {
            #     "processor": TrOCRProcessor.from_pretrained(f"{local_fold}/microsoft/trocr-base-printed"),
            #     "model": VisionEncoderDecoderModel.from_pretrained(f"{local_fold}/microsoft/trocr-base-printed"),
            #     "device": device
            # },
            # "microsoft/trocr-base-handwritten": {
            #     "processor": TrOCRProcessor.from_pretrained(f"{local_fold}/microsoft/trocr-base-handwritten"),
            #     "model": VisionEncoderDecoderModel.from_pretrained(f"{local_fold}/microsoft/trocr-base-handwritten"),
            #     "device": device
            # },
            "JorisCos/DCCRNet_Libri1Mix_enhsingle_16k": {
                "model": BaseModel.from_pretrained("JorisCos/DCCRNet_Libri1Mix_enhsingle_16k"),
                "device": device
            },
            "espnet/kan-bayashi_ljspeech_vits": {
                "model": Text2Speech.from_pretrained(f"espnet/kan-bayashi_ljspeech_vits"),
                "device": device
            },
            "lambdalabs/sd-image-variations-diffusers": {
                "model": DiffusionPipeline.from_pretrained(f"{local_fold}/lambdalabs/sd-image-variations-diffusers"), #torch_dtype=torch.float16
                "device": device
            },
            # "CompVis/stable-diffusion-v1-4": {
            #     "model": DiffusionPipeline.from_pretrained(f"{local_fold}/CompVis/stable-diffusion-v1-4"),
            #     "device": device
            # },
            # "stabilityai/stable-diffusion-2-1": {
            #     "model": DiffusionPipeline.from_pretrained(f"{local_fold}/stabilityai/stable-diffusion-2-1"),
            #     "device": device
            # },
            "runwayml/stable-diffusion-v1-5": {
                "model": DiffusionPipeline.from_pretrained(f"{local_fold}/runwayml/stable-diffusion-v1-5"),
                "device": device
            },
            # "microsoft/speecht5_tts":{
            #     "processor": SpeechT5Processor.from_pretrained(f"{local_fold}/microsoft/speecht5_tts"),
            #     "model": SpeechT5ForTextToSpeech.from_pretrained(f"{local_fold}/microsoft/speecht5_tts"),
            #     "vocoder":  SpeechT5HifiGan.from_pretrained(f"{local_fold}/microsoft/speecht5_hifigan"),
            #     "embeddings_dataset": load_dataset(f"{local_fold}/Matthijs/cmu-arctic-xvectors", split="validation"),
            #     "device": device
            # },
            # "speechbrain/mtl-mimic-voicebank": {
            #     "model": WaveformEnhancement.from_hparams(source="speechbrain/mtl-mimic-voicebank", savedir="models/mtl-mimic-voicebank"),
            #     "device": device
            # },
            "microsoft/speecht5_vc":{
                "processor": SpeechT5Processor.from_pretrained(f"{local_fold}/microsoft/speecht5_vc"),
                "model": SpeechT5ForSpeechToSpeech.from_pretrained(f"{local_fold}/microsoft/speecht5_vc"),
                "vocoder": SpeechT5HifiGan.from_pretrained(f"{local_fold}/microsoft/speecht5_hifigan"),
                "embeddings_dataset": load_dataset(f"{local_fold}/Matthijs/cmu-arctic-xvectors", split="validation"),
                "device": device
            },
            # "julien-c/wine-quality": {
            #     "model": joblib.load(cached_download(hf_hub_url("julien-c/wine-quality", "sklearn_model.joblib")))
            # },
            # "facebook/timesformer-base-finetuned-k400": {
            #     "processor": AutoImageProcessor.from_pretrained(f"{local_fold}/facebook/timesformer-base-finetuned-k400"),
            #     "model": TimesformerForVideoClassification.from_pretrained(f"{local_fold}/facebook/timesformer-base-finetuned-k400"),
            #     "device": device
            # },
            "facebook/maskformer-swin-base-coco": {
                "feature_extractor": MaskFormerFeatureExtractor.from_pretrained(f"{local_fold}/facebook/maskformer-swin-base-coco"),
                "model": MaskFormerForInstanceSegmentation.from_pretrained(f"{local_fold}/facebook/maskformer-swin-base-coco"),
                "device": device
            },
            "Intel/dpt-hybrid-midas": {
                "model": DPTForDepthEstimation.from_pretrained(f"{local_fold}/Intel/dpt-hybrid-midas", low_cpu_mem_usage=True),
                "feature_extractor": DPTFeatureExtractor.from_pretrained(f"{local_fold}/Intel/dpt-hybrid-midas"),
                "device": device
            }
        }

    if local_deployment in ["full", "standard"]:
        standard_pipes = {
            # "superb/wav2vec2-base-superb-ks": {
            #     "model": pipeline(task="audio-classification", model=f"{local_fold}/superb/wav2vec2-base-superb-ks"), 
            #     "device": device
            # },
            "openai/whisper-base": {
                "model": pipeline(task="automatic-speech-recognition", model=f"{local_fold}/openai/whisper-base"), 
                "device": device
            },
            "microsoft/speecht5_asr": {
                "model": pipeline(task="automatic-speech-recognition", model=f"{local_fold}/microsoft/speecht5_asr"), 
                "device": device
            },
            "Intel/dpt-large": {
                "model": pipeline(task="depth-estimation", model=f"{local_fold}/Intel/dpt-large"), 
                "device": device
            },
            # "microsoft/beit-base-patch16-224-pt22k-ft22k": {
            #     "model": pipeline(task="image-classification", model=f"{local_fold}/microsoft/beit-base-patch16-224-pt22k-ft22k"), 
            #     "device": device
            # },
            "facebook/detr-resnet-50-panoptic": {
                "model": pipeline(task="image-segmentation", model=f"{local_fold}/facebook/detr-resnet-50-panoptic"), 
                "device": device
            },
            "facebook/detr-resnet-101": {
                "model": pipeline(task="object-detection", model=f"{local_fold}/facebook/detr-resnet-101"), 
                "device": device
            },
            # "openai/clip-vit-large-patch14": {
            #     "model": pipeline(task="zero-shot-image-classification", model=f"{local_fold}/openai/clip-vit-large-patch14"), 
            #     "device": device
            # },
            "google/owlvit-base-patch32": {
                "model": pipeline(task="zero-shot-object-detection", model=f"{local_fold}/google/owlvit-base-patch32"), 
                "device": device
            },
            # "microsoft/DialoGPT-medium": {
            #     "model": pipeline(task="conversational", model=f"{local_fold}/microsoft/DialoGPT-medium"), 
            #     "device": device
            # },
            # "bert-base-uncased": {
            #     "model": pipeline(task="fill-mask", model=f"{local_fold}/bert-base-uncased"), 
            #     "device": device
            # },
            # "deepset/roberta-base-squad2": {
            #     "model": pipeline(task = "question-answering", model=f"{local_fold}/deepset/roberta-base-squad2"), 
            #     "device": device
            # },
            # "facebook/bart-large-cnn": {
            #     "model": pipeline(task="summarization", model=f"{local_fold}/facebook/bart-large-cnn"), 
            #     "device": device
            # },
            # "google/tapas-base-finetuned-wtq": {
            #     "model": pipeline(task="table-question-answering", model=f"{local_fold}/google/tapas-base-finetuned-wtq"), 
            #     "device": device
            # },
            # "distilbert-base-uncased-finetuned-sst-2-english": {
            #     "model": pipeline(task="text-classification", model=f"{local_fold}/distilbert-base-uncased-finetuned-sst-2-english"), 
            #     "device": device
            # },
            # "gpt2": {
            #     "model": pipeline(task="text-generation", model="gpt2"), 
            #     "device": device
            # },
            # "mrm8488/t5-base-finetuned-question-generation-ap": {
            #     "model": pipeline(task="text2text-generation", model=f"{local_fold}/mrm8488/t5-base-finetuned-question-generation-ap"), 
            #     "device": device
            # },
            # "Jean-Baptiste/camembert-ner": {
            #     "model": pipeline(task="token-classification", model=f"{local_fold}/Jean-Baptiste/camembert-ner", aggregation_strategy="simple"), 
            #     "device": device
            # },
            # "t5-base": {
            #     "model": pipeline(task="translation", model=f"{local_fold}/t5-base"), 
            #     "device": device
            # },
            "impira/layoutlm-document-qa": {
                "model": pipeline(task="document-question-answering", model=f"{local_fold}/impira/layoutlm-document-qa"), 
                "device": device
            },
            "ydshieh/vit-gpt2-coco-en": {
                "model": pipeline(task="image-to-text", model=f"{local_fold}/ydshieh/vit-gpt2-coco-en"), 
                "device": device
            },
            "dandelin/vilt-b32-finetuned-vqa": {
                "model": pipeline(task="visual-question-answering", model=f"{local_fold}/dandelin/vilt-b32-finetuned-vqa"), 
                "device": device
            }
        }

    if local_deployment in ["full", "standard", "minimal"]:
        controlnet = ControlNetModel.from_pretrained(f"{local_fold}/lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
        controlnetpipe = StableDiffusionControlNetPipeline.from_pretrained(
            f"{local_fold}/runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
        )

        def mlsd_control_network():
            model = MobileV2_MLSD_Large()
            model.load_state_dict(torch.load(f"{local_fold}/lllyasviel/ControlNet/annotator/ckpts/mlsd_large_512_fp32.pth"), strict=True)
            return MLSDdetector(model)


        hed_network = Network(f"{local_fold}/lllyasviel/ControlNet/annotator/ckpts/network-bsds500.pth")

        controlnet_sd_pipes = {
            "openpose-control": {
                "model": OpenposeDetector(Body(f"{local_fold}/lllyasviel/ControlNet/annotator/ckpts/body_pose_model.pth"))
            },
            "mlsd-control": {
                "model": mlsd_control_network()
            },
            "hed-control": {
                "model": HEDdetector(hed_network)
            },
            "scribble-control": {
                "model": HEDdetector(hed_network)
            },
            "midas-control": {
                "model": MidasDetector(model_path=f"{local_fold}/lllyasviel/ControlNet/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt")
            },
            "canny-control": {
                "model": CannyDetector()
            },
            "lllyasviel/sd-controlnet-canny":{
                "control": controlnet, 
                "model": controlnetpipe,
                "device": device
            },
            "lllyasviel/sd-controlnet-depth":{
                "control": ControlNetModel.from_pretrained(f"{local_fold}/lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16),
                "model": controlnetpipe,
                "device": device
            },
            "lllyasviel/sd-controlnet-hed":{
                "control": ControlNetModel.from_pretrained(f"{local_fold}/lllyasviel/sd-controlnet-hed", torch_dtype=torch.float16), 
                "model": controlnetpipe,
                "device": device
            },
            "lllyasviel/sd-controlnet-mlsd":{
                "control": ControlNetModel.from_pretrained(f"{local_fold}/lllyasviel/sd-controlnet-mlsd", torch_dtype=torch.float16), 
                "model": controlnetpipe,
                "device": device
            },
            "lllyasviel/sd-controlnet-openpose":{
                "control": ControlNetModel.from_pretrained(f"{local_fold}/lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16), 
                "model": controlnetpipe,
                "device": device
            },
            "lllyasviel/sd-controlnet-scribble":{
                "control": ControlNetModel.from_pretrained(f"{local_fold}/lllyasviel/sd-controlnet-scribble", torch_dtype=torch.float16), 
                "model": controlnetpipe,
                "device": device
            },
            "lllyasviel/sd-controlnet-seg":{
                "control": ControlNetModel.from_pretrained(f"{local_fold}/lllyasviel/sd-controlnet-seg", torch_dtype=torch.float16), 
                "model": controlnetpipe,
                "device": device
            }    
        }
    pipes = {**standard_pipes, **other_pipes, **controlnet_sd_pipes}
    return pipes

pipes = load_pipes(local_deployment)

end = time.time()
during = end - start

print(f"[ ready ] {during}s")

@app.route('/running', methods=['GET'])
def running():
    return jsonify({"running": True})

@app.route('/status/<path:model_id>', methods=['GET'])
def status(model_id):
    disabled_models = ["microsoft/trocr-base-printed", "microsoft/trocr-base-handwritten"]
    if model_id in pipes.keys() and model_id not in disabled_models:
        print(f"[ check {model_id} ] success")
        return jsonify({"loaded": True})
    else:
        print(f"[ check {model_id} ] failed")
        return jsonify({"loaded": False})

@app.route('/models/<path:model_id>', methods=['POST'])
def models(model_id):
    while "using" in pipes[model_id] and pipes[model_id]["using"]:
        print(f"[ inference {model_id} ] waiting")
        time.sleep(0.1)
    pipes[model_id]["using"] = True
    print(f"[ inference {model_id} ] start")

    start = time.time()

    pipe = pipes[model_id]["model"]
    
    if "device" in pipes[model_id]:
        try:
            pipe.to(pipes[model_id]["device"])
        except:
            pipe.device = torch.device(pipes[model_id]["device"])
            pipe.model.to(pipes[model_id]["device"])
    
    result = None
    try:
        # text to video
        if model_id == "damo-vilab/text-to-video-ms-1.7b":
            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
            # pipe.enable_model_cpu_offload()
            prompt = request.get_json()["text"]
            video_frames = pipe(prompt, num_inference_steps=50, num_frames=40).frames
            video_path = export_to_video(video_frames)
            file_name = str(uuid.uuid4())[:4]
            os.system(f"LD_LIBRARY_PATH=/usr/local/lib /usr/local/bin/ffmpeg -i {video_path} -vcodec libx264 public/videos/{file_name}.mp4")
            result = {"path": f"/videos/{file_name}.mp4"}

        # controlnet
        if model_id.startswith("lllyasviel/sd-controlnet-"):
            pipe.controlnet.to('cpu')
            pipe.controlnet = pipes[model_id]["control"].to(pipes[model_id]["device"])
            pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
            control_image = load_image(request.get_json()["img_url"])
            # generator = torch.manual_seed(66)
            out_image: Image = pipe(request.get_json()["text"], num_inference_steps=20, image=control_image).images[0]
            file_name = str(uuid.uuid4())[:4]
            out_image.save(f"public/images/{file_name}.png")
            result = {"path": f"/images/{file_name}.png"}

        if model_id.endswith("-control"):
            image = load_image(request.get_json()["img_url"])
            if "scribble" in model_id:
                control = pipe(image, scribble = True)
            elif "canny" in model_id:
                control = pipe(image, low_threshold=100, high_threshold=200)
            else:
                control = pipe(image)
            file_name = str(uuid.uuid4())[:4]
            control.save(f"public/images/{file_name}.png")
            result = {"path": f"/images/{file_name}.png"}

        # image to image
        if model_id == "lambdalabs/sd-image-variations-diffusers":
            im = load_image(request.get_json()["img_url"])
            file_name = str(uuid.uuid4())[:4]
            with open(f"public/images/{file_name}.png", "wb") as f:
                f.write(request.data)
            tform = transforms.Compose([
                transforms.ToTensor(),
                transforms.Resize(
                    (224, 224),
                    interpolation=transforms.InterpolationMode.BICUBIC,
                    antialias=False,
                    ),
                transforms.Normalize(
                [0.48145466, 0.4578275, 0.40821073],
                [0.26862954, 0.26130258, 0.27577711]),
            ])
            inp = tform(im).to(pipes[model_id]["device"]).unsqueeze(0)
            out = pipe(inp, guidance_scale=3)
            out["images"][0].save(f"public/images/{file_name}.jpg")
            result = {"path": f"/images/{file_name}.jpg"}

        # image to text
        if model_id == "Salesforce/blip-image-captioning-large":
            raw_image = load_image(request.get_json()["img_url"]).convert('RGB')
            text = request.get_json()["text"]
            inputs = pipes[model_id]["processor"](raw_image, return_tensors="pt").to(pipes[model_id]["device"])
            out = pipe.generate(**inputs)
            caption = pipes[model_id]["processor"].decode(out[0], skip_special_tokens=True)
            result = {"generated text": caption}
        if model_id == "ydshieh/vit-gpt2-coco-en":
            img_url = request.get_json()["img_url"]
            generated_text = pipe(img_url)[0]['generated_text']
            result = {"generated text": generated_text}
        if model_id == "nlpconnect/vit-gpt2-image-captioning":
            image = load_image(request.get_json()["img_url"]).convert("RGB")
            pixel_values = pipes[model_id]["feature_extractor"](images=image, return_tensors="pt").pixel_values
            pixel_values = pixel_values.to(pipes[model_id]["device"])
            generated_ids = pipe.generate(pixel_values, **{"max_length": 200, "num_beams": 1})
            generated_text = pipes[model_id]["tokenizer"].batch_decode(generated_ids, skip_special_tokens=True)[0]
            result = {"generated text": generated_text}
        # image to text: OCR
        if model_id == "microsoft/trocr-base-printed" or  model_id == "microsoft/trocr-base-handwritten":
            image = load_image(request.get_json()["img_url"]).convert("RGB")
            pixel_values = pipes[model_id]["processor"](image, return_tensors="pt").pixel_values
            pixel_values = pixel_values.to(pipes[model_id]["device"])
            generated_ids = pipe.generate(pixel_values)
            generated_text = pipes[model_id]["processor"].batch_decode(generated_ids, skip_special_tokens=True)[0]
            result = {"generated text": generated_text}

        # text to image
        if model_id == "runwayml/stable-diffusion-v1-5":
            file_name = str(uuid.uuid4())[:4]
            text = request.get_json()["text"]
            out = pipe(prompt=text)
            out["images"][0].save(f"public/images/{file_name}.jpg")
            result = {"path": f"/images/{file_name}.jpg"}

        # object detection
        if model_id == "google/owlvit-base-patch32" or model_id == "facebook/detr-resnet-101":
            img_url = request.get_json()["img_url"]
            open_types = ["cat", "couch", "person", "car", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird"]
            result = pipe(img_url, candidate_labels=open_types)
        
        # VQA
        if model_id == "dandelin/vilt-b32-finetuned-vqa":
            question = request.get_json()["text"]
            img_url = request.get_json()["img_url"]
            result = pipe(question=question, image=img_url)
        
        #DQA
        if model_id == "impira/layoutlm-document-qa":
            question = request.get_json()["text"]
            img_url = request.get_json()["img_url"]
            result = pipe(img_url, question)

        # depth-estimation
        if model_id == "Intel/dpt-large":
            output = pipe(request.get_json()["img_url"])
            image = output['depth']
            name = str(uuid.uuid4())[:4]
            image.save(f"public/images/{name}.jpg")
            result = {"path": f"/images/{name}.jpg"}

        if model_id == "Intel/dpt-hybrid-midas" and model_id == "Intel/dpt-large":
            image = load_image(request.get_json()["img_url"])
            inputs = pipes[model_id]["feature_extractor"](images=image, return_tensors="pt")
            with torch.no_grad():
                outputs = pipe(**inputs)
                predicted_depth = outputs.predicted_depth
            prediction = torch.nn.functional.interpolate(
                predicted_depth.unsqueeze(1),
                size=image.size[::-1],
                mode="bicubic",
                align_corners=False,
            )
            output = prediction.squeeze().cpu().numpy()
            formatted = (output * 255 / np.max(output)).astype("uint8")
            image = Image.fromarray(formatted)
            name = str(uuid.uuid4())[:4]
            image.save(f"public/images/{name}.jpg")
            result = {"path": f"/images/{name}.jpg"}

        # TTS
        if model_id == "espnet/kan-bayashi_ljspeech_vits":
            text = request.get_json()["text"]
            wav = pipe(text)["wav"]
            name = str(uuid.uuid4())[:4]
            sf.write(f"public/audios/{name}.wav", wav.cpu().numpy(), pipe.fs, "PCM_16")
            result = {"path": f"/audios/{name}.wav"}

        if model_id == "microsoft/speecht5_tts":
            text = request.get_json()["text"]
            inputs = pipes[model_id]["processor"](text=text, return_tensors="pt")
            embeddings_dataset = pipes[model_id]["embeddings_dataset"]
            speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0).to(pipes[model_id]["device"])
            pipes[model_id]["vocoder"].to(pipes[model_id]["device"])
            speech = pipe.generate_speech(inputs["input_ids"].to(pipes[model_id]["device"]), speaker_embeddings, vocoder=pipes[model_id]["vocoder"])
            name = str(uuid.uuid4())[:4]
            sf.write(f"public/audios/{name}.wav", speech.cpu().numpy(), samplerate=16000)
            result = {"path": f"/audios/{name}.wav"}

        # ASR
        if model_id == "openai/whisper-base" or model_id == "microsoft/speecht5_asr":
            audio_url = request.get_json()["audio_url"]
            result = { "text": pipe(audio_url)["text"]}

        # audio to audio
        if model_id == "JorisCos/DCCRNet_Libri1Mix_enhsingle_16k":
            audio_url = request.get_json()["audio_url"]
            wav, sr = torchaudio.load(audio_url)
            with torch.no_grad():
                result_wav = pipe(wav.to(pipes[model_id]["device"]))
            name = str(uuid.uuid4())[:4]
            sf.write(f"public/audios/{name}.wav", result_wav.cpu().squeeze().numpy(), sr)
            result = {"path": f"/audios/{name}.wav"}
        
        if model_id == "microsoft/speecht5_vc":
            audio_url = request.get_json()["audio_url"]
            wav, sr = torchaudio.load(audio_url)
            inputs = pipes[model_id]["processor"](audio=wav, sampling_rate=sr, return_tensors="pt")
            embeddings_dataset = pipes[model_id]["embeddings_dataset"]
            speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
            pipes[model_id]["vocoder"].to(pipes[model_id]["device"])
            speech = pipe.generate_speech(inputs["input_ids"].to(pipes[model_id]["device"]), speaker_embeddings, vocoder=pipes[model_id]["vocoder"])
            name = str(uuid.uuid4())[:4]
            sf.write(f"public/audios/{name}.wav", speech.cpu().numpy(), samplerate=16000)
            result = {"path": f"/audios/{name}.wav"}
        
        # segmentation
        if model_id == "facebook/detr-resnet-50-panoptic":
            result = []
            segments = pipe(request.get_json()["img_url"])
            image = load_image(request.get_json()["img_url"])

            colors = []
            for i in range(len(segments)):
                colors.append((random.randint(100, 255), random.randint(100, 255), random.randint(100, 255), 50))

            for segment in segments:
                mask = segment["mask"]
                mask = mask.convert('L')
                layer = Image.new('RGBA', mask.size, colors[i])
                image.paste(layer, (0, 0), mask)
            name = str(uuid.uuid4())[:4]
            image.save(f"public/images/{name}.jpg")
            result = {"path": f"/images/{name}.jpg"}

        if model_id == "facebook/maskformer-swin-base-coco" or model_id == "facebook/maskformer-swin-large-ade":
            image = load_image(request.get_json()["img_url"])
            inputs = pipes[model_id]["feature_extractor"](images=image, return_tensors="pt").to(pipes[model_id]["device"])
            outputs = pipe(**inputs)
            result = pipes[model_id]["feature_extractor"].post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
            predicted_panoptic_map = result["segmentation"].cpu().numpy()
            predicted_panoptic_map = Image.fromarray(predicted_panoptic_map.astype(np.uint8))
            name = str(uuid.uuid4())[:4]
            predicted_panoptic_map.save(f"public/images/{name}.jpg")
            result = {"path": f"/images/{name}.jpg"}

    except Exception as e:
        print(e)
        traceback.print_exc()
        result = {"error": {"message": "Error when running the model inference."}}

    if "device" in pipes[model_id]:
        try:
            pipe.to("cpu")
            torch.cuda.empty_cache()
        except:
            pipe.device = torch.device("cpu")
            pipe.model.to("cpu")
            torch.cuda.empty_cache()

    pipes[model_id]["using"] = False

    if result is None:
        result = {"error": {"message": "model not found"}}
    
    end = time.time()
    during = end - start
    print(f"[ complete {model_id} ] {during}s")
    print(f"[ result {model_id} ] {result}")

    return jsonify(result)


if __name__ == '__main__':
    # temp folders
    if not os.path.exists("public/audios"):
        os.makedirs("public/audios")
    if not os.path.exists("public/images"):
        os.makedirs("public/images")
    if not os.path.exists("public/videos"):
        os.makedirs("public/videos")
        
    waitress.serve(app, host="0.0.0.0", port=port)


File: hugginggpt/server/run_gradio_demo.py

import uuid
import gradio as gr
import re
from diffusers.utils import load_image
import requests
from awesome_chat import chat_huggingface

all_messages = []
OPENAI_KEY = ""

def add_message(content, role):
    message = {"role":role, "content":content}
    all_messages.append(message)

def extract_medias(message):
    image_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(jpg|jpeg|tiff|gif|png)")
    image_urls = []
    for match in image_pattern.finditer(message):
        if match.group(0) not in image_urls:
            image_urls.append(match.group(0))

    audio_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(flac|wav)")
    audio_urls = []
    for match in audio_pattern.finditer(message):
        if match.group(0) not in audio_urls:
            audio_urls.append(match.group(0))

    video_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(mp4)")
    video_urls = []
    for match in video_pattern.finditer(message):
        if match.group(0) not in video_urls:
            video_urls.append(match.group(0))

    return image_urls, audio_urls, video_urls

def set_openai_key(openai_key):
    global OPENAI_KEY
    OPENAI_KEY = openai_key
    return OPENAI_KEY

def add_text(messages, message):
    if len(OPENAI_KEY) == 0 or not OPENAI_KEY.startswith("sk-"):
        return messages, "Please set your OpenAI API key first."
    add_message(message, "user")
    messages = messages + [(message, None)]
    image_urls, audio_urls, video_urls = extract_medias(message)

    for image_url in image_urls:
        if not image_url.startswith("http"):
            image_url = "public/" + image_url
        image = load_image(image_url)
        name = f"public/images/{str(uuid.uuid4())[:4]}.jpg" 
        image.save(name)
        messages = messages + [((f"{name}",), None)]
    for audio_url in audio_urls:
        if not audio_url.startswith("http"):
            audio_url = "public/" + audio_url
        ext = audio_url.split(".")[-1]
        name = f"public/audios/{str(uuid.uuid4()[:4])}.{ext}"
        response = requests.get(audio_url)
        with open(name, "wb") as f:
            f.write(response.content)
        messages = messages + [((f"{name}",), None)]
    for video_url in video_urls:
        if not video_url.startswith("http"):
            video_url = "public/" + video_url
        ext = video_url.split(".")[-1]
        name = f"public/audios/{str(uuid.uuid4()[:4])}.{ext}"
        response = requests.get(video_url)
        with open(name, "wb") as f:
            f.write(response.content)
        messages = messages + [((f"{name}",), None)]
    return messages, ""

def bot(messages):
    if len(OPENAI_KEY) == 0 or not OPENAI_KEY.startswith("sk-"):
        return messages
    message = chat_huggingface(all_messages, OPENAI_KEY, "openai", "https://api.openai.com/v1/completions")["message"]
    image_urls, audio_urls, video_urls = extract_medias(message)
    add_message(message, "assistant")
    messages[-1][1] = message
    for image_url in image_urls:
        if not image_url.startswith("http"):
            image_url = image_url.replace("public/", "")
            messages = messages + [((None, (f"public/{image_url}",)))]
    for audio_url in audio_urls:
        if not audio_url.startswith("http"):
            audio_url = audio_url.replace("public/", "")
            messages = messages + [((None, (f"public/{audio_url}",)))]
    for video_url in video_urls:
        if not video_url.startswith("http"):
            video_url = video_url.replace("public/", "")
            messages = messages + [((None, (f"public/{video_url}",)))]
    return messages

with gr.Blocks() as demo:
    gr.Markdown("<h2><center>HuggingGPT (Dev)</center></h2>")
    with gr.Row():
        openai_api_key = gr.Textbox(
            show_label=False,
            placeholder="Set your OpenAI API key here and press Enter",
            lines=1,
            type="password",
        )

    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=500)

    with gr.Row():
        txt = gr.Textbox(
            show_label=False,
            placeholder="Enter text and press enter. The url of the multimedia resource must contain the extension name.",
        ).style(container=False)

    txt.submit(add_text, [chatbot, txt], [chatbot, txt]).then(
        bot, chatbot, chatbot
    )
    openai_api_key.submit(set_openai_key, [openai_api_key], [openai_api_key])

    gr.Examples(
        examples=["Given a collection of image A: /examples/a.jpg, B: /examples/b.jpg, C: /examples/c.jpg, please tell me how many zebras in these picture?",
                    "Please generate a canny image based on /examples/f.jpg",
                    "show me a joke and an image of cat",
                    "what is in the /examples/a.jpg",
                    "generate a video and audio about a dog is running on the grass",
                    "based on the /examples/a.jpg, please generate a video and audio",
                    "based on pose of /examples/d.jpg and content of /examples/e.jpg, please show me a new image",
                    ],
        inputs=txt
    )

demo.launch()


File: taskbench/generate_graph.py

import json
import click

def generate_graph_resource(tool_file):
    with open(tool_file) as f:
        data = json.load(f)
    data = data["nodes"]
    assert "input-type" in data[0] and "output-type" in data[0], "Input and output types are not defined"
    nodes = []
    for i in range(len(data)):
        nodes.append({"id": data[i]["id"], "desc": data[i]["desc"], "input-type": data[i]["input-type"], "output-type": data[i]["output-type"]})
    links = []
    for i in range(len(nodes)):
        for j in range(len(nodes)):
            if i != j:
                if len(set(nodes[i]["output-type"]).intersection(set(nodes[j]["input-type"]))) > 0:
                    links.append({"source": nodes[i]["id"], "target": nodes[j]["id"], "type": list(set(nodes[i]["output-type"]).intersection(set(nodes[j]["input-type"])))[0]})
    graph = {"nodes": nodes, "links": links}
    with open(tool_file.replace("tools", "graph"), 'w') as f:
        json.dump(graph, f, indent=2)

def generate_graph_temporal(tool_file):
    with open(tool_file) as f:
        data = json.load(f)
    nodes = []
    data = data["nodes"]
    if "parameters" not in data[0] and "input-type" not in data[0]:
        for i in range(len(data)):
            nodes.append({"id": data[i]["id"], "desc": data[i]["desc"]})
    elif "input-type" not in data[0]:
        for i in range(len(data)):
            nodes.append({"id": data[i]["id"], "desc": data[i]["desc"], "parameters": data[i]["parameters"]})
    else:
        for i in range(len(data)):
            nodes.append({"id": data[i]["id"], "desc": data[i]["desc"], "parameters": data[i]["parameters"], "input-type": data[i]["input-type"], "output-type": data[i]["output-type"]})
    links = []
    for i in range(len(nodes)):
        for j in range(len(nodes)):
            if i != j:
                links.append({"source": nodes[i]["id"], "target": nodes[j]["id"], "type": "complete"})
    graph = {"nodes": nodes, "links": links}
    with open(tool_file.replace("tools", "graph"), 'w') as f:
        json.dump(graph, f, indent=2)

@click.command()
@click.option('--data_dir')
@click.option('--tool_desc', default=None, help='Path to the tool description file')
@click.option('--dependency_type', default='resource', help='Type of graph to generate')
def generate_graph(tool_desc, data_dir, dependency_type):
    if tool_desc:
        tool_file = tool_desc
    else:
        tool_file = f"{data_dir}/graph_desc.json"
    if dependency_type == "temporal":
        generate_graph_temporal(tool_file)
    elif dependency_type == "resource":
        generate_graph_resource(tool_file)
    else:
        print("Type not supported")

if __name__ == "__main__":
    generate_graph()


File: taskbench/visualize_graph.py

import json
import networkx as nx
import matplotlib.pyplot as plt
import click

@click.command()
@click.option('--data_dir')
def visialize_graph(data_dir):
    graph_file = f"{data_dir}/graph_desc.json"
    with open(graph_file, "r") as f:
        data = json.load(f)

    G = nx.DiGraph()

    for node in data["nodes"]:
        G.add_node(node["id"])

    for link in data["links"]:
        G.add_edge(link["source"], link["target"])

    pos = nx.spring_layout(G)
    pos = nx.random_layout(G)    
    pos = nx.kamada_kawai_layout(G)
    
    # Show the visualization
    plt.figure(figsize=(60, 60), dpi=80)
    plt.tight_layout()
    plt.axis("off")
    plt.show()

    nx.draw_networkx_nodes(G, pos, node_color="skyblue", node_size=1200)
    nx.draw_networkx_edges(G, pos, arrows=True, arrowsize=40)
    nx.draw_networkx_labels(G, pos, font_size=50, font_color="green", font_weight="bold")
    plt.savefig(graph_file.replace(".json", ".pdf"))

if __name__ == "__main__":
    visialize_graph()


File: taskbench/data_engine.py

import os
import random
import traceback
import uuid
import json
import networkx as nx
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import shutil
import warnings

from graph_sampler import GraphSampler
import click
import matplotlib.pyplot as plt
import asyncio
import aiohttp
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.handlers = []

console_handler = logging.StreamHandler()
console_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
console_handler.setLevel(logging.INFO)
logger.addHandler(console_handler)

@click.command()
@click.option("--temperature", type=float, default=0.7)
@click.option("--top_p", type=float, default=1)
@click.option("--check", type=bool, default=False)
@click.option("--data_dir", type=str, default=None)
@click.option("--graph_desc", type=str, default=None)
@click.option("--tool_desc", type=str, default=None)
@click.option("--api_addr", type=str, default="localhost")
@click.option("--api_port", type=int, default=4000)
@click.option("--api_key", type=str, default="your api key")
@click.option("--play", type=bool, default=False)
@click.option("--method", type=str, default=None)
@click.option("--tool_number", type=int, default=None)
@click.option("--number_of_samples", type=int, default=5000)
@click.option("--seed", type=int, default=-1)
@click.option("--save_figure", type=bool, default=False)
@click.option("--multiworker", type=int, default=1)
@click.option("--llm", type=str, default="gpt-4")
@click.option("--use_async", type=bool, default=False)
@click.option("--dependency_type", type=str, default="resource")
def main(temperature, top_p, check, graph_desc, tool_desc, api_addr, api_port, api_key, play, method, tool_number, number_of_samples, seed, data_dir, save_figure, multiworker, llm, use_async, dependency_type):
    args = locals()
    url = f"http://{api_addr}:{api_port}/v1/chat/completions"
    header = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    now = datetime.now()

    if data_dir:
        if os.path.exists(data_dir):
            if graph_desc or tool_desc:
                warnings.warn(f"Data directory {data_dir} already exists, tool graph and tool desc will be ignored.")
            graph_desc = f"{data_dir}/graph_desc.json"
            tool_desc = f"{data_dir}/tool_desc.json"
    else:
        data_dir = f"result_{now.strftime('%Y%m%d%H%M%S')}_{llm}_t{temperature}_p{top_p}{'_check' if check else ''}".replace(".", "_")
    
    assert data_dir and graph_desc and tool_desc

    if seed > -1:
        random.seed(seed)
    
    tool_list = json.load(open(tool_desc, "r"))["nodes"]
    tools = {}
    if dependency_type == "resource":
        assert "input-type" in tool_list[0] and "output-type" in tool_list[0], "Input and output types are not defined"
        for tool in tool_list:
            tools[tool["id"]] = {"id": tool["id"], "desc": tool["desc"], "input-type": tool["input-type"], "output-type": tool["output-type"]}
    elif dependency_type == "temporal":
        assert "parameters" in tool_list[0], "Parameters are not defined"
        for tool in tool_list:
            tools[tool["id"]] = {"id": tool["id"], "desc": tool["desc"], "parameters": tool["parameters"]}
    else:
        raise ValueError(f"Unsupported dependency type: {dependency_type}")

    sampler = GraphSampler(file_name=graph_desc)

    if play:
        assert method is not None
        assert tool_number is not None
        result = asyncio.run(sample(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, "./", None, dependency_type))
        logger.info(json.dumps(result, indent=2))
        return

    figure_dir = None

    if not os.path.exists(data_dir):
        os.makedirs(data_dir, exist_ok=True)
        shutil.copy(graph_desc, f"{data_dir}/graph_desc.json")
        shutil.copy(tool_desc, f"{data_dir}/tool_desc.json")

    output = f"{data_dir}/data_raw.json"
    statistics_output = f"{data_dir}/statistics.json"

    file_handler = logging.FileHandler(f"{data_dir}/data_engine.log")
    file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
    file_handler.setLevel(logging.INFO)
    logger.addHandler(file_handler)

    logger.info(json.dumps(args))

    if save_figure:
        figure_dir = f"{data_dir}/task_graphs"
        os.makedirs(figure_dir, exist_ok=True)

    wf = open(output, "a")
    statistics_wf = open(f"{statistics_output}", "a")
    args["start_time"] = now.strftime("%Y-%m-%d %H:%M:%S")
    statistics_wf.write(json.dumps(args) + "\n")

    method_weights = {
        "single": 3,
        "chain": 7,
        "dag": 8,
    }

    number_weights = {
        2: 0.1,
        3: 0.2,
        4: 0.3,
        5: 0.2,
        6: 0.1,
        7: 0.05,
        8: 0.025,
        9: 0.025,
        10: 0.025,
    }

    statistics = {"total": 0, "avg_time_per_sample": 0, "success": 0, "fail": 0}
    done, failed = [], []
    if use_async:
        # coroutine with Semaphore
        sem = asyncio.Semaphore(multiworker)
        async def sample_with_statistics(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, figure_dir, wf, statistics, now, dependency_type):
            async with sem:  # semaphore limits num of simultaneous sampling
                if statistics["total"] % 100 == 0 and statistics["total"] != 0:
                    logger.info(json.dumps(statistics, indent=2))
                    statistics_wf.write(json.dumps(statistics) + "\n")
                try:
                    await sample(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, figure_dir, wf, dependency_type)
                except Exception as e:
                    statistics["total"] += 1
                    statistics["fail"] += 1
                    if str(type(e)) not in statistics:
                        statistics[str(type(e))] = 0
                    statistics[str(type(e))] += 1
                    raise e
                statistics["total"] += 1
                statistics["success"] += 1
                statistics["avg_time_per_sample"] = str((datetime.now() - now) / statistics["success"])

        async def run(url, header, llm, temperature, top_p, check, sampler, tools, figure_dir, wf, statistics, now, dependency_type):
            method = random.choices(list(method_weights.keys()), weights=list(method_weights.values()))[0]
            if method == "single":
                tool_number = 1
            else:
                tool_number = random.choices(list(number_weights.keys()), weights=list(number_weights.values()))[0]
                if method == "dag":
                    tool_number = max(tool_number, 3)
            await sample_with_statistics(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, figure_dir, wf, statistics, now, dependency_type)

        tasks = []
        for _ in range(number_of_samples):
            tasks.append(run(url, header, llm, temperature, top_p, check, sampler, tools, figure_dir, wf, statistics, now, dependency_type))

        loop = asyncio.get_event_loop()
        results = loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))

        for result in results:
            if isinstance(result, Exception):
                failed.append(result)
            else:
                done.append(result)
    else:
        # multi-thread with ThreadPoolExecutor
        executor = ThreadPoolExecutor(max_workers=multiworker)
        def sample_with_statistics(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, figure_dir, wf, statistics, now, dependency_type):
            if statistics["total"] % 100 == 0 and statistics["total"] != 0:
                logger.info(json.dumps(statistics, indent=2))
                statistics_wf.write(json.dumps(statistics) + "\n")
            try:
                asyncio.run(sample(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, figure_dir, wf, dependency_type))
            except Exception as e:
                statistics["total"] += 1
                statistics["fail"] += 1
                if str(type(e)) not in statistics:
                    statistics[str(type(e))] = 0
                statistics[str(type(e))] += 1
                raise e
            statistics["total"] += 1
            statistics["success"] += 1
            statistics["avg_time_per_sample"] = str((datetime.now() - now) / statistics["success"])

        def run(url, header, llm, temperature, top_p, check, sampler, tools, figure_dir, wf, statistics, now, dependency_type):
            method = random.choices(list(method_weights.keys()), weights=list(method_weights.values()))[0]
            if method == "single":
                tool_number = 1
            else:
                tool_number = random.choices(list(number_weights.keys()), weights=list(number_weights.values()))[0]
                if method == "dag":
                    tool_number = max(tool_number, 3)
            sample_with_statistics(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, figure_dir, wf, statistics, now, dependency_type)

        tasks = []
        for _ in range(number_of_samples):
            tasks.append(executor.submit(run, url, header, llm, temperature, top_p, check, sampler, tools, figure_dir, wf, statistics, now, dependency_type))
        for future in as_completed(tasks):
            try:
                future.result()
                done.append(future)
            except Exception as e:
                failed.append(future)

    statistics_wf.write(json.dumps(statistics) + "\n")
    logger.info(f"Done: {len(done)}, Failed: {len(failed)}")

class RateLimitError(Exception):
    def __init__(self, message):
        super().__init__(message)

class ContentFormatError(Exception):
    def __init__(self, message):
        super().__init__(message)

async def sample(url, header, llm, temperature, top_p, check, tool_number, sampler, tools, method, figure_dir, wf, dependency_type):
    start_time = datetime.now()
    sample_id = str(uuid.uuid4().int)[:8]
    sub_G = sampler.sample_subgraph(tool_number, sample_method=method)

    tool_list = list(sub_G.nodes)
    tool_edge = list(sub_G.edges)  
    seed = random.randint(0, 1000000)
    sampled_tools_string = "Given a tool graph with tools as nodes, and invoking chains between tools as edges. The following tools (nodes) are available with their corresponding descriptions and input/outputs types:\n"
    for k, tool in enumerate(tool_list):
        sampled_tools_string += f"Node {k+1}:" + json.dumps(tools[tool]) + "\n"

    sampled_links_string = "These tools can be connected as follows (the directed edges are invoking chains among tools):\n"
    for k, edge in enumerate(tool_edge):
        sampled_links_string += f"Edge: " + edge[0] + " -> " + edge[1] + "\n"
    prompt = """\nBased on the above tool graph, please be skillful to generate the according task steps, user request and tool invoking graph. \nRequirements: \n1. the generated user request should be somewhat clear, self-contained (user-specified text, image, video, audio, content should be contained in the request) and practical (help users solve a practical problem); \n2. the task steps must be strictly aligned with the tool graph (nodes and edges) and reasonable, the tool invoking graph must align with task steps, also with the given tool graph; \n3. the user request just can be decomposed into task steps solved by the tool invoking graph; \n4. each task step corresponds to a tool node in the tool graph and tool invoking graph, and the number of task steps must be same with the nodes. Each tool node can only be used once; \n5. if need image/audio/video resources in user request, please use files 'example.[jpg/mp4/wav/png]'; \n6. the dependencies among task steps must align with the edges of tool graph and tool invoking graph; \n7. the number and types of tool parameters in the generated tool invoking graph need to be consistent with the pre-defined input/outputs types of the tools. \nNow please generate your result (with random seed {""" + f"{seed}"+ """}) in a compact JSON format"""
    if dependency_type == "resource":
        prompt += """{"task_steps": [ step description of one or more steps ], "user_request": "your high-quality and self-contained synthesized request", "invoking_graph": {"nodes": [{"id": "tool name", "input": [ either user-specified text or resource file 'example.[jpg/mp4/wav/png' ] in the above user request, or the dependent tool name whose output is required by this node ]}], "links": [{"source": "tool name i", "target": "tool name j"}]}"""
    else:
        prompt += """{"task_steps": [ "concrete steps, format as Step x: Call xxx tool with xxx: 'xxx' and xxx: 'xxx'" ], "user_request": "your high-quality, concrete and self-contained synthesized request, with explicit parameter values", "invoking_graph": {"nodes": [{"id": "tool name", "arguments": [ {"name": "parameter name", "value": "parameter value, either user-specified text or the specific name of the tool whose result is required by this node"} ]}], "links": [{"source": "tool name i", "target": "tool name j"}]}"""
    if check:
        prompt += """, "check_by_teacher": "This field is filled by your strict and well-trained teacher, minor mistakes are complete intolerable to him. He evaluated whether your synthesized user request, tool invoking graph are valid and whether they are aligned with the given tool graph (strictly checked step by step according to the above requirements). Some comments from him place here (start with 'Let me check your result step by step, and evaluate the 'Executable' and 'Correct' of the tool invoking graph (Executable means that the tool invoking graph executed successfully, regardless of alignment with the given tool graph. While Correct implies that the tool invoking graph are not only 'Executable' but also strictly consistent (with strictly same nodes and same edges) with the given tool graph). After carefully evaluating, found some mistakes:' and end with a conclusion: 'Conclusion: Executable: no/yes, Correct: no/yes'.)"""
    prompt += "}:"

    final_prompt = sampled_tools_string + sampled_links_string + prompt

    if dependency_type == "temporal":
        final_prompt = final_prompt.replace("tool", "API")

    payload = json.dumps({
        "model": f"{llm}",
        "messages": [
            {
            "role": "user",
            "content":  final_prompt
            }
        ],
        "temperature": temperature,
        "top_p": top_p,
        "frequency_penalty": 0,
        "presence_penalty": 0,
        "max_tokens": 2500,
        "stream": False,
        "stop": None
    })
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(url, headers=header, data=payload, timeout=120) as response:
                resp = await response.json()

        if response.status == 429:
            raise RateLimitError(f"{resp}")
        if response.status != 200:
            raise Exception(f"{resp}")

        content = resp["choices"][0]["message"]["content"]
        content = content.replace("\n", "")
        json_start = 0
        json_end = len(content)
        
        if "```json" in content:
            json_start = content.find("```json") + 7
        if "```" in content:
            json_end = content.rfind("```")
        content = content[json_start:json_end]
        logger.info(content)
        try:
            content = json.loads(content)
        except json.JSONDecodeError as e:
            raise ContentFormatError(f"{content}")

        if dependency_type == "resource":
            sampled_nodes = [{"id": tool, "input-type": tools[tool]["input-type"], "output-type": tools[tool]["output-type"]} for tool in tool_list]
        else:
            sampled_nodes = [{"id": tool, "parameters": tools[tool]["parameters"]} for tool in tool_list]

        sampled_links = [{"source": edge[0], "target": edge[1]} for edge in tool_edge]
        sampled_nodes = sorted(sampled_nodes, key=lambda x: x["id"])
        sampled_links = sorted(sampled_links, key=lambda x: (x["source"], x["target"]))

        content["invoking_graph"]["nodes"] = sorted(content["invoking_graph"]["nodes"], key=lambda x: x["id"])
        content["invoking_graph"]["links"] = sorted(content["invoking_graph"]["links"], key=lambda x: (x["source"], x["target"]))
            
        result = {"id": sample_id, "seed": seed, "method": method, "number_of_tools": tool_number, "sampled_nodes": sampled_nodes, "sampled_links": sampled_links, "result": content}
        
        if wf:
            wf.write(json.dumps(result) + "\n")

        if figure_dir:
            plt.figure()
            pos = nx.circular_layout(sub_G)
            nx.draw_networkx_nodes(sub_G, pos, node_color="skyblue", node_size=300)
            nx.draw_networkx_edges(sub_G, pos, arrows=True)
            nx.draw_networkx_labels(sub_G, pos, font_size=8)
            plt.tight_layout()
            plt.savefig(f"{figure_dir}/{sample_id}.jpg")
            plt.close()

        sampled_nodes_ids = [node["id"] for node in sampled_nodes]
        generated_nodes_ids = [node["id"] for node in content["invoking_graph"]["nodes"]]

        end_time = datetime.now()
        logger.info(f"Sample {sample_id} finished, time cost: {end_time - start_time}")
        if sampled_links == content["invoking_graph"]["links"] and sampled_nodes_ids == generated_nodes_ids:
            logger.info("Check success: invoking graph and sampled graph are aligned.")
        elif sampled_nodes_ids != generated_nodes_ids:
            logger.info("Check fail: mismatched nodes")
            logger.info("Sampled node:\n" + json.dumps(sampled_nodes_ids, indent=2))
            logger.info("Generated node:\n" + json.dumps(generated_nodes_ids, indent=2))
            logger.info(f"Sample {sample_id}:\n{json.dumps(result, indent=2)}")
        else:
            logger.info("Check fail: mismatched links")
            logger.info("Sampled link:\n" + json.dumps(sampled_links, indent=2))
            logger.info("Generated link:\n" + json.dumps(content["invoking_graph"]["links"], indent=2))
            logger.info(f"Sample {sample_id}:\n{json.dumps(result, indent=2)}")
    except Exception as e:
        logger.info(f"Failed: {type(e)}")
        print(traceback.format_exc())
        raise e
    return result

if __name__ == "__main__":
    main()


File: taskbench/graph_sampler.py

import networkx as nx
import random
import json
import matplotlib.pyplot as plt
import click

random.seed(0)
class GraphSampler:
    def __init__(self, graph: nx.Graph = None, file_name = None):
        if file_name:
            with open(file_name, "r") as f:
                data = json.load(f)

            # Represent your graph in NetworkX
            graph = nx.DiGraph()

            # Add nodes to the graph
            if "input-type" in data["nodes"][0]:
                for node in data["nodes"]:
                    graph.add_node(node["id"], desc=node["desc"], input_type=node["input-type"], output_type=node["output-type"])
            else:
                for node in data["nodes"]:
                    graph.add_node(node["id"], desc=node["desc"], parameters=node["parameters"])

            # Add edges to the graph
            for link in data["links"]:
                graph.add_edge(link["source"], link["target"], type=link["type"])

        self.graph = graph

    def sample_subgraph_by_weight(self, number_weights, method_weights):
        method = random.choices(list(method_weights.keys()), weights=list(method_weights.values()))[0]
        if method == "single":
            tool_number = 1
        else:
            tool_number = random.choices(list(number_weights.keys()), weights=list(number_weights.values()))[0]
        return self.sample_subgraph(tool_number, sample_method=method)

    def sample_subgraph(self, num_nodes=3, sample_method="chain"):
        seed_node = random.choice(list(self.graph.nodes))
        if sample_method == "single":
            sub_G = nx.DiGraph()
            sub_G.add_node(seed_node)
            return sub_G
        elif sample_method == "chain":
            return self.sample_subgraph_chain(seed_node, num_nodes)
        elif sample_method == "dag":
            return self.sample_subgraph_dag(seed_node, num_nodes)
        else:
            raise ValueError("Invalid sample method")

    def sample_subgraph_chain(self, seed_node, num_nodes):
        # Create a list to store the sub-graph nodes
        sub_graph_nodes = [seed_node]
        head_node = seed_node
        tail_node = seed_node
        edges = []

        # Keep adding nodes until we reach the desired number
        while len(sub_graph_nodes) < num_nodes:
            # Get the neighbors of the last node in the sub-graph
            head_node_neighbors = list(self.graph.predecessors(head_node))
            tail_node_neighbors = list(self.graph.successors(tail_node))
            neighbors = head_node_neighbors + tail_node_neighbors

            # If the node has neighbors, randomly select one and add it to the sub-graph
            if len(neighbors) > 0:
                neighbor = random.choice(neighbors)
                if neighbor not in sub_graph_nodes:
                    if neighbor in head_node_neighbors:
                        sub_graph_nodes.insert(0, neighbor)
                        edges.insert(0, (neighbor, head_node))
                        head_node = neighbor
                    else:
                        sub_graph_nodes.append(neighbor)
                        edges.append((tail_node, neighbor))
                        tail_node = neighbor
            else:
                break

        # Create the sub-graph
        sub_G = nx.DiGraph()
        sub_G.add_nodes_from(sub_graph_nodes)
        sub_G.add_edges_from(edges)

        return sub_G

    def sample_subgraph_dag(self, seed_node, num_nodes):
        # Create a list to store the sub-graph nodes
        sub_graph_nodes = [seed_node]
        edges = []

        # Keep adding nodes until we reach the desired number
        while len(sub_graph_nodes) < num_nodes:
            # Randomly select a node from the current sub-graph
            node = random.choice(sub_graph_nodes)
            # prec_neighbors = list(self.graph.predecessors(node))
            succ_neighbors = list(self.graph.successors(node))

            if "input_type" in self.graph.nodes[node]:
                # filter exisiting income edge type
                prec_neighbors = []
                input_type = list(self.graph.nodes[node]["input_type"])
                all_in_edges = list(self.graph.in_edges(node, data=True))
                for edge in edges:
                    for ref_edge in all_in_edges:
                        if edge[0] == ref_edge[0] and edge[1] == ref_edge[1]:
                            input_type.remove(ref_edge[2]["type"])
                for edge in all_in_edges:
                    if edge[2]["type"] in input_type:
                        prec_neighbors.append(edge[0])
            else:
                prec_neighbors = list(self.graph.predecessors(node))

            neighbors = prec_neighbors + succ_neighbors

            # If the node has neighbors, randomly select one and add it to the sub-graph
            if neighbors:
                neighbor = random.choice(neighbors)
                if neighbor not in sub_graph_nodes:
                    if neighbor in prec_neighbors:
                        edges.append((neighbor, node))
                    else:
                        edges.append((node, neighbor))
                    sub_graph_nodes.append(neighbor)
            # If the node has no neighbors, select a new node from the original graph
            else:
                node = random.choice(list(self.graph.nodes))
                if node not in sub_graph_nodes:
                    sub_graph_nodes.append(node)

        # Create the sub-graph
        sub_G = nx.DiGraph()
        sub_G.add_nodes_from(sub_graph_nodes)
        sub_G.add_edges_from(edges)

        return sub_G
    
    def sample_subgraph_random_walk(self, seed_node, num_nodes):
        # Create a list to store the sub-graph nodes
        sub_graph_nodes = [seed_node]
        edges = []

        # Keep adding nodes until we reach the desired number
        while len(sub_graph_nodes) < num_nodes:
            # Randomly select a node from the current sub-graph
            node = random.choice(sub_graph_nodes)
            neighbors = list(self.graph.successors(node))

            # If the node has neighbors, randomly select one and add it to the sub-graph
            if neighbors:
                neighbor = random.choice(neighbors)
                if neighbor not in sub_graph_nodes:
                    edges.append((node, neighbor))
                    sub_graph_nodes.append(neighbor)
            # If the node has no neighbors, select a new node from the original graph
            else:
                node = random.choice(list(self.graph.nodes))
                if node not in sub_graph_nodes:
                    sub_graph_nodes.append(node)

        # Create the sub-graph
        sub_G = nx.DiGraph()
        sub_G.add_nodes_from(sub_graph_nodes)
        sub_G.add_edges_from(edges)

        return sub_G
    
    def sample_subgraph_random_walk_with_restart(self, seed_node, num_nodes, restart_prob=0.15):
        # Create a list to store the sub-graph nodes
        sub_graph_nodes = [seed_node]
        edges = []

        # Keep adding nodes until we reach the desired number
        while len(sub_graph_nodes) < num_nodes:
            # Randomly select a node from the current sub-graph
            node = random.choice(sub_graph_nodes)
            neighbors = list(self.graph.successors(node))

            # If the node has neighbors, randomly select one and add it to the sub-graph
            if neighbors:
                neighbor = random.choice(neighbors)
                if neighbor not in sub_graph_nodes:
                    edges.append((node, neighbor))
                    sub_graph_nodes.append(neighbor)
            # If the node has no neighbors, select a new node from the original graph
            else:
                node = random.choice(list(self.graph.nodes))
                if node not in sub_graph_nodes:
                    sub_graph_nodes.append(node)
            
            # Randomly restart the walk
            if random.random() < restart_prob:
                node = random.choice(list(self.graph.nodes))
                if node not in sub_graph_nodes:
                    sub_graph_nodes.append(node)

        # Create the sub-graph
        sub_G = nx.DiGraph()
        sub_G.add_nodes_from(sub_graph_nodes)
        sub_G.add_edges_from(edges)

        return sub_G

@click.command()
@click.option('--file_name', default='graph_desc_original.json', help='Path to the json file')
@click.option('--sample_method', default='chain', help='Type of graph to generate')
@click.option('--num_nodes', default=3, help='Number of nodes in the subgraph')
@click.option('--save_figure', default=False, help='Save the figure')
def sample_subgraph(file_name, sample_method, num_nodes, save_figure):
    # Create a graph sampler
    random.seed(0)
    sampler = GraphSampler(file_name=file_name)

    # Sample a sub-graph
    sub_G = sampler.sample_subgraph(num_nodes, sample_method=sample_method)
    print("Sub-graph nodes:", sub_G.nodes)
    print("Sub-graph edges:", sub_G.edges)

    # Visualize the sub-graph
    if save_figure:
        pos = nx.circular_layout(sub_G)
        nx.draw_networkx_nodes(sub_G, pos, node_color="skyblue", node_size=300)
        nx.draw_networkx_edges(sub_G, pos, arrows=True)
        nx.draw_networkx_labels(sub_G, pos, font_size=8)
        plt.axis("off")
        plt.tight_layout()
        plt.savefig("test.png")



if __name__ == "__main__":
    sample_subgraph()


File: taskbench/inference.py


import os
import json
import click
import asyncio
import aiohttp
import logging
import emoji

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.handlers = []

class RateLimitError(Exception):
    def __init__(self, message):
        super().__init__(message)

class ContentFormatError(Exception):
    def __init__(self, message):
        super().__init__(message)

@click.command()
@click.option("--data_dir", default="data_huggingface", help="The directory of the data.")
@click.option("--temperature", type=float, default=0.2)
@click.option("--top_p", type=float, default=0.1)
@click.option("--api_addr", type=str, default="localhost")
@click.option("--api_port", type=int, default=4000)
@click.option("--api_key", type=str, default="your api key")
@click.option("--multiworker", type=int, default=1)
@click.option("--llm", type=str, default="gpt-4")
@click.option("--use_demos", type=int, default=0)
@click.option("--reformat", type=bool, default=False)
@click.option("--reformat_by", type=str, default="self")
@click.option("--tag", type=bool, default=False)
@click.option("--dependency_type", type=str, default="resource")
@click.option("--log_first_detail", type=bool, default=False)
def main(data_dir, temperature, top_p, api_addr, api_key, api_port, multiworker, llm, use_demos, reformat, reformat_by, tag, dependency_type, log_first_detail):
    assert dependency_type in ["resource", "temporal"], "Dependency type not supported"
    arguments = locals()
    url = f"http://{api_addr}:{api_port}/v1/chat/completions"
    header = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    prediction_dir = f"{data_dir}/predictions{f'_use_demos_{use_demos}' if use_demos and tag else ''}{f'_reformat_by_{ reformat_by}' if reformat and tag else ''}"
    wf_name = f"{prediction_dir}/{llm}.json"
    
    if not os.path.exists(prediction_dir):
        os.makedirs(prediction_dir, exist_ok=True)

    has_inferenced = []
    if os.path.exists(wf_name):
        rf = open(wf_name, "r")
        for line in rf:
            data = json.loads(line)
            has_inferenced.append(data["id"])
        rf.close()

    rf_ur = open(f"{data_dir}/user_requests.json", "r")
    inputs = []
    for line in rf_ur:
        input = json.loads(line)
        if input["id"] not in has_inferenced:
            inputs.append(input)
    rf_ur.close()

    wf = open(wf_name, "a")
    
    tool_list = json.load(open(f"{data_dir}/tool_desc.json", "r"))["nodes"]
    if "input-type" not in tool_list[0]:
        assert dependency_type == "temporal", "Tool type is not ignored, but the tool list does not contain input-type and output-type"
    if dependency_type == "temporal":
        for tool in tool_list:
            parameter_list = []
            for parameter in tool["parameters"]:
                parameter_list.append(parameter["name"])
            tool["parameters"] = parameter_list

    # log llm name in format
    formatter = logging.Formatter(f"%(asctime)s - [ {llm} ] - %(levelname)s - %(message)s")
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    console_handler.setLevel(logging.INFO)
    logger.addHandler(console_handler)

    file_handler = logging.FileHandler(f"{prediction_dir}/{llm}.log")
    file_handler.setFormatter(formatter)
    file_handler.setLevel(logging.INFO)
    logger.addHandler(file_handler)

    # logging all args
    logger.info(f"Arguments: {arguments}")

    demos = []
    if use_demos:
        if dependency_type == "temporal":
            demos_id = [ "38563456", "27267145", "91005535"]
        else:
            if "huggingface" in data_dir: 
                demos_id = [ "10523150", "14611002", "22067492"]
            elif "multimedia" in data_dir:
                demos_id = [ "30934207", "20566230", "19003517"]
        demos_id = demos_id[:use_demos]
        logger.info(f"Use {len(demos_id)} demos: {demos_id}")
        demos_rf = open(f"{data_dir}/data.json", "r")
        for line in demos_rf:
            data = json.loads(line)
            if data["id"] in demos_id:
                if dependency_type == "temporal":
                    demo = {
                        "user_request": data["user_request"],
                        "result":{
                            "task_steps": data["task_steps"],
                            "task_nodes": data["task_nodes"],
                            "task_links": data["task_links"]
                        }
                    }
                else:
                    demo = {
                        "user_request": data["user_request"],
                        "result":{
                            "task_steps": data["task_steps"],
                            "task_nodes": data["task_nodes"]
                        }
                    }
                demos.append(demo)
        demos_rf.close()

    tool_string = "# TASK LIST #:\n"
    for k, tool in enumerate(tool_list):
        tool_string += json.dumps(tool) + "\n"
    
    sem = asyncio.Semaphore(multiworker)

    async def inference_wrapper(input, url, header, temperature, top_p, tool_string, wf, llm, demos, reformat, reformat_by, dependency_type, log_detail = False):
        async with sem:
            await inference(input, url, header, temperature, top_p, tool_string, wf, llm, demos, reformat, reformat_by, dependency_type, log_detail)

    if len(inputs) == 0:
        logger.info("All Completed!")
        return
    else:
        logger.info(f"Detected {len(has_inferenced)} has been inferenced,")
        logger.info(f"Start inferencing {len(inputs)} tasks...")
    
    loop = asyncio.get_event_loop()

    if log_first_detail:
        tasks = [inference_wrapper(inputs[0], url, header, temperature, top_p, tool_string, wf, llm, demos, reformat, reformat_by, dependency_type, log_detail=True)]
        results = loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
        inputs = inputs[1:]

    tasks = []
    for input in inputs:
        tasks.append(inference_wrapper(input, url, header, temperature, top_p, tool_string, wf, llm, demos, reformat, reformat_by, dependency_type))

    results += loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
    failed = []
    done = []
    for result in results:
        if isinstance(result, Exception):
            failed.append(result)
        else:
            done.append(result)
    logger.info(f"Completed: {len(done)}")
    logger.info(f"Failed: {len(failed)}")
    loop.close()

async def inference(input, url, header, temperature, top_p, tool_string, wf, llm, demos, reformat, reformat_by, dependency_type, log_detail = False):
    user_request = input["user_request"]
    if dependency_type == "resource":
        prompt = """\n# GOAL #: Based on the above tools, I want you generate task steps and task nodes to solve the # USER REQUEST #. The format must in a strict JSON format, like: {"task_steps": [ step description of one or more steps ], "task_nodes": [{"task": "tool name must be from # TOOL LIST #", "arguments": [ a concise list of arguments for the tool. Either original text, or user-mentioned filename, or tag '<node-j>' (start from 0) to refer to the output of the j-th node. ]}]} """
        prompt += """\n\n# REQUIREMENTS #: \n1. the generated task steps and task nodes can resolve the given user request # USER REQUEST # perfectly. Task name must be selected from # TASK LIST #; \n2. the task steps should strictly aligned with the task nodes, and the number of task steps should be same with the task nodes; \n3. the dependencies among task steps should align with the argument dependencies of the task nodes; \n4. the tool arguments should be align with the input-type field of # TASK LIST #;"""
    else:
        prompt = """\n# GOAL #:\nBased on the above tools, I want you generate task steps and task nodes to solve the # USER REQUEST #. The format must in a strict JSON format, like: {"task_steps": [ "concrete steps, format as Step x: Call xxx tool with xxx: 'xxx' and xxx: 'xxx'" ], "task_nodes": [{"task": "task name must be from # TASK LIST #", "arguments": [ {"name": "parameter name", "value": "parameter value, either user-specified text or the specific name of the tool whose result is required by this node"} ]}], "task_links": [{"source": "task name i", "target": "task name j"}]}"""
        prompt += """\n\n# REQUIREMENTS #: \n1. the generated task steps and task nodes can resolve the given user request # USER REQUEST # perfectly. Task name must be selected from # TASK LIST #; \n2. the task steps should strictly aligned with the task nodes, and the number of task steps should be same with the task nodes; \n3. The task links (task_links) should reflect the temporal dependencies among task nodes, i.e. the order in which the APIs are invoked;"""

    if len(demos) > 0:
        prompt += "\n"
        for demo in demos:
            prompt += f"""\n# EXAMPLE #:\n# USER REQUEST #: {demo["user_request"]}\n# RESULT #: {json.dumps(demo["result"])}"""
    
    prompt += """\n\n# USER REQUEST #: {{user_request}}\nnow please generate your result in a strict JSON format:\n# RESULT #:"""

    final_prompt = tool_string + prompt.replace("{{user_request}}", user_request)
    payload = json.dumps({
        "model": f"{llm}",
        "messages": [
            {
            "role": "user",
            "content":  final_prompt
            }
        ],
        "temperature": temperature,
        "top_p": top_p,
        "frequency_penalty": 0,
        "presence_penalty": 1.05,
        "max_tokens": 2000,
        "stream": False,
        "stop": None
    })
    try:
        result = await get_response(url, header, payload, input['id'], reformat, reformat_by, dependency_type, log_detail)
    except Exception as e:
        logger.info(f"Failed #id {input['id']}: {type(e)} {e}")
        raise e
    logger.info(f"Success #id {input['id']}")
    input["result"] = result
    wf.write(json.dumps(input) + "\n")
    wf.flush()

async def get_response(url, header, payload, id, reformat, reformat_by, dependency_type, log_detail=False):
    async with aiohttp.ClientSession() as session:
        async with session.post(url, headers=header, data=payload, timeout=300) as response:
            resp = await response.json()

    if response.status == 429:
        raise RateLimitError(f"{resp}")
    if response.status != 200:
        raise Exception(f"{resp}")
    
    if log_detail:
        logger.info(json.loads(payload)["messages"][0]["content"])
        logger.info(resp["choices"][0]["message"]["content"])

    oring_content = resp["choices"][0]["message"]["content"]
    oring_content = oring_content.replace("\n", "")
    oring_content = oring_content.replace("\_", "_")
    content = oring_content.replace("\\", "")

    start_pos = content.find("RESULT #:")
    if start_pos!=-1:
        content = content[start_pos+len("RESULT #:"):]
        
    content = content[content.find("{"):content.rfind("}")+1]
    try:
        content = json.loads(content)
        if isinstance(content, list) and len(content):
            merge_content = {}
            for c in content:
                for k, v in c.items():
                    merge_content[k].extend(v) if k in merge_content else merge_content.update({k: v})
        return content
    except json.JSONDecodeError as e:
        if reformat:
            if dependency_type == "resource":
                prompt = """Please format the result # RESULT # to a strict JSON format # STRICT JSON FORMAT #. \nRequirements:\n1. Do not change the meaning of task steps and task nodes;\n2. Don't tolerate any possible irregular formatting to ensure that the generated content can be converted by json.loads();\n3. You must output the result in this schema: {"task_steps": [ step description of one or more steps ], "task_nodes": [{"task": "tool name must be from # TOOL LIST #", "arguments": [ a concise list of arguments for the tool. Either original text, or user-mentioned filename, or tag '<node-j>' (start from 0) to refer to the output of the j-th node. ]}]}\n# RESULT #:{{illegal_result}}\n# STRICT JSON FORMAT #:"""
            else:
                prompt = """Please format the result # RESULT # to a strict JSON format # STRICT JSON FORMAT #. \nRequirements:\n1. Do not change the meaning of task steps, task nodes and task links;\n2. Don't tolerate any possible irregular formatting to ensure that the generated content can be converted by json.loads();\n3. Pay attention to the matching of brackets. Write in a compact format and avoid using too many space formatting controls;\n4. You must output the result in this schema: {"task_steps": [ "concrete steps, format as Step x: Call xxx tool with xxx: 'xxx' and xxx: 'xxx'" ], "task_nodes": [{"task": "task name must be from # TASK LIST #", "arguments": [ {"name": "parameter name", "value": "parameter value, either user-specified text or the specific name of the tool whose result is required by this node"} ]}], "task_links": [{"source": "task name i", "target": "task name j"}]}\n# RESULT #:{{illegal_result}}\n# STRICT JSON FORMAT #:"""
            prompt = prompt.replace("{{illegal_result}}", oring_content)
            payload = json.loads(payload)
            if reformat_by != "self":
                payload["model"] = reformat_by

            if log_detail:
                logger.info(f"{emoji.emojize(':warning:')}  #id {id} Illegal JSON format: {content}")
                logger.info(f"{emoji.emojize(':sparkles:')} #id {id} Detected illegal JSON format, try to reformat by {payload['model']}...")

            payload["messages"][0]["content"] = prompt
            payload = json.dumps(payload)
            
            async with aiohttp.ClientSession() as session:
                async with session.post(url, headers=header, data=payload, timeout=120) as response:
                    resp = await response.json()

            if response.status == 429:
                raise RateLimitError(f"{resp}")
            if response.status != 200:
                raise Exception(f"{resp}")
            
            if log_detail:
                logger.info(json.loads(payload)["messages"][0]["content"])
                logger.info(resp["choices"][0]["message"]["content"])

            content = resp["choices"][0]["message"]["content"]
            content = content.replace("\n", "")
            content = content.replace("\_", "_")
            start_pos = content.find("STRICT JSON FORMAT #:")
            if start_pos!=-1:
                content = content[start_pos+len("STRICT JSON FORMAT #:"):]

            content = content[content.find("{"):content.rfind("}")+1]
            try:
                content = json.loads(content)
                return content
            except json.JSONDecodeError as e:
                raise ContentFormatError(f"{content}")
        else:
            raise ContentFormatError(f"{content}")

if __name__ == "__main__":
    main()


File: taskbench/evaluate.py

import traceback
import numpy as np
from scipy.optimize import linear_sum_assignment
import json
import click
from datasets import load_metric
import Levenshtein
from sklearn.metrics import precision_recall_fscore_support as prfs
import warnings
import logging
import os
import itertools

warnings.filterwarnings("ignore")

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.handlers = []

def sim(name_1, name_2):
    if name_1 == "<PAD>" or name_2 == "<PAD>":
        return 0
    return 1 if name_1 == name_2 else 0

def create_cost_matrix(graph_1, graph_2):
    nodes_1 = graph_1["nodes"]
    nodes_2 = graph_2["nodes"]

    num_nodes_1 = len(nodes_1)
    num_nodes_2 = len(nodes_2)

    nodes_similarity_matrix = np.zeros((num_nodes_1, num_nodes_2))
    
    for i, node_1 in enumerate(graph_1["nodes"]):
        for j, node_2 in enumerate(graph_2["nodes"]):
            nodes_similarity_matrix[i, j] = sim(node_1, node_2)  # 

    links_similarity_matrix = np.zeros((num_nodes_1, num_nodes_2))
    for link_1 in graph_1["links"]:
        for link_2 in graph_2["links"]:
            if link_1["source"] == link_2["source"] and link_1["target"] == link_2["target"]:
                try:
                    i_index_1 = nodes_1.index(link_1["source"])
                    i_index_2 = nodes_2.index(link_2["source"])
                    j_index_1 = nodes_1.index(link_1["target"])
                    j_index_2 = nodes_2.index(link_2["target"])
                except ValueError:
                    continue
                links_similarity_matrix[i_index_1, i_index_2] += 1
                links_similarity_matrix[j_index_1, j_index_2] += 1
    
    cost_matrix = 2 - nodes_similarity_matrix - 0.5 * links_similarity_matrix
    return cost_matrix

def compute_assignment_matrix(graph_1, graph_2):
    cost_matrix = create_cost_matrix(graph_1, graph_2)
    row_ind, col_ind = linear_sum_assignment(cost_matrix) 
    return row_ind, col_ind, cost_matrix[row_ind, col_ind].sum()

def matching(graph_1, graph_2):
    indices_1, indices_2, total_cost = compute_assignment_matrix(graph_1, graph_2)
    return indices_1, indices_2

def ratio_levenshtein(x, y):
    assert len(x) == len(y)
    n = len(x)
    total = 0
    for i in range(n):
        total += Levenshtein.ratio(x[i], y[i])
    return total / n


def flatten(gt, pred, types = None):
    assert len(gt) == len(pred)

    gt_flat = []
    pred_flat = []

    for (sample_gt, sample_pred) in zip(gt, pred):
        union = set()

        union.update(sample_gt)
        union.update(sample_pred)

        for s in union:
            if types: 
                if s in types:
                    if s in sample_gt:
                        gt_flat.append(types.index(s)+1)
                    else:
                        gt_flat.append(0)

                    if s in sample_pred:
                        pred_flat.append(types.index(s)+1)
                    else:
                        pred_flat.append(0)
                else:
                    gt_flat.append(0)
                    pred_flat.append(0)
            else:
                if s in sample_gt:
                    gt_flat.append(1)
                else:
                    gt_flat.append(0)

                if s in sample_pred:
                    pred_flat.append(1)
                else:
                    pred_flat.append(0)
    return gt_flat, pred_flat

def print_results(per_type, micro, macro, types, result_dict = None):
    columns = ('type', 'precision', 'recall', 'f1-score', 'support')

    row_fmt = "%30s" + (" %12s" * (len(columns) - 1))
    logger.info(row_fmt % columns)

    metrics_per_type = []
    for i, t in enumerate(types):
        metrics = []
        for j in range(len(per_type)):
            metrics.append(per_type[j][i])
        metrics_per_type.append(metrics)

    for m, t in zip(metrics_per_type, types):
        logger.info(row_fmt % get_row(m, t))
        if result_dict is not None:
            result_dict[t] = {}
            result_dict[t]["precision"] = m[0]
            result_dict[t]["recall"] = m[1]
            result_dict[t]["f1-score"] = m[2]
            result_dict[t]["support"] = int(m[3])

    logger.info('')

    # micro
    logger.info(row_fmt % get_row(micro, 'micro'))

    # macro
    logger.info(row_fmt % get_row(macro, 'macro'))

def get_row(data, label):
    row = [label]
    for i in range(len(data) - 1):
        row.append("%.2f" % (data[i] * 100))
    row.append(data[3])
    return tuple(row)

def get_content_type(content):
    content = content.strip('\'')
    assert isinstance(content, str), content
    # image
    for ext in ["jpg", "png", "jpeg", "gif", "bmp", "tiff", "svg", "ico"]:
        if "."+ext in content:
            return "image"
    # audio
    for ext in ["mp3", "wav", "wma", "ogg", "aac", "flac", "aiff", "au"]:
        if "."+ext in content:
            return "audio"
    # video
    for ext in ["mp4", "avi", "mov", "flv", "wmv", "mkv", "webm", "m4v", "mpg", "mpeg"]:
        if "."+ext in content:
            return "video"
    return "text"

@click.command()
@click.option("--data_dir", default="data_huggingface", help="The directory of the data.")
@click.option("--prediction_dir", default="predictions", help="The directory of the data.")
@click.option("--save_dir", default=None, help="The directory to save the evaluation results")
@click.option("--alignment", default=None)
@click.option("--splits", "-s", multiple=True, default=["overall"])
@click.option("--n_tools", "-n", multiple=True, default=["overall"])
@click.option("--mode", default="add")
@click.option("--metric", "-m", multiple=True, default=["all"])
@click.option("--llm", default="gpt-3.5-turbo")
@click.option("--dependency_type", type=str, default="resource")
@click.option("--prompting", default="cot")
def main(data_dir, prediction_dir, save_dir, splits, n_tools, mode, metric, llm, dependency_type, alignment, prompting):
    assert dependency_type in ["resource", "temporal"], "Dependency type not supported"
    args = locals()
    
    if save_dir is None:
        save_dir = prediction_dir.replace("predictions", "metrics") 
        save_dir = save_dir + f"_alignment_{alignment}" if alignment is not None else save_dir

    formatter = logging.Formatter(f'%(asctime)s - [ {llm} ] - %(levelname)s - %(message)s')
    if not os.path.exists(f'{data_dir}/{save_dir}'):
        os.makedirs(f'{data_dir}/{save_dir}')
    
    metric_file = f'{data_dir}/{save_dir}/{llm}.json'
    if os.path.exists(metric_file):
        all_metric_dict = json.load(open(metric_file, "r"))
    else:
        all_metric_dict = {}
    
    file_handler = logging.FileHandler(f'{data_dir}/{save_dir}/{llm}.log')
    stream_handler = logging.StreamHandler()
    
    file_handler.setFormatter(formatter)
    stream_handler.setFormatter(formatter)

    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)

    if "all" in metric:
        metric = ["f1", "ed", "link", "argument", "rouge", "bertscore"]
        if prompting != "cot":
            metric = ["f1", "ed", "link", "argument"]

    logger.info(f"Starts with: {args}")

    tool_desc = json.load(open(f"{data_dir}/tool_desc.json", "r"))
    tool_map = {tool["id"]: i+1 for i, tool in enumerate(tool_desc["nodes"])}
    tool_map_reverse = {i+1: tool["id"] for i, tool in enumerate(tool_desc["nodes"])}
    tool_map_reverse[0] = "NEGATIVE"
    tool_map["<PAD>"] = -1

    tool_output_type_map = None
    if dependency_type == "resource":
        tool_output_type_map = {tool["id"]: tool["output-type"][0] if len(tool["output-type"]) else "none" for tool in tool_desc["nodes"]}

    splits = list(splits)
    n_tools = list(n_tools)

    if "all" in splits:
        splits = ["overall", "single", "chain", "dag", ]
    if "all" in n_tools:
        n_tools = ["overall"] + [str(i) for i in range(1, 11)]
    
    group = []

    if mode == "mul":
        for s in splits:
            for n in n_tools:
                if (s, n) not in group:
                    group.append((s, n))
    elif mode == "add":
        for s in splits:
            if (s, "overall") not in group:
                group.append((s, "overall"))
        for n in n_tools:
            if ("overall", n) not in group:
                group.append(("overall", n))
    else:
        assert False, "mode should be mul or add"

    for s, n in group:
        logger.info("-"*15)
        logger.info(f"Tools Number: {n}, Task Split: {s}")
        evaluate(data_dir, prediction_dir, llm, s, n, metric, tool_desc, tool_map, tool_output_type_map, tool_map_reverse, all_metric_dict, dependency_type=dependency_type, alignment=alignment)

    metric_json = open(metric_file, "w")
    metric_json.write(json.dumps(all_metric_dict, indent=2))

def evaluate(data_dir, prediction_dir, llm, split, n_tool, metric, tool_desc, tool_map, tool_output_type_map, tool_map_reverse, all_metric_dict, dependency_type, alignment = None):
    if f"{split}_{n_tool}" in all_metric_dict:
        metric_dict = all_metric_dict[f"{split}_{n_tool}"]
    else:
        metric_dict = {}
        all_metric_dict[f"{split}_{n_tool}"] = metric_dict

    label_rf = open(f"{data_dir}/data.json", "r")
    
    alignment_ids = None
    if alignment is not None:
        if alignment == "human":
            label_rf = open(f"{data_dir}/data.json", "r")
            logger.info(f"Alignment Mode: {alignment} ({len(label_rf.readlines())})")
        else:
            alignment_file = open(f"{data_dir}/alignment_ids.json", "r")
            alignment_ids = json.load(alignment_file)
            alignment_ids = list(itertools.chain(*alignment_ids[f"{alignment}_alignment_id"].values()))
            logger.info(f"Alignment Mode: {alignment} ({len(alignment_ids)})")
        
    predcition_rf = open(f"{data_dir}/{prediction_dir}/{llm}.json", "r")

    predcitions = {}
    labels = {}
    label_rf = open(f"{data_dir}/data.json", "r")
    for line in label_rf:
        data = json.loads(line)
        real_tool_num = len(data["task_nodes"])
        if alignment_ids is None or data["id"] in alignment_ids:
            if split == "overall" or data["type"] == split:
                if n_tool == "overall" or str(real_tool_num) == n_tool:
                    id = data["id"]
                    labels[id] = data

    for line in predcition_rf:
        try:
            data = json.loads(line)
        except Exception as e:
            print(e)
            print(line)
            exit()
        id = data["id"]
        predcitions[id] = data

    ids = set(labels.keys()).intersection(set(predcitions.keys()))
    labels = {id: labels[id] for id in ids}
    predcitions = {id: predcitions[id] for id in ids}

    predcition_task_steps = []
    label_task_steps = []
    predcition_names = []
    label_names = []
    label_graphs = []
    predcition_graphs = []
    label_links = []
    predcition_links = []
    label_task_arg_names = []
    predcition_task_arg_names = []
    label_task_arg_name_values = []
    predcition_task_arg_name_values = []

    for id in ids:
        try:
            label = labels[id]
            predcition = predcitions[id]

            if "rouge" in metric or "bertscore" in metric:
                predcition_task_step = predcition["result"]["task_steps"]
                label_task_step = label["task_steps"]
                
                try:
                    if isinstance(predcition_task_step[0], str):
                        predcition_task_steps.append("\n".join(predcition_task_step))
                    else:
                        if "task" in predcition_task_step[0]:
                            predcition_task_steps.append("\n".join([step["task"] for step in predcition_task_step]))
                        elif "step" in predcition_task_step[0]:
                            predcition_task_steps.append("\n".join([step["step"] for step in predcition_task_step]))
                        elif "id" in predcition_task_step[0]:
                            predcition_task_steps.append("\n".join([step["id"] for step in predcition_task_step]))
                        elif "step_name" in predcition_task_step[0]:
                            predcition_task_steps.append("\n".join([step["step_name"] for step in predcition_task_step]))
                        else:
                            predcition_task_steps.append("\n".join([step["description"] for step in predcition_task_step]))
                except Exception as e:
                    predcition_task_steps.append(str(predcition_task_step))

                label_task_steps.append("\n".join(label_task_step))

            label_nodes = label["task_nodes"]
            predcition_nodes = predcition["result"]["task_nodes"] 

            label_node_name = [node["task"] for node in label_nodes]
            predcition_node_name = [node["task"] for node in predcition_nodes]

            label_task_arg_name = []
            predcition_task_arg_name = []

            label_task_arg_name_value = []
            predcition_task_arg_name_value = []
                
            if dependency_type == "resource":
                predcition_node_name = [name.replace("_", " ") for name in predcition_node_name]
                label_node_name = [name.replace("_", " ") for name in label_node_name]
                label_link = []
                predcition_link = []
                for inx, node in enumerate(label_nodes):
                    new_arguments = []
                    for i, argument in enumerate(node["arguments"]):
                        try:
                            if isinstance(argument, dict):
                                argument = list(argument.values())[0]
                            if isinstance(argument, list):
                                argument = " ".join(argument)
                            if "<node-" in argument:
                                index_start = argument.index("<node-") + 6
                                index_end = argument.index(">")
                                if int(argument[index_start: index_end]) == inx:
                                    continue
                                argument_tool_name = label_node_name[int(argument[index_start: index_end])]
                                label_link.append({"source": argument_tool_name, "target": node["task"]})
                                new_argument = {"name": tool_output_type_map.get(argument_tool_name, "other"), "value": argument_tool_name}
                            else:
                                new_argument = {"name": get_content_type(argument), "value": argument}
                        except Exception as e:
                            pass
                        new_arguments.append(new_argument)
                    node["arguments"] = new_arguments
                    
                for inx, node in enumerate(predcition_nodes):
                    new_arguments = []
                    for i, argument in enumerate(node.get("arguments", [])):
                        try:
                            if isinstance(argument, dict):
                                argument = list(argument.values())[0]
                            if isinstance(argument, list):
                                argument = " ".join(argument)
                            if isinstance(argument, str) and "<node-" in argument:
                                index_start = argument.index("<node-") + 6
                                index_end = argument.index(">")
                            
                                if int(argument[index_start: index_end]) == inx:
                                    continue
                                prediction_tool_name = predcition_node_name[int(argument[index_start: index_end])]
                                predcition_link.append({"source": prediction_tool_name, "target": node["task"]})
                                new_argument = {"name": tool_output_type_map.get(prediction_tool_name, "other"), "value": prediction_tool_name}
                            else:
                                new_argument = {"name": get_content_type(argument), "value": argument}

                        except Exception as e:
                            pass
                        new_arguments.append(new_argument)
                    node["arguments"] = new_arguments
            else:
                predcition_link = predcition["result"]["task_links"]
                label_link = label["task_links"]

            predcition_node_argument = [node.get("arguments", []) for node in predcition_nodes]
            label_node_argument = [node["arguments"] for node in label_nodes]

            for task, arguments in zip (predcition_node_name, predcition_node_argument):
                for argument in arguments:
                    label_task_arg_name.append(f"{task}-{argument['name']}")
                    label_task_arg_name_value.append(f"{task}-{argument['name']}-{argument['value']}")
            
            for task, arguments in zip (label_node_name, label_node_argument):
                for argument in arguments:
                    predcition_task_arg_name.append(f"{task}-{argument['name']}")
                    predcition_task_arg_name_value.append(f"{task}-{argument['name']}-{argument['value']}")

            label_graph = {
                "nodes": label_node_name,
                "links": label_link,
                "arguments": label_node_argument
            }
            predcition_graph = {
                "nodes": predcition_node_name,
                "links": predcition_link,
                "arguments": predcition_node_argument
            }

            label_graphs.append(label_graph)
            predcition_graphs.append(predcition_graph)

            for node_name in predcition_node_name:
                assert isinstance(node_name, str), node_name

            predcition_names.append(predcition_node_name)
            label_names.append(label_node_name)

            predcition_task_arg_names.append(predcition_task_arg_name)
            label_task_arg_names.append(label_task_arg_name)
        
            predcition_task_arg_name_values.append(predcition_task_arg_name_value)
            label_task_arg_name_values.append(label_task_arg_name_value)

            label_links.append(label_link)
            predcition_links.append(predcition_link)

        except Exception as e:
            logger.info(f"Parsing Error: {e}, Ignore #id {id}")
            logger.info(traceback.format_exc())
            
    logger.info(f"Step Supports: {len(label_task_steps)} / {len(ids)}")
    logger.info(f"Node Support: {len(label_names)} / {len(ids)}")
    logger.info(f"Link Support: {len(label_links)} / {len(ids)}")
    logger.info(f"Argument Support: {len(label_graphs)} / {len(ids)}")

    metric_dict["all_samples"] = len(ids)
    metric_dict["step_supports"] = len(label_task_steps)
    metric_dict["node_supports"] = len(label_names)
    metric_dict["link_supports"] = len(label_links)
    metric_dict["argument_supports"] = len(label_graphs)

    if len(label_graphs) == 0 or len(label_names) == 0 or len(label_links) == 0:
        logger.info("No supports, skip")
        return

    if "rouge" in metric:
        rouge = load_metric("rouge")
        rouge_scores = rouge.compute(predictions=predcition_task_steps, references=label_task_steps, use_aggregator=True)
        for key in rouge_scores:
            logger.info(f"Step {key}: {rouge_scores[key].mid.fmeasure}")
            metric_dict[f"step_{key}"] = rouge_scores[key].mid.fmeasure

    if "bertscore" in metric:
        bertscore = load_metric("bertscore")
        bertscore_scores = bertscore.compute(predictions=predcition_task_steps, references=label_task_steps, model_type="roberta-large")
        for key in bertscore_scores:
            if key in ["precision", "recall", "f1"]:
                bertscore_scores[key] = np.mean(bertscore_scores[key])
                logger.info(f"Step BERTScore {key}: {bertscore_scores[key]}")
                metric_dict[f"step_bertscore_{key}"] = bertscore_scores[key]
    
    if "f1" in metric or "argument" in metric:
        types = list(range(1, len(tool_desc["nodes"])+1))
        types_name = [tool_map_reverse[i] for i in types]
        gt_flat, pred_flat = flatten(label_names, predcition_names, types = types_name)

        per_type = prfs(gt_flat, pred_flat, labels=types, average=None)
        micro = prfs(gt_flat, pred_flat, labels=types, average='micro')[:-1]
        macro = prfs(gt_flat, pred_flat, labels=types, average='macro')[:-1]
        total_support = sum(per_type[-1])

        logger.info(f"Node Micro Precision [ No Matching ]: {micro[0]}")
        logger.info(f"Node Micro Recall [ No Matching ]: {micro[1]}")
        logger.info(f"Node Micro F1 [ No Matching ]: {micro[2]}")
        logger.info(f"Node Macro Precision [ No Matching ]: {macro[0]}")
        logger.info(f"Node Macro Recall [ No Matching ]: {macro[1]}")
        logger.info(f"Node Macro F1 [ No Matching ]: {macro[2]}")
        logger.info("Node Detailed Report [ No Matching ]: ")
        metric_dict["node_micro_precision_no_matching"] = micro[0]
        metric_dict["node_micro_recall_no_matching"] = micro[1]
        metric_dict["node_micro_f1_no_matching"] = micro[2]
        metric_dict["node_macro_precision_no_matching"] = macro[0]
        metric_dict["node_macro_recall_no_matching"] = macro[1]
        metric_dict["node_macro_f1_no_matching"] = macro[2]

        per_type_metric = {}
        metric_dict["node_per_type_no_matchcing"] = per_type_metric
        print_results(per_type, list(micro) + [total_support], list(macro) + [total_support], types_name, result_dict = per_type_metric)


        gt_flat, pred_flat = flatten(label_task_arg_names, predcition_task_arg_names)
        micro = prfs(gt_flat, pred_flat, average="binary")[:-1]
        logger.info(f"Argument Task-ArgName Binary F1: [ No Matching ]: {micro[-1]}")
        metric_dict["argument_task_argname_binary_f1_no_matching"] = micro[-1]

        gt_flat, pred_flat = flatten(label_task_arg_name_values, predcition_task_arg_name_values)
        micro = prfs(gt_flat, pred_flat, average="binary")[:-1]
        logger.info(f"Argument Task-ArgName-Value Binary F1 [ No Matching ]: {micro[-1]}")
        metric_dict["argument_task_argname_value_binary_f1_no_matching"] = micro[-1]

    if "ed" in metric:
        labels = []
        predcitions = []
        for label_name, predcition_name in zip(label_names, predcition_names):
            labels.append([tool_map.get(name, 0) for name in label_name])
            predcitions.append([tool_map.get(name, 0) for name in predcition_name])
        ed = ratio_levenshtein(predcitions, labels)
        logger.info(f"Edit Distance: {1-ed}")
        metric_dict["edit_distance"] = 1-ed
    
    if "link" in metric:
        tuple_label_links = []
        tuple_predcition_links = []
        for label_link, predcition_link in zip(label_links, predcition_links):
            tuple_label_links.append([(link["source"], link["target"]) for link in label_link])
            tuple_predcition_links.append([(link["source"], link["target"]) for link in predcition_link])
        
        gt_flat, pred_flat = flatten(tuple_label_links, tuple_predcition_links)


        micro = prfs(gt_flat, pred_flat, average="binary")[:-1]
        logger.info(f"Link Binary F1: {micro[-1]}")
        metric_dict["link_binary_f1"] = micro[-1]

if __name__ == "__main__":
    main()


File: taskbench/format_data.py

import json
import click
import traceback

def formulate_sample(data, dependency_type):
    try:
        user_request = data["result"]["user_request"]
        invoking_graph = data["result"]["invoking_graph"]
        task_steps = data["result"]["task_steps"]
        nodes = invoking_graph["nodes"]
        links = invoking_graph["links"]
        user_request = data["result"]["user_request"]
        if "check_by_teacher" in data["result"]:
            check_by_teacher = data["result"]["check_by_teacher"]
        else:
            check_by_teacher = invoking_graph["check_by_teacher"]
        for node in nodes:
            node["task"] = node["id"]
            node.pop("id")
            if dependency_type == "resource":
                node["task"] = node["task"].replace("_", " ")
                node["arguments"] = node["input"]
                node.pop("input")

        for node in nodes:
            assert isinstance(node, dict)
            assert "task" in node
            assert "arguments" in node
            if isinstance(node["arguments"], str) and node["arguments"].startswith("<node-"):
                node["arguments"] = [node["arguments"]]
            assert isinstance(node["arguments"], list), node["arguments"]
            if dependency_type == "resource":
                assert len(node["arguments"]) <= 2
                for i, argument in enumerate(node["arguments"]):
                    for j, n in enumerate(nodes):
                        if n["task"] in argument:
                            node["arguments"][i] = f"<node-{j}>"
                            break
        for link in links:
            assert isinstance(link, dict)
            assert "source" in link
            assert "target" in link
            if dependency_type == "resource":
                link["source"] = link["source"].replace("_", " ")
                link["target"] = link["target"].replace("_", " ")
        assert isinstance(task_steps, list)
        assert isinstance(nodes, list)
        assert len(nodes) == len(task_steps)
        assert isinstance(user_request, str)
        assert isinstance(check_by_teacher, str)
    except Exception as e:
        print(e)
        traceback.print_exc()
        return None, None, None, None, None
    return user_request, task_steps, links, nodes, check_by_teacher

@click.command()
@click.option('--data_dir', default='data_huggingface', help='Path to the data directory')
@click.option('--dependency_type', default='resource')
def formulate(data_dir, dependency_type):
    rf = open(f"{data_dir}/data_raw.json", "r")
    wf_format = open(f"{data_dir}/data.json", "w")
    wf_error = open(f"{data_dir}/data_error.json", "w")
    wf_ur = open(f"{data_dir}/user_requests.json", "w")

    all = 0
    format = 0
    for line in rf:
        all += 1
        data = json.loads(line)
        method = data["method"]
        n_tools = data["number_of_tools"]
        seed = data["seed"]
        _id = data["id"]
        sampled_nodes = data["sampled_nodes"]
        sampled_links = data["sampled_links"]
        for sampled_node in sampled_nodes:
            sampled_node["task"] = sampled_node["id"]
            sampled_node.pop("id")
            if dependency_type == "resource":
                sampled_node["task"] = sampled_node["task"].replace("_", " ")
                if "input" in sampled_node:
                    sampled_node["input-type"] = sampled_node["input"]
                    sampled_node.pop("input")
                    sampled_node["output-type"] = sampled_node["output"]
                    sampled_node.pop("output")
            else:
                sampled_node["arguments"] = sampled_node["parameters"]
                sampled_node.pop("parameters")
        user_request, task_steps, links, nodes, check_by_teacher = formulate_sample(data, dependency_type)
        if user_request is None:
            wf_error.write(line)
            continue
        format += 1
        result = {
            "id": _id,
            "seed": seed,
            "type": method,
            "n_tools": n_tools,
            "sampled_nodes": sampled_nodes,
            "sampled_links": sampled_links,
            "user_request": user_request,
            "task_steps": task_steps,
            "task_nodes": nodes,
            "task_links": links,
            "check_by_teacher": check_by_teacher,
        }
        wf_format.write(json.dumps(result)+"\n")
        ur_result = {
            "id": _id,
            "user_request": user_request,
        }
        wf_ur.write(json.dumps(ur_result)+"\n")
    wf_format.close()
    wf_error.close()
    wf_ur.close()
    rf.close()
    print(f"Format {format} out of {all}")

if __name__ == "__main__":
    formulate()


# JARVIS


[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2303.17580)
[![Open in Spaces](https://img.shields.io/badge/%F0%9F%A4%97-Open%20in%20Spaces-blue)](https://huggingface.co/spaces/microsoft/HuggingGPT)

The mission of JARVIS is to explore artificial general intelligence (AGI) and deliver cutting-edge research to the whole community.

## What's New

+  [2024.01.15] We release Easytool for easier tool usage.
   + The code and datasets are available at [EasyTool](/easytool).
   + The paper is available at [EasyTool: Enhancing LLM-based Agents with Concise Tool Instruction](https://arxiv.org/abs/2401.06201).
+  [2023.11.30] We release TaskBench for evaluating task automation capability of LLMs.
   + The code and datasets are avaliable at [TaskBench](/taskbench).
   + The paper is avaliable at [TaskBench: Benchmarking Large Language Models for Task Automation](https://arxiv.org/abs/2311.18760).
+  [2023.07.28] We are now in the process of planning evaluation and project rebuilding. We will release a new version of Jarvis in the near future.
+  [2023.07.24] We released a light langchain version of Jarvis. See <a href="https://github.com/langchain-ai/langchain/tree/master/libs/experimental/langchain_experimental/autonomous_agents/hugginggpt">here</a>.
+  [2023.04.16] Jarvis now supports the OpenAI service on the Azure platform and the GPT-4 model.
+  [2023.04.06] We added the Gradio demo and built the web API for `/tasks` and `/results` in `server` mode.
   +  The Gradio demo is now hosted on Hugging Face Space. (Build with `inference_mode=hybrid` and `local_deployment=standard`)
   +  The Web API `/tasks` and `/results` access intermediate results for `Stage #1`: task planning and `Stage #1-3`: model selection with execution results. See <a href="#Server">here</a>.
+  [2023.04.03] We added the CLI mode and provided parameters for configuring the scale of local endpoints.
   +  You can enjoy a lightweight experience with Jarvis without deploying the models locally. See <a href="#Configuration">here</a>.
   +  Just run `python awesome_chat.py --config configs/config.lite.yaml` to experience it.
+  [2023.04.01] We updated a version of code for building.


### Overview

Language serves as an interface for LLMs to connect numerous AI models for solving complicated AI tasks!

<p align="center">
<img width="100%" alt="image" src="./hugginggpt/assets/intro.png">    
</p>


See our paper: [HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in HuggingFace](http://arxiv.org/abs/2303.17580), Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu and Yueting Zhuang (the first two authors contribute equally)

We introduce a collaborative system that consists of **an LLM as the controller** and **numerous expert models as collaborative executors** (from HuggingFace Hub). The workflow of our system consists of four stages:
+ **Task Planning**: Using ChatGPT to analyze the requests of users to understand their intention, and disassemble them into possible solvable tasks.
+ **Model Selection**: To solve the planned tasks, ChatGPT selects expert models hosted on Hugging Face based on their descriptions.
+ **Task Execution**: Invokes and executes each selected model, and return the results to ChatGPT.
+ **Response Generation**: Finally, using ChatGPT to integrate the prediction of all models, and generate responses.

<p align="center"><img src="./hugginggpt/assets/overview.jpg"></p>

### System Requirements

#### Default (Recommended)

For `configs/config.default.yaml`:

+ Ubuntu 16.04 LTS
+ VRAM >= 24GB
+ RAM > 12GB (minimal), 16GB (standard), 80GB (full)
+ Disk > 284GB 
  + 42GB for `damo-vilab/text-to-video-ms-1.7b`
  + 126GB for `ControlNet`
  + 66GB for `stable-diffusion-v1-5`
  + 50GB for others
  
#### Minimum (Lite)

For `configs/config.lite.yaml`:

+ Ubuntu 16.04 LTS
+ Nothing else

The configuration `configs/config.lite.yaml` does not require any expert models to be downloaded and deployed locally. However, it means that Jarvis is restricted to models running stably on HuggingFace Inference Endpoints.

### Quick Start

First replace `openai.key` and `huggingface.token` in `server/configs/config.default.yaml` with **your personal OpenAI Key** and **your Hugging Face Token**, or put them in the environment variables `OPENAI_API_KEY` and `HUGGINGFACE_ACCESS_TOKEN` respectively. Then run the following commands:

<span id="Server"></span>

#### For Server:

```bash
# setup env
cd server
conda create -n jarvis python=3.8
conda activate jarvis
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
pip install -r requirements.txt

# download models. Make sure that `git-lfs` is installed.
cd models
bash download.sh # required when `inference_mode` is `local` or `hybrid`. 

# run server
cd ..
python models_server.py --config configs/config.default.yaml # required when `inference_mode` is `local` or `hybrid`
python awesome_chat.py --config configs/config.default.yaml --mode server # for text-davinci-003
```

Now you can access Jarvis' services by the Web API. 

+ `/hugginggpt` --method `POST`, access the full service.
+ `/tasks` --method `POST`, access intermediate results for Stage #1.
+ `/results` --method `POST`, access intermediate results for Stage #1-3.

For example:

```bash
# request
curl --location 'http://localhost:8004/tasks' \
--header 'Content-Type: application/json' \
--data '{
    "messages": [
        {
            "role": "user",
            "content": "based on pose of /examples/d.jpg and content of /examples/e.jpg, please show me a new image"
        }
    ]
}'

# response
[{"args":{"image":"/examples/d.jpg"},"dep":[-1],"id":0,"task":"openpose-control"},{"args":{"image":"/examples/e.jpg"},"dep":[-1],"id":1,"task":"image-to-text"},{"args":{"image":"<GENERATED>-0","text":"<GENERATED>-1"},"dep":[1,0],"id":2,"task":"openpose-text-to-image"}]
```


#### For Web:

We provide a user-friendly web page. After starting `awesome_chat.py` in a server mode, you can run the commands to communicate with Jarvis in your browser:
 
- you need to install `nodejs` and `npm` first.
- [ IMPORTANT ] if you are running the web client on another machine, you need set `http://{LAN_IP_of_the_server}:{port}/` to `HUGGINGGPT_BASE_URL` of `web/src/config/index.ts`.
- if you want to use the video generation feature, you need to compile `ffmpeg` manually with H.264.
- you can switch to ChatGPT by `double click` on the setting icon!

```bash
cd web
npm install
npm run dev
```

```bash
# Optional: Install ffmpeg
# This command need be executed without errors.
LD_LIBRARY_PATH=/usr/local/lib /usr/local/bin/ffmpeg -i input.mp4 -vcodec libx264 output.mp4
```

<span id="Gradio"></span>

#### For Gradio

The Gradio demo is now hosted on Hugging Face Space. You can also run the following commands to start the demo locally:

```bash
python models_server.py --config configs/config.gradio.yaml
python run_gradio_demo.py --config configs/config.gradio.yaml

# or run the HF Space as a Docker image (Build with `inference_mode=hibrid` and `local_deployment=standard`)
docker run -it -p 7860:7860 --platform=linux/amd64 registry.hf.space/microsoft-hugginggpt:latest python app.py
```

#### For CLI:

You can also run Jarvis more easily in CLI mode:

```bash
cd server
python awesome_chat.py --config configs/config.default.yaml --mode cli
```

Examples of CLI mode:

```
Welcome to Jarvis! A collaborative system that consists of an LLM as the controller and numerous expert models as collaborative executors. Jarvis can plan tasks, schedule Hugging Face models, generate friendly responses based on your requests, and help you with many things. Please enter your request (`exit` to exit).
[ User ]: Given a collection of image A: /examples/a.jpg, B: /examples/b.jpg, C: /examples/c.jpg, please tell me how many zebras in these picture?
[ Jarvis ]: Based on your request, I have used several models to process the images you provided. For the image A: /examples/a.jpg, I used the image-to-text model nlpconnect/vit-gpt2-image-captioning to generate the text "a cat sitting on a window sill looking out". Then I used the object-detection model facebook/detr-resnet-50 to generate the image with predicted box '/images/a3dc.jpg', which contains two objects with labels 'potted plant' and 'cat'. For the image B: /examples/b.jpg, I used the image-to-text model nlpconnect/vit-gpt2-image-captioning to generate the text "two zebras standing in a field of dry grass". Then I used the object-detection model facebook/detr-resnet-50 to generate the image with predicted box '/images/f5df.jpg', which contains three objects with labels 'zebra'. For the image C: /examples/c.jpg, I used the image-to-text model nlpconnect/vit-gpt2-image-captioning to generate the text "a zebra standing in a field of tall grass". Then I used the object-detection model facebook/detr-resnet-50 to generate the image with predicted box '/images/69c7.jpg', which contains one object with label 'zebra'. Finally, I used the canny-control model canny-control to generate the image '/images/7e18.png' based on the image /examples/f.jpg. Therefore, there are four zebras in these pictures. Is there anything else I can help you with?
[ User ]: Please answer all the named entities in the sentence: Iron Man is a superhero appearing in American comic books published by Marvel Comics. The character was co-created by writer and editor Stan Lee, developed by scripter Larry Lieber, and designed by artists Don Heck and Jack Kirby.
[ Jarvis ]: Sure, I understand your request. Based on the inference results of the models, I will answer your request step by step. Firstly, I used the token-classification model dslim/bert-base-NER to recognize the named entities in the sentence. The inference result is that the named entities are Iron Man, Stan Lee, Larry Lieber, Don Heck and Jack Kirby. Then, I used the question-answering model deepset/roberta-base-squad2 to answer your request. The inference result is that there is no output since the context cannot be empty. Therefore, I cannot make it. I hope this answer your request. Is there anything else I can help you with?
[ User ]: Please generate a canny image based on /examples/f.jpg
[ Jarvis ]: Sure. I have generated a canny image based on /examples/f.jpg. To do this, I first used the image-to-text model nlpconnect/vit-gpt2-image-captioning to generate the text description of the image, which is "a herd of giraffes and zebras grazing in a field". Then I used the canny-control model to generate the canny image of the original image, which is stored in the path "/images/1ea2.png". Finally, I used the canny-text-to-image model lllyasviel/sd-controlnet-canny to generate the canny image based on the text description and the canny image, which is stored in the path "/images/ba66.png". Do you need anything else?
```


### Configuration

The server-side configuration file is `server/configs/config.default.yaml`, and some parameters are presented as follows:

+ `model`: LLM, currently supports `text-davinci-003`. We are working on integrating more open-source LLMs.
+ `inference_mode`: mode of inference endpoints
  + `local`: only use the local inference endpoints
  + `huggingface`: only use the Hugging Face Inference Endpoints **(free of local inference endpoints)**
  + `hybrid`: both of `local` and `huggingface`
+ `local_deployment`: scale of locally deployed models, works under `local` or `hybrid` inference mode:
  +  `minimal` (RAM>12GB, ControlNet only)
  +  `standard` (RAM>16GB, ControlNet + Standard Pipelines)
  +  `full` (RAM>42GB, All registered models)

On a personal laptop, we recommend the configuration of `inference_mode: hybrid `and `local_deployment: minimal`. But the available models under this setting may be limited due to the instability of remote Hugging Face Inference Endpoints.

### NVIDIA Jetson Embedded Device Support
A [Dockerfile](./Dockerfile.jetson) is included that provides experimental support for [NVIDIA Jetson embedded devices](https://developer.nvidia.com/embedded-computing).  This image provides accelerated ffmpeg, pytorch, torchaudio, and torchvision dependencies.  To build the docker image, [ensure that the default docker runtime is set to 'nvidia'](https://github.com/NVIDIA/nvidia-docker/wiki/Advanced-topics#default-runtime).  A pre-built image is provided at https://hub.docker.com/r/toolboc/nv-jarvis.

```bash
#Build the docker image
docker build --pull --rm -f "Dockerfile.jetson" -t toolboc/nv-jarvis:r35.2.1 
```

Due to to memory requirements, JARVIS is required to run on Jetson AGX Orin family devices (64G on-board RAM device preferred) with config options set to:
* `inference_mode: local` 
* `local_deployment: standard`  

Models and configs are recommended to be provided through a volume mount from the host to the container as shown in the `docker run` step below.  It is possible to uncomment the `# Download local models` section of the [Dockerfile](./Dockerfile.jetson) to build a container with models included.

#### Start the model server, awesomechat, and web app on Jetson Orin AGX

```bash
# run the container which will automatically start the model server
docker run --name jarvis --net=host --gpus all -v ~/jarvis/configs:/app/server/configs -v ~/src/JARVIS/server/models:/app/server/models toolboc/nv-jarvis:r35.2.1

# (wait for model server to complete initialization)

# start awesome_chat.py 
docker exec jarvis python3 awesome_chat.py --config configs/config.default.yaml --mode server

#start the web application (application will be acessible at http://localhost:9999)
docker exec jarvis npm run dev --prefix=/app/web
```

### Screenshots

<p align="center"><img src="./hugginggpt/assets/screenshot_q.jpg"><img src="./hugginggpt/assets/screenshot_a.jpg"></p>




## Citation
If you find this work useful in your method, you can cite the paper as below:

    @inproceedings{shen2023hugginggpt,
      author = {Shen, Yongliang and Song, Kaitao and Tan, Xu and Li, Dongsheng and Lu, Weiming and Zhuang, Yueting},
      booktitle = {Advances in Neural Information Processing Systems},
      title = {HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in HuggingFace},
      year = {2023}
    }

>
    @article{shen2023taskbench,
      title   = {TaskBench: Benchmarking Large Language Models for Task Automation},
      author  = {Shen, Yongliang and Song, Kaitao and Tan, Xu and Zhang, Wenqi and Ren, Kan and Yuan, Siyu and Lu, Weiming and Li, Dongsheng and Zhuang, Yueting},
      journal = {arXiv preprint arXiv:2311.18760},
      year    = {2023}
    }

>
    @article{yuan2024easytool,
      title   = {EASYTOOL: Enhancing LLM-based Agents with Concise Tool Instruction},
      author  = {Siyu Yuan and Kaitao Song and Jiangjie Chen and Xu Tan and Yongliang Shen and Ren Kan and Dongsheng Li and Deqing Yang},
      journal = {arXiv preprint arXiv:2401.06201},
      year    = {2024}
    }
