feat: add queue to celery task (#688)

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.
+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.

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

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

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.

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.

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.

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

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

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

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

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

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

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

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

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

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

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

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

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