rag-gateway/app/service/user.py

import time
from datetime import datetime

from app.api import pwd_context
from app.api.dialog import dialog_list
from app.config.config import settings
from app.config.const import RAGFLOW, BISHENG, DIFY, USER_STATSU_DELETE, ROLE_STATUS_ON, DEPT_STATUS_ON
from app.models import RoleModel, GroupModel, AgentType, role_resource_table, DialogModel, OrganizationModel, \
    ResourceModel
from app.models.menu_model import WebMenuModel, MenuCapacityModel
from app.models.user_model import UserModel, UserAppModel
from Log import logger
from app.service.auth import UserAppDao
from app.service.bisheng import BishengService
from app.service.v2.app_register import AppRegisterDao
from app.service.difyService import DifyService
from app.service.ragflow import RagflowService
from app.service.service_token import get_ragflow_token, get_bisheng_token, get_new_token, get_dify_token
from app.utils.password_handle import generate_password


async def get_user_list(db, page_index: int, page_size: int, keyword: str, role_key: str, user_id):
    query = db.query(UserModel).filter(UserModel.permission != "admin", UserModel.status != USER_STATSU_DELETE)
    # if role_key != "admin":
    #     query.filter(UserModel.creator==user_id)
    if keyword:
        query = query.filter(UserModel.username.like('%{}%'.format(keyword)))
    total = query.count()
    users = query.order_by(UserModel.id.desc()).limit(page_size).offset(
        (page_index - 1) * page_size).all()
    return {"total": total, "rows": [user.to_json() for user in users]}


async def edit_user_status(db, status: str, user_id: int):
    try:
        db.query(UserModel).filter(UserModel.id == user_id, UserModel.status != USER_STATSU_DELETE).update({"status": status})
        db.commit()
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True


async def delete_user_data(db, user_id: str):
    try:
        db.query(UserModel).filter(UserModel.id == user_id, UserModel.permission != "admin").update({"status": USER_STATSU_DELETE})
        db.commit()
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True


async def create_user(db, user_name, email, phone, login_name, password, roles, groups, user_id):
    try:
        app_register = AppRegisterDao(db).get_apps()
        register_dict = {}
        token = ""
        app_password = generate_password()
        for app in app_register:
            if app["id"] == RAGFLOW:
                service = RagflowService(settings.fwr_base_url)
            elif app["id"] == BISHENG:
                service = BishengService(settings.sgb_base_url)
            elif app["id"] == DIFY:
                token = await get_dify_token(db, user_id)
                service = DifyService(settings.dify_base_url)
            else:
                logger.error("未知注册应用---")
                continue
            try:
                name = app["id"] + str(int(time.time()))
                register_info = await service.register(name, app_password, token)
                # print(register_info)
                register_dict[app['id']] = {"id": register_info.get("id"), "name": name,
                                            "email": register_info.get("email")}
            except Exception as e:
                logger.error(e)
                return False

        # 存储用户信息
        hashed_password = pwd_context.hash(password)

        user_model = UserModel(username=user_name, hashed_password=hashed_password, email=email,
                               ## ragflow_id=ragflow_info.get("id"),bisheng_id=bisheng_info.get("user_id"),
                               phone=phone, login_name=login_name)
        pwd = user_model.encrypted_password(app_password)
        user_model.roles = [db.get(RoleModel, roleId) for roleId in roles]
        user_model.password = pwd
        if groups:
            user_model.groups = [db.get(GroupModel, groupId) for groupId in groups]
        user_model.creator = user_id
        db.add(user_model)
        db.commit()
        db.refresh(user_model)
        u_id = user_model.id
        for k, v in register_dict.items():
            await UserAppDao(db).update_and_insert_data(v.get("name"), pwd, v.get("email"), u_id, str(v.get("id")), k)
    except Exception as e:
        logger.error(e)
        # db.rollback()
        return False
    return True


async def edit_user_data(db, user_id, email, phone, login_name, roles, groups):
    try:
        user = db.query(UserModel).filter(UserModel.id == user_id).first()
        user.email = email
        user.phone = phone
        user.login_name = login_name
        user.email = email
        user.updated_at = datetime.now()
        user.roles = [db.get(RoleModel, roleId) for roleId in roles]
        # user.groups = [db.get(GroupModel, groupId) for groupId in groups]
        db.commit()
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True


async def edit_user_pwd(db, user_id, current_user_id, new_password="basic123456"):
    try:
        user = db.query(UserModel).filter(UserModel.id == user_id).first()
        # pwd = user.decrypted_password()
        # for app in AppRegisterDao(db).get_apps():
        #     if app.get("id") == RAGFLOW:
        #         token = await get_new_token(db, user_id, app.get("id"))
        #         ragflow_service = RagflowService(settings.fwr_base_url)
        #         await ragflow_service.set_user_password(token, pwd, new_password)
        #     elif app.get("id") == BISHENG:
        #         token = await get_bisheng_token(db, current_user_id)
        #         bisheng_service = BishengService(settings.sgb_base_url)
        #         await bisheng_service.change_password_public(token, user.username, pwd, new_password)
        #     else:
        #         logger.error("注册未知应用：{}".format(app.get("id")))
        user.hashed_password = pwd_context.hash(new_password)
        # user.password = user.encrypted_password(new_password)
        db.commit()
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True

async def change_user_pwd(db, user_id, new_password):
    try:
        user = db.query(UserModel).filter(UserModel.id == user_id).first()
        user.hashed_password = pwd_context.hash(new_password)
        db.commit()
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True



async def get_user_info(db, user_id):
    permissions = set()
    roles = []
    dept = []
    role_set = set()
    dept_set = set()
    user = db.query(UserModel).filter(UserModel.id==user_id,UserModel.status !=USER_STATSU_DELETE).first()
    await role_resource(role_set, roles, permissions, user.roles)
    for ogt in user.organizations:
        if ogt.status != DEPT_STATUS_ON:
            continue
        dept.append(ogt.to_json())
        if ogt.roles:
            await role_resource(role_set, roles, permissions, user.roles)
        parent_ogt = ogt.parent
        while parent_ogt:
            if parent_ogt.id not in dept_set:
                await role_resource(role_set, roles, permissions, parent_ogt.roles)
                dept_set.add(parent_ogt.id)
            parent_ogt = parent_ogt.parent
    if user.permission == "admin":
        permissions = ["*:*:*"]

    return {"permissions": list(permissions), "dept": dept, "roles": roles, "user": user.to_dict()}


async def role_resource(role_set, role_list, permissions, roles):
    for role in roles:
        if role.id not in role_set and role.status == ROLE_STATUS_ON:
            role_set.add(role.id)
            role_list.append(role.to_dict())
            for r in role.resources:
                if r.resource_type_id == "1" and r.status==DEPT_STATUS_ON:
                    permissions.add(r.perms)


async def get_user_routers(db, user_id):
    permissions = {}
    role_set = set()
    dept_set = set()
    user = db.query(UserModel).filter_by(id=user_id).first()
    parent_id = ""
    tmp_dit = {}
    if user.permission == "admin":
        dept_list = db.query(ResourceModel).filter(ResourceModel.status==DEPT_STATUS_ON, ResourceModel.resource_type_id != "1").all()
        for dept in dept_list:
            if not dept.resource_id:
                parent_id = dept.id
                continue
            permissions[dept.id] = dept.to_router_dict()
    else:
        async def role_resource(role_set, permissions, roles):
            nonlocal parent_id
            for role in roles:
                if role.id not in role_set and role.status == ROLE_STATUS_ON:
                    role_set.add(role.id)

                    for r in role.resources:
                        if r.resource_type_id != "1" and r.status==DEPT_STATUS_ON:
                            if not r.resource_id:
                                parent_id = r.id
                                continue
                            permissions[r.id] = r.to_router_dict()

        await role_resource(role_set, permissions, user.roles)
        for ogt in user.organizations:
            if ogt.status != DEPT_STATUS_ON:
                continue
            if ogt.roles:
                await role_resource(role_set, permissions, user.roles)
            parent_ogt = ogt.parent
            while parent_ogt:
                if parent_ogt.id not in dept_set:
                    await role_resource(role_set, permissions, parent_ogt.roles)
                    dept_set.add(parent_ogt.id)
                    parent_ogt = parent_ogt.parent
                else:
                    break

    for permission in permissions.values():
        tmp_dit[permission["parentId"]] = tmp_dit.get(permission["parentId"], []) + [permission]

    def get_child(parent_id):
        res = permissions[parent_id]
        res["children"] = [get_child(i["id"]) for i in tmp_dit.get(parent_id, [])]
        return res

    return {"routers": [get_child(i["id"]) for i in tmp_dit.get(parent_id, [])]}


async def get_user_menus(db, user_id):
    dialog_list = []
    menu_dict = {}
    res = []
    # print(user_id)
    user = db.query(UserModel).filter_by(id=user_id).first()
    for group in user.groups:
        for dialog in group.dialogs:
            dialog_list.append(dialog.id)
    # print(dialog_list)
    menu_list = db.query(WebMenuModel.id, WebMenuModel.title, WebMenuModel.describe, WebMenuModel.icon, WebMenuModel.desc,WebMenuModel.rank,
                         WebMenuModel.img, MenuCapacityModel.capacity_id, MenuCapacityModel.capacity_type, MenuCapacityModel.chat_id.label("agentId")).outerjoin(
        MenuCapacityModel, WebMenuModel.id == MenuCapacityModel.menu_id).outerjoin(
        DialogModel, MenuCapacityModel.capacity_id == DialogModel.id).filter(DialogModel.status=="1").all()

    for menu in menu_list:
        menu_dict[menu.id] = menu_dict.get(menu.id, []) + [menu]
    # print(dialog_list)
    for menus in menu_dict.values():

        for m in menus:
            if user.permission == "admin":
                continue
            elif m.capacity_id not in dialog_list:
                break
        else:
            res.append({
            'id': menus[0].id,
            'title': menus[0].title,
            'icon': menus[0].icon,
            'img': menus[0].img,
            'desc': menus[0].desc,
            'dialog': menus[0].describe,
            'agentId': menus[0].agentId,
            'rank': menus[0].rank
        })
    return sorted(res, key=lambda x: x['rank'], reverse=True)


async def get_user_permission(db,  user_id):
    res = {}
    permissions = {}
    role_set = set()
    dept_set = set()
    agent_dict = {}
    knowledge_dict = {}
    user = db.query(UserModel).filter_by(id=user_id).first()
    parent_id = ""
    # print(111111111111111)
    async def role_resource(role_set, permissions, roles):
        nonlocal parent_id
        for role in roles:
            if role.id not in role_set:
                role_set.add(role.id)

                for r in role.resources:
                    if r.resource_type_id != "1":
                        if not r.resource_id:
                            parent_id = r.id
                            continue
                        permissions[r.id] = r.to_router_dict()

    await role_resource(role_set, permissions, user.roles)
    for ogt in user.organizations:
        if ogt.roles:
            await role_resource(role_set, permissions, user.roles)
        parent_ogt = ogt.parent
        while parent_ogt:

            if parent_ogt.id not in dept_set:
                await role_resource(role_set, permissions, parent_ogt.roles)
                dept_set.add(parent_ogt.id)

                parent_ogt = parent_ogt.parent
            else:
                break

    tmp_dit = {}
    for permission in permissions.values():
        tmp_dit[permission["parentId"]] = tmp_dit.get(permission["parentId"], []) + [permission]

    def get_child(parent_id):
        res = permissions[parent_id]
        res["children"] = [get_child(i["id"]) for i in tmp_dit.get(parent_id, [])]
        return res

    res["menus"] = [get_child(i["id"]) for i in tmp_dit.get(parent_id, [])]
    for group in user.groups:
        for klg in group.knowledges:
            knowledge_dict[klg.id] = klg.to_json()
        for dialog in group.dialogs:
            agent_dict[dialog.id] = dialog.to_json()
    res["knowledge"] = list(knowledge_dict.values())
    res["dialog"] = list(agent_dict.values())
    return res


async def get_user_dept(db, user_id):
    res = {}
    user = db.query(UserModel).filter_by(id=user_id).first()
    res["rows"] = [i.to_dict() for i in user.organizations]
    return res


async def user_data_service(db, user_id):
    user = db.query(UserModel).filter_by(id=user_id).first()


    return {"roles": [i.to_dict() for i in user.roles], "user": user.to_dict()}


async def edit_user_dept(db, user_id, dept_list):
    try:
        user = db.query(UserModel).filter(UserModel.id == user_id).first()
        user.organizations = [db.get(OrganizationModel, deptId) for deptId in dept_list]
        db.commit()
    except Exception as e:
        logger.error(e)
        db.rollback()
        return False
    return True