send knowledge base auto disable notification

langgenius · Dec 26, 2024 · b872628 · b872628
1 parent 8339d2c
commit b872628
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 11 deletions.
diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py
@@ -69,6 +69,7 @@ def __call__(self, *args: object, **kwargs: object) -> object:
         "schedule.create_tidb_serverless_task",
         "schedule.update_tidb_serverless_status_task",
         "schedule.clean_messages",
+        "schedule.mail_clean_document_notify_task",
     ]
     day = dify_config.CELERY_BEAT_SCHEDULER_TIME
     beat_schedule = {
@@ -92,6 +93,11 @@ def __call__(self, *args: object, **kwargs: object) -> object:
             "task": "schedule.clean_messages.clean_messages",
             "schedule": timedelta(days=day),
         },
+        # every Monday
+        "mail_clean_document_notify_task": {
+            "task": "schedule.mail_clean_document_notify_task.mail_clean_document_notify_task",
+            "schedule": crontab(minute="0", hour="10", day_of_week="1"),
+        },
     }
     celery_app.conf.update(beat_schedule=beat_schedule, imports=imports)
 

diff --git a/api/schedule/mail_clean_document_notify_task.py b/api/schedule/mail_clean_document_notify_task.py
@@ -3,14 +3,17 @@
 from collections import defaultdict
 
 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 models.account import Account, Tenant, TenantAccountJoin
 from models.dataset import Dataset, DatasetAutoDisableLog
 
 
-@shared_task(queue="mail")
+@app.celery.task(queue="dataset")
 def send_document_clean_notify_task():
     """
     Async Send document clean notify mail
@@ -29,13 +32,17 @@ def send_document_clean_notify_task():
         # group by tenant_id
         dataset_auto_disable_logs_map: dict[str, list[DatasetAutoDisableLog]] = defaultdict(list)
         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)
-
+        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
@@ -45,6 +52,8 @@ def send_document_clean_notify_task():
 
             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
                 )
@@ -53,11 +62,23 @@ def send_document_clean_notify_task():
                 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>")
+                    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()
         logging.info(
             click.style("Send document clean notify mail succeeded: latency: {}".format(end_at - start_at), fg="green")
         )
     except Exception:
-        logging.exception("Send invite member mail to failed")
+        logging.exception("Send document clean notify mail failed")
diff --git a/api/templates/clean_document_job_mail_template-US.html b/api/templates/clean_document_job_mail_template-US.html
@@ -45,14 +45,14 @@
     .content ul li {
       margin-bottom: 10px;
     }
-    .cta-button {
+    .cta-button, .cta-button:hover, .cta-button:active, .cta-button:visited, .cta-button:focus {
       display: block;
       margin: 20px auto;
       padding: 10px 20px;
       background-color: #4e89f9;
-      color: #ffffff;
+      color: #ffffff !important;
       text-align: center;
-      text-decoration: none;
+      text-decoration: none !important;
       border-radius: 5px;
       width: fit-content;
     }
@@ -69,7 +69,7 @@
   <div class="email-container">
     <!-- Header -->
     <div class="header">
-      <img src="https://via.placeholder.com/150x40?text=Dify" alt="Dify Logo">
+      <img src="https://img.mailinblue.com/6365111/images/content_library/original/64cb67ca60532312c211dc72.png" alt="Dify Logo">
     </div>
 
     <!-- Content -->
@@ -78,11 +78,13 @@ <h1>Some Documents in Your Knowledge Base Have Been Disabled</h1>
       <p>Dear {{userName}},</p>
       <p>
         We're sorry for the inconvenience. To ensure optimal performance, documents 
-        that haven’t been updated or accessed in the past 7 days have been disabled in 
+        that haven’t been updated or accessed in the past 30 days have been disabled in
         your knowledge bases:
       </p>
       <ul>
-        {{knowledge_details}}
+          {% for item in knowledge_details %}
+            <li>{{ item }}</li>
+          {% endfor %}
       </ul>
       <p>You can re-enable them anytime.</p>
       <a href={{url}} class="cta-button">Re-enable in Dify</a>