rewired login functionality

ChitChat/__init__.py

@@ -3,17 +3,18 @@ from flask_sqlalchemy import SQLAlchemy
 from flask_bcrypt import Bcrypt
 from flask_cors import CORS 
 from ChitChat.config import Config
-
+from flask_login import LoginManager
 
 db = SQLAlchemy()
 bcrypt = Bcrypt()
-
+login_manager = LoginManager()
 def create_app(config_class = Config):
     app = Flask(__name__)
     CORS(app)
     app.config.from_object(Config)
     db.init_app(app)
     bcrypt.init_app(app)
+    login_manager.init_app(app)
 
     from ChitChat.resources.routes import resources
     app.register_blueprint(resources)

ChitChat/common/files/.gitignore

@@ -0,0 +1 @@
+**/**

ChitChat/common/utils.py

@@ -3,7 +3,9 @@ from transformers import T5Tokenizer, T5ForConditionalGeneration
 import torch
 from flask import current_app
 from ChitChat import db
-from ChitChat.models import ChatHistory
+from ChitChat.models import User
+# from flask_login import current_user
+
 
 small_model_name = 'Th3BossC/DialoGPT-medium-AICLUB_NITC'
 default_model = 'microsoft/DialoGPT-medium'
@@ -24,7 +26,7 @@ def getChatHistory(user):
 
 def saveChatHistory(user, chat_history_ids):
     location = current_app.config.get('SAVE_FOLDER')
-    file_name = location + user.user_id + '.pt'
+    file_name = location + str(user.id) + '.pt'
     if chat_history_ids.shape[-1] > 100:
         user.history = None 
         db.session.commit()

ChitChat/models.py

@@ -1,15 +1,19 @@
 from ChitChat import db 
 from datetime import datetime 
 import uuid
+from flask_login import UserMixin
+from ChitChat import login_manager
 
-class ChatHistory(db.Model):
+@login_manager.user_loader
+def load_user(user_id):
+    return User.query.get(int(user_id)) 
+
+class User(db.Model, UserMixin):
     id = db.Column(db.Integer, primary_key = True)
-    user_id = db.Column(db.String(), unique = True, nullable = False, default = str(uuid.uuid4()))
     username  = db.Column(db.String(10), unique = True, nullable = False)
     email = db.Column(db.String(100), nullable = False)
     password = db.Column(db.String(60), nullable = False)
     date = db.Column(db.DateTime, nullable = False, default = datetime.utcnow)
     history = db.Column(db.String(100), nullable = True)
-
     def __repr__(self):
-        return f"ChatHistory({self.username}, {self.email})"
+        return f"User({self.username}, {self.email})"

ChitChat/resources/routes.py

@@ -1,9 +1,9 @@
 from flask import Blueprint, request, current_app
 from flask_restful import Api, Resource 
-from ChitChat.models import ChatHistory
+from ChitChat.models import User
 from ChitChat import bcrypt, db
 from ChitChat.common.utils import conversation, complexChat
-
+# from flask_login import login_user, logout_user, current_user
 
 resources = Blueprint('resources', __name__)
 api = Api(resources)
@@ -12,13 +12,12 @@ api = Api(resources)
 class UserLogin(Resource):
     def post(self):
         userInfo = request.json
-        
-        user = ChatHistory.query.filter_by(email = userInfo['email']).first()
+        user = User.query.filter_by(username = userInfo['username']).first()
 
         if user is None:
             return {'status' : "Account doesn't exist", 'user_id' : None}
         elif bcrypt.check_password_hash(pw_hash = user.password, password = userInfo['password']):
-            return {'status' : "login successful", 'user_id' : user.user_id}
+            return {'status' : "login successful", 'user_id' : str(user.id)}
         else:
             return {'status' : "Invalid password", 'user_id' : None}
 api.add_resource(UserLogin, '/login/')
@@ -27,23 +26,24 @@ class RegisterUser(Resource):
     def post(self):
         userInfo = request.json 
 
-        user = ChatHistory.query.filter_by(email = userInfo['email']).first()
+        user = User.query.filter_by(email = userInfo['email']).first()
         if user is not None:
             return {'status' : 'email already registered'}
         else:
-            user = ChatHistory.query.filter_by(username = userInfo['username']).first()
+            user = User.query.filter_by(username = userInfo['username']).first()
             if user is not None:
                 return {'status' : 'Username already exists'} 
         
-        newUser = ChatHistory(username = userInfo['username'], email = userInfo['email'], password = bcrypt.generate_password_hash(password = userInfo['password']))
+        newUser = User(username = userInfo['username'], email = userInfo['email'], password = bcrypt.generate_password_hash(password = userInfo['password']))
         db.session.add(newUser)
         db.session.commit()
         return {'status' : 'User created successfully'} 
 api.add_resource(RegisterUser, '/register/')
 
+
 class ChatBot(Resource):
     def post(self, user_id):
-        user = ChatHistory.query.filter_by(user_id = user_id).first()
+        user = User.query.filter_by(id = user_id).first()
         userInput = request.json['user']
         if user is None:
             return {'error' : "User doesn't exist"}, 400

app.py

@@ -3,8 +3,8 @@ from ChitChat import create_app
 
 app = create_app()
 
-# if __name__ == '__main__':
-#     app.run(debug = True, port = 5000)
-
 if __name__ == '__main__':
-    app.run(debug = False, host = "0.0.0.0", port = 7860)
+    app.run(debug = True, port = 5000)
+
+# if __name__ == '__main__':
+#     app.run(debug = False, host = "0.0.0.0", port = 7860)