browser-chatgptapi/run.py at master · heshen/browser-chatgptapi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.templating import Jinja2Templates

app = FastAPI(
    docs_url='/doc'
)

app.add_middleware(
    CORSMiddleware,

    allow_origins=[
        "*",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

socket_manager = dict()

wait = False


async def sendto_server(data):
    # global wait
    # if not wait:
    #     wait = True
    #     await socket_manager['server'].send_json(data)
    # else:
    #     return True

    await socket_manager['server'].send_json(data)


async def sendto_user(data):
    global wait
    await socket_manager['user'].send_json(data)
    wait = False


@app.websocket("/server/{server}")
async def websocket_endpoint(websocket: WebSocket, server: str):
    ws = websocket
    await ws.accept()

    socket_manager['server'] = ws
    # if not websocket.headers.get('x-user-type') == 'server':
    #     await ws.send_json({"message": "这个接口只供服务端连接"})
    #     await ws.close(code=882)

    if socket_manager.get('user'):
        await sendto_user({'msg': '服务端加入成功'})

    try:
        while True:
            data = await websocket.receive_json()

            print(data)
            if not socket_manager.get('user'):
                await ws.send_json({'msg': '客户端未加入'})
            else:
                await sendto_user(data)

    except WebSocketDisconnect:
        # await ws.close()
        socket_manager.pop('server')


@app.websocket("/user/{user}")
async def websocket_endpoint(ws: WebSocket, user: str):

    await ws.accept()
    socket_manager['user'] = ws
    await ws.send_json({'msg': '连接成功'})

    try:
        while True:
            data = await ws.receive_json()
            print(data)
            if not socket_manager.get('server'):
                await ws.send_json({'msg': '服务端未加入'})
            else:
                result = await sendto_server(data)
                if result:
                    await ws.send_json({'msg': '请等待上一条消息返回'})

    except WebSocketDisconnect:

        print('user disconnected')
        socket_manager.pop('user')


templates = Jinja2Templates(directory='./')


@app.get("/")
def index(request: Request):
    api_url = 'ws://localhost:8010/user/user1'
    return templates.TemplateResponse("index2.html",
                                      {"request": request, 'api_url': api_url})


if __name__ == "__main__":
    import uvicorn

    # 官方推荐是用命令后启动 uvicorn main:app --host=127.0.0.1 --port=8010 --reload
    uvicorn.run('run:app', host='localhost', port=8010, reload=True,debug=True)