feat: add queue to celery task (#688)

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