Spaces:

Moonfanz
/

Gemini

Running

App Files Files Community

Moonfanz commited on Dec 15, 2024

Commit

48c1fc0

verified ·

1 Parent(s): 39bdd55

Upload 4 files

Browse files

Files changed (4) hide show

app.py +47 -9
func.py +1 -2
templates/login.html +20 -0
templates/manage.html +33 -0

app.py CHANGED Viewed

@@ -1,5 +1,6 @@
-from flask import Flask, request, jsonify, Response, stream_with_context
 import google.generativeai as genai
 import json
 from datetime import datetime
@@ -10,15 +11,16 @@ import func
 os.environ['TZ'] = 'Asia/Shanghai'
 app = Flask(__name__)
-if 'API_KEYS' not in os.environ:
-    raise EnvironmentError("API_KEYS environment variable not set")
-if 'HF_API_KEY' not in os.environ:
-    raise EnvironmentError("HF_API_KEY environment variable not set")
-formatter = logging.Formatter('%(message)s')  # 包含更多信息
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO) # 或者 logging.DEBUG
 handler = logging.StreamHandler()
 handler.setFormatter(formatter)
@@ -73,11 +75,47 @@ GEMINI_MODELS = [
     {"id": "gemini-2.0-exp"},
     {"id": "gemini-2.0-pro-exp"},
 ]
 @app.route('/hf/v1/chat/completions', methods=['POST'])
 def chat_completions():
     global current_api_key
-    is_authenticated, auth_error, status_code = func.authenticate_request(request)
     if not is_authenticated:
         return auth_error if auth_error else jsonify({'error': 'Unauthorized'}), status_code if status_code else 401
     try:
@@ -217,7 +255,7 @@ def chat_completions():
 @app.route('/hf/v1/models', methods=['GET'])
 def list_models():
-    is_authenticated, auth_error, status_code = func.authenticate_request(request)
     if not is_authenticated:
         return auth_error if auth_error else jsonify({'error': 'Unauthorized'}), status_code if status_code else 401
     response = {"object": "list", "data": GEMINI_MODELS}

+from flask import Flask, request, jsonify, Response, stream_with_context, render_template, session, redirect, url_for
 import google.generativeai as genai
 import json
 from datetime import datetime
 os.environ['TZ'] = 'Asia/Shanghai'
 app = Flask(__name__)
+app.secret_key = os.urandom(24)
+hf_api_keys = os.environ['HF_API_KEY'].split(',')
+formatter = logging.Formatter('%(message)s')
+ADMIN_PASSWORD = os.environ.get('AD_PASSWORD')
 logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
 handler = logging.StreamHandler()
 handler.setFormatter(formatter)
     {"id": "gemini-2.0-exp"},
     {"id": "gemini-2.0-pro-exp"},
 ]
+# Login route
+@app.route("/", methods=["GET", "POST"])
+def login():
+    if request.method == "POST":
+        password = request.form.get("password")
+        if password == ADMIN_PASSWORD:
+            session["logged_in"] = True
+            return redirect(url_for("manage_keys"))
+        else:
+            return render_template("login.html", error="Incorrect password")
+    return render_template("login.html", error=None)
+# Manage API keys route (requires authentication)
+@app.route("/manage", methods=["GET", "POST"])
+def manage_keys():
+    if not session.get("logged_in"):
+        return redirect(url_for("login"))
+    if request.method == "POST":
+        action = request.form.get("action")
+        if action == "add":
+            new_key = request.form.get("new_key")
+            if new_key:
+                hf_api_keys.append(new_key)
+        elif action == "delete":
+            key_to_delete = request.form.get("key_to_delete")
+            if key_to_delete in hf_api_keys:
+                hf_api_keys.remove(key_to_delete)
+    return render_template("manage.html", keys=hf_api_keys)
+# Logout route
+@app.route("/logout")
+def logout():
+    session.pop("logged_in", None)
+    return redirect(url_for("login"))
 @app.route('/hf/v1/chat/completions', methods=['POST'])
 def chat_completions():
     global current_api_key
+    is_authenticated, auth_error, status_code = func.authenticate_request(hf_api_keys, request)
     if not is_authenticated:
         return auth_error if auth_error else jsonify({'error': 'Unauthorized'}), status_code if status_code else 401
     try:
 @app.route('/hf/v1/models', methods=['GET'])
 def list_models():
+    is_authenticated, auth_error, status_code = func.authenticate_request(hf_api_keys, request)
     if not is_authenticated:
         return auth_error if auth_error else jsonify({'error': 'Unauthorized'}), status_code if status_code else 401
     response = {"object": "list", "data": GEMINI_MODELS}

func.py CHANGED Viewed

@@ -22,9 +22,8 @@ GEMINI_MODELS = [
 logger = logging.getLogger(__name__)
-def authenticate_request(request):
     auth_header = request.headers.get('Authorization')
-    hf_api_key = os.environ.get('HF_API_KEY').split(',')
     if not auth_header:
         return False, jsonify({'error': 'Authorization header is missing'}), 401

 logger = logging.getLogger(__name__)
+def authenticate_request(hf_api_key, request):
     auth_header = request.headers.get('Authorization')
     if not auth_header:
         return False, jsonify({'error': 'Authorization header is missing'}), 401

templates/login.html ADDED Viewed

	@@ -0,0 +1,20 @@

+<!DOCTYPE html>
+<html>
+<head>
+    <title>Login</title>
+</head>
+<body>
+    <h1>Admin Login</h1>
+    {% if error %}
+    <p style="color: red;">{{ error }}</p>
+    {% endif %}
+    <form method="POST">
+        <label for="password">Password:</label>
+        <input type="password" name="password" id="password" required>
+        <button type="submit">Login</button>
+    </form>
+</body>
+</html>

templates/manage.html ADDED Viewed

	@@ -0,0 +1,33 @@

+<!DOCTYPE html>
+<html>
+<head>
+    <title>Manage API Keys</title>
+</head>
+<body>
+    <h1>Manage Hugging Face API Keys</h1>
+    <h2>Current Keys:</h2>
+    <ul>
+        {% for key in keys %}
+        <li>
+            {{ key }}
+            <form method="POST" style="display: inline;">
+                <input type="hidden" name="action" value="delete">
+                <input type="hidden" name="key_to_delete" value="{{ key }}">
+                <button type="submit">Delete</button>
+            </form>
+        </li>
+        {% endfor %}
+    </ul>
+    <h2>Add New Key:</h2>
+    <form method="POST">
+        <input type="hidden" name="action" value="add">
+        <input type="text" name="new_key" placeholder="Enter new API key">
+        <button type="submit">Add</button>
+    </form>
+    <br>
+    <a href="{{ url_for('logout') }}">Logout</a>
+</body>
+</html>