users.py

from pydantic import BaseModel, ConfigDict, parse_obj_as
from typing import List, Union, Optional
import time

from sqlalchemy import String, Column, BigInteger, Text

from utils.misc import get_gravatar_url

from apps.webui.internal.db import Base, JSONField, Session
from apps.webui.models.chats import Chats

####################
# User DB Schema
####################


class User(Base):
    __tablename__ = "user"

    id = Column(String, primary_key=True)
    name = Column(String)
    email = Column(String)
    role = Column(String)
    profile_image_url = Column(String)

    last_active_at = Column(BigInteger)
    updated_at = Column(BigInteger)
    created_at = Column(BigInteger)

    api_key = Column(String, nullable=True, unique=True)
    settings = Column(JSONField, nullable=True)
    info = Column(JSONField, nullable=True)

    oauth_sub = Column(Text, unique=True)


class UserSettings(BaseModel):
    ui: Optional[dict] = {}
    model_config = ConfigDict(extra="allow")
    pass


class UserModel(BaseModel):
    id: str
    name: str
    email: str
    role: str = "pending"
    profile_image_url: str

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

    api_key: Optional[str] = None
    settings: Optional[UserSettings] = None
    info: Optional[dict] = None

    oauth_sub: Optional[str] = None

    model_config = ConfigDict(from_attributes=True)


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


class UserRoleUpdateForm(BaseModel):
    id: str
    role: str


class UserUpdateForm(BaseModel):
    name: str
    email: str
    profile_image_url: str
    password: Optional[str] = None


class UsersTable:

    def insert_new_user(
        self,
        id: str,
        name: str,
        email: str,
        profile_image_url: str = "/user.png",
        role: str = "pending",
        oauth_sub: Optional[str] = None,
    ) -> Optional[UserModel]:
        user = UserModel(
            **{
                "id": id,
                "name": name,
                "email": email,
                "role": role,
                "profile_image_url": profile_image_url,
                "last_active_at": int(time.time()),
                "created_at": int(time.time()),
                "updated_at": int(time.time()),
                "oauth_sub": oauth_sub,
            }
        )
        result = User(**user.model_dump())
        Session.add(result)
        Session.commit()
        Session.refresh(result)
        if result:
            return user
        else:
            return None

    def get_user_by_id(self, id: str) -> Optional[UserModel]:
        try:
            user = Session.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
        except Exception as e:
            return None

    def get_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
        try:
            user = Session.query(User).filter_by(api_key=api_key).first()
            return UserModel.model_validate(user)
        except:
            return None

    def get_user_by_email(self, email: str) -> Optional[UserModel]:
        try:
            user = Session.query(User).filter_by(email=email).first()
            return UserModel.model_validate(user)
        except:
            return None

    def get_user_by_oauth_sub(self, sub: str) -> Optional[UserModel]:
        try:
            user = Session.query(User).filter_by(oauth_sub=sub).first()
            return UserModel.model_validate(user)
        except:
            return None

    def get_users(self, skip: int = 0, limit: int = 50) -> List[UserModel]:
        users = (
            Session.query(User)
            # .offset(skip).limit(limit)
            .all()
        )
        return [UserModel.model_validate(user) for user in users]

    def get_num_users(self) -> Optional[int]:
        return Session.query(User).count()

    def get_first_user(self) -> UserModel:
        try:
            user = Session.query(User).order_by(User.created_at).first()
            return UserModel.model_validate(user)
        except:
            return None

    def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]:
        try:
            Session.query(User).filter_by(id=id).update({"role": role})
            Session.commit()

            user = Session.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
        except:
            return None

    def update_user_profile_image_url_by_id(
        self, id: str, profile_image_url: str
    ) -> Optional[UserModel]:
        try:
            Session.query(User).filter_by(id=id).update(
                {"profile_image_url": profile_image_url}
            )
            Session.commit()

            user = Session.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
        except:
            return None

    def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]:
        try:
            Session.query(User).filter_by(id=id).update(
                {"last_active_at": int(time.time())}
            )

            user = Session.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
        except:
            return None

    def update_user_oauth_sub_by_id(
        self, id: str, oauth_sub: str
    ) -> Optional[UserModel]:
        try:
            Session.query(User).filter_by(id=id).update({"oauth_sub": oauth_sub})

            user = Session.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
        except:
            return None

    def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]:
        try:
            Session.query(User).filter_by(id=id).update(updated)
            Session.commit()

            user = Session.query(User).filter_by(id=id).first()
            return UserModel.model_validate(user)
            # return UserModel(**user.dict())
        except Exception as e:
            return None

    def delete_user_by_id(self, id: str) -> bool:
        try:
            # Delete User Chats
            result = Chats.delete_chats_by_user_id(id)

            if result:
                # Delete User
                Session.query(User).filter_by(id=id).delete()
                Session.commit()

                return True
            else:
                return False
        except:
            return False

    def update_user_api_key_by_id(self, id: str, api_key: str) -> str:
        try:
            result = Session.query(User).filter_by(id=id).update({"api_key": api_key})
            Session.commit()
            return True if result == 1 else False
        except:
            return False

    def get_user_api_key_by_id(self, id: str) -> Optional[str]:
        try:
            user = Session.query(User).filter_by(id=id).first()
            return user.api_key
        except Exception as e:
            return None


Users = UsersTable()