Update main.py

main.py

@@ -1,6 +1,12 @@
-from flask import Flask, request, jsonify
 from huggingface_hub import InferenceClient
 
+import random
+
+from flask import Flask, request, jsonify, make_response
+from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
+from werkzeug.security import generate_password_hash, check_password_hash
+import sqlite3
+
 client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")
 
 app = Flask(__name__)
@@ -13,7 +19,49 @@ with open(file_path, "r") as file:
 def home():
     return jsonify({"message": "Welcome to the Recommendation API!"})
 
-import random
+
+
+app.config['JWT_SECRET_KEY'] = 123456
+jwt = JWTManager(app)
+
+# Create SQLite database
+conn = sqlite3.connect('users.db', check_same_thread=False)
+c = conn.cursor()
+c.execute('''CREATE TABLE IF NOT EXISTS users
+             (id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT UNIQUE, password TEXT)''')
+conn.commit()
+
+# Endpoint for user registration
+@app.route('/register', methods=['POST'])
+def register():
+    data = request.get_json()
+    username = data.get('username')
+    password = data.get('password')
+    if not username or not password:
+        return jsonify({"message": "Missing username or password"}), 400
+    hashed_password = generate_password_hash(password)
+    try:
+        c.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, hashed_password))
+        conn.commit()
+        return jsonify({"message": "User created successfully"}), 201
+    except sqlite3.IntegrityError:
+        return jsonify({"message": "Username already exists"}), 400
+
+# Endpoint for user login
+@app.route('/login', methods=['POST'])
+def login():
+    data = request.get_json()
+    username = data.get('username')
+    password = data.get('password')
+    if not username or not password:
+        return jsonify({"message": "Missing username or password"}), 400
+    user = c.execute("SELECT * FROM users WHERE username=?", (username,)).fetchone()
+    if user and check_password_hash(user[2], password):
+        access_token = create_access_token(identity=username, expires_delta=False)
+        return jsonify({"access_token": access_token}), 200
+    else:
+        return jsonify({"message": "Invalid username or password"}), 401
+
 
 def format_prompt(message):
     # Generate a random user prompt and bot response pair
@@ -25,7 +73,9 @@ def format_prompt(message):
 
 
 @app.route('/get_course', methods=['POST'])
+@jwt_required()
 def recommend():
+    current_user = get_jwt_identity()
     temperature = 0.9
     max_new_tokens = 256
     top_p = 0.95
@@ -60,8 +110,11 @@ def recommend():
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=False, details=False, return_full_text=False)
     return jsonify({"ans": stream})
 
+
 @app.route('/get_mentor', methods=['POST'])
+@jwt_required()
 def mentor():
+    current_user = get_jwt_identity()
     temperature = 0.9
     max_new_tokens = 256
     top_p = 0.95