rag-gateway/app/service/user.py

import time
from datetime import datetime

from app.api import pwd_context
from app.config.config import settings
from app.config.const import RAGFLOW, BISHENG, DIFY
from app.models import RoleModel, GroupModel, AgentType, role_resource_table
from app.models.menu_model import WebMenuModel, MenuCapacityModel
from app.models.user_model import UserModel
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_ragflow_new_token


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")
    # if role_key != "admin":
    #     query.filter(UserModel.creator==user_id)
    if keyword:
        query = query.filter(UserModel.group_name.like('%{}%'.format(keyword)))
    users = query.order_by(UserModel.id.desc()).limit(page_size).offset(
        (page_index - 1) * page_size).all()
    return {"total": query.count(), "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).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").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:
        # bisheng_service = BishengService(settings.sgb_base_url)
        # ragflow_service = RagflowService(settings.fwr_base_url)
        #
        # # 注册到毕昇
        # try:
        #     bisheng_info = await bisheng_service.register(user_name, password)
        # except Exception as e:
        #     logger.error(f"Failed to register with Bisheng: {str(e)}")
        #     return False
        #
        # # 注册到ragflow
        # try:
        #     ragflow_info = await ragflow_service.register(user_name, password)
        # except Exception as e:
        #     logger.error(f"Failed to register with Ragflow: {str(e)}")
        #     return False
        app_register = AppRegisterDao(db).get_apps()
        register_dict = {}
        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:
                service = DifyService(settings.dify_base_url)
            else:
                logger.error("未知注册应用---")
                continue
            try:
                name = app["id"] + str(int(time.time()))
                register_info = await service.register(name, password)
                # 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(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="000000"):
    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_ragflow_new_token(db, user_id, 1)
                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")))
        # hashed_password = pwd_context.hash(password)
        hashed_password = user.encrypted_password(new_password)

        user.password = hashed_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_by(id=user_id).first()
    await role_resource(role_set, roles, permissions, user.roles)
    for ogt in user.organizations:
        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:
            role_set.add(role.id)
            role_list.append(role.to_dict())
            for r in role.resources:
                if r.resource_type_id == "1":
                    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 = ""

    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
    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

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


async def get_user_menus(db, user_id):
    dialog_list = []
    agent_list = []
    menu_dict = {}
    res = []
    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)
        for agent in group.agents:
            agent_list.append(agent.id)
    menu_list = db.query(WebMenuModel.id, WebMenuModel.title, WebMenuModel.describe, WebMenuModel.icon, WebMenuModel.desc,
                         WebMenuModel.img, MenuCapacityModel.capacity_id, MenuCapacityModel.capacity_type).outerjoin(
        MenuCapacityModel, WebMenuModel.id == MenuCapacityModel.menu_id).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 not m.capacity_type or m.capacity_type == 1 and m.capacity_id not in dialog_list:
                break
            elif not m.capacity_type or m.capacity_type == 2 and m.capacity_id not in agent_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
        })
    return res