users.py 6.97 KB
Newer Older
1
from pydantic import BaseModel, ConfigDict, parse_obj_as
2
3
from typing import List, Union, Optional
import time
4
5
6
7

from sqlalchemy import String, Column, BigInteger, Text
from sqlalchemy.orm import Session

Timothy J. Baek's avatar
Timothy J. Baek committed
8
9
from utils.misc import get_gravatar_url

10
from apps.webui.internal.db import Base, JSONField
11
from apps.webui.models.chats import Chats
12

13
14
15
16
17
####################
# User DB Schema
####################


18
19
class User(Base):
    __tablename__ = "user"
Timothy J. Baek's avatar
Timothy J. Baek committed
20

21
22
23
24
25
    id = Column(String, primary_key=True)
    name = Column(String)
    email = Column(String)
    role = Column(String)
    profile_image_url = Column(String)
Timothy J. Baek's avatar
Timothy J. Baek committed
26

27
28
29
    last_active_at = Column(BigInteger)
    updated_at = Column(BigInteger)
    created_at = Column(BigInteger)
Timothy J. Baek's avatar
Timothy J. Baek committed
30

31
32
33
    api_key = Column(String, nullable=True, unique=True)
    settings = Column(JSONField, nullable=True)
    info = Column(JSONField, nullable=True)
34

35
    oauth_sub = Column(Text, unique=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
36
37


38
39
40
41
42
43
class UserSettings(BaseModel):
    ui: Optional[dict] = {}
    model_config = ConfigDict(extra="allow")
    pass


44
class UserModel(BaseModel):
45
46
    model_config = ConfigDict(from_attributes=True)

47
48
49
    id: str
    name: str
    email: str
Timothy J. Baek's avatar
Timothy J. Baek committed
50
    role: str = "pending"
51
    profile_image_url: str
Timothy J. Baek's avatar
Timothy J. Baek committed
52
53
54
55
56

    last_active_at: int  # timestamp in epoch
    updated_at: int  # timestamp in epoch
    created_at: int  # timestamp in epoch

57
    api_key: Optional[str] = None
58
    settings: Optional[UserSettings] = None
Timothy J. Baek's avatar
Timothy J. Baek committed
59
    info: Optional[dict] = None
60

61
62
    oauth_sub: Optional[str] = None

63
64
65
66
67
68

####################
# Forms
####################


Timothy J. Baek's avatar
Timothy J. Baek committed
69
70
71
72
73
class UserRoleUpdateForm(BaseModel):
    id: str
    role: str


Timothy J. Baek's avatar
Timothy J. Baek committed
74
75
76
77
78
79
class UserUpdateForm(BaseModel):
    name: str
    email: str
    profile_image_url: str
    password: Optional[str] = None

80

Timothy J. Baek's avatar
Timothy J. Baek committed
81
class UsersTable:
82

Timothy J. Baek's avatar
Timothy J. Baek committed
83
    def insert_new_user(
Danny Liu's avatar
Danny Liu committed
84
        self,
85
        db: Session,
Danny Liu's avatar
Danny Liu committed
86
87
88
        id: str,
        name: str,
        email: str,
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
89
        profile_image_url: str = "/user.png",
Danny Liu's avatar
Danny Liu committed
90
        role: str = "pending",
91
        oauth_sub: Optional[str] = None,
Timothy J. Baek's avatar
Timothy J. Baek committed
92
    ) -> Optional[UserModel]:
93
94
95
96
97
98
        user = UserModel(
            **{
                "id": id,
                "name": name,
                "email": email,
                "role": role,
99
                "profile_image_url": profile_image_url,
Timothy J. Baek's avatar
Timothy J. Baek committed
100
101
102
                "last_active_at": int(time.time()),
                "created_at": int(time.time()),
                "updated_at": int(time.time()),
103
                "oauth_sub": oauth_sub,
Timothy J. Baek's avatar
Timothy J. Baek committed
104
105
            }
        )
106
107
108
109
        result = User(**user.model_dump())
        db.add(result)
        db.commit()
        db.refresh(result)
110
111
112
113
114
        if result:
            return user
        else:
            return None

115
    def get_user_by_id(self, db: Session, id: str) -> Optional[UserModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
116
        try:
117
118
119
            user = db.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
        except Exception as e:
Timothy J. Baek's avatar
Timothy J. Baek committed
120
            return None
121

122
    def get_user_by_api_key(self, db: Session, api_key: str) -> Optional[UserModel]:
123
        try:
124
125
            user = db.query(User).filter_by(api_key=api_key).first()
            return UserModel.model_validate(user)
126
127
128
        except:
            return None

129
    def get_user_by_email(self, db: Session, email: str) -> Optional[UserModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
130
        try:
131
132
            user = db.query(User).filter_by(email=email).first()
            return UserModel.model_validate(user)
Timothy J. Baek's avatar
Timothy J. Baek committed
133
        except:
134
135
            return None

136
137
138
139
140
141
142
    def get_user_by_oauth_sub(self, sub: str) -> Optional[UserModel]:
        try:
            user = User.get(User.oauth_sub == sub)
            return UserModel(**model_to_dict(user))
        except:
            return None

143
144
145
146
147
148
149
    def get_users(self, db: Session, skip: int = 0, limit: int = 50) -> List[UserModel]:
        users = (
            db.query(User)
            # .offset(skip).limit(limit)
            .all()
        )
        return [UserModel.model_validate(user) for user in users]
150

151
152
    def get_num_users(self, db: Session) -> Optional[int]:
        return db.query(User).count()
Timothy J. Baek's avatar
Timothy J. Baek committed
153

154
    def get_first_user(self, db: Session) -> UserModel:
155
        try:
156
157
            user = db.query(User).order_by(User.created_at).first()
            return UserModel.model_validate(user)
158
159
160
        except:
            return None

161
162
163
    def update_user_role_by_id(
        self, db: Session, id: str, role: str
    ) -> Optional[UserModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
164
        try:
165
166
            db.query(User).filter_by(id=id).update({"role": role})
            db.commit()
Timothy J. Baek's avatar
Timothy J. Baek committed
167

168
169
            user = db.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
Timothy J. Baek's avatar
Timothy J. Baek committed
170
171
        except:
            return None
Timothy J. Baek's avatar
Timothy J. Baek committed
172

173
    def update_user_profile_image_url_by_id(
174
        self, db: Session, id: str, profile_image_url: str
175
176
    ) -> Optional[UserModel]:
        try:
177
178
            db.query(User).filter_by(id=id).update(
                {"profile_image_url": profile_image_url}
179
            )
180
            db.commit()
181

182
183
            user = db.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
184
185
186
        except:
            return None

187
188
189
    def update_user_last_active_by_id(
        self, db: Session, id: str
    ) -> Optional[UserModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
190
        try:
191
            db.query(User).filter_by(id=id).update({"last_active_at": int(time.time())})
Timothy J. Baek's avatar
Timothy J. Baek committed
192

193
194
            user = db.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
Timothy J. Baek's avatar
Timothy J. Baek committed
195
196
197
        except:
            return None

198
    def update_user_oauth_sub_by_id(
199
        self, db: Session, id: str, oauth_sub: str
200
201
    ) -> Optional[UserModel]:
        try:
202
            db.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub})
203

204
205
            user = db.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
206
207
208
        except:
            return None

209
210
211
    def update_user_by_id(
        self, db: Session, id: str, updated: dict
    ) -> Optional[UserModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
212
        try:
213
214
            db.query(User).filter_by(id=id).update(updated)
            db.commit()
Timothy J. Baek's avatar
Timothy J. Baek committed
215

216
217
218
219
            user = db.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
            # return UserModel(**user.dict())
        except Exception as e:
Timothy J. Baek's avatar
Timothy J. Baek committed
220
221
            return None

222
    def delete_user_by_id(self, db: Session, id: str) -> bool:
223
224
        try:
            # Delete User Chats
225
            result = Chats.delete_chats_by_user_id(db, id)
226

227
228
            if result:
                # Delete User
229
230
                db.query(User).filter_by(id=id).delete()
                db.commit()
231

232
233
234
                return True
            else:
                return False
235
236
237
        except:
            return False

238
    def update_user_api_key_by_id(self, db: Session, id: str, api_key: str) -> str:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
239
        try:
240
241
            result = db.query(User).filter_by(id=id).update({"api_key": api_key})
            db.commit()
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
242
243
244
245
            return True if result == 1 else False
        except:
            return False

246
    def get_user_api_key_by_id(self, db: Session, id: str) -> Optional[str]:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
247
        try:
248
            user = db.query(User).filter_by(id=id).first()
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
249
            return user.api_key
250
        except Exception as e:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
251
252
            return None

253

254
Users = UsersTable()