Update awesome_chat.py

awesome_chat.py

@@ -1,6 +1,5 @@
 import base64
 import copy
-import datetime
 from io import BytesIO
 import io
 import os
@@ -19,50 +18,43 @@ 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
-from models_server import models, status
-from functools import partial
-from huggingface_hub import Repository
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--config", type=str, default="config.yaml.dev")
+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 = "config.gradio.yaml"
+    args.config = "configs/config.gradio.yaml"
+    args.mode = "gradio"
 
 config = yaml.load(open(args.config, "r"), Loader=yaml.FullLoader)
 
-if not os.path.exists("logs"):
-    os.mkdir("logs")
+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)
 
-now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
-
-DATASET_REPO_URL = "https://huggingface.co/datasets/deepcode-ai/HuggingSpace_logs"
-LOG_HF_TOKEN = os.environ.get("LOG_HF_TOKEN")
-if LOG_HF_TOKEN:
-    repo = Repository(
-        local_dir="logs", clone_from=DATASET_REPO_URL, use_auth_token=LOG_HF_TOKEN
-    )
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-logger.handlers = []
-logger.propagate = False
+logger.setLevel(logging.DEBUG)
 
 handler = logging.StreamHandler()
 formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 handler.setFormatter(formatter)
-if config["debug"]:
-    handler.setLevel(logging.DEBUG)
+if not config["debug"]:
+    handler.setLevel(logging.CRITICAL)
 logger.addHandler(handler)
 
 log_file = config["log_file"]
 if log_file:
-    log_file = log_file.replace("TIMESTAMP", now)
     filehandler = logging.FileHandler(log_file)
     filehandler.setLevel(logging.DEBUG)
     filehandler.setFormatter(formatter)
@@ -73,7 +65,7 @@ use_completion = config["use_completion"]
 
 # consistent: wrong msra model name 
 LLM_encoding = LLM
-if LLM == "gpt-3.5-turbo":
+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)
@@ -87,20 +79,35 @@ if use_completion:
 else:
     api_name = "chat/completions"
 
-if not config["dev"]:
-    if not config["openai"]["key"].startswith("sk-") and not config["openai"]["key"]=="gradio":
-        raise ValueError("Incrorrect OpenAI key. Please check your config.yaml file.")
-    OPENAI_KEY = config["openai"]["key"]
-    endpoint = f"https://api.openai.com/v1/{api_name}"
-    if OPENAI_KEY.startswith("sk-"):
-        HEADER = {
-            "Authorization": f"Bearer {OPENAI_KEY}"
-        }
-    else:
-        HEADER = None
+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:
-    endpoint = f"{config['local']['endpoint']}/v1/{api_name}"
-    HEADER = None
+    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"]:
@@ -110,6 +117,19 @@ if 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()
@@ -133,6 +153,18 @@ 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 = ""
@@ -155,19 +187,26 @@ def convert_chat_to_completion(data):
     return data
 
 def send_request(data):
-    global HEADER
-    openaikey = data.pop("openaikey")
+    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 openaikey and openaikey.startswith("sk-"):
+    if api_type == "openai":
         HEADER = {
-            "Authorization": f"Bearer {openaikey}"
+            "Authorization": f"Bearer {api_key}"
         }
-    
-    response = requests.post(endpoint, json=data, headers=HEADER, proxies=PROXY)
-    logger.debug(response.text.strip())
-    if "choices" not in response.json():
+    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:
@@ -177,7 +216,7 @@ 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', "").replace('\\', '\\\\'))
+        text = text.replace("{{" + key +"}}", value.replace('"', "'").replace('\n', ""))
     return text
 
 def find_json(s):
@@ -204,14 +243,13 @@ def get_id_reason(choose_str):
     return id.strip(), reason.strip(), choose
 
 def record_case(success, **args):
-    if not success:
-        return
-    f = open(f"logs/log_success_{now}.jsonl", "a")
+    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()
-    if LOG_HF_TOKEN:
-        commit_url = repo.push_to_hub(blocking=False)
 
 def image_to_bytes(img_url):
     img_byte = io.BytesIO()
@@ -266,20 +304,19 @@ def unfold(tasks):
         
     return tasks
 
-def chitchat(messages, openaikey=None):
+def chitchat(messages, api_key, api_type, api_endpoint):
     data = {
         "model": LLM,
         "messages": messages,
-        "openaikey": openaikey
+        "api_key": api_key,
+        "api_type": api_type,
+        "api_endpoint": api_endpoint
     }
     return send_request(data)
 
-def parse_task(context, input, openaikey=None):
+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)
-    for message in messages:
-        if not isinstance(message["content"], str):
-            message["content"] = json.dumps(message["content"], ensure_ascii=False)
     messages.insert(0, {"role": "system", "content": parse_task_tprompt})
 
     # cut chat logs
@@ -304,11 +341,13 @@ def parse_task(context, input, openaikey=None):
         "messages": messages,
         "temperature": 0,
         "logit_bias": {item: config["logit_bias"]["parse_task"] for item in task_parsing_highlight_ids},
-        "openaikey": openaikey
+        "api_key": api_key,
+        "api_type": api_type,
+        "api_endpoint": api_endpoint
     }
     return send_request(data)
 
-def choose_model(input, task, metas, openaikey = None):
+def choose_model(input, task, metas, api_key, api_type, api_endpoint):
     prompt = replace_slot(choose_model_prompt, {
         "input": input,
         "task": task,
@@ -328,12 +367,14 @@ def choose_model(input, task, metas, openaikey = None):
         "messages": messages,
         "temperature": 0,
         "logit_bias": {item: config["logit_bias"]["choose_model"] for item in choose_model_highlight_ids}, # 5
-        "openaikey": openaikey
+        "api_key": api_key,
+        "api_type": api_type,
+        "api_endpoint": api_endpoint
     }
     return send_request(data)
 
 
-def response_results(input, results, openaikey=None):
+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,
@@ -342,7 +383,7 @@ def response_results(input, results, openaikey=None):
         "input": input,
         "processes": results
     })
-    messages = json.loads(demos_or_presteps, strict=False)
+    messages = json.loads(demos_or_presteps)
     messages.insert(0, {"role": "system", "content": response_results_tprompt})
     messages.append({"role": "user", "content": prompt})
     logger.debug(messages)
@@ -350,19 +391,15 @@ def response_results(input, results, openaikey=None):
         "model": LLM,
         "messages": messages,
         "temperature": 0,
-        "openaikey": openaikey
+        "api_key": api_key,
+        "api_type": api_type,
+        "api_endpoint": api_endpoint
     }
     return send_request(data)
 
-def huggingface_model_inference(model_id, data, task, huggingfacetoken=None):
-    if huggingfacetoken is None:
-        HUGGINGFACE_HEADERS = {}
-    else:
-        HUGGINGFACE_HEADERS = {
-            "Authorization": f"Bearer {huggingfacetoken}",
-    }
+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=huggingfacetoken)
+    inference = InferenceApi(repo_id=model_id, token=config["huggingface"]["token"])
     
     # NLP tasks
     if task == "question-answering":
@@ -426,7 +463,7 @@ def huggingface_model_inference(model_id, data, task, huggingfacetoken=None):
         name = str(uuid.uuid4())[:4]
         image.save(f"public/images/{name}.jpg")
         result = {}
-        result["generated image with segmentation mask"] = f"/images/{name}.jpg"
+        result["generated image"] = f"/images/{name}.jpg"
         result["predicted"] = predicted
 
     if task == "object-detection":
@@ -447,7 +484,7 @@ def huggingface_model_inference(model_id, data, task, huggingfacetoken=None):
         name = str(uuid.uuid4())[:4]
         image.save(f"public/images/{name}.jpg")
         result = {}
-        result["generated image with predicted box"] = f"/images/{name}.jpg"
+        result["generated image"] = f"/images/{name}.jpg"
         result["predicted"] = predicted
 
     if task in ["image-classification"]:
@@ -459,7 +496,7 @@ def huggingface_model_inference(model_id, data, task, huggingfacetoken=None):
         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)
+        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")
@@ -493,62 +530,71 @@ def huggingface_model_inference(model_id, data, task, huggingfacetoken=None):
     return result
 
 def local_model_inference(model_id, data, task):
-    inference = partial(models, model_id)
+    task_url = f"{Model_Server}/models/{model_id}"
+    
     # contronlet
     if model_id.startswith("lllyasviel/sd-controlnet-"):
         img_url = data["image"]
         text = data["text"]
-        results = inference({"img_url": img_url, "text": 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"]
-        results = inference({"img_url": img_url})
+        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":
-        results = inference(data)
+        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":
-        results = inference(json=data)
-        return results
+        response = requests.post(task_url, json=data)
+        return response.json()
     if task in ["text-classification",  "token-classification", "text2text-generation", "summarization", "translation", "conversational", "text-generation"]:
-        results = inference(json=data)
-        return results
+        response = requests.post(task_url, json=data)
+        return response.json()
 
     # CV tasks
     if task == "depth-estimation":
         img_url = data["image"]
-        results = inference({"img_url": img_url})
+        response = requests.post(task_url, json={"img_url": img_url})
+        results = response.json()
         if "path" in results:
-            results["generated depth image"] = results.pop("path")
+            results["generated image"] = results.pop("path")
         return results
     if task == "image-segmentation":
         img_url = data["image"]
-        results = inference({"img_url": img_url})
-        results["generated image with segmentation mask"] = results.pop("path")
+        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"]
-        results = inference({"img_url": img_url})
+        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":
-        results = inference(data)
+        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"]
-        predicted = inference({"img_url": img_url})
+        response = requests.post(task_url, json={"img_url": img_url})
+        predicted = response.json()
         if "error" in predicted:
             return predicted
         image = load_image(img_url)
@@ -565,7 +611,7 @@ def local_model_inference(model_id, data, task):
         name = str(uuid.uuid4())[:4]
         image.save(f"public/images/{name}.jpg")
         results = {}
-        results["generated image with predicted box"] = f"/images/{name}.jpg"
+        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"]:
@@ -573,43 +619,40 @@ def local_model_inference(model_id, data, task):
         text = None
         if "text" in data:
             text = data["text"]
-        results = inference({"img_url": img_url, "text": 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":
-        results = inference(data)
+        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"]
-        results = inference({"audio_url": audio_url})
-        return results
+        response = requests.post(task_url, json={"audio_url": audio_url})
+        return response.json()
 
 
-def model_inference(model_id, data, hosted_on, task, huggingfacetoken=None):
-    if huggingfacetoken:
-        HUGGINGFACE_HEADERS = {
-            "Authorization": f"Bearer {huggingfacetoken}",
-        }
-    else:
-        HUGGINGFACE_HEADERS = None
+def model_inference(model_id, data, hosted_on, task):
     if hosted_on == "unknown":
-        r = status(model_id)
-        logger.debug("Local Server Status: " + str(r))
-        if "loaded" in r and r["loaded"]:
+        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 "loaded" in r and r["loaded"]:
+            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, huggingfacetoken)
+            inference_result = huggingface_model_inference(model_id, data, task)
     except Exception as e:
         print(e)
         traceback.print_exc()
@@ -622,8 +665,8 @@ def get_model_status(model_id, url, headers, queue = None):
     if "huggingface" in url:
         r = requests.get(url, headers=headers, proxies=PROXY)
     else:
-        r = status(model_id)
-    if "loaded" in r and r["loaded"]:
+        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
@@ -632,13 +675,11 @@ def get_model_status(model_id, url, headers, queue = None):
             queue.put((model_id, False, None))
         return False
 
-def get_avaliable_models(candidates, topk=10, huggingfacetoken = None):
+def get_avaliable_models(candidates, topk=5):
     all_available_models = {"local": [], "huggingface": []}
     threads = []
     result_queue = Queue()
-    HUGGINGFACE_HEADERS = {
-        "Authorization": f"Bearer {huggingfacetoken}",
-    }
+
     for candidate in candidates:
         model_id = candidate["id"]
 
@@ -649,7 +690,8 @@ def get_avaliable_models(candidates, topk=10, huggingfacetoken = None):
             thread.start()
         
         if inference_mode != "huggingface" and config["local_deployment"] != "minimal":
-            thread = threading.Thread(target=get_model_status, args=(model_id, "", {}, result_queue))
+            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()
         
@@ -675,7 +717,7 @@ def collect_result(command, choose, inference_result):
     return result
 
 
-def run_task(input, command, results, openaikey = None, huggingfacetoken = None):
+def run_task(input, command, results, api_key, api_type, api_endpoint):
     id = command["id"]
     args = command["args"]
     task = command["task"]
@@ -776,9 +818,9 @@ def run_task(input, command, results, openaikey = None, huggingfacetoken = None)
         choose = {"id": best_model_id, "reason": reason}
         messages = [{
             "role": "user",
-            "content": f"[ {input} ] contains a task in JSON format {command}, 'task' indicates the task type and 'args' indicates the arguments required for the task. Don't explain the task to me, just help me do it and give me the result. The result must be in text form without any urls."
+            "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, openaikey)
+        response = chitchat(messages, api_key, api_type, api_endpoint)
         results[id] = collect_result(command, choose, {"response": response})
         return True
     else:
@@ -789,8 +831,8 @@ def run_task(input, command, results, openaikey = None, huggingfacetoken = None)
             results[id] = collect_result(command, "", inference_result)
             return False
 
-        candidates = MODELS_MAP[task][:20]
-        all_avaliable_models = get_avaliable_models(candidates, config["num_candidate_models"], huggingfacetoken)
+        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}")
 
@@ -800,8 +842,7 @@ def run_task(input, command, results, openaikey = None, huggingfacetoken = None)
             inference_result = {"error": f"no available models on {command['task']} task."}
             results[id] = collect_result(command, "", inference_result)
             return False
-        
-        all_avaliable_model_ids = all_avaliable_model_ids[:1]
+            
         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"
@@ -817,14 +858,14 @@ def run_task(input, command, results, openaikey = None, huggingfacetoken = None)
                     ),
                     "likes": model.get("likes"),
                     "description": model.get("description", "")[:config["max_description_length"]],
-                    "language": model.get("language"),
-                    "tags": model.get("tags"),
+                    # "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, openaikey)
+            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)
@@ -836,7 +877,7 @@ def run_task(input, command, results, openaikey = None, huggingfacetoken = None)
                 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'], huggingfacetoken)
+    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']}")
@@ -847,42 +888,39 @@ def run_task(input, command, results, openaikey = None, huggingfacetoken = None)
     results[id] = collect_result(command, choose, inference_result)
     return True
 
-def chat_huggingface(messages, openaikey = None, huggingfacetoken = None, return_planning = False, return_results = False):
+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, openaikey)
-    logger.info(task_str)
+    task_str = parse_task(context, input, api_key, api_type, api_endpoint)
 
     if "error" in task_str:
-        return str(task_str), {}
-    else:
-        task_str = task_str.strip()
+        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, openaikey)
+        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 response, {}
-
+        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, openaikey)
-        return response, {}
+        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": "task parsing fail: empty", "op": "chitchat"})
-        response = chitchat(messages, openaikey)
-        best_model_id = "ChatGPT"
-        reason = "ChatGPT performs well on some NLP tasks as well."
-        choose = {"id": best_model_id, "reason": reason}
-        return response, collect_result(tasks[0], choose, {"response": 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)
@@ -897,24 +935,23 @@ def chat_huggingface(messages, openaikey = None, huggingfacetoken = None, return
     d = dict()
     retry = 0
     while True:
-        num_threads = len(threads)
+        num_thread = len(threads)
         for task in tasks:
-            dep = task["dep"]
             # logger.debug(f"d.keys(): {d.keys()}, dep: {dep}")
-            for dep_id in dep:
+            for dep_id in task["dep"]:
                 if dep_id >= task["id"]:
                     task["dep"] = [-1]
-                    dep = [-1]
                     break
-            if len(list(set(dep).intersection(d.keys()))) == len(dep) or dep[0] == -1:
+            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, openaikey, huggingfacetoken))
+                thread = threading.Thread(target=run_task, args=(input, task, d, api_key, api_type, api_endpoint))
                 thread.start()
                 threads.append(thread)
-        if num_threads == len(threads):
+        if num_thread == len(threads):
             time.sleep(0.5)
             retry += 1
-        if retry > 80:
+        if retry > 160:
             logger.debug("User has waited too long, Loop break.")
             break
         if len(tasks) == 0:
@@ -928,7 +965,7 @@ def chat_huggingface(messages, openaikey = None, huggingfacetoken = None, return
     if return_results:
         return results
     
-    response = response_results(input, results, openaikey).strip()
+    response = response_results(input, results, api_key, api_type, api_endpoint).strip()
 
     end = time.time()
     during = end - start
@@ -936,4 +973,98 @@ def chat_huggingface(messages, openaikey = None, huggingfacetoken = None, return
     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 response, results
+    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()