From 0c18cab111fc41c2e8b2d2f06eb34adeea2ec014 Mon Sep 17 00:00:00 2001
From: takatost <takatost@users.noreply.github.com>
Date: Mon, 31 Jul 2023 13:13:08 +0800
Subject: [PATCH] feat: add queue to celery task (#688)

---
 api/README.md                                   | 2 +-
 api/docker/entrypoint.sh                        | 3 ++-
 api/tasks/add_document_to_index_task.py         | 2 +-
 api/tasks/clean_dataset_task.py                 | 2 +-
 api/tasks/clean_document_task.py                | 2 +-
 api/tasks/clean_notion_document_task.py         | 2 +-
 api/tasks/create_segment_to_index_task.py       | 2 +-
 api/tasks/deal_dataset_vector_index_task.py     | 2 +-
 api/tasks/document_indexing_sync_task.py        | 2 +-
 api/tasks/document_indexing_task.py             | 2 +-
 api/tasks/document_indexing_update_task.py      | 2 +-
 api/tasks/enable_segment_to_index_task.py       | 2 +-
 api/tasks/generate_conversation_summary_task.py | 2 +-
 api/tasks/mail_invite_member_task.py            | 2 +-
 api/tasks/recover_document_indexing_task.py     | 2 +-
 api/tasks/remove_document_from_index_task.py    | 2 +-
 api/tasks/remove_segment_from_index_task.py     | 2 +-
 api/tasks/update_segment_index_task.py          | 2 +-
 api/tasks/update_segment_keyword_index_task.py  | 2 +-
 19 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/api/README.md b/api/README.md
index 924374e5d3..0f487b2248 100644
--- a/api/README.md
+++ b/api/README.md
@@ -38,4 +38,4 @@
    flask run --host 0.0.0.0 --port=5001 --debug
    ```
 7. Setup your application by visiting http://localhost:5001/console/api/setup or other apis...
-8. If you need to debug local async processing, you can run `celery -A app.celery worker`, celery can do dataset importing and other async tasks.
\ No newline at end of file
+8. If you need to debug local async processing, you can run `celery -A app.celery worker -Q dataset,generation,mail`, celery can do dataset importing and other async tasks.
\ No newline at end of file
diff --git a/api/docker/entrypoint.sh b/api/docker/entrypoint.sh
index 50b8cbd86a..6d60e3d792 100644
--- a/api/docker/entrypoint.sh
+++ b/api/docker/entrypoint.sh
@@ -8,7 +8,8 @@ if [[ "${MIGRATION_ENABLED}" == "true" ]]; then
 fi
 
 if [[ "${MODE}" == "worker" ]]; then
-  celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} -c ${CELERY_WORKER_AMOUNT:-1} --loglevel INFO
+  celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} -c ${CELERY_WORKER_AMOUNT:-1} --loglevel INFO \
+    -Q ${CELERY_QUEUES:-dataset,generation,mail}
 else
   if [[ "${DEBUG}" == "true" ]]; then
     flask run --host=${DIFY_BIND_ADDRESS:-0.0.0.0} --port=${DIFY_PORT:-5001} --debug
diff --git a/api/tasks/add_document_to_index_task.py b/api/tasks/add_document_to_index_task.py
index db802d208a..20d54d97d4 100644
--- a/api/tasks/add_document_to_index_task.py
+++ b/api/tasks/add_document_to_index_task.py
@@ -14,7 +14,7 @@ from models.dataset import DocumentSegment
 from models.dataset import Document as DatasetDocument
 
 
-@shared_task
+@shared_task(queue='dataset')
 def add_document_to_index_task(dataset_document_id: str):
     """
     Async Add document to index
diff --git a/api/tasks/clean_dataset_task.py b/api/tasks/clean_dataset_task.py
index 14e17220b4..bd40d20c4e 100644
--- a/api/tasks/clean_dataset_task.py
+++ b/api/tasks/clean_dataset_task.py
@@ -10,7 +10,7 @@ from models.dataset import DocumentSegment, Dataset, DatasetKeywordTable, Datase
     AppDatasetJoin, Document
 
 
-@shared_task
+@shared_task(queue='dataset')
 def clean_dataset_task(dataset_id: str, tenant_id: str, indexing_technique: str, index_struct: str):
     """
     Clean dataset when dataset deleted.
diff --git a/api/tasks/clean_document_task.py b/api/tasks/clean_document_task.py
index 650037b6eb..2f37e5ff6d 100644
--- a/api/tasks/clean_document_task.py
+++ b/api/tasks/clean_document_task.py
@@ -9,7 +9,7 @@ from extensions.ext_database import db
 from models.dataset import DocumentSegment, Dataset
 
 
-@shared_task
+@shared_task(queue='dataset')
 def clean_document_task(document_id: str, dataset_id: str):
     """
     Clean document when document deleted.
diff --git a/api/tasks/clean_notion_document_task.py b/api/tasks/clean_notion_document_task.py
index a516fd23c2..d68b257c4c 100644
--- a/api/tasks/clean_notion_document_task.py
+++ b/api/tasks/clean_notion_document_task.py
@@ -10,7 +10,7 @@ from extensions.ext_database import db
 from models.dataset import DocumentSegment, Dataset, Document
 
 
-@shared_task
+@shared_task(queue='dataset')
 def clean_notion_document_task(document_ids: List[str], dataset_id: str):
     """
     Clean document when document deleted.
diff --git a/api/tasks/create_segment_to_index_task.py b/api/tasks/create_segment_to_index_task.py
index 0f7bf6ebaa..6e0af2d933 100644
--- a/api/tasks/create_segment_to_index_task.py
+++ b/api/tasks/create_segment_to_index_task.py
@@ -14,7 +14,7 @@ from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment
 
 
-@shared_task
+@shared_task(queue='dataset')
 def create_segment_to_index_task(segment_id: str, keywords: Optional[List[str]] = None):
     """
     Async create segment to index
diff --git a/api/tasks/deal_dataset_vector_index_task.py b/api/tasks/deal_dataset_vector_index_task.py
index 7eac3b3271..a8690647c2 100644
--- a/api/tasks/deal_dataset_vector_index_task.py
+++ b/api/tasks/deal_dataset_vector_index_task.py
@@ -11,7 +11,7 @@ from models.dataset import DocumentSegment, Dataset
 from models.dataset import Document as DatasetDocument
 
 
-@shared_task
+@shared_task(queue='dataset')
 def deal_dataset_vector_index_task(dataset_id: str, action: str):
     """
     Async deal dataset from index
diff --git a/api/tasks/document_indexing_sync_task.py b/api/tasks/document_indexing_sync_task.py
index dd50284f32..7cb388c667 100644
--- a/api/tasks/document_indexing_sync_task.py
+++ b/api/tasks/document_indexing_sync_task.py
@@ -14,7 +14,7 @@ from models.dataset import Document, Dataset, DocumentSegment
 from models.source import DataSourceBinding
 
 
-@shared_task
+@shared_task(queue='dataset')
 def document_indexing_sync_task(dataset_id: str, document_id: str):
     """
     Async update document
diff --git a/api/tasks/document_indexing_task.py b/api/tasks/document_indexing_task.py
index 328dfdfc45..451e6f9b35 100644
--- a/api/tasks/document_indexing_task.py
+++ b/api/tasks/document_indexing_task.py
@@ -11,7 +11,7 @@ from extensions.ext_database import db
 from models.dataset import Document
 
 
-@shared_task
+@shared_task(queue='dataset')
 def document_indexing_task(dataset_id: str, document_ids: list):
     """
     Async process document
diff --git a/api/tasks/document_indexing_update_task.py b/api/tasks/document_indexing_update_task.py
index ae9e2c3cda..c3fbe7172c 100644
--- a/api/tasks/document_indexing_update_task.py
+++ b/api/tasks/document_indexing_update_task.py
@@ -12,7 +12,7 @@ from extensions.ext_database import db
 from models.dataset import Document, Dataset, DocumentSegment
 
 
-@shared_task
+@shared_task(queue='dataset')
 def document_indexing_update_task(dataset_id: str, document_id: str):
     """
     Async update document
diff --git a/api/tasks/enable_segment_to_index_task.py b/api/tasks/enable_segment_to_index_task.py
index 5553420196..76fa3de447 100644
--- a/api/tasks/enable_segment_to_index_task.py
+++ b/api/tasks/enable_segment_to_index_task.py
@@ -13,7 +13,7 @@ from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment
 
 
-@shared_task
+@shared_task(queue='dataset')
 def enable_segment_to_index_task(segment_id: str):
     """
     Async enable segment to index
diff --git a/api/tasks/generate_conversation_summary_task.py b/api/tasks/generate_conversation_summary_task.py
index fcdd1ef6d8..2862a871d7 100644
--- a/api/tasks/generate_conversation_summary_task.py
+++ b/api/tasks/generate_conversation_summary_task.py
@@ -10,7 +10,7 @@ from extensions.ext_database import db
 from models.model import Conversation, Message
 
 
-@shared_task
+@shared_task(queue='generation')
 def generate_conversation_summary_task(conversation_id: str):
     """
     Async Generate conversation summary
diff --git a/api/tasks/mail_invite_member_task.py b/api/tasks/mail_invite_member_task.py
index 94145b9144..907e62731e 100644
--- a/api/tasks/mail_invite_member_task.py
+++ b/api/tasks/mail_invite_member_task.py
@@ -8,7 +8,7 @@ from flask import current_app
 from extensions.ext_mail import mail
 
 
-@shared_task
+@shared_task(queue='mail')
 def send_invite_member_mail_task(to: str, token: str, inviter_name: str, workspace_id: str, workspace_name: str):
     """
     Async Send invite member mail
diff --git a/api/tasks/recover_document_indexing_task.py b/api/tasks/recover_document_indexing_task.py
index bde8541bea..a9917da9b2 100644
--- a/api/tasks/recover_document_indexing_task.py
+++ b/api/tasks/recover_document_indexing_task.py
@@ -10,7 +10,7 @@ from extensions.ext_database import db
 from models.dataset import Document
 
 
-@shared_task
+@shared_task(queue='dataset')
 def recover_document_indexing_task(dataset_id: str, document_id: str):
     """
     Async recover document
diff --git a/api/tasks/remove_document_from_index_task.py b/api/tasks/remove_document_from_index_task.py
index 1635d38bad..36ec02e48e 100644
--- a/api/tasks/remove_document_from_index_task.py
+++ b/api/tasks/remove_document_from_index_task.py
@@ -11,7 +11,7 @@ from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment, Document
 
 
-@shared_task
+@shared_task(queue='dataset')
 def remove_document_from_index_task(document_id: str):
     """
     Async Remove document from index
diff --git a/api/tasks/remove_segment_from_index_task.py b/api/tasks/remove_segment_from_index_task.py
index 821d7dc3a7..b0f118649d 100644
--- a/api/tasks/remove_segment_from_index_task.py
+++ b/api/tasks/remove_segment_from_index_task.py
@@ -11,7 +11,7 @@ from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment
 
 
-@shared_task
+@shared_task(queue='dataset')
 def remove_segment_from_index_task(segment_id: str):
     """
     Async Remove segment from index
diff --git a/api/tasks/update_segment_index_task.py b/api/tasks/update_segment_index_task.py
index cf36793919..40089ad3e4 100644
--- a/api/tasks/update_segment_index_task.py
+++ b/api/tasks/update_segment_index_task.py
@@ -14,7 +14,7 @@ from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment
 
 
-@shared_task
+@shared_task(queue='dataset')
 def update_segment_index_task(segment_id: str, keywords: Optional[List[str]] = None):
     """
     Async update segment index
diff --git a/api/tasks/update_segment_keyword_index_task.py b/api/tasks/update_segment_keyword_index_task.py
index 831076e708..de3dfd10df 100644
--- a/api/tasks/update_segment_keyword_index_task.py
+++ b/api/tasks/update_segment_keyword_index_task.py
@@ -14,7 +14,7 @@ from extensions.ext_redis import redis_client
 from models.dataset import DocumentSegment
 
 
-@shared_task
+@shared_task(queue='dataset')
 def update_segment_keyword_index_task(segment_id: str):
     """
     Async update segment index