rag-gateway/app/service/service_token.py

from Log import logger
from app.config.config import settings
from app.models import UserModel
from app.models.token_model import TokenModel
from app.service.bisheng import BishengService
from app.service.ragflow import RagflowService


def get_bisheng_token(db, user_id: int):
    token = db.query(TokenModel).filter(TokenModel.user_id == user_id).first()
    if not token:
        return None
    return token.bisheng_token


def get_ragflow_token(db, user_id: int):
    token = db.query(TokenModel).filter(TokenModel.user_id == user_id).first()
    if not token:
        return None
    return token.ragflow_token


async def get_ragflow_new_token(db, user_id: int, app_type):
    user = db.query(UserModel).filter(UserModel.id == user_id).first()
    if not user:
        return None
    if app_type == 1:
        ragflow_service = RagflowService(settings.fwr_base_url)
        # 登录到ragflow
        try:
            ragflow_token = await ragflow_service.login(user.username, user.decrypted_password())
            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 == 2:
        bisheng_service = BishengService(settings.sgb_base_url)
        # 登录到毕昇
        try:
            bisheng_token = await bisheng_service.login(user.username, user.decrypted_password())
            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

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