#該函數(shù)段由RdFast小程序自動生成
from pydantic import BaseModel
class User(BaseModel):
    prompt: str
    history: list
@app.post("/http/noasync")
def create_item(request: User):
    global model, tokenizer
    json_post_raw = request.dict()
    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')
    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()

import json
import time
import requests
import threading
def get_ans(id, prompt):
    t0 = time.time()
    headers = {'Content-Type': 'application/json'}
    url = 'http://IP:Port/http/noasync'
    data = {'prompt': prompt, 'history': []}
    data = json.dumps(data)
    reponse = requests.post(url=url, data=data, headers=headers)
    print(id, '耗時為：', round(time.time() - t0, 2), 's，結(jié)果為：', reponse .text)
if __name__ == '__main__':
    t1 = threading.Thread(target=get_ans, args=('線程1', '你好'))
    t2 = threading.Thread(target=get_ans, args=('線程2', '你好'))
    t3 = threading.Thread(target=get_ans, args=('線程3', '你好'))
    t1.start()
    t2.start()
    t3.start()

app = FastAPI()
tokenizer = AutoTokenizer.from_pretrained("chatglm-6b-int4-qe", trust_remote_code=True)
model = AutoModel.from_pretrained("chatglm-6b-int4-qe", trust_remote_code=True).half().cuda()
model.eval()

#該示例程序來源于RdFast小程序自動生成。
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
app = FastAPI()
connected_websockets = {}
@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await websocket.accept()
    connected_websockets[client_id] = websocket
    try:
        while True:
            # 接收客戶端websocket發(fā)送過來的信息
            data = await websocket.receive_text()
            # 將接收到的信息通過所有已連接的websocket廣播給其他客戶端
            for ws in connected_websockets.values():
                await ws.send_text(f"Client {client_id}: {data}")
    except WebSocketDisconnect:
        # 連接斷開時，從已連接的客戶端中移除
        del connected_websockets[client_id]

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(ws: WebSocket, client_id: str):
    await ws.accept()
    print('已連接')
    try:
        while True:
            # 接收客戶端websocket發(fā)送過來的信息
            data = await ws.receive_text()
            print('收到消息：', data)
            resp0 = ''
            for response, history in model.stream_chat(tokenizer, data, [], max_length= 2048,top_p= 0.7, temperature= 0.95):
                print('response：', response)
                res = response.replace(resp0, '')
                resp0 = response
                await ws.send_text(res)
            await ws.send_text('<rdfast stop>')#自定義結(jié)束符
    except WebSocketDisconnect:
        print('連接已斷開')

from websocket import create_connection
def connect_node(ques):
    ans = ''
    url = "ws://IP:Port/ws/2"
    ws = create_connection(url)
    ws.send(ques)
    while True:
        try:
            recv_text = ws.recv()
            print(recv_text)
            if '<rdfast stop>' in recv_text:
                print('break')
                break
            ans += recv_text
        except Exception as e:
            print('except： ', str(e))
            recv_text = ws.recv()
            break
    print(ans)
    ws.close()
connect_node('你好')

demo.queue(
    concurrency_count=5,
    max_size=500).launch(share=False,
                inbrowser=True,
                server_name="0.0.0.0",
                server_port=8000)