Spaces:
Running
Running
File size: 4,924 Bytes
e06f93a c64b502 e06f93a c64b502 964dddf c64b502 e06f93a c64b502 e06f93a c64b502 964dddf c64b502 964dddf c64b502 964dddf c64b502 964dddf c64b502 e06f93a c64b502 964dddf c64b502 964dddf c64b502 964dddf c64b502 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | from flask_login import UserMixin
from werkzeug.security import generate_password_hash, check_password_hash
from cryptography.fernet import Fernet
import pymongo
from bson.objectid import ObjectId
from config import ENCRYPTION_KEY, MONGO_URI
# Connect to MongoDB
mongo_client = pymongo.MongoClient(MONGO_URI)
try:
db = mongo_client.get_default_database()
except pymongo.errors.ConfigurationError:
db = mongo_client["rag_app"]
users_collection = db["users"]
cipher_suite = Fernet(ENCRYPTION_KEY)
class User(UserMixin):
def __init__(self, username, email, password=None, _id=None, google_id=None,
profile_pic=None, groq_api_key=None, gemini_api_key=None,
pinecone_api_key=None, pinecone_index_name=None,
preferred_model="groq", is_admin=False):
self.username = username
self.email = email
self.password = password
self.google_id = google_id
self.profile_pic = profile_pic
self.groq_api_key = groq_api_key
self.gemini_api_key = gemini_api_key
self.pinecone_api_key = pinecone_api_key
self.pinecone_index_name = pinecone_index_name or ""
self.preferred_model = preferred_model
self.is_admin = is_admin
if _id:
self.id = str(_id)
else:
self.id = None
def get_id(self):
return self.id or self.username
def save(self):
user_data = {
"username": self.username,
"email": self.email,
"password": self.password,
"google_id": self.google_id,
"profile_pic": self.profile_pic,
"groq_api_key": self.groq_api_key,
"gemini_api_key": self.gemini_api_key,
"pinecone_api_key": self.pinecone_api_key,
"pinecone_index_name": self.pinecone_index_name,
"preferred_model": self.preferred_model,
"is_admin": self.is_admin
}
if self.id:
users_collection.update_one({"_id": ObjectId(self.id)}, {"$set": user_data})
else:
result = users_collection.insert_one(user_data)
self.id = str(result.inserted_id)
def set_password(self, password):
self.password = generate_password_hash(password)
def check_password(self, password):
return check_password_hash(self.password, password)
# ── Groq Key ─────────────────────────────────────
def set_groq_key(self, api_key):
if api_key:
self.groq_api_key = cipher_suite.encrypt(api_key.encode('utf-8')).decode('utf-8')
else:
self.groq_api_key = None
def get_groq_key(self):
if self.groq_api_key:
try:
return cipher_suite.decrypt(self.groq_api_key.encode('utf-8')).decode('utf-8')
except Exception:
return None
return None
# ── Gemini Key ───────────────────────────────────
def set_gemini_key(self, api_key):
if api_key:
self.gemini_api_key = cipher_suite.encrypt(api_key.encode('utf-8')).decode('utf-8')
else:
self.gemini_api_key = None
def get_gemini_key(self):
if self.gemini_api_key:
try:
return cipher_suite.decrypt(self.gemini_api_key.encode('utf-8')).decode('utf-8')
except Exception:
return None
return None
# ── Pinecone Key ─────────────────────────────────
def set_pinecone_key(self, api_key):
if api_key:
self.pinecone_api_key = cipher_suite.encrypt(api_key.encode('utf-8')).decode('utf-8')
else:
self.pinecone_api_key = None
def get_pinecone_key(self):
if self.pinecone_api_key:
try:
return cipher_suite.decrypt(self.pinecone_api_key.encode('utf-8')).decode('utf-8')
except Exception:
return None
return None
@classmethod
def get(cls, user_id):
try:
data = users_collection.find_one({"_id": ObjectId(user_id)})
if data:
return cls(**data)
return None
except Exception:
return None
@classmethod
def find_by_username(cls, username):
data = users_collection.find_one({"username": username})
if data:
return cls(**data)
return None
@classmethod
def find_by_email(cls, email):
data = users_collection.find_one({"email": email})
if data:
return cls(**data)
return None
@classmethod
def get_all(cls):
return [cls(**data) for data in users_collection.find()] |