style: Refactor import statements and improve code formatting in embedding.py

Author: liuruibin

apps/knowledge/task/embedding.py

@@ -8,8 +8,12 @@
 from django.utils.translation import gettext_lazy as _
 
 from common.config.embedding_config import ModelManage
-from common.event.listener_manage import ListenerManagement, UpdateProblemArgs, UpdateEmbeddingKnowledgeIdArgs, \
-    UpdateEmbeddingDocumentIdArgs
+from common.event.listener_manage import (
+    ListenerManagement,
+    UpdateProblemArgs,
+    UpdateEmbeddingKnowledgeIdArgs,
+    UpdateEmbeddingDocumentIdArgs,
+)
 from common.utils.logger import maxkb_logger
 from knowledge.models import Document, TaskType, State
 from knowledge.serializers.common import drop_knowledge_index
@@ -18,11 +22,12 @@
 from ops import celery_app
 
 
-def get_embedding_model(model_id, exception_handler=lambda e: maxkb_logger.error(
-    _('Failed to obtain vector model: {error} {traceback}').format(
-        error=str(e),
-        traceback=traceback.format_exc()
-    ))):
+def get_embedding_model(
+    model_id,
+    exception_handler=lambda e: maxkb_logger.error(
+        _("Failed to obtain vector model: {error} {traceback}").format(error=str(e), traceback=traceback.format_exc())
+    ),
+):
     try:
         model = QuerySet(Model).filter(id=model_id).first()
 
@@ -35,25 +40,25 @@ def get_embedding_model(model_id, exception_handler=lambda e: maxkb_logger.error
     return embedding_model
 
 
-@celery_app.task(base=QueueOnce, once={'keys': ['paragraph_id']}, name='celery:embedding_by_paragraph')
+@celery_app.task(base=QueueOnce, once={"keys": ["paragraph_id"]}, name="celery:embedding_by_paragraph")
 def embedding_by_paragraph(paragraph_id, model_id):
     embedding_model = get_embedding_model(model_id)
     ListenerManagement.embedding_by_paragraph(paragraph_id, embedding_model)
 
 
-@celery_app.task(base=QueueOnce, once={'keys': ['paragraph_id_list']}, name='celery:embedding_by_paragraph_data_list')
+@celery_app.task(base=QueueOnce, once={"keys": ["paragraph_id_list"]}, name="celery:embedding_by_paragraph_data_list")
 def embedding_by_paragraph_data_list(data_list, paragraph_id_list, model_id):
     embedding_model = get_embedding_model(model_id)
     ListenerManagement.embedding_by_paragraph_data_list(data_list, paragraph_id_list, embedding_model)
 
 
-@celery_app.task(base=QueueOnce, once={'keys': ['paragraph_id_list']}, name='celery:embedding_by_paragraph_list')
+@celery_app.task(base=QueueOnce, once={"keys": ["paragraph_id_list"]}, name="celery:embedding_by_paragraph_list")
 def embedding_by_paragraph_list(paragraph_id_list, model_id):
     embedding_model = get_embedding_model(model_id)
     ListenerManagement.embedding_by_paragraph_list(paragraph_id_list, embedding_model)
 
 
-@celery_app.task(base=QueueOnce, once={'keys': ['document_id']}, name='celery:embedding_by_document')
+@celery_app.task(base=QueueOnce, once={"keys": ["document_id"]}, name="celery:embedding_by_document")
 def embedding_by_document(document_id, model_id, state_list=None):
     """
     向量化文档
@@ -64,25 +69,30 @@ def embedding_by_document(document_id, model_id, state_list=None):
     """
 
     if state_list is None:
-        state_list = [State.PENDING.value, State.STARTED.value, State.SUCCESS.value, State.FAILURE.value,
-                      State.REVOKE.value,
-                      State.REVOKED.value, State.IGNORED.value]
+        state_list = [
+            State.PENDING.value,
+            State.STARTED.value,
+            State.SUCCESS.value,
+            State.FAILURE.value,
+            State.REVOKE.value,
+            State.REVOKED.value,
+            State.IGNORED.value,
+        ]
 
     def exception_handler(e):
-        ListenerManagement.update_status(QuerySet(Document).filter(id=document_id), TaskType.EMBEDDING,
-                                         State.FAILURE)
+        ListenerManagement.update_status(QuerySet(Document).filter(id=document_id), TaskType.EMBEDDING, State.FAILURE)
         maxkb_logger.error(
-            _('Failed to obtain vector model: {error} {traceback}').format(
-                error=str(e),
-                traceback=traceback.format_exc()
-            ))
+            _("Failed to obtain vector model: {error} {traceback}").format(
+                error=str(e), traceback=traceback.format_exc()
+            )
+        )
 
     embedding_model = get_embedding_model(model_id, exception_handler)
     #
     ListenerManagement.embedding_by_document(document_id, embedding_model, state_list)
 
 
-@celery_app.task(name='celery:embedding_by_document_list')
+@celery_app.task(name="celery:embedding_by_document_list")
 def embedding_by_document_list(document_id_list, model_id):
     """
     向量化文档
@@ -94,33 +104,37 @@ def embedding_by_document_list(document_id_list, model_id):
         embedding_by_document.delay(document_id, model_id)
 
 
-@celery_app.task(base=QueueOnce, once={'keys': ['knowledge_id']}, name='celery:embedding_by_knowledge')
+@celery_app.task(base=QueueOnce, once={"keys": ["knowledge_id"]}, name="celery:embedding_by_knowledge")
 def embedding_by_knowledge(knowledge_id, model_id):
     """
-          向量化知识库
-          @param knowledge_id: 知识库id
-          @param model_id 向量模型
-          :return: None
-          """
-    maxkb_logger.info(_('Start--->Vectorized knowledge: {knowledge_id}').format(knowledge_id=knowledge_id))
+    向量化知识库
+    @param knowledge_id: 知识库id
+    @param model_id 向量模型
+    :return: None
+    """
+    maxkb_logger.info(_("Start--->Vectorized knowledge: {knowledge_id}").format(knowledge_id=knowledge_id))
     try:
         ListenerManagement.delete_embedding_by_knowledge(knowledge_id)
         drop_knowledge_index(knowledge_id=knowledge_id)
         document_list = QuerySet(Document).filter(knowledge_id=knowledge_id)
-        maxkb_logger.info(_('Knowledge documentation: {document_names}').format(
-            document_names=", ".join([d.name for d in document_list])))
+        maxkb_logger.info(
+            _("Knowledge documentation: {document_names}").format(
+                document_names=", ".join([d.name for d in document_list])
+            )
+        )
         for document in document_list:
             try:
                 embedding_by_document.delay(document.id, model_id)
             except Exception as e:
                 pass
     except Exception as e:
         maxkb_logger.error(
-            _('Vectorized knowledge: {knowledge_id} error {error} {traceback}').format(knowledge_id=knowledge_id,
-                                                                                      error=str(e),
-                                                                                      traceback=traceback.format_exc()))
+            _("Vectorized knowledge: {knowledge_id} error {error} {traceback}").format(
+                knowledge_id=knowledge_id, error=str(e), traceback=traceback.format_exc()
+            )
+        )
     finally:
-        maxkb_logger.info(_('End--->Vectorized knowledge: {knowledge_id}').format(knowledge_id=knowledge_id))
+        maxkb_logger.info(_("End--->Vectorized knowledge: {knowledge_id}").format(knowledge_id=knowledge_id))
 
 
 def embedding_by_problem(args, model_id):
@@ -233,7 +247,8 @@ def update_embedding_knowledge_id(paragraph_id_list, target_knowledge_id):
     """
 
     ListenerManagement.update_embedding_knowledge_id(
-        UpdateEmbeddingKnowledgeIdArgs(paragraph_id_list, target_knowledge_id))
+        UpdateEmbeddingKnowledgeIdArgs(paragraph_id_list, target_knowledge_id)
+    )
 
 
 def delete_embedding_by_paragraph_ids(paragraph_ids: List[str]):
@@ -245,13 +260,17 @@ def delete_embedding_by_paragraph_ids(paragraph_ids: List[str]):
     ListenerManagement.delete_embedding_by_paragraph_ids(paragraph_ids)
 
 
-def update_embedding_document_id(paragraph_id_list, target_document_id, target_knowledge_id,
-                                 target_embedding_model_id=None):
-    target_embedding_model = get_embedding_model(
-        target_embedding_model_id) if target_embedding_model_id is not None else None
+def update_embedding_document_id(
+    paragraph_id_list, target_document_id, target_knowledge_id, target_embedding_model_id=None
+):
+    target_embedding_model = (
+        get_embedding_model(target_embedding_model_id) if target_embedding_model_id is not None else None
+    )
     ListenerManagement.update_embedding_document_id(
-        UpdateEmbeddingDocumentIdArgs(paragraph_id_list, target_document_id, target_knowledge_id,
-                                      target_embedding_model))
+        UpdateEmbeddingDocumentIdArgs(
+            paragraph_id_list, target_document_id, target_knowledge_id, target_embedding_model
+        )
+    )
 
 
 def delete_embedding_by_knowledge_id_list(knowledge_id_list):