|
@@ -3,14 +3,18 @@ import time
|
|
from collections import defaultdict
|
|
from collections import defaultdict
|
|
|
|
|
|
import click
|
|
import click
|
|
-from celery import shared_task # type: ignore
|
|
|
|
|
|
+from flask import render_template # type: ignore
|
|
|
|
|
|
|
|
+import app
|
|
|
|
+from configs import dify_config
|
|
|
|
+from extensions.ext_database import db
|
|
from extensions.ext_mail import mail
|
|
from extensions.ext_mail import mail
|
|
from models.account import Account, Tenant, TenantAccountJoin
|
|
from models.account import Account, Tenant, TenantAccountJoin
|
|
from models.dataset import Dataset, DatasetAutoDisableLog
|
|
from models.dataset import Dataset, DatasetAutoDisableLog
|
|
|
|
+from services.feature_service import FeatureService
|
|
|
|
|
|
|
|
|
|
-@shared_task(queue="mail")
|
|
|
|
|
|
+@app.celery.task(queue="dataset")
|
|
def send_document_clean_notify_task():
|
|
def send_document_clean_notify_task():
|
|
"""
|
|
"""
|
|
Async Send document clean notify mail
|
|
Async Send document clean notify mail
|
|
@@ -29,35 +33,58 @@ def send_document_clean_notify_task():
|
|
# group by tenant_id
|
|
# group by tenant_id
|
|
dataset_auto_disable_logs_map: dict[str, list[DatasetAutoDisableLog]] = defaultdict(list)
|
|
dataset_auto_disable_logs_map: dict[str, list[DatasetAutoDisableLog]] = defaultdict(list)
|
|
for dataset_auto_disable_log in dataset_auto_disable_logs:
|
|
for dataset_auto_disable_log in dataset_auto_disable_logs:
|
|
|
|
+ if dataset_auto_disable_log.tenant_id not in dataset_auto_disable_logs_map:
|
|
|
|
+ dataset_auto_disable_logs_map[dataset_auto_disable_log.tenant_id] = []
|
|
dataset_auto_disable_logs_map[dataset_auto_disable_log.tenant_id].append(dataset_auto_disable_log)
|
|
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():
|
|
for tenant_id, tenant_dataset_auto_disable_logs in dataset_auto_disable_logs_map.items():
|
|
- knowledge_details = []
|
|
|
|
- tenant = Tenant.query.filter(Tenant.id == tenant_id).first()
|
|
|
|
- if not tenant:
|
|
|
|
- continue
|
|
|
|
- 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:
|
|
|
|
- dataset_auto_dataset_map[dataset_auto_disable_log.dataset_id].append(
|
|
|
|
- dataset_auto_disable_log.document_id
|
|
|
|
- )
|
|
|
|
|
|
+ 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
|
|
|
|
|
|
- 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(f"<li>Knowledge base {dataset.name}: {document_count} documents</li>")
|
|
|
|
|
|
+ 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
|
|
|
|
+ )
|
|
|
|
+
|
|
|
|
+ # update notified to True
|
|
|
|
+ for dataset_auto_disable_log in tenant_dataset_auto_disable_logs:
|
|
|
|
+ dataset_auto_disable_log.notified = True
|
|
|
|
+ db.session.commit()
|
|
end_at = time.perf_counter()
|
|
end_at = time.perf_counter()
|
|
logging.info(
|
|
logging.info(
|
|
click.style("Send document clean notify mail succeeded: latency: {}".format(end_at - start_at), fg="green")
|
|
click.style("Send document clean notify mail succeeded: latency: {}".format(end_at - start_at), fg="green")
|
|
)
|
|
)
|
|
except Exception:
|
|
except Exception:
|
|
- logging.exception("Send invite member mail to failed")
|
|
|
|
|
|
+ logging.exception("Send document clean notify mail failed")
|