From db6fbe83662c4bef74d11dab49f664a2b5fbb32a Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Tue, 7 Nov 2023 22:22:29 +0800
Subject: [PATCH] add python webui server

---
 backend-python/main.py         | 45 +++++++++++++++++++++++-----------
 backend-python/webui_server.py | 14 +++++++++++
 2 files changed, 45 insertions(+), 14 deletions(-)
 create mode 100644 backend-python/webui_server.py

diff --git a/backend-python/main.py b/backend-python/main.py
index 58e472a..3f8e9d9 100644
--- a/backend-python/main.py
+++ b/backend-python/main.py
@@ -22,6 +22,11 @@ def get_args(args: Union[Sequence[str], None] = None):
         help="host to run the server on (default: 127.0.0.1)",
     )
     group = parser.add_argument_group(title="mode arguments")
+    group.add_argument(
+        "--webui",
+        action="store_true",
+        help="whether to enable WebUI (default: False)",
+    )
     group.add_argument(
         "--rwkv-beta",
         action="store_true",
@@ -79,6 +84,32 @@ app.include_router(misc.router)
 app.include_router(state_cache.router)
 
 
+@app.post("/exit", tags=["Root"])
+def exit():
+    parent_pid = os.getpid()
+    parent = psutil.Process(parent_pid)
+    for child in parent.children(recursive=True):
+        child.kill()
+    parent.kill()
+
+
+try:
+    if (
+        "RWKV_RUNNER_PARAMS" in os.environ
+        and "--webui" in os.environ["RWKV_RUNNER_PARAMS"].split(" ")
+    ) or args.webui:
+        from webui_server import webui_server
+
+        app.mount("/", webui_server)
+except NameError:
+    pass
+
+
+@app.get("/", tags=["Root"])
+def read_root():
+    return {"Hello": "World!"}
+
+
 def init():
     global_var.init()
     cmd_params = os.environ["RWKV_RUNNER_PARAMS"]
@@ -94,20 +125,6 @@ def init():
         ngrok_connect()
 
 
-@app.get("/", tags=["Root"])
-def read_root():
-    return {"Hello": "World!"}
-
-
-@app.post("/exit", tags=["Root"])
-def exit():
-    parent_pid = os.getpid()
-    parent = psutil.Process(parent_pid)
-    for child in parent.children(recursive=True):
-        child.kill()
-    parent.kill()
-
-
 if __name__ == "__main__":
     os.environ["RWKV_RUNNER_PARAMS"] = " ".join(sys.argv[1:])
     print("--- %s seconds ---" % (time.time() - start_time))
diff --git a/backend-python/webui_server.py b/backend-python/webui_server.py
new file mode 100644
index 0000000..5d58d35
--- /dev/null
+++ b/backend-python/webui_server.py
@@ -0,0 +1,14 @@
+from fastapi import FastAPI
+from fastapi.middleware.gzip import GZipMiddleware
+from fastapi.staticfiles import StaticFiles
+import uvicorn
+
+webui_server = FastAPI()
+
+webui_server.add_middleware(GZipMiddleware, minimum_size=1000)
+webui_server.mount(
+    "/", StaticFiles(directory="frontend/dist", html=True), name="static"
+)
+
+if __name__ == "__main__":
+    uvicorn.run("webui_server:webui_server")