tjdify/api/services/robot_account_service.py

import logging
from datetime import UTC, datetime, timedelta
from typing import Optional

from configs import dify_config
from extensions.ext_redis import redis_client
from libs.passport import PassportService
from models import Account, AccountStatus, db
from services.account_service import AccountService

ACCESS_TOKEN_PREFIX = "access_token:"
ACCESS_TOKEN_EXPIRY = timedelta(minutes=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES)

class RobotAccountService:

    @staticmethod
    def login(email: str, password: str, *, ip_address: Optional[str] = None) -> str:
        account = AccountService.authenticate(email, password)

        if ip_address:
            AccountService.update_login_info(account=account, ip_address=ip_address)

        if account.status == AccountStatus.PENDING.value:
            account.status = AccountStatus.ACTIVE.value
            db.session.commit()

        access_token = RobotAccountService.get_account_jwt_token(account=account)
        access_token_key = RobotAccountService._get_access_token_key(email)
        redis_client.setex(access_token_key, ACCESS_TOKEN_EXPIRY, access_token)

        return access_token

    @staticmethod
    def get_account_jwt_token(account: Account) -> str:
        exp_dt = datetime.now(UTC) + ACCESS_TOKEN_EXPIRY
        exp = int(exp_dt.timestamp())
        payload = {
            "user_id": account.id,
            "exp": exp,
            "iss": dify_config.EDITION,
            "sub": "Console API Passport",
        }

        token: str = PassportService().issue(payload)
        return token

    @staticmethod
    def get_account_access_token(email) -> str:
        access_token =  redis_client.get(RobotAccountService._get_access_token_key(email))
        logging.info(f"access_token: {access_token}")
        return access_token

    @staticmethod
    def _get_access_token_key(access_token: str) -> str:
        return f"{ACCESS_TOKEN_PREFIX}{access_token}"