first commit

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-16 11:12:17 +08:00
commit 4003624b8c
80 changed files with 9990 additions and 0 deletions
@@ -0,0 +1,80 @@
+"""Tests for spreadsheet layout detection and row splitting."""
+
+from __future__ import annotations
+
+import unittest
+
+from rag_cut.parsers.pdf.tables import detect_spreadsheet_layout, is_qa_style_table
+from rag_cut.pipeline import chunk_document
+from rag_cut.splitters.by_row import split_by_row
+from rag_cut.models import Block, BlockType, SplitConfig, SplitMode
+
+
+BACKMAN_HEADER = [
+    [
+        "Session会话（必填）：用于标识1个对话",
+        "query 用户输入（必填）：消息内容",
+        "用户ID（必填）",
+        "使用的大语言模型（必填）",
+        "要求AI回复的语言（必填）",
+        "reference_output 标准答案（可选）",
+    ],
+    ["session", "query", "userid", "model", "lang", "reference_output"],
+    ["1", "账户余额是多少？", "SUPPORT", "gpt-4o-mini", "zh", "SELECT 1"],
+    ["2", "今日有哪些账户透支？", "SUPPORT", "gpt-4o-mini", "zh", "SELECT 2"],
+]
+
+
+class SpreadsheetLayoutTest(unittest.TestCase):
+    def test_detect_template_description_and_header_rows(self) -> None:
+        layout = detect_spreadsheet_layout(BACKMAN_HEADER)
+        self.assertEqual(layout["preamble_rows"], 1)
+        self.assertEqual(layout["header_row_start"], 2)
+        self.assertEqual(layout["header_row_end"], 2)
+        self.assertEqual(layout["data_start_row"], 3)
+        self.assertTrue(is_qa_style_table(BACKMAN_HEADER, layout))
+
+    def test_split_by_row_uses_real_header_and_one_row_per_chunk(self) -> None:
+        layout = detect_spreadsheet_layout(BACKMAN_HEADER)
+        block = Block(
+            type=BlockType.TABLE,
+            markdown="",
+            meta={"rows": BACKMAN_HEADER, "table_title": "testset", **layout},
+        )
+        groups = split_by_row(
+            [block],
+            SplitConfig(
+                mode=SplitMode.BY_ROW,
+                header_row_start=layout["header_row_start"],
+                header_row_end=layout["header_row_end"],
+                start_row=layout["data_start_row"],
+                rows_per_chunk=1,
+            ),
+        )
+        self.assertEqual(len(groups), 2)
+        first_md = groups[0][0].markdown
+        self.assertIn("| session | query | userid | model | lang | reference_output |", first_md)
+        self.assertIn("账户余额是多少？", first_md)
+        self.assertNotIn("Session会话（必填）", first_md)
+
+
+class BackmanFixtureTest(unittest.TestCase):
+    def test_backman_upload_chunks_one_row_each(self) -> None:
+        from pathlib import Path
+
+        uploads = list((Path(__file__).resolve().parents[2] / "storage" / "uploads").rglob("BackmanAI*.xlsx"))
+        if not uploads:
+            self.skipTest("BackmanAI fixture not uploaded")
+        result = chunk_document(uploads[0])
+        self.assertEqual(result.split_config["rows_per_chunk"], 1)
+        self.assertEqual(result.split_config["header_row_start"], 2)
+        self.assertEqual(result.split_config["start_row"], 3)
+        self.assertGreaterEqual(result.chunk_count, 40)
+        first = result.chunks[0].content
+        self.assertIn("| session | query |", first)
+        self.assertNotIn("Session会话（必填）", first)
+        self.assertLess(result.chunks[0].char_count, 1200)
+
+
+if __name__ == "__main__":
+    unittest.main()