feat: 文档状态70%

Author: shaohuzhang1

apps/common/db/search.py

@@ -12,7 +12,7 @@
 from django.db.models import QuerySet
 
 from common.db.compiler import AppSQLCompiler
-from common.db.sql_execute import select_one, select_list
+from common.db.sql_execute import select_one, select_list, update_execute
 from common.response.result import Page
 
 
@@ -109,6 +109,24 @@ def native_search(queryset: QuerySet | Dict[str, QuerySet], select_string: str,
         return select_list(exec_sql, exec_params)
 
 
+def native_update(queryset: QuerySet | Dict[str, QuerySet], select_string: str,
+                  field_replace_dict: None | Dict[str, Dict[str, str]] | Dict[str, str] = None,
+                  with_table_name=False):
+    """
+    复杂查询
+    :param with_table_name:     生成sql是否包含表名
+    :param queryset:            查询条件构造器
+    :param select_string:       查询前缀 不包括 where limit 等信息
+    :param field_replace_dict:  需要替换的字段
+    :return: 查询结果
+    """
+    if isinstance(queryset, Dict):
+        exec_sql, exec_params = generate_sql_by_query_dict(queryset, select_string, field_replace_dict, with_table_name)
+    else:
+        exec_sql, exec_params = generate_sql_by_query(queryset, select_string, field_replace_dict, with_table_name)
+    return update_execute(exec_sql, exec_params)
+
+
 def page_search(current_page: int, page_size: int, queryset: QuerySet, post_records_handler):
     """
     分页查询

apps/common/event/listener_manage.py

@@ -13,15 +13,18 @@
 from typing import List
 
 import django.db.models
+from django.db import models
 from django.db.models import QuerySet
+from django.db.models.functions import Substr, Reverse
 from langchain_core.embeddings import Embeddings
 
 from common.config.embedding_config import VectorStore
-from common.db.search import native_search, get_dynamics_model
-from common.event.common import embedding_poxy
+from common.db.search import native_search, get_dynamics_model, native_update
+from common.db.sql_execute import sql_execute, update_execute
 from common.util.file_util import get_file_content
 from common.util.lock import try_lock, un_lock
-from dataset.models import Paragraph, Status, Document, ProblemParagraphMapping
+from common.util.page_utils import page
+from dataset.models import Paragraph, Status, Document, ProblemParagraphMapping, TaskType, State
 from embedding.models import SourceType, SearchMode
 from smartdoc.conf import PROJECT_DIR
 
@@ -114,7 +117,8 @@ def embedding_by_paragraph(paragraph_id, embedding_model: Embeddings):
         @param embedding_model:  向量模型
         """
         max_kb.info(f"开始--->向量化段落:{paragraph_id}")
-        status = Status.success
+        # 更新到开始状态
+        ListenerManagement.update_status(QuerySet(Paragraph).filter(id=paragraph_id), TaskType.EMBEDDING, State.STARTED)
         try:
             data_list = native_search(
                 {'problem': QuerySet(get_dynamics_model({'paragraph.id': django.db.models.CharField()})).filter(
@@ -125,23 +129,89 @@ def embedding_by_paragraph(paragraph_id, embedding_model: Embeddings):
             # 删除段落
             VectorStore.get_embedding_vector().delete_by_paragraph_id(paragraph_id)
 
-            def is_save_function():
-                return QuerySet(Paragraph).filter(id=paragraph_id).exists()
+            def is_the_task_interrupted():
+                _paragraph = QuerySet(Paragraph).filter(id=paragraph_id).first()
+                if _paragraph is None or Status(_paragraph.status)[TaskType.EMBEDDING] == State.REVOKE:
+                    return True
+                return False
 
             # 批量向量化
-            VectorStore.get_embedding_vector().batch_save(data_list, embedding_model, is_save_function)
+            VectorStore.get_embedding_vector().batch_save(data_list, embedding_model, is_the_task_interrupted)
+            # 更新到开始状态
+            ListenerManagement.update_status(QuerySet(Paragraph).filter(id=paragraph_id), TaskType.EMBEDDING,
+                                             State.SUCCESS)
         except Exception as e:
             max_kb_error.error(f'向量化段落:{paragraph_id}出现错误{str(e)}{traceback.format_exc()}')
-            status = Status.error
+            ListenerManagement.update_status(QuerySet(Paragraph).filter(id=paragraph_id), TaskType.EMBEDDING,
+                                             State.FAILURE)
         finally:
-            QuerySet(Paragraph).filter(id=paragraph_id).update(**{'status': status})
             max_kb.info(f'结束--->向量化段落:{paragraph_id}')
 
     @staticmethod
     def embedding_by_data_list(data_list: List, embedding_model: Embeddings):
         # 批量向量化
         VectorStore.get_embedding_vector().batch_save(data_list, embedding_model, lambda: True)
 
+    @staticmethod
+    def get_embedding_paragraph_apply(embedding_model, is_the_task_interrupted, post_apply=lambda: None):
+        def embedding_paragraph_apply(paragraph_list):
+            for paragraph in paragraph_list:
+                if is_the_task_interrupted():
+                    break
+                ListenerManagement.embedding_by_paragraph(str(paragraph.get('id')), embedding_model)
+            post_apply()
+
+        return embedding_paragraph_apply
+
+    @staticmethod
+    def get_aggregation_document_status(document_id):
+        def aggregation_document_status():
+            sql = get_file_content(
+                os.path.join(PROJECT_DIR, "apps", "dataset", 'sql', 'update_document_status_meta.sql'))
+            update_execute(sql, [document_id, document_id])
+
+        return aggregation_document_status
+
+    @staticmethod
+    def post_update_document_status(document_id, task_type: TaskType):
+        _document = QuerySet(Document).filter(id=document_id).first()
+
+        status = Status(_document.status)
+        if status[task_type] == State.REVOKE:
+            status[task_type] = State.REVOKED
+        else:
+            status[task_type] = State.SUCCESS
+        for item in _document.status_meta.get('aggs', []):
+            agg_status = item.get('status')
+            agg_count = item.get('count')
+            if Status(agg_status)[task_type] == State.FAILURE and agg_count > 0:
+                status[task_type] = State.FAILURE
+        _document.status = status.__str__()
+        _document.save()
+        ListenerManagement.update_status(QuerySet(Paragraph).annotate(
+            reversed_status=Reverse('status'),
+            task_type_status=Substr('reversed_status', task_type.value,
+                                    task_type.value),
+        ).filter(task_type_status=State.REVOKE.value).filter(document_id=document_id).values('id'),
+                                         task_type,
+                                         State.REVOKED)
+
+    @staticmethod
+    def update_status(query_set: QuerySet, taskType: TaskType, state: State):
+        exec_sql = get_file_content(
+            os.path.join(PROJECT_DIR, "apps", "dataset", 'sql', 'update_paragraph_status.sql'))
+        bit_number = len(TaskType)
+        up_index = taskType.value - 1
+        next_index = taskType.value + 1
+        status_number = state.value
+        params_dict = {'${bit_number}': bit_number, '${up_index}': up_index,
+                       '${status_number}': status_number, '${next_index}': next_index,
+                       '${table_name}': query_set.model._meta.db_table}
+        for key in params_dict:
+            _value_ = params_dict[key]
+            exec_sql = exec_sql.replace(key, str(_value_))
+        native_update(query_set, exec_sql)
+
     @staticmethod
     def embedding_by_document(document_id, embedding_model: Embeddings):
         """
@@ -153,33 +223,28 @@ def embedding_by_document(document_id, embedding_model: Embeddings):
         if not try_lock('embedding' + str(document_id)):
             return
         max_kb.info(f"开始--->向量化文档:{document_id}")
-        QuerySet(Document).filter(id=document_id).update(**{'status': Status.embedding})
-        QuerySet(Paragraph).filter(document_id=document_id).update(**{'status': Status.embedding})
-        status = Status.success
+        # 批量修改状态为PADDING
+        ListenerManagement.update_status(QuerySet(Document).filter(id=document_id), TaskType.EMBEDDING, State.STARTED)
         try:
-            data_list = native_search(
-                {'problem': QuerySet(
-                    get_dynamics_model({'paragraph.document_id': django.db.models.CharField()})).filter(
-                    **{'paragraph.document_id': document_id}),
-                    'paragraph': QuerySet(Paragraph).filter(document_id=document_id)},
-                select_string=get_file_content(
-                    os.path.join(PROJECT_DIR, "apps", "common", 'sql', 'list_embedding_text.sql')))
             # 删除文档向量数据
             VectorStore.get_embedding_vector().delete_by_document_id(document_id)
 
-            def is_save_function():
-                return QuerySet(Document).filter(id=document_id).exists()
-
-            # 批量向量化
-            VectorStore.get_embedding_vector().batch_save(data_list, embedding_model, is_save_function)
+            def is_the_task_interrupted():
+                document = QuerySet(Document).filter(id=document_id).first()
+                if document is None or Status(document.status)[TaskType.EMBEDDING] == State.REVOKE:
+                    return True
+                return False
+
+            # 根据段落进行向量化处理
+            page(QuerySet(Paragraph).filter(document_id=document_id).values('id'), 10,
+                 ListenerManagement.get_embedding_paragraph_apply(embedding_model, is_the_task_interrupted,
+                                                                  ListenerManagement.get_aggregation_document_status(
+                                                                      document_id)),
+                 is_the_task_interrupted)
         except Exception as e:
             max_kb_error.error(f'向量化文档:{document_id}出现错误{str(e)}{traceback.format_exc()}')
-            status = Status.error
         finally:
-            # 修改状态
-            QuerySet(Document).filter(id=document_id).update(
-                **{'status': status, 'update_time': datetime.datetime.now()})
-            QuerySet(Paragraph).filter(document_id=document_id).update(**{'status': status})
+            ListenerManagement.post_update_document_status(document_id, TaskType.EMBEDDING)
             max_kb.info(f"结束--->向量化文档:{document_id}")
             un_lock('embedding' + str(document_id))
 

apps/common/util/page_utils.py

@@ -0,0 +1,27 @@
+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author：虎
+    @file： page_utils.py
+    @date：2024/11/21 10:32
+    @desc:
+"""
+from math import ceil
+
+
+def page(query_set, page_size, handler, is_the_task_interrupted=lambda: False):
+    """
+
+    @param query_set: 查询query_set
+    @param page_size: 每次查询大小
+    @param handler:   数据处理器
+    @param is_the_task_interrupted: 任务是否被中断
+    @return:
+    """
+    count = query_set.count()
+    for i in range(0, ceil(count / page_size)):
+        if is_the_task_interrupted():
+            return
+        offset = i * page_size
+        paragraph_list = query_set[offset: offset + page_size]
+        handler(paragraph_list)

apps/dataset/migrations/0011_document_status_meta_paragraph_status_meta_and_more.py

@@ -0,0 +1,34 @@
+# Generated by Django 4.2.15 on 2024-11-22 14:44
+
+import dataset.models.data_set
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dataset', '0010_file_meta'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='document',
+            name='status_meta',
+            field=models.JSONField(default=dataset.models.data_set.default_status_meta, verbose_name='状态统计数据'),
+        ),
+        migrations.AddField(
+            model_name='paragraph',
+            name='status_meta',
+            field=models.JSONField(default=dataset.models.data_set.default_status_meta, verbose_name='状态数据'),
+        ),
+        migrations.AlterField(
+            model_name='document',
+            name='status',
+            field=models.CharField(default=dataset.models.data_set.Status.__str__, max_length=20, verbose_name='状态'),
+        ),
+        migrations.AlterField(
+            model_name='paragraph',
+            name='status',
+            field=models.CharField(default=dataset.models.data_set.Status.__str__, max_length=20, verbose_name='状态'),
+        ),
+    ]

apps/dataset/models/data_set.py

@@ -7,6 +7,7 @@
     @desc: 数据集
 """
 import uuid
+from enum import Enum
 
 from django.db import models
 from django.db.models.signals import pre_delete
@@ -18,13 +19,60 @@
 from users.models import User
 
 
-class Status(models.TextChoices):
-    """订单类型"""
-    embedding = 0, '导入中'
-    success = 1, '已完成'
-    error = 2, '导入失败'
-    queue_up = 3, '排队中'
-    generating = 4, '生成问题中'
+class TaskType(Enum):
+    # 向量
+    EMBEDDING = 1
+    # 生成问题
+    GENERATE_PROBLEM = 2
+    # 同步
+    SYNC = 3
+
+
+class State(Enum):
+    # 等待
+    PENDING = '0'
+    # 执行中
+    STARTED = '1'
+    # 成功
+    SUCCESS = '2'
+    # 失败
+    FAILURE = '3'
+    # 取消任务
+    REVOKE = '4'
+    # 取消成功
+    REVOKED = '5'
+
+
+class Status:
+    type_cls = TaskType
+    state_cls = State
+
+    def __init__(self, status: str = None):
+        self.task_status = {}
+        status_list = list(status[::-1] if status is not None else '')
+        for _type in self.type_cls:
+            index = _type.value - 1
+            _state = self.state_cls(status_list[index] if len(status_list) > index else '2')
+            self.task_status[_type] = _state
+
+    @staticmethod
+    def of(status: str):
+        return Status(status)
+
+    def __str__(self):
+        result = []
+        for _type in sorted(self.type_cls, key=lambda item: item.value, reverse=True):
+            result.insert(len(self.type_cls) - _type.value, self.task_status[_type].value)
+        return ''.join(result)
+
+    def __setitem__(self, key, value):
+        self.task_status[key] = value
+
+    def __getitem__(self, item):
+        return self.task_status[item]
+
+    def update_status(self, task_type: TaskType, state: State):
+        self.task_status[task_type] = state
 
 
 class Type(models.TextChoices):
@@ -42,6 +90,10 @@ def default_model():
     return uuid.UUID('42f63a3d-427e-11ef-b3ec-a8a1595801ab')
 
 
+def default_status_meta():
+    return {"state_time": {}}
+
+
 class DataSet(AppModelMixin):
     """
     数据集表
@@ -68,8 +120,8 @@ class Document(AppModelMixin):
     dataset = models.ForeignKey(DataSet, on_delete=models.DO_NOTHING)
     name = models.CharField(max_length=150, verbose_name="文档名称")
     char_length = models.IntegerField(verbose_name="文档字符数 冗余字段")
-    status = models.CharField(verbose_name='状态', max_length=1, choices=Status.choices,
-                              default=Status.queue_up)
+    status = models.CharField(verbose_name='状态', max_length=20, default=Status('').__str__)
+    status_meta = models.JSONField(verbose_name="状态统计数据", default=default_status_meta)
     is_active = models.BooleanField(default=True)
 
     type = models.CharField(verbose_name='类型', max_length=1, choices=Type.choices,
@@ -94,8 +146,8 @@ class Paragraph(AppModelMixin):
     dataset = models.ForeignKey(DataSet, on_delete=models.DO_NOTHING)
     content = models.CharField(max_length=102400, verbose_name="段落内容")
     title = models.CharField(max_length=256, verbose_name="标题", default="")
-    status = models.CharField(verbose_name='状态', max_length=1, choices=Status.choices,
-                              default=Status.embedding)
+    status = models.CharField(verbose_name='状态', max_length=20, default=Status('').__str__)
+    status_meta = models.JSONField(verbose_name="状态数据", default=default_status_meta)
     hit_num = models.IntegerField(verbose_name="命中次数", default=0)
     is_active = models.BooleanField(default=True)
 
@@ -145,7 +197,6 @@ class File(AppModelMixin):
 
     meta = models.JSONField(verbose_name="文件关联数据", default=dict)
 
-
     class Meta:
         db_table = "file"
 
@@ -161,7 +212,6 @@ def get_byte(self):
         return result['data']
 
 
-
 @receiver(pre_delete, sender=File)
 def on_delete_file(sender, instance, **kwargs):
     select_one(f'SELECT lo_unlink({instance.loid})', [])