feat: add initial implementations of various file handling classes for CSV, XLS, and XLSX formats

Author: liuruibin

apps/common/constants/permission_constants.py

@@ -218,6 +218,14 @@ class PermissionConstants(Enum):
                                                                                             RoleConstants.USER])
     KNOWLEDGE_DELETE = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
                                                                                             RoleConstants.USER])
+    DOCUMENT_READ = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
+                                                                                            RoleConstants.USER])
+    DOCUMENT_CREATE = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
+                                                                                            RoleConstants.USER])
+    DOCUMENT_EDIT = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
+                                                                                            RoleConstants.USER])
+    DOCUMENT_DELETE = Permission(group=Group.KNOWLEDGE, operate=Operate.DELETE, role_list=[RoleConstants.ADMIN,
+                                                                                            RoleConstants.USER])
 
     def get_workspace_application_permission(self):
         return lambda r, kwargs: Permission(group=self.value.group, operate=self.value.operate,

apps/common/handle/__init__.py

@@ -0,0 +1,8 @@
+# coding=utf-8
+"""
+    @project: qabot
+    @Author：虎
+    @file： __init__.py.py
+    @date：2023/9/6 10:09
+    @desc:
+"""

apps/common/handle/base_parse_qa_handle.py

@@ -0,0 +1,52 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： base_parse_qa_handle.py
+    @date：2024/5/21 14:56
+    @desc:
+"""
+from abc import ABC, abstractmethod
+
+
+def get_row_value(row, title_row_index_dict, field):
+    index = title_row_index_dict.get(field)
+    if index is None:
+        return None
+    if (len(row) - 1) >= index:
+        return row[index]
+    return None
+
+
+def get_title_row_index_dict(title_row_list):
+    title_row_index_dict = {}
+    if len(title_row_list) == 1:
+        title_row_index_dict['content'] = 0
+    elif len(title_row_list) == 1:
+        title_row_index_dict['title'] = 0
+        title_row_index_dict['content'] = 1
+    else:
+        title_row_index_dict['title'] = 0
+        title_row_index_dict['content'] = 1
+        title_row_index_dict['problem_list'] = 2
+    for index in range(len(title_row_list)):
+        title_row = title_row_list[index]
+        if title_row is None:
+            title_row = ''
+        if title_row.startswith('分段标题'):
+            title_row_index_dict['title'] = index
+        if title_row.startswith('分段内容'):
+            title_row_index_dict['content'] = index
+        if title_row.startswith('问题'):
+            title_row_index_dict['problem_list'] = index
+    return title_row_index_dict
+
+
+class BaseParseQAHandle(ABC):
+    @abstractmethod
+    def support(self, file, get_buffer):
+        pass
+
+    @abstractmethod
+    def handle(self, file, get_buffer, save_image):
+        pass

apps/common/handle/base_parse_table_handle.py

@@ -0,0 +1,23 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： base_parse_qa_handle.py
+    @date：2024/5/21 14:56
+    @desc:
+"""
+from abc import ABC, abstractmethod
+
+
+class BaseParseTableHandle(ABC):
+    @abstractmethod
+    def support(self, file, get_buffer):
+        pass
+
+    @abstractmethod
+    def handle(self, file, get_buffer,save_image):
+        pass
+
+    @abstractmethod
+    def get_content(self, file, save_image):
+        pass

apps/common/handle/base_split_handle.py

@@ -0,0 +1,24 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： base_split_handle.py
+    @date：2024/3/27 18:13
+    @desc:
+"""
+from abc import ABC, abstractmethod
+from typing import List
+
+
+class BaseSplitHandle(ABC):
+    @abstractmethod
+    def support(self, file, get_buffer):
+        pass
+
+    @abstractmethod
+    def handle(self, file, pattern_list: List, with_filter: bool, limit: int, get_buffer, save_image):
+        pass
+
+    @abstractmethod
+    def get_content(self, file, save_image):
+        pass

apps/common/handle/base_to_response.py

@@ -0,0 +1,30 @@
+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author：虎
+    @file： base_to_response.py
+    @date：2024/9/6 16:04
+    @desc:
+"""
+from abc import ABC, abstractmethod
+
+from rest_framework import status
+
+
+class BaseToResponse(ABC):
+
+    @abstractmethod
+    def to_block_response(self, chat_id, chat_record_id, content, is_end, completion_tokens,
+                          prompt_tokens, other_params: dict = None,
+                          _status=status.HTTP_200_OK):
+        pass
+
+    @abstractmethod
+    def to_stream_chunk_response(self, chat_id, chat_record_id, node_id, up_node_id_list, content, is_end,
+                                 completion_tokens,
+                                 prompt_tokens, other_params: dict = None):
+        pass
+
+    @staticmethod
+    def format_stream_chunk(response_str):
+        return 'data: ' + response_str + '\n\n'

apps/common/handle/handle_exception.py

@@ -0,0 +1,94 @@
+# coding=utf-8
+"""
+    @project: qabot
+    @Author：虎
+    @file： handle_exception.py
+    @date：2023/9/5 19:29
+    @desc:
+"""
+import logging
+import traceback
+
+from rest_framework.exceptions import ValidationError, ErrorDetail, APIException
+from rest_framework.views import exception_handler
+
+from common.exception.app_exception import AppApiException
+
+from django.utils.translation import gettext_lazy as _
+
+from common.result import result
+
+
+def to_result(key, args, parent_key=None):
+    """
+    将校验异常 args转换为统一数据
+    :param key:       校验key
+    :param args:      校验异常参数
+    :param parent_key 父key
+    :return: 接口响应对象
+    """
+    error_detail = list(filter(
+        lambda d: True if isinstance(d, ErrorDetail) else True if isinstance(d, dict) and len(
+            d.keys()) > 0 else False,
+        (args[0] if len(args) > 0 else {key: [ErrorDetail(_('Unknown exception'), code='unknown')]}).get(key)))[0]
+
+    if isinstance(error_detail, dict):
+        return list(map(lambda k: to_result(k, args=[error_detail],
+                                            parent_key=key if parent_key is None else parent_key + '.' + key),
+                        error_detail.keys() if len(error_detail) > 0 else []))[0]
+
+    return result.Result(500 if isinstance(error_detail.code, str) else error_detail.code,
+                         message=f"【{key if parent_key is None else parent_key + '.' + key}】为必填参数" if str(
+                             error_detail) == "This field is required." else error_detail)
+
+
+def validation_error_to_result(exc: ValidationError):
+    """
+    校验异常转响应对象
+    :param exc: 校验异常
+    :return: 接口响应对象
+    """
+    try:
+        v = find_err_detail(exc.detail)
+        if v is None:
+            return result.error(str(exc.detail))
+        return result.error(str(v))
+    except Exception as e:
+        return result.error(str(exc.detail))
+
+
+def find_err_detail(exc_detail):
+    if isinstance(exc_detail, ErrorDetail):
+        return exc_detail
+    if isinstance(exc_detail, dict):
+        keys = exc_detail.keys()
+        for key in keys:
+            _value = exc_detail[key]
+            if isinstance(_value, list):
+                return find_err_detail(_value)
+            if isinstance(_value, ErrorDetail):
+                return _value
+            if isinstance(_value, dict) and len(_value.keys()) > 0:
+                return find_err_detail(_value)
+    if isinstance(exc_detail, list):
+        for v in exc_detail:
+            r = find_err_detail(v)
+            if r is not None:
+                return r
+
+
+def handle_exception(exc, context):
+    exception_class = exc.__class__
+    # 先调用REST framework默认的异常处理方法获得标准错误响应对象
+    response = exception_handler(exc, context)
+    # 在此处补充自定义的异常处理
+    if issubclass(exception_class, ValidationError):
+        return validation_error_to_result(exc)
+    if issubclass(exception_class, AppApiException):
+        return result.Result(exc.code, exc.message, response_status=exc.status_code)
+    if issubclass(exception_class, APIException):
+        return result.error(exc.detail)
+    if response is None:
+        logging.getLogger("max_kb_error").error(f'{str(exc)}:{traceback.format_exc()}')
+        return result.error(str(exc))
+    return response

apps/common/handle/impl/__init__.py

@@ -0,0 +1,8 @@
+# coding=utf-8
+"""
+    @project: qabot
+    @Author：虎
+    @file： __init__.py.py
+    @date：2023/9/6 10:09
+    @desc:
+"""

apps/common/handle/impl/common_handle.py

@@ -0,0 +1,116 @@
+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author：虎
+    @file： tools.py
+    @date：2024/9/11 16:41
+    @desc:
+"""
+import io
+import uuid_utils.compat as uuid
+from functools import reduce
+from io import BytesIO
+from xml.etree.ElementTree import fromstring
+from zipfile import ZipFile
+
+from PIL import Image as PILImage
+from openpyxl.drawing.image import Image as openpyxl_Image
+from openpyxl.packaging.relationship import get_rels_path, get_dependents
+from openpyxl.xml.constants import SHEET_DRAWING_NS, REL_NS, SHEET_MAIN_NS
+
+from common.handle.base_parse_qa_handle import get_title_row_index_dict, get_row_value
+from knowledge.models import File
+
+
+def parse_element(element) -> {}:
+    data = {}
+    xdr_namespace = "{%s}" % SHEET_DRAWING_NS
+    targets = level_order_traversal(element, xdr_namespace + "nvPicPr")
+    for target in targets:
+        cNvPr = embed = ""
+        for child in target:
+            if child.tag == xdr_namespace + "nvPicPr":
+                cNvPr = child[0].attrib["name"]
+            elif child.tag == xdr_namespace + "blipFill":
+                _rel_embed = "{%s}embed" % REL_NS
+                embed = child[0].attrib[_rel_embed]
+        if cNvPr:
+            data[cNvPr] = embed
+    return data
+
+
+def parse_element_sheet_xml(element) -> []:
+    data = []
+    xdr_namespace = "{%s}" % SHEET_MAIN_NS
+    targets = level_order_traversal(element, xdr_namespace + "f")
+    for target in targets:
+        for child in target:
+            if child.tag == xdr_namespace + "f":
+                data.append(child.text)
+    return data
+
+
+def level_order_traversal(root, flag: str) -> []:
+    queue = [root]
+    targets = []
+    while queue:
+        node = queue.pop(0)
+        children = [child.tag for child in node]
+        if flag in children:
+            targets.append(node)
+            continue
+        for child in node:
+            queue.append(child)
+    return targets
+
+
+def handle_images(deps, archive: ZipFile) -> []:
+    images = []
+    if not PILImage:  # Pillow not installed, drop images
+        return images
+    for dep in deps:
+        try:
+            image_io = archive.read(dep.target)
+            image = openpyxl_Image(BytesIO(image_io))
+        except Exception as e:
+            print(e)
+            continue
+        image.embed = dep.id  # 文件rId
+        image.target = dep.target  # 文件地址
+        images.append(image)
+    return images
+
+
+def xlsx_embed_cells_images(buffer) -> {}:
+    archive = ZipFile(buffer)
+    # 解析cellImage.xml文件
+    deps = get_dependents(archive, get_rels_path("xl/cellimages.xml"))
+    image_rel = handle_images(deps=deps, archive=archive)
+    # 工作表及其中图片ID
+    sheet_list = {}
+    for item in archive.namelist():
+        if not item.startswith('xl/worksheets/sheet'):
+            continue
+        key = item.split('/')[-1].split('.')[0].split('sheet')[-1]
+        sheet_list[key] = parse_element_sheet_xml(fromstring(archive.read(item)))
+    cell_images_xml = parse_element(fromstring(archive.read("xl/cellimages.xml")))
+    cell_images_rel = {}
+    for image in image_rel:
+        cell_images_rel[image.embed] = image
+    for cnv, embed in cell_images_xml.items():
+        cell_images_xml[cnv] = cell_images_rel.get(embed)
+    result = {}
+    for key, img in cell_images_xml.items():
+        image_excel_id_list = [_xl for _xl in
+                               reduce(lambda x, y: [*x, *y], [sheet for sheet_id, sheet in sheet_list.items()], []) if
+                               key in _xl]
+        if len(image_excel_id_list) > 0:
+            image_excel_id = image_excel_id_list[-1]
+            f = archive.open(img.target)
+            img_byte = io.BytesIO()
+            im = PILImage.open(f).convert('RGB')
+            im.save(img_byte, format='JPEG')
+            image = File(id=uuid.uuid7(), file_name=img.path, meta={'debug': False, 'content': img_byte.getvalue()})
+            result['=' + image_excel_id] = image
+    archive.close()
+    return result