from typing import Union
from fastapi import FastAPI, Header, BackgroundTasks,Request,Body
from urllib.parse import unquote
import re
import pickle
import os
from fastapi.responses import JSONResponse
import asyncio
from pydantic import BaseModel
import time
from transformers import AutoTokenizer, AutoModel
import uvicorn, json, datetime
import torch

DEVICE = "cuda"
DEVICE_ID = "0"
CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE


def torch_gc():
    if torch.cuda.is_available():
        with torch.cuda.device(CUDA_DEVICE):
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()



app = FastAPI()


class Item(BaseModel):
    query: Union[str, None] = None


@app.get("/")
def index():
    return "hello"




@app.get("/home")
# def read_root():
# return {"Hello": "World"}
def read_items(
    appid: Union[str, None] = Header(default=None),
    bid: Union[str, None] = Header(default=None),
    requestid: Union[str, None] = Header(default=None),
    uid: Union[str, None] = Header(default=None),
):
    print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid)
    return {"appid": appid, "bid": bid, "requestid": requestid, "uid": uid}


@app.post("/glm")
async def create_item(request: Request):
    # 获取请求参数
    global model, tokenizer
    json_post_raw = await request.json()
    json_post = json.dumps(json_post_raw)
    json_post_list = json.loads(json_post)
    prompt = json_post_list.get('prompt') #问题
    history = json_post_list.get('history')  #默认空数组
    max_length = json_post_list.get('max_length')
    top_p = json_post_list.get('top_p')
    temperature = json_post_list.get('temperature')

    # 请求ai模型
    response, history = model.chat(tokenizer,
                                   prompt,
                                   history=history,
                                   max_length=max_length if max_length else 2048,
                                   top_p=top_p if top_p else 0.7,
                                   temperature=temperature if temperature else 0.95)
    now = datetime.datetime.now()
    time = now.strftime("%Y-%m-%d %H:%M:%S")
    answer = {
        "response": response,
        "history": history,
        "status": 200,
        "time": time
    }
    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
    print(log)
    torch_gc()
    return answer


# 通过网址发送闪信
@app.post("/callback")
async def callback(
    # item: Item,
    # inputArguments: Union[str, None] = None,
    query=Body(None),
    appid: Union[str, None] = Header(default=None),
    bid: Union[str, None] = Header(default=None),
    requestid: Union[str, None] = Header(default=None),
    uid: Union[str, None] = Header(default=None),
):
    print('query',query)
    print('query[query]',query["query"])

    # 获取请求参数
    global model, tokenizer
    # time.sleep(10)
    # print("item", item)
    history = []  #默认空数组
    prompt=''
    try:
        query = unquote(query["query"], "utf-8")
        prompt = query #问题
        
    except:
        prompt = ""
   
    max_length=None
    top_p=None
    temperature=None
    # 请求ai模型
    response, history = model.chat(tokenizer,
                                   prompt,
                                   history=history,
                                   max_length=max_length if max_length else 2048,
                                   top_p=top_p if top_p else 0.7,
                                   temperature=temperature if temperature else 0.95)
    # log日志
    now = datetime.datetime.now()
    time = now.strftime("%Y-%m-%d %H:%M:%S")
    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
    print(log)
    torch_gc()
    # return answer
    content = {
        "answer": response,
        "answer_type": "text",
    }
    headers = {"Content-Type": "application/json"}
    return JSONResponse(content=content, headers=headers)

if __name__ == '__main__':
    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
    model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).cuda()
    # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
    # model_path = "THUDM/chatglm2-6b"
    # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    # model = load_model_on_gpus(model_path, num_gpus=2)
    model.eval()
    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)