add token file

qwen_server/assistant_server.py

@@ -25,6 +25,7 @@ mem = Memory(llm=llm, stream=False)
 
 cache_file = os.path.join(server_config.path.cache_root, 'browse.jsonl')
 cache_file_popup_url = os.path.join(server_config.path.cache_root, 'popup_url.jsonl')
+access_token_file = os.path.join(server_config.path.cache_root, 'access_token.jsonl')
 
 PAGE_URL = {}
 
@@ -70,6 +71,18 @@ def set_page_url(access_token):
 
 def initialize(request: gr.Request):
     access_token = request.query_params["access_token"] or request.session["access_token"]
+    is_valid = False
+    if access_token:
+        if os.getenv("ACCESS_TOKEN") == access_token:
+            is_valid = True
+        else:
+            for line in jsonlines.open(access_token_file):
+                if line['access_token'] == access_token:
+                    is_valid = True
+                    break
+    if not is_valid:
+        gr.Info("The token is not valid, Please reset!")
+        return
     set_page_url(access_token)
     return access_token
 

qwen_server/database_server.py

@@ -32,7 +32,7 @@ origins = [
     'http://' + get_local_ip() + ':' +
     str(server_config.server.workstation_port),
 ]
-
+access_token_file = os.path.join(server_config.path.cache_root, 'access_token.jsonl')
 app.add_middleware(
     CORSMiddleware,
     # allow_origins=origins,
@@ -80,10 +80,19 @@ def change_checkbox_state(text, cache_file, access_token):
 
 
 @app.middleware("http")
-async def add_process_time_header(request: Request, call_next):
+async def access_token_auth(request: Request, call_next):
     # print(f"Request URL path: {request.url}")
     access_token: str = request.headers.get("Authorization") or request.query_params.get("access_token") or request.session.get("access_token")
-    if not access_token or os.getenv("ACCESS_TOKEN") != access_token:
+    is_valid = False
+    if access_token:
+        if os.getenv("ACCESS_TOKEN") == access_token:
+            is_valid = True
+        else:
+            for line in jsonlines.open(access_token_file):
+                if line['access_token'] == access_token:
+                    is_valid = True
+                    break
+    if not is_valid:
         return Response(status_code=401, content="the token is not valid")
     request.session.setdefault("access_token", access_token)
     return await call_next(request)
@@ -94,8 +103,27 @@ async def healthz(request: Request):
     return JSONResponse({"healthz": True})
 
 
+@app.post('/token/add')
+async def add_token(request: Request):
+    access_token: str = request.headers.get("Authorization") or request.query_params.get("access_token") or request.session.get("access_token")
+    if not access_token or os.getenv("ACCESS_TOKEN") != access_token:
+        return Response(status_code=401, content="the token is not valid")
+    data = await request.json()
+    lines = []
+    for line in jsonlines.open(access_token_file):
+        lines.append(line)
+    lines.append(data)
+    with jsonlines.open(access_token_file, mode='w') as writer:
+        for new_line in lines:
+            writer.write(new_line)
+    return JSONResponse({"success": True})
+
+
 @app.get('/cachedata/{file_name}')
-async def cache_data(file_name: str):
+async def cache_data(request: Request, file_name: str):
+    access_token: str = request.headers.get("Authorization") or request.query_params.get("access_token") or request.session.get("access_token")
+    if not access_token or os.getenv("ACCESS_TOKEN") != access_token:
+        return Response(status_code=401, content="the token is not valid")
     cache_file = os.path.join(server_config.path.cache_root, file_name)
     lines = []
     for line in jsonlines.open(cache_file):