refactor: reorganize imports and update field name in EducationApi

api/controllers/console/workspace/account.py

@@ -1,38 +1,29 @@
-import pytz
-from flask import request
-from flask_login import current_user
-from flask_restful import Resource, fields, marshal_with, reqparse
-from sqlalchemy import select
-from sqlalchemy.orm import Session
+from datetime import datetime
 
+import pytz
 from configs import dify_config
 from constants.languages import supported_language
 from controllers.console import api
-from controllers.console.auth.error import (
-    EmailAlreadyInUseError,
-    EmailChangeLimitError,
-    EmailCodeError,
-    InvalidEmailError,
-    InvalidTokenError,
-)
-from controllers.console.error import AccountInFreezeError, AccountNotFound, EmailSendIpLimitError
+from controllers.console.auth.error import (EmailAlreadyInUseError,
+                                            EmailChangeLimitError,
+                                            EmailCodeError, InvalidEmailError,
+                                            InvalidTokenError)
+from controllers.console.error import (AccountInFreezeError, AccountNotFound,
+                                       EmailSendIpLimitError)
 from controllers.console.workspace.error import (
-    AccountAlreadyInitedError,
-    CurrentPasswordIncorrectError,
-    InvalidAccountDeletionCodeError,
-    InvalidInvitationCodeError,
-    RepeatPasswordNotMatchError,
-)
-from controllers.console.wraps import (
-    account_initialization_required,
-    cloud_edition_billing_enabled,
-    enable_change_email,
-    enterprise_license_required,
-    only_edition_cloud,
-    setup_required,
-)
+    AccountAlreadyInitedError, CurrentPasswordIncorrectError,
+    InvalidAccountDeletionCodeError, InvalidInvitationCodeError,
+    RepeatPasswordNotMatchError)
+from controllers.console.wraps import (account_initialization_required,
+                                       cloud_edition_billing_enabled,
+                                       enable_change_email,
+                                       enterprise_license_required,
+                                       only_edition_cloud, setup_required)
 from extensions.ext_database import db
 from fields.member_fields import account_fields
+from flask import request
+from flask_login import current_user
+from flask_restful import Resource, fields, marshal_with, reqparse
 from libs.datetime_utils import naive_utc_now
 from libs.helper import TimestampField, email, extract_remote_ip, timezone
 from libs.login import login_required
@@ -40,7 +31,10 @@ from models import AccountIntegrate, InvitationCode
 from models.account import Account
 from services.account_service import AccountService
 from services.billing_service import BillingService
-from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError
+from services.errors.account import \
+    CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError
+from sqlalchemy import select
+from sqlalchemy.orm import Session
 
 
 class AccountInitApi(Resource):
@@ -327,6 +321,9 @@ class EducationVerifyApi(Resource):
 class EducationApi(Resource):
     status_fields = {
         "result": fields.Boolean,
+        "is_student": fields.Boolean,
+        "expire_at": TimestampField,
+        "allow_refresh": fields.Boolean,
     }
 
     @setup_required
@@ -354,7 +351,11 @@ class EducationApi(Resource):
     def get(self):
         account = current_user
 
-        return BillingService.EducationIdentity.is_active(account.id)
+        res = BillingService.EducationIdentity.status(account.id)
+        # convert expire_at to UTC timestamp from isoformat
+        if res and "expire_at" in res:
+            res["expire_at"] = datetime.fromisoformat(res["expire_at"]).astimezone(pytz.utc)
+        return res
 
 
 class EducationAutoCompleteApi(Resource):
@@ -540,3 +541,4 @@ api.add_resource(ChangeEmailResetApi, "/account/change-email/reset")
 api.add_resource(CheckEmailUnique, "/account/change-email/check-email-unique")
 # api.add_resource(AccountEmailApi, '/account/email')
 # api.add_resource(AccountEmailVerifyApi, '/account/email-verify')
+# api.add_resource(AccountEmailVerifyApi, '/account/email-verify')

api/controllers/inner_api/mail.py

@@ -1,27 +1,38 @@
-from flask_restful import (
-    Resource,  # type: ignore
-    reqparse,
-)
+from flask_restful import Resource, reqparse
 
 from controllers.console.wraps import setup_required
 from controllers.inner_api import api
-from controllers.inner_api.wraps import enterprise_inner_api_only
-from services.enterprise.mail_service import DifyMail, EnterpriseMailService
+from controllers.inner_api.wraps import billing_inner_api_only, enterprise_inner_api_only
+from tasks.mail_inner_task import send_inner_email_task
+
+_mail_parser = reqparse.RequestParser()
+_mail_parser.add_argument("to", type=str, action="append", required=True)
+_mail_parser.add_argument("subject", type=str, required=True)
+_mail_parser.add_argument("body", type=str, required=True)
+_mail_parser.add_argument("substitutions", type=dict, required=False)
 
 
-class EnterpriseMail(Resource):
-    @setup_required
-    @enterprise_inner_api_only
+class BaseMail(Resource):
+    """Shared logic for sending an inner email."""
+
     def post(self):
-        parser = reqparse.RequestParser()
-        parser.add_argument("to", type=str, action="append", required=True)
-        parser.add_argument("subject", type=str, required=True)
-        parser.add_argument("body", type=str, required=True)
-        parser.add_argument("substitutions", type=dict, required=False)
-        args = parser.parse_args()
-
-        EnterpriseMailService.send_mail(DifyMail(**args))
+        args = _mail_parser.parse_args()
+        send_inner_email_task.delay(
+            to=args["to"],
+            subject=args["subject"],
+            body=args["body"],
+            substitutions=args["substitutions"],
+        )
         return {"message": "success"}, 200
 
 
+class EnterpriseMail(BaseMail):
+    method_decorators = [setup_required, enterprise_inner_api_only]
+
+
+class BillingMail(BaseMail):
+    method_decorators = [setup_required, billing_inner_api_only]
+
+
 api.add_resource(EnterpriseMail, "/enterprise/mail")
+api.add_resource(BillingMail, "/billing/mail")

api/controllers/inner_api/wraps.py

@@ -10,6 +10,22 @@ from extensions.ext_database import db
 from models.model import EndUser
 
 
+def billing_inner_api_only(view):
+    @wraps(view)
+    def decorated(*args, **kwargs):
+        if not dify_config.INNER_API:
+            abort(404)
+
+        # get header 'X-Inner-Api-Key'
+        inner_api_key = request.headers.get("X-Inner-Api-Key")
+        if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY:
+            abort(401)
+
+        return view(*args, **kwargs)
+
+    return decorated
+
+
 def enterprise_inner_api_only(view):
     @wraps(view)
     def decorated(*args, **kwargs):

api/services/billing_service.py

@@ -123,7 +123,7 @@ class BillingService:
             return BillingService._send_request("GET", "/education/verify", params=params)
 
         @classmethod
-        def is_active(cls, account_id: str):
+        def status(cls, account_id: str):
             params = {"account_id": account_id}
             return BillingService._send_request("GET", "/education/status", params=params)
 

api/services/enterprise/mail_service.py

@@ -1,18 +0,0 @@
-from pydantic import BaseModel
-
-from tasks.mail_enterprise_task import send_enterprise_email_task
-
-
-class DifyMail(BaseModel):
-    to: list[str]
-    subject: str
-    body: str
-    substitutions: dict[str, str] = {}
-
-
-class EnterpriseMailService:
-    @classmethod
-    def send_mail(cls, mail: DifyMail):
-        send_enterprise_email_task.delay(
-            to=mail.to, subject=mail.subject, body=mail.body, substitutions=mail.substitutions
-        )

api/tasks/mail_inner_task.py

@@ -11,7 +11,7 @@ from libs.email_i18n import get_email_i18n_service
 
 
 @shared_task(queue="mail")
-def send_enterprise_email_task(to: list[str], subject: str, body: str, substitutions: Mapping[str, str]):
+def send_inner_email_task(to: list[str], subject: str, body: str, substitutions: Mapping[str, str]):
     if not mail.is_inited():
         return