rag-gateway/app/service/service_token.py

from Log import logger
from app.config.config import settings
from app.config.const import BISHENG, RAGFLOW, DIFY, chat_server
from app.models import UserModel, UserAppModel, UserTokenModel
from app.models.token_model import TokenModel
from app.service.auth import UserAppDao
from app.service.bisheng import BishengService
from app.service.difyService import DifyService
from app.service.ragflow import RagflowService
from app.task.fetch_agent import update_ragflow_user_tenant


async def get_bisheng_token(db, user_id: int):

    token = db.query(TokenModel).filter(TokenModel.user_id == user_id).first()
    if not token:
        token = await UserAppDao(db).get_data_by_id(user_id, BISHENG)
        if not token:
            return None
        return token.access_token
    else:
        return token.bisheng_token


async def get_ragflow_token(db, user_id: int):
    # token = await UserAppDao(db).get_data_by_id(user_id, RAGFLOW)
    token = db.query(TokenModel).filter(TokenModel.user_id == user_id).first()
    token = db.query(UserTokenModel).filter(UserTokenModel.id == chat_server).first()
    if not token:
        token = await UserAppDao(db).get_data_by_id(user_id, RAGFLOW)
        if not token:
            return None
        return token.access_token
    else:
        return token.access_token


async def get_dify_token(db, user_id: int):

    token = await UserAppDao(db).get_data_by_id(user_id, DIFY)
    if not token:
        return None
    return token.access_token

async def get_new_token(db, user_id: int, app_type):
    # user = db.query(UserModel).filter(UserModel.id == user_id).first()
    user = db.query(UserAppModel).filter(UserAppModel.user_id == user_id, UserAppModel.app_type==app_type).first()
    if not user:
        return None
    pwd = user.password
    if app_type == RAGFLOW:
        ragflow_service = RagflowService(settings.fwr_base_url)
        # 登录到ragflow
        try:
            ragflow_token = await ragflow_service.login(user.username, user.decrypted_password(pwd))
            return ragflow_token
        except Exception as e:
            logger.error(e)
            # return Response(code=500, msg=f"Failed to login with Ragflow: {str(e)}")
            return None
    elif app_type == BISHENG:
        bisheng_service = BishengService(settings.sgb_base_url)
        # 登录到毕昇
        try:
            bisheng_token = await bisheng_service.login(user.username, user.decrypted_password(pwd))
            return bisheng_token
        except Exception as e:
            logger.error(e)
            # return Response(code=500, msg=f"Failed to login with Ragflow: {str(e)}")
            return None

    elif app_type == DIFY:
        dify_service = DifyService(settings.dify_base_url)
        # 登录到毕昇
        try:
            dify_token = await dify_service.login(user.username, user.decrypted_password(pwd))
            return dify_token
        except Exception as e:
            logger.error(e)
            # return Response(code=500, msg=f"Failed to login with Ragflow: {str(e)}")
            return None

    else:
        logger.error("错误的应用类型{}".format(app_type))


async def get_token(db, user_id: int):
    res = {i.app_type.replace("app", "token"): i.access_token for i in await UserAppDao(db).get_user_datas(user_id)}
    if not res:
        token = db.query(TokenModel).filter_by(user_id=user_id).first()
        res = {"ragflow_token": token.ragflow_token, "bisheng_token": token.bisheng_token}
    return res


async def update_user_group(db, sync_flag):
    admin_user = db.query(UserModel).filter(UserModel.permission == "admin").first()
    sync_user = db.query(UserModel).filter(UserModel.sync_flag == sync_flag).first()
    if not admin_user or not sync_user:
        logger.error("注册用户失败！--------")
        return False
    token = await get_new_token(db, admin_user.id, RAGFLOW)
    # print(token)
    if not token:
        logger.error("注册用户获取token失败！")
        return False
    user_list = db.query(UserAppModel).filter(UserAppModel.app_type==RAGFLOW).all()
    user = db.query(UserAppModel).filter(UserAppModel.user_id==sync_user.id, UserAppModel.app_type == RAGFLOW).first()
    if not user:
        logger.error("注册用户获取信息失败！")
        return False
    ragflow_service = RagflowService(settings.fwr_base_url)
    for u in user_list:
        if u.id and user.id != u.id:
            # print(user.id, u.id)
            await ragflow_service.add_user_tenant(token,u.app_id,
                                                      user.email,
                                                      user.app_id)
            await ragflow_service.add_user_tenant(token, user.app_id,
                                                  u.email,
                                                  u.app_id)
    update_ragflow_user_tenant(user.app_id)