Add task moduel, and pipline the task and every parser (#49)

api/db/__init__.py

@@ -59,3 +59,14 @@ class ChatStyle(StrEnum):
     PRECISE = 'Precise'
     EVENLY = 'Evenly'
     CUSTOM = 'Custom'
+
+
+class ParserType(StrEnum):
+    GENERAL = "general"
+    PRESENTATION = "presentation"
+    LAWS = "laws"
+    MANUAL = "manual"
+    PAPER = "paper"
+    RESUME = ""
+    BOOK = ""
+    QA = ""

api/db/db_models.py

@@ -496,15 +496,27 @@ class Document(DataBaseModel):
     token_num = IntegerField(default=0)
     chunk_num = IntegerField(default=0)
     progress = FloatField(default=0)
-    progress_msg = CharField(max_length=255, null=True, help_text="process message", default="")
+    progress_msg = CharField(max_length=512, null=True, help_text="process message", default="")
     process_begin_at = DateTimeField(null=True)
     process_duation = FloatField(default=0)
+    run = CharField(max_length=1, null=True, help_text="start to run processing or cancel.(1: run it; 2: cancel)", default="0")
     status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted，1: validate)", default="1")
 
     class Meta:
         db_table = "document"
 
 
+class Task(DataBaseModel):
+    id = CharField(max_length=32, primary_key=True)
+    doc_id = CharField(max_length=32, null=False, index=True)
+    from_page = IntegerField(default=0)
+    to_page = IntegerField(default=-1)
+    begin_at = DateTimeField(null=True)
+    process_duation = FloatField(default=0)
+    progress = FloatField(default=0)
+    progress_msg = CharField(max_length=255, null=True, help_text="process message", default="")
+
+
 class Dialog(DataBaseModel):
     id = CharField(max_length=32, primary_key=True)
     tenant_id = CharField(max_length=32, null=False)
@@ -553,72 +565,6 @@ class Conversation(DataBaseModel):
 
 
 """
-class Job(DataBaseModel):
-    # multi-party common configuration
-    f_user_id = CharField(max_length=25, null=True)
-    f_job_id = CharField(max_length=25, index=True)
-    f_name = CharField(max_length=500, null=True, default='')
-    f_description = TextField(null=True, default='')
-    f_tag = CharField(max_length=50, null=True, default='')
-    f_dsl = JSONField()
-    f_runtime_conf = JSONField()
-    f_runtime_conf_on_party = JSONField()
-    f_train_runtime_conf = JSONField(null=True)
-    f_roles = JSONField()
-    f_initiator_role = CharField(max_length=50)
-    f_initiator_party_id = CharField(max_length=50)
-    f_status = CharField(max_length=50)
-    f_status_code = IntegerField(null=True)
-    f_user = JSONField()
-    # this party configuration
-    f_role = CharField(max_length=50, index=True)
-    f_party_id = CharField(max_length=10, index=True)
-    f_is_initiator = BooleanField(null=True, default=False)
-    f_progress = IntegerField(null=True, default=0)
-    f_ready_signal = BooleanField(default=False)
-    f_ready_time = BigIntegerField(null=True)
-    f_cancel_signal = BooleanField(default=False)
-    f_cancel_time = BigIntegerField(null=True)
-    f_rerun_signal = BooleanField(default=False)
-    f_end_scheduling_updates = IntegerField(null=True, default=0)
-
-    f_engine_name = CharField(max_length=50, null=True)
-    f_engine_type = CharField(max_length=10, null=True)
-    f_cores = IntegerField(default=0)
-    f_memory = IntegerField(default=0)  # MB
-    f_remaining_cores = IntegerField(default=0)
-    f_remaining_memory = IntegerField(default=0)  # MB
-    f_resource_in_use = BooleanField(default=False)
-    f_apply_resource_time = BigIntegerField(null=True)
-    f_return_resource_time = BigIntegerField(null=True)
-
-    f_inheritance_info = JSONField(null=True)
-    f_inheritance_status = CharField(max_length=50, null=True)
-
-    f_start_time = BigIntegerField(null=True)
-    f_start_date = DateTimeField(null=True)
-    f_end_time = BigIntegerField(null=True)
-    f_end_date = DateTimeField(null=True)
-    f_elapsed = BigIntegerField(null=True)
-
-    class Meta:
-        db_table = "t_job"
-        primary_key = CompositeKey('f_job_id', 'f_role', 'f_party_id')
-
-
-
-class PipelineComponentMeta(DataBaseModel):
-    f_model_id = CharField(max_length=100, index=True)
-    f_model_version = CharField(max_length=100, index=True)
-    f_role = CharField(max_length=50, index=True)
-    f_party_id = CharField(max_length=10, index=True)
-    f_component_name = CharField(max_length=100, index=True)
-    f_component_module_name = CharField(max_length=100)
-    f_model_alias = CharField(max_length=100, index=True)
-    f_model_proto_index = JSONField(null=True)
-    f_run_parameters = JSONField(null=True)
-    f_archive_sha256 = CharField(max_length=100, null=True)
-    f_archive_from_ip = CharField(max_length=100, null=True)
 
     class Meta:
         db_table = 't_pipeline_component_meta'

api/db/db_utils.py

@@ -32,19 +32,19 @@ LOGGER = getLogger()
 def bulk_insert_into_db(model, data_source, replace_on_conflict=False):
     DB.create_tables([model])
 
-    current_time = current_timestamp()
-    current_date = timestamp_to_date(current_time)
 
     for data in data_source:
-        if 'f_create_time' not in data:
-            data['f_create_time'] = current_time
-        data['f_create_date'] = timestamp_to_date(data['f_create_time'])
-        data['f_update_time'] = current_time
-        data['f_update_date'] = current_date
+        current_time = current_timestamp()
+        current_date = timestamp_to_date(current_time)
+        if 'create_time' not in data:
+            data['create_time'] = current_time
+        data['create_date'] = timestamp_to_date(data['create_time'])
+        data['update_time'] = current_time
+        data['update_date'] = current_date
 
-    preserve = tuple(data_source[0].keys() - {'f_create_time', 'f_create_date'})
+    preserve = tuple(data_source[0].keys() - {'create_time', 'create_date'})
 
-    batch_size = 50 if RuntimeConfig.USE_LOCAL_DATABASE else 1000
+    batch_size = 1000
 
     for i in range(0, len(data_source), batch_size):
         with DB.atomic():

api/db/services/common_service.py

@@ -70,6 +70,7 @@ class CommonService:
     @DB.connection_context()
     def insert_many(cls, data_list, batch_size=100):
         with DB.atomic():
+            for d in data_list: d["create_time"] = datetime_format(datetime.now())
             for i in range(0, len(data_list), batch_size):
                 cls.model.insert_many(data_list[i:i + batch_size]).execute()
 

api/db/services/document_service.py

@@ -61,8 +61,8 @@ class DocumentService(CommonService):
 
     @classmethod
     @DB.connection_context()
-    def get_newly_uploaded(cls, tm, mod, comm, items_per_page=64):
-        fields = [cls.model.id, cls.model.kb_id, cls.model.parser_id, cls.model.name, cls.model.location, cls.model.size, Knowledgebase.tenant_id, Tenant.embd_id, Tenant.img2txt_id, cls.model.update_time]
+    def get_newly_uploaded(cls, tm, mod=0, comm=1, items_per_page=64):
+        fields = [cls.model.id, cls.model.kb_id, cls.model.parser_id, cls.model.name, cls.model.type, cls.model.location, cls.model.size, Knowledgebase.tenant_id, Tenant.embd_id, Tenant.img2txt_id, Tenant.asr_id, cls.model.update_time]
         docs = cls.model.select(*fields) \
             .join(Knowledgebase, on=(cls.model.kb_id == Knowledgebase.id)) \
             .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))\
@@ -76,6 +76,18 @@ class DocumentService(CommonService):
             .paginate(1, items_per_page)
         return list(docs.dicts())
 
+    @classmethod
+    @DB.connection_context()
+    def get_unfinished_docs(cls):
+        fields = [cls.model.id, cls.model.process_begin_at]
+        docs = cls.model.select(*fields) \
+            .where(
+                cls.model.status == StatusEnum.VALID.value,
+                ~(cls.model.type == FileType.VIRTUAL.value),
+                cls.model.progress < 1,
+                cls.model.progress > 0)
+        return list(docs.dicts())
+
     @classmethod
     @DB.connection_context()
     def increment_chunk_num(cls, doc_id, kb_id, token_num, chunk_num, duation):

api/db/services/task_service.py

@@ -0,0 +1,53 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+from peewee import Expression
+from api.db.db_models import DB
+from api.db import StatusEnum, FileType
+from api.db.db_models import Task, Document, Knowledgebase, Tenant
+from api.db.services.common_service import CommonService
+
+
+class TaskService(CommonService):
+    model = Task
+
+    @classmethod
+    @DB.connection_context()
+    def get_tasks(cls, tm, mod=0, comm=1, items_per_page=64):
+        fields = [cls.model.id, cls.model.doc_id, cls.model.from_page,cls.model.to_page, Document.kb_id, Document.parser_id, Document.name, Document.type, Document.location, Document.size, Knowledgebase.tenant_id, Tenant.embd_id, Tenant.img2txt_id, Tenant.asr_id, cls.model.update_time]
+        docs = cls.model.select(*fields) \
+            .join(Document, on=(cls.model.doc_id == Document.id)) \
+            .join(Knowledgebase, on=(Document.kb_id == Knowledgebase.id)) \
+            .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))\
+            .where(
+                Document.status == StatusEnum.VALID.value,
+                ~(Document.type == FileType.VIRTUAL.value),
+                cls.model.progress == 0,
+                cls.model.update_time >= tm,
+                (Expression(cls.model.create_time, "%%", comm) == mod))\
+            .order_by(cls.model.update_time.asc())\
+            .paginate(1, items_per_page)
+        return list(docs.dicts())
+
+
+    @classmethod
+    @DB.connection_context()
+    def do_cancel(cls, id):
+        try:
+            cls.model.get_by_id(id)
+            return False
+        except Exception as e:
+            pass
+        return True