fix: refine document data model

Author: liningping

parsers/base_models.py

@@ -1,29 +1,42 @@
 import logging
 from abc import ABC, abstractmethod
 from enum import Enum
+from typing import Any
 
 from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 
+
 class ChunkType(str, Enum):
     """块类型"""
     TEXT = "text"
     IMAGE = "image"
     TABLE = "table"
     FORMULA = "formula"
 
+class TableDataItem(BaseModel):
+    """表格数据类"""
+    rows: int  # 行数
+    columns: int  # 列数
+    row_headers: list[Any] = []  # 行头
+    column_headers: list[Any] = []  # 列头
+    data: list[list[str]] = []  # 数据
+
 class ChunkData(BaseModel):
     """块数据类"""
     type: ChunkType
     name: str
-    content: str = ""
+    content: str|TableDataItem = ""
     description: str = ""
 
 class DocumentData(BaseModel):
     """解析结果类"""
     title: str = ""
-    chunks: list[ChunkData] = []
+    texts: list[ChunkData] = []
+    tables: list[ChunkData] = []
+    images: list[ChunkData] = []
+    formulas: list[ChunkData] = []
     processing_time: float = 0
     success: bool
     error_message: str | None = None

parsers/excel_parser.py

@@ -19,7 +19,13 @@
 from openpyxl.workbook.workbook import Workbook  # type: ignore
 from openpyxl.worksheet.worksheet import Worksheet  # type: ignore
 
-from parsers.base_models import ChunkData, ChunkType, DocumentData, DocumentParser
+from parsers.base_models import (
+    ChunkData,
+    ChunkType,
+    DocumentData,
+    DocumentParser,
+    TableDataItem,
+)
 
 # 忽略 openpyxl 的特定警告
 warnings.filterwarnings('ignore', category=UserWarning, module='openpyxl')
@@ -57,105 +63,77 @@ def __init__(self, config: ExcelParseConfig | None = None):
         self.image_index: int = 0
         self.supported_formats: list[str] = ['.xlsx', '.xls']
 
-    async def parse(self, excel_path: str) -> DocumentData:
+    def can_parse(self, file_path: str) -> bool:
+        """
+        验证输入文件
+        Args:
+            file_path: 文件路径
+        Returns:
+            bool: 是否支持解析
         """
-        解析Excel文件并保存结果
+        return any(file_path.lower().endswith(fmt) for fmt in self.supported_formats)
 
+    async def parse(self, excel_path: str) -> DocumentData:
+        """
+        将Excel文件转换为JSON格式
         Args:
             excel_path: Excel文件路径
-            output_dir: 输出目录路径
         Returns:
-            ParseResult: 解析结果对象
-        Raises:
-            ExcelParseError: 当解析失败时抛出
+            DocumentData: 文档数据
         """
+        # 获取文件名作为标题（不带扩展名）
         start_time = time.time()
 
         try:
-            # 转换Excel到JSON格式
-            title, document_data = self._excel_to_json(excel_path)
-
-            # 计算处理时间
+            # 初始化内容列表和图片列表
+            texts: list[ChunkData] = []
+            tables: list[ChunkData] = []
+            images: list[ChunkData] = []
+
+            # 加载工作簿
+            workbook = self._load_workbook(excel_path)
+
+            # 处理每个工作表
+            for sheet_index, sheet_name in enumerate(workbook.sheetnames):
+                sheet = workbook[sheet_name]
+
+                # 添加工作表标题
+                texts.append(ChunkData(
+                    type=ChunkType.TEXT,
+                    name=sheet_name,
+                    content=f"工作表 {sheet_index + 1}: {sheet_name}",
+                    description="工作表标题"
+                ))
+
+                # 处理图片
+                sheet_images = self._extract_sheet_images(sheet)
+                images.extend(sheet_images)
+
+                # 处理表格数据
+                table_content = self._extract_table_data(sheet)
+                tables.append(ChunkData(
+                    type=ChunkType.TABLE,
+                    name="表格",
+                    content=table_content,
+                    description="表格"
+                ))
             processing_time = time.time() - start_time
-
-
             return DocumentData(
-                title=title,
-                chunks=document_data,
+                title=Path(excel_path).stem,
+                texts=texts,
+                tables=tables,
+                images=images,
                 processing_time=processing_time,
                 success=True
             )
-
         except Exception as e:
             processing_time = time.time() - start_time
             return DocumentData(
                 success=False,
-                error_message=str(e)
+                error_message=str(e),
+                processing_time=processing_time
             )
 
-    def can_parse(self, file_path: str) -> bool:
-        """
-        验证输入文件
-        Args:
-            file_path: 文件路径
-        Returns:
-            bool: 是否支持解析
-        """
-        return any(file_path.lower().endswith(fmt) for fmt in self.supported_formats)
-
-    def _excel_to_json(self, excel_path: str) -> tuple[str, list[ChunkData]]:
-        """
-        将Excel文件转换为JSON格式
-        Args:
-            excel_path: Excel文件路径
-        Returns:
-            DocumentData: 文档数据
-        """
-        # 获取文件名作为标题（不带扩展名）
-        title = Path(excel_path).stem
-
-        # 初始化内容列表和图片列表
-        content: list[ChunkData] = []
-        self.image_index = 0
-
-        # 加载工作簿
-        workbook = self._load_workbook(excel_path)
-
-        # 处理每个工作表
-        for sheet_index, sheet_name in enumerate(workbook.sheetnames):
-            sheet = workbook[sheet_name]
-
-            # 添加工作表标题
-            content.append(ChunkData(
-                type=ChunkType.TEXT,
-                name=sheet_name,
-                content=f"工作表 {sheet_index + 1}: {sheet_name}",
-                description="工作表标题"
-            ))
-
-            # 处理图片
-            sheet_images = self._extract_sheet_images(sheet)
-            content.extend(sheet_images)
-
-            # 处理表格数据
-            table_content = self._extract_table_data(sheet)
-            content.append(ChunkData(
-                type=ChunkType.TABLE,
-                name="表格",
-                content=json.dumps(table_content),
-                description="表格"
-            ))
-
-        # 添加结束文本
-        content.append(ChunkData(
-            type=ChunkType.TEXT,
-            name="结束文本",
-            content="",
-            description="结束文本"
-        ))
-
-        return title, content
-
     def _load_workbook(self, excel_path: str) -> Workbook:
         """
         加载Excel工作簿
@@ -250,13 +228,13 @@ def _get_image_format(self, img_obj: Image) -> str:
             return img_format
         return self.config.default_image_format
 
-    def _process_cell_value(self, cell_value: Any) -> CellValue:
+    def _process_cell_value(self, cell_value: Any) -> str:
         """
         预处理单元格值，将datetime对象转换为字符串
         Args:
             cell_value: 原始单元格值
         Returns:
-            CellValue: 处理后的单元格值
+            str: 处理后的单元格值
         """
         if cell_value is None:
             return ""
@@ -269,14 +247,10 @@ def _process_cell_value(self, cell_value: Any) -> CellValue:
         if isinstance(cell_value, date):
             return cell_value.strftime("%Y-%m-%d")
 
-        # 处理其他类型
-        if isinstance(cell_value, str|int|float|bool):
-            return cell_value
-
         # 对于其他类型，转换为字符串
         return str(cell_value)
 
-    def _extract_table_data(self, sheet: Worksheet) -> dict[str, Any]:
+    def _extract_table_data(self, sheet: Worksheet) -> TableDataItem:
         """
         提取表格数据
         Args:
@@ -295,16 +269,14 @@ def _extract_table_data(self, sheet: Worksheet) -> dict[str, Any]:
         # 提取所有数据
         all_rows = self._extract_all_rows(sheet, max_row, max_col, merged_map)
 
-        return {
-            "dimensions": {
-                "rows": len(all_rows),
-                "columns": max_col
-            },
-            "headers": all_rows[0] if all_rows else [],
-            "data": all_rows[1:] if len(all_rows) > 1 else []
-        }
+        return TableDataItem(
+            rows=len(all_rows),
+            columns=max_col,
+            row_headers=all_rows[0] if all_rows else [],
+            data=all_rows[1:] if len(all_rows) > 1 else []
+        )
 
-    def _get_merged_cells(self, sheet: Worksheet) -> dict[tuple[int, int, int, int], CellValue]:
+    def _get_merged_cells(self, sheet: Worksheet) -> dict[tuple[int, int, int, int], str]:
         """
         获取合并单元格信息
         Args:
@@ -323,7 +295,7 @@ def _get_merged_cells(self, sheet: Worksheet) -> dict[tuple[int, int, int, int],
                 merged_ranges[(min_row, min_col, max_row, max_col)] = merged_value
         return merged_ranges
 
-    def _create_merged_cell_map(self, merged_ranges: dict, sheet: Worksheet) -> dict[tuple[int, int], CellValue]:
+    def _create_merged_cell_map(self, merged_ranges: dict, sheet: Worksheet) -> dict[tuple[int, int], str]:
         """
         创建合并单元格映射
         Args:
@@ -342,7 +314,7 @@ def _create_merged_cell_map(self, merged_ranges: dict, sheet: Worksheet) -> dict
         return merged_map
 
     def _extract_all_rows(self, sheet: Worksheet, max_row: int, max_col: int,
-                          merged_map: dict[tuple[int, int], CellValue]) -> TableData:
+                          merged_map: dict[tuple[int, int], str]) -> list[list[str]]:
         """
         提取所有行数据
         Args:

tests/test_excel_parser.py

@@ -48,20 +48,22 @@ async def test_parse_real_basic_and_image():
             result = await parser.parse(xlsx_path)
 
             assert result.success is True
-            # 内容：Sheet1标题、Sheet1图片、Sheet1表格、Sheet2标题、Sheet2表格、结束文本
-            content = result.chunks
-            assert len(content) == 6
+            # 内容：Sheet1标题、Sheet1图片、Sheet1表格、Sheet2标题、Sheet2表格
+            content = result.tables
+            assert len(content) == 2
+
+            assert len(result.images) == 1
+            assert len(result.texts) == 2
 
             # 校验顺序与关键字段
-            assert content[0].type == "text" and content[0].name == "Sheet1"
-            assert content[1].type == "image"
-            assert content[1].name == "#/pictures/0"
-            assert content[1].content.startswith("data:image/")
-
-            assert content[2].type == "table"
-            assert content[3].type == "text" and content[3].name == "Sheet2"
-            assert content[4].type == "table"
-            assert content[5].type == "text" and content[5].name == "结束文本"
+            assert result.texts[0].type == "text" and result.texts[0].name == "Sheet1"
+            assert result.images[0].type == "image"
+            assert result.images[0].name == "#/pictures/0"
+            assert result.images[0].content.startswith("data:image/")
+
+            assert result.tables[0].type == "table"
+            assert result.texts[1].type == "text" and result.texts[1].name == "Sheet2"
+            assert result.tables[1].type == "table"
         finally:
             os.remove(xlsx_path)
     finally:
@@ -91,19 +93,19 @@ async def test_parse_real_merged_cells():
         result = await parser.parse(xlsx_path)
 
         assert result.success is True
-        content = result.chunks
-        # 结构：标题、表格、结束文本
-        assert len(content) == 3
+        # 结构：标题、表格
+        assert len(result.tables) == 1
+        assert len(result.texts) == 1
 
         # 表格在索引1
-        table_chunk: ChunkData = content[1]
+        table_chunk: ChunkData = result.tables[0]
         assert table_chunk.type == "table"
 
-        import json as _json
-        payload = _json.loads(table_chunk.content)
-        assert payload["headers"] == ["Merged Header", "Merged Header"]
-        assert payload["dimensions"]["rows"] == 2
-        assert payload["dimensions"]["columns"] == 2
+        payload = table_chunk.content
+        assert payload.row_headers == ["Merged Header", "Merged Header"]
+        assert payload.data == [["Value1", "Value2"]]
+        assert payload.rows == 2
+        assert payload.columns == 2
     finally:
         os.remove(xlsx_path)
 

worker.py

@@ -22,7 +22,7 @@ async def worker(app: Sanic) -> dict[str, Any]:
         parse_result = await parser_factory.parse_document(file_path)
         if not parse_result.success:
             continue
-        chunk_list = parse_result.chunks
+        chunk_list = parse_result.texts + parse_result.tables + parse_result.images + parse_result.formulas
         # 控制并发数量，防止访问量过大导致失败
         SEMAPHORE_LIMIT = 10  # 可根据实际情况调整
         semaphore = asyncio.Semaphore(SEMAPHORE_LIMIT)
@@ -35,5 +35,8 @@ async def enhance_with_semaphore(chunk: ChunkData, semaphore: asyncio.Semaphore)
         enhanced_chunk_list = await asyncio.gather(
             *(enhance_with_semaphore(chunk, semaphore) for chunk in chunk_list)
         )
-        parse_result.chunks = enhanced_chunk_list
+        parse_result.texts = enhanced_chunk_list[:len(parse_result.texts)]
+        parse_result.tables = enhanced_chunk_list[len(parse_result.texts):len(parse_result.texts) + len(parse_result.tables)]
+        parse_result.images = enhanced_chunk_list[len(parse_result.texts) + len(parse_result.tables):len(parse_result.texts) + len(parse_result.tables) + len(parse_result.images)]
+        parse_result.formulas = enhanced_chunk_list[len(parse_result.texts) + len(parse_result.tables) + len(parse_result.images):]
         return parse_result.model_dump(mode="json")