Add set api key page (#45)

* add html page for set key * change the gif howto set key
siliconflow · Jul 22, 2024 · 0bdbb2b · 0bdbb2b
1 parent 816a76f
commit 0bdbb2b
Show file tree

Hide file tree

Showing 5 changed files with 150 additions and 29 deletions.
diff --git a/auth.py b/auth.py
@@ -4,9 +4,89 @@
 import server
 from aiohttp import web
 
-from .utils import validate_api_key
 
 API_KEY = None
+set_api_key_html = """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>BizyAir - Set API Key</title>
+    <style>
+        body {
+            background-color: #121212;
+            color: #ffffff;
+            font-family: Arial, sans-serif;
+            display: flex;
+            justify-content: center;
+            align-items: center;
+            height: 100vh;
+            margin: 0;
+        }
+        .container {
+            text-align: center;
+        }
+        input[type="password"] {
+            padding: 10px;
+            margin: 10px;
+            border: 2px solid rgb(130, 88, 245);
+            border-radius: 5px;
+            background-color: #1e1e1e;
+            color: #ffffff;
+            width: 200px;
+        }
+        button {
+            padding: 10px 20px;
+            border: none;
+            border-radius: 5px;
+            background-color: rgb(130, 88, 245);
+            color: #ffffff;
+            cursor: pointer;
+        }
+        button:hover {
+            background-color: #0056b3;
+        }
+    </style>
+    <script>
+        async function setApiKey() {
+            const apiKey = document.getElementById('apiKey').value;
+            const response = await fetch('/bizyair/set_api_key', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/x-www-form-urlencoded'
+                },
+                body: `api_key=${encodeURIComponent(apiKey)}`
+            });
+            if (response.ok) {
+                alert('API Key set successfully!');
+                if (window.opener) {
+                    window.close();
+                }
+            } else {
+                alert('Failed to set API Key: ' + await response.text());
+            }
+        }
+    </script>
+</head>
+<body>
+    <div class="container">
+        <h1>Set API Key</h1>
+        <input type="password" id="apiKey" placeholder="Enter API Key">
+        <button onclick="setApiKey()">Set API Key</button>
+        <p>To get your key, visit <a href="https://cloud.siliconflow.cn" target="_blank">https://cloud.siliconflow.cn</a></p>
+    </div>
+</body>
+</html>
+    </div>
+</body>
+</html>
+"""
+
+
+@server.PromptServer.instance.routes.get("/bizyair/set-api-key")
+async def set_api_key_page(request):
+    return web.Response(text=set_api_key_html, content_type="text/html")
 
 
 @server.PromptServer.instance.routes.post("/bizyair/set_api_key")
@@ -22,33 +102,37 @@ async def set_api_key(request):
             API_KEY = api_key
             return response
         else:
-            return web.Response(text="No token provided", status=400)
+            return web.Response(
+                text="No token provided, please refer to cloud.siliconflow.cn to get the key",
+                status=400,
+            )
     except Exception as e:
         return web.Response(text=str(e), status=500)
 
+
 @server.PromptServer.instance.routes.get("/bizyair/get_api_key")
 async def get_api_key(request):
     global API_KEY
-    # 检查是否存在 api_key.ini，文件中是否存在 key
-    api_key = ''
+    api_key = ""
     current_directory = os.path.dirname(os.path.abspath(__file__))
-    file_path = os.path.join(current_directory, 'api_key.ini')
+    file_path = os.path.join(current_directory, "api_key.ini")
     if os.path.exists(file_path):
         config = configparser.ConfigParser()
         config.read(file_path)
-        api_key = config.get('auth', 'api_key', fallback='').strip()
-        # print(f"api_key from file = {api_key}")
-    if api_key=='':
+        api_key = config.get("auth", "api_key", fallback="").strip()
+    if api_key == "":
         api_key = request.cookies.get("api_key")
-    # print(f"api_key = {api_key}")
     try:
         if api_key:
             API_KEY = api_key
             response = web.Response(text="ok")
             response.set_cookie("api_key", api_key)
             return response
         else:
-            return web.Response(text="No api key found in cookie", status=404)
+            return web.Response(
+                text="No api key found in cookie, please refer to cloud.siliconflow.cn to get the key",
+                status=404,
+            )
     except Exception as e:
         return web.Response(text="str(e)", status=500)
 
@@ -65,7 +149,6 @@ def INPUT_TYPES(s):
     OUTPUT_NODE = True
 
     def set_api_key(self, API_KEY="YOUR_API_KEY"):
-        validate_api_key(API_KEY)
         return {"ui": {"api_key": (API_KEY,)}, "result": ()}
 
     @classmethod
@@ -77,5 +160,5 @@ def IS_CHANGED(s, latent):
     "BizyAirSetAPIKey": SetAPIKey,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "BizyAirSetAPIKey": "☁️Set SiliconCloud API Key",
+    "BizyAirSetAPIKey": "☁️Set SiliconCloud API Key(deprecated)",
 }
diff --git a/docs/docs/getting-started/imgs/how-to-set-key.gif b/docs/docs/getting-started/imgs/how-to-set-key.gif
diff --git a/docs/docs/getting-started/quick-start.md b/docs/docs/getting-started/quick-start.md
@@ -5,13 +5,10 @@
 !!! note
     You can get your API KEY from [SiliconCloud](https://cloud.siliconflow.cn/) for free.
 
-For the first use, you need to set your API key using "Set SiliconCloud API Key". Add the "Set SiliconCloud API Key" node and run it(click "Queue Prompt" button). The key changes to "**********", indicating successful operation.
+For the first use, you need to set your API key. Click the "BizyAir Key" button to launch the key setting page and set the key.
 
 ![](./imgs/how-to-set-key.gif)
 
-!!! note
-    You **should** remove this node after you set the key.
-
 !!! note
     You don't need to set the key every time you start the program. The key is cached in your local browser, and as long as you don't get a pop-up prompt asking you to enter the key, you can continue to use BizyAir Nodes.
 

diff --git a/js/set_api_key.js b/js/set_api_key.js
@@ -2,6 +2,50 @@ import { app } from "../../scripts/app.js";
 import { api } from "../../../../scripts/api.js";
 
 
+app.registerExtension({
+    name: "BizyAir_SetAPIKey",
+
+    async setup() {
+        const menu = document.querySelector(".comfy-menu");
+
+        // 添加分隔线
+        const separator = document.createElement("div");
+        separator.style.margin = "20px 0";
+        separator.style.width = "100%";
+        separator.style.height = "2px";
+        separator.style.backgroundColor = "gray";
+        menu.append(separator);
+
+        // 创建按钮
+        const BizyAir_SetAPIKey = document.createElement("button");
+        BizyAir_SetAPIKey.textContent = "BizyAir Key";
+        BizyAir_SetAPIKey.style.backgroundColor = "rgb(130, 88, 245)"; // 设置按钮背景颜色
+        BizyAir_SetAPIKey.style.border = "none";
+        BizyAir_SetAPIKey.style.color = "white";
+        BizyAir_SetAPIKey.style.padding = "10px 20px";
+        BizyAir_SetAPIKey.style.textAlign = "center";
+        BizyAir_SetAPIKey.style.textDecoration = "none";
+        BizyAir_SetAPIKey.style.display = "inline-block";
+        BizyAir_SetAPIKey.style.fontSize = "16px";
+        BizyAir_SetAPIKey.style.margin = "4px 2px";
+        BizyAir_SetAPIKey.style.cursor = "pointer";
+        BizyAir_SetAPIKey.style.borderRadius = "12px";
+
+        // 动态获取当前页面的协议、地址和端口
+        const currentUrl = window.location.protocol + "//" + window.location.host;
+        const apiKeyUrl = currentUrl + "/bizyair/set-api-key";
+
+        // 设置按钮点击事件
+        BizyAir_SetAPIKey.onclick = () => {
+            window.open(apiKeyUrl, "Set API Key");
+        };
+
+        menu.append(BizyAir_SetAPIKey);
+    }
+});
+
+
+
 app.registerExtension({
     name: "bizyair.set.api.key",
     async setup() {
@@ -49,4 +93,4 @@ or you can only use nodes locally.`);
             };
         }
     },
-})
+});
diff --git a/utils.py b/utils.py
@@ -31,7 +31,15 @@ def send_post_request(api_url, payload, headers):
             response_data = response.read().decode("utf-8")
         return response_data
     except urllib.error.URLError as e:
-        raise Exception(f"Failed to connect to the server: {e}")
+        if "Unauthorized" in str(e):
+            raise Exception(
+                "Key is invalid, please refer to https://cloud.siliconflow.cn to get the API key.\n"
+                "If you have the key, please click the 'BizyAir Key' button at the bottom right to set the key."
+            )
+        else:
+            raise Exception(
+                f"Failed to connect to the server: {e}, if you have no key, "
+            )
 
 
 def serialize_and_encode(obj: Union[np.ndarray], compress=True) -> Tuple[str, bool]:
@@ -114,19 +122,9 @@ def format_bytes(num_bytes: int) -> str:
         return f"{num_bytes / (1024 * 1024):.2f} MB"
 
 
-def validate_api_key(key: str):
-    assert isinstance(
-        key, str
-    ), f"API_KEY must be a string, got {type(key)}. Please refer to https://cloud.siliconflow.cn to get your API_KEY"
-    assert key.startswith(
-        "sk"
-    ), f"invalid key, please refer to https://cloud.siliconflow.cn to get your API_KEY, got {key}"
-
-
 def get_api_key():
     from .auth import API_KEY
 
-    validate_api_key(API_KEY)
     return API_KEY
 
 
@@ -139,7 +137,6 @@ def get_llm_response(
 ):
     api_url = "https://api.siliconflow.cn/v1/chat/completions"
     API_KEY = get_api_key()
-    validate_api_key(API_KEY)
     headers = {
         "accept": "application/json",
         "content-type": "application/json",