diff --git a/api/.env.example b/api/.env.example index e9d34d4c4..64f6b5661 100644 --- a/api/.env.example +++ b/api/.env.example @@ -495,5 +495,5 @@ RMB_TO_USD_RATE=7.26 # 默认语言 DEFAULT_LANGUAGE=zh-Hans -# Bedrock Proxy -BEDROCK_PROXY= \ No newline at end of file +# 用户初始额度限制(二开新增配置) +ACCOUNT_TOTAL_QUOTA=30 \ No newline at end of file diff --git a/api/configs/extend/__init__.py b/api/configs/extend/__init__.py index 9d13628b9..859f2da4e 100644 --- a/api/configs/extend/__init__.py +++ b/api/configs/extend/__init__.py @@ -47,6 +47,11 @@ class ExtendInfo(BaseSettings): default="7.26", ) + ACCOUNT_TOTAL_QUOTA: Optional[decimal.Decimal] = Field( + description="用户额度初始总额度", + default="15", + ) + DEFAULT_LANGUAGE: Optional[str] = Field( description="默认语言", default="zh-Hans", diff --git a/api/controllers/console/auth/register_extend.py b/api/controllers/console/auth/register_extend.py index 425d857e6..06b7701ab 100644 --- a/api/controllers/console/auth/register_extend.py +++ b/api/controllers/console/auth/register_extend.py @@ -11,6 +11,7 @@ from extensions.ext_database import db from libs.login import login_required from models import Account from models.account import AccountStatus +from models.account_money_extend import AccountMoneyExtend from services.account_service import AccountService, TenantService from services.account_service_extend import TenantExtendService @@ -49,6 +50,15 @@ class AdminRegisterApi(Resource): is_setup=True, ) + # extend begin:初始化用户额度数据 + account_money_add = AccountMoneyExtend( + account_id=account.id, + used_quota=0, + total_quota=dify_config.ACCOUNT_TOTAL_QUOTA, + ) + db.session.add(account_money_add) + # extend end:初始化用户额度数据 + account.last_login_ip = "" account.status = AccountStatus.ACTIVE.value account.initialized_at = datetime.now(UTC).replace(tzinfo=None) diff --git a/api/controllers/console/workspace/account_extend.py b/api/controllers/console/workspace/account_extend.py index d3a3b0252..0a028f33d 100644 --- a/api/controllers/console/workspace/account_extend.py +++ b/api/controllers/console/workspace/account_extend.py @@ -1,6 +1,7 @@ from flask_login import current_user from flask_restful import Resource, marshal_with +from configs import dify_config from controllers.console import api from controllers.console.workspace.workspace import account_initialization_required, setup_required from extensions.ext_database import db @@ -19,7 +20,7 @@ class AccountMoneyApi(Resource): db.session.query(AccountMoneyExtend).filter(AccountMoneyExtend.account_id == current_user.id).first() ) if not account_money: - return {"total_quota": "15", "used_quota": "0"}, 200 + return {"total_quota": dify_config.ACCOUNT_TOTAL_QUOTA, "used_quota": "0"}, 200 return account_money diff --git a/api/events/event_handlers/update_account_money_when_messaeg_created_extend.py b/api/events/event_handlers/update_account_money_when_messaeg_created_extend.py index 601de7eec..d7e236b13 100644 --- a/api/events/event_handlers/update_account_money_when_messaeg_created_extend.py +++ b/api/events/event_handlers/update_account_money_when_messaeg_created_extend.py @@ -40,7 +40,7 @@ def handle(sender, **kwargs): account_money_add = AccountMoneyExtend( account_id=payerId, used_quota=price, - total_quota=15, # TODO 初始总额度这里到时候默认15要改 + total_quota=dify_config.ACCOUNT_TOTAL_QUOTA, ) db.session.add(account_money_add) diff --git a/api/services/account_service.py b/api/services/account_service.py index 8329ef364..e9ee342b7 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -31,6 +31,7 @@ from models.account import ( TenantAccountRole, TenantStatus, ) +from models.account_money_extend import AccountMoneyExtend from models.model import DifySetup from services.billing_service import BillingService from services.errors.account import ( @@ -857,6 +858,15 @@ class RegisterService: is_setup=True, ) + # extend begin:初始化用户额度数据 + account_money_add = AccountMoneyExtend( + account_id=account.id, + used_quota=0, + total_quota=dify_config.ACCOUNT_TOTAL_QUOTA, + ) + db.session.add(account_money_add) + # extend end:初始化用户额度数据 + account.last_login_ip = ip_address account.initialized_at = datetime.now(UTC).replace(tzinfo=None) diff --git a/api/tasks/extend/update_account_money_when_workflow_node_execution_created_extend.py b/api/tasks/extend/update_account_money_when_workflow_node_execution_created_extend.py index 10dfbd6b9..da3fa3d9c 100644 --- a/api/tasks/extend/update_account_money_when_workflow_node_execution_created_extend.py +++ b/api/tasks/extend/update_account_money_when_workflow_node_execution_created_extend.py @@ -64,7 +64,7 @@ def update_account_money_when_workflow_node_execution_created_extend(self, workf account_money_add = AccountMoneyExtend( account_id=payerId, used_quota=price, - total_quota=15, # TODO 初始总额度这里到时候默认15要改 + total_quota=dify_config.ACCOUNT_TOTAL_QUOTA, ) db.session.add(account_money_add)