send knowledge base auto disable notification

langgenius · Dec 26, 2024 · dc40069 · dc40069
1 parent b872628
commit dc40069
Showing 1 changed file with 36 additions and 32 deletions.
diff --git a/api/schedule/mail_clean_document_notify_task.py b/api/schedule/mail_clean_document_notify_task.py
@@ -11,6 +11,7 @@
 from extensions.ext_mail import mail
 from models.account import Account, Tenant, TenantAccountJoin
 from models.dataset import Dataset, DatasetAutoDisableLog
+from services.feature_service import FeatureService
 
 
 @app.celery.task(queue="dataset")
@@ -37,40 +38,43 @@ def send_document_clean_notify_task():
             dataset_auto_disable_logs_map[dataset_auto_disable_log.tenant_id].append(dataset_auto_disable_log)
         url = f"{dify_config.CONSOLE_WEB_URL}/datasets"
         for tenant_id, tenant_dataset_auto_disable_logs in dataset_auto_disable_logs_map.items():
-            knowledge_details = []
-            # check tenant
-            tenant = Tenant.query.filter(Tenant.id == tenant_id).first()
-            if not tenant:
-                continue
-            # check current owner
-            current_owner_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, role="owner").first()
-            if not current_owner_join:
-                continue
-            account = Account.query.filter(Account.id == current_owner_join.account_id).first()
-            if not account:
-                continue
+            features = FeatureService.get_features(tenant_id)
+            plan = features.billing.subscription.plan
+            if plan != "sandbox":
+                knowledge_details = []
+                # check tenant
+                tenant = Tenant.query.filter(Tenant.id == tenant_id).first()
+                if not tenant:
+                    continue
+                # check current owner
+                current_owner_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, role="owner").first()
+                if not current_owner_join:
+                    continue
+                account = Account.query.filter(Account.id == current_owner_join.account_id).first()
+                if not account:
+                    continue
 
-            dataset_auto_dataset_map = {}  # type: ignore
-            for dataset_auto_disable_log in tenant_dataset_auto_disable_logs:
-                if dataset_auto_disable_log.dataset_id not in dataset_auto_dataset_map:
-                    dataset_auto_dataset_map[dataset_auto_disable_log.dataset_id] = []
-                dataset_auto_dataset_map[dataset_auto_disable_log.dataset_id].append(
-                    dataset_auto_disable_log.document_id
-                )
+                dataset_auto_dataset_map = {}  # type: ignore
+                for dataset_auto_disable_log in tenant_dataset_auto_disable_logs:
+                    if dataset_auto_disable_log.dataset_id not in dataset_auto_dataset_map:
+                        dataset_auto_dataset_map[dataset_auto_disable_log.dataset_id] = []
+                    dataset_auto_dataset_map[dataset_auto_disable_log.dataset_id].append(
+                        dataset_auto_disable_log.document_id
+                    )
 
-            for dataset_id, document_ids in dataset_auto_dataset_map.items():
-                dataset = Dataset.query.filter(Dataset.id == dataset_id).first()
-                if dataset:
-                    document_count = len(document_ids)
-                    knowledge_details.append(rf"Knowledge base {dataset.name}: {document_count} documents")
-            if knowledge_details:
-                html_content = render_template(
-                    "clean_document_job_mail_template-US.html",
-                    userName=account.email,
-                    knowledge_details=knowledge_details,
-                    url=url,
-                )
-                mail.send(to=account.email, subject="Dify Knowledge base auto disable notification", html=html_content)
+                for dataset_id, document_ids in dataset_auto_dataset_map.items():
+                    dataset = Dataset.query.filter(Dataset.id == dataset_id).first()
+                    if dataset:
+                        document_count = len(document_ids)
+                        knowledge_details.append(rf"Knowledge base {dataset.name}: {document_count} documents")
+                if knowledge_details:
+                    html_content = render_template(
+                        "clean_document_job_mail_template-US.html",
+                        userName=account.email,
+                        knowledge_details=knowledge_details,
+                        url=url,
+                    )
+                    mail.send(to=account.email, subject="Dify Knowledge base auto disable notification", html=html_content)
 
             # update notified to True
             for dataset_auto_disable_log in tenant_dataset_auto_disable_logs: