import json
import argparse
import requests

parse = argparse.ArgumentParser()
parse.add_argument('--query', default='请写一首诗')
parse.add_argument('--use_hf', action='store_true')
args = parse.parse_args()

print(args.query)
headers = {"Content-Type": "application/json"}
data = {
        "query": args.query,
        "history": []
        }

json_str = json.dumps(data)
if args.use_hf:
    response = requests.post("http://localhost:8888/hf_inference", headers=headers, data=json_str.encode("utf-8"), verify=False)
else:
    response = requests.post("http://localhost:8888/vllm_inference", headers=headers, data=json_str.encode("utf-8"), verify=False)

str_response = response.content.decode("utf-8")
print(json.loads(str_response))