diff --git a/api/README.md b/api/README.md index 924374e5d..0f487b224 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 50b8cbd86..6d60e3d79 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 db802d208..20d54d97d 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 14e17220b..bd40d20c4 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 650037b6e..2f37e5ff6 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 a516fd23c..d68b257c4 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 0f7bf6eba..6e0af2d93 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 7eac3b327..a8690647c 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 dd50284f3..7cb388c66 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 328dfdfc4..451e6f9b3 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 ae9e2c3cd..c3fbe7172 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 555342019..76fa3de44 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 fcdd1ef6d..2862a871d 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 94145b914..907e62731 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 bde8541be..a9917da9b 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 1635d38ba..36ec02e48 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 821d7dc3a..b0f118649 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 cf3679391..40089ad3e 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 831076e70..de3dfd10d 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