YG-Soft/YG_FT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(data-process): 完善文件解析与切分存储链路

Browse Source
This commit is contained in:
caoxiaozhu
2026-07-24 11:27:51 +08:00
parent 6d4bf85284
commit 33d0ed2e01
12 changed files with 4813 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files

546
backend/tests/test_data_process_algorithms.py
View File

@@ -1,23 +1,151 @@
from __future__ import annotations
import io
import json
import xml.etree.ElementTree as ET
import zipfile
from datetime import datetime
import pytest
from docx import Document
from openpyxl import Workbook
from pptx import Presentation
from pptx.util import Inches
from pypdf import PdfWriter
from app.modules.data_process.algorithms import (
chunk_unstructured,
content_quality_flags,
desensitize_pii,
desensitize_structured_record,
detect_document_structure,
detect_text_format,
estimate_token_count,
extract_pdf_page_texts,
extract_structured_records,
generate_standard_records,
is_near_duplicate,
merge_short_blocks,
normalize_text,
parse_text_content,
preprocess_structured_records,
record_fingerprint,
score_quality,
stable_split,
)
def _minimal_pdf(text: str = "Hello PDF") -> bytes:
stream = f"BT /F1 12 Tf 72 720 Td ({text}) Tj ET".encode("ascii")
objects = [
b"<< /Type /Catalog /Pages 2 0 R >>",
b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
(
b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] "
b"/Resources << /Font << /F1 5 0 R >> >> /Contents 4 0 R >>"
),
b"<< /Length " + str(len(stream)).encode("ascii") + b" >>\nstream\n"
+ stream
+ b"\nendstream",
b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
]
result = bytearray(b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n")
offsets = [0]
for object_number, value in enumerate(objects, start=1):
offsets.append(len(result))
result.extend(f"{object_number} 0 obj\n".encode("ascii"))
result.extend(value)
result.extend(b"\nendobj\n")
xref_offset = len(result)
result.extend(f"xref\n0 {len(objects) + 1}\n".encode("ascii"))
result.extend(b"0000000000 65535 f \n")
for offset in offsets[1:]:
result.extend(f"{offset:010d} 00000 n \n".encode("ascii"))
result.extend(
(
f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n"
f"startxref\n{xref_offset}\n%%EOF\n"
).encode("ascii")
)
return bytes(result)
def _aes_encrypted_pdf(*, user_password: str) -> bytes:
writer = PdfWriter(clone_from=io.BytesIO(_minimal_pdf()))
writer.encrypt(
user_password=user_password,
owner_password="owner-secret",
algorithm="AES-256",
)
output = io.BytesIO()
writer.write(output)
return output.getvalue()
def _docx_bytes() -> bytes:
document = Document()
document.add_heading("服务说明", level=1)
document.add_paragraph("这是 DOCX 正文。")
table = document.add_table(rows=1, cols=2)
table.cell(0, 0).text = "字段"
table.cell(0, 1).text = "内容"
output = io.BytesIO()
document.save(output)
return output.getvalue()
def _xlsx_bytes() -> bytes:
workbook = Workbook()
worksheet = workbook.active
worksheet.title = "数据"
worksheet.append(["name", "score", "created_at"])
worksheet.append(["Alice", 95, datetime(2026, 7, 23, 10, 30)])
worksheet.append(["Bob", 88, datetime(2026, 7, 24, 9, 0)])
output = io.BytesIO()
workbook.save(output)
workbook.close()
return output.getvalue()
def _xlsx_with_worksheet_relationship(
raw: bytes,
target: str,
*,
target_mode: str | None = None,
) -> bytes:
member_name = "xl/_rels/workbook.xml.rels"
source = io.BytesIO(raw)
output = io.BytesIO()
with zipfile.ZipFile(source) as original, zipfile.ZipFile(output, "w") as rewritten:
for member in original.infolist():
content = original.read(member.filename)
if member.filename == member_name:
root = ET.fromstring(content)
worksheet_relationship = next(
element
for element in root
if element.attrib.get("Type", "").endswith("/worksheet")
)
worksheet_relationship.set("Target", target)
if target_mode is None:
worksheet_relationship.attrib.pop("TargetMode", None)
else:
worksheet_relationship.set("TargetMode", target_mode)
content = ET.tostring(root, encoding="utf-8", xml_declaration=True)
rewritten.writestr(member, content)
return output.getvalue()
def _pptx_bytes() -> bytes:
presentation = Presentation()
slide = presentation.slides.add_slide(presentation.slide_layouts[6])
text_box = slide.shapes.add_textbox(Inches(1), Inches(1), Inches(6), Inches(1))
text_box.text = "PPTX 页面正文"
output = io.BytesIO()
presentation.save(output)
return output.getvalue()
def test_parse_utf8_json_jsonl_csv_markdown_and_txt() -> None:
parsed_json = parse_text_content(
b'\xef\xbb\xbf{"data":[{"name":"\xe5\xbc\xa0\xe4\xb8\x89"}]}',
@@ -44,6 +172,269 @@ def test_parse_utf8_json_jsonl_csv_markdown_and_txt() -> None:
assert parsed_txt.text == "普通文本"
def test_parse_pdf_docx_xlsx_and_pptx() -> None:
parsed_pdf = parse_text_content(_minimal_pdf(), filename="manual.pdf")
assert parsed_pdf.format == "pdf"
assert "Hello PDF" in parsed_pdf.text
assert parsed_pdf.records == ()
pdf_pages = extract_pdf_page_texts(_minimal_pdf())
assert len(pdf_pages) == 1
assert pdf_pages[0].page_number == 1
assert pdf_pages[0].text == "Hello PDF"
assert pdf_pages[0].source_start == 0
assert pdf_pages[0].source_end == len(parsed_pdf.text)
parsed_docx = parse_text_content(_docx_bytes(), filename="manual.docx")
assert parsed_docx.format == "docx"
assert "服务说明" in parsed_docx.text
assert "这是 DOCX 正文。" in parsed_docx.text
assert "字段\t内容" in parsed_docx.text
assert parsed_docx.records == ()
parsed_xlsx = parse_text_content(_xlsx_bytes(), filename="records.xlsx")
assert parsed_xlsx.format == "xlsx"
assert parsed_xlsx.records == (
{"name": "Alice", "score": 95, "created_at": "2026-07-23T10:30:00"},
{"name": "Bob", "score": 88, "created_at": "2026-07-24T09:00:00"},
)
assert json.loads(parsed_xlsx.text.splitlines()[0]) == parsed_xlsx.records[0]
parsed_pptx = parse_text_content(_pptx_bytes(), filename="slides.pptx")
assert parsed_pptx.format == "pptx"
assert parsed_pptx.text == "PPTX 页面正文"
assert parsed_pptx.records == ()
def test_xlsx_merged_multilevel_headers_are_flattened_without_losing_columns() -> None:
workbook = Workbook()
worksheet = workbook.active
worksheet.merge_cells("A1:A2")
worksheet.merge_cells("B1:C1")
worksheet["A1"] = "地区"
worksheet["B1"] = "销售"
worksheet["B2"] = "Q1"
worksheet["C2"] = "Q2"
worksheet.append(["华东", 100, 120])
output = io.BytesIO()
workbook.save(output)
workbook.close()
parsed = parse_text_content(output.getvalue(), filename="sales.xlsx")
assert parsed.records == ({"地区": "华东", "销售.Q1": 100, "销售.Q2": 120},)
def test_xlsx_header_inference_skips_more_than_eight_merged_report_titles() -> None:
workbook = Workbook()
worksheet = workbook.active
for row_number in range(1, 13):
worksheet.merge_cells(
start_row=row_number,
start_column=1,
end_row=row_number,
end_column=4,
)
worksheet.cell(row_number, 1, f"报表说明 {row_number}")
worksheet.append(["姓名", "部门", "得分", "日期"])
worksheet.append(["张三", "研发", 95, "2026-07-23"])
output = io.BytesIO()
workbook.save(output)
workbook.close()
parsed = parse_text_content(output.getvalue(), filename="report.xlsx")
assert parsed.records == (
{"姓名": "张三", "部门": "研发", "得分": 95, "日期": "2026-07-23"},
)
def test_xlsx_header_inference_ignores_continuous_body_merges() -> None:
workbook = Workbook()
worksheet = workbook.active
worksheet.append(["类别", "名称", "数量"])
worksheet.append(["水果", "苹果", 10])
worksheet.append([None, "香蕉", 12])
worksheet.append(["蔬菜", "白菜", 8])
worksheet.append([None, "萝卜", 9])
worksheet.merge_cells("A2:A3")
worksheet.merge_cells("A4:A5")
output = io.BytesIO()
workbook.save(output)
workbook.close()
parsed = parse_text_content(output.getvalue(), filename="inventory.xlsx")
assert parsed.records == (
{"类别": "水果", "名称": "苹果", "数量": 10},
{"类别": "", "名称": "香蕉", "数量": 12},
{"类别": "蔬菜", "名称": "白菜", "数量": 8},
{"类别": "", "名称": "萝卜", "数量": 9},
)
def test_xlsx_header_inference_supports_title_and_two_header_levels() -> None:
workbook = Workbook()
worksheet = workbook.active
worksheet.merge_cells("A1:C1")
worksheet["A1"] = "区域销售报表"
worksheet["A2"] = "统计日期"
worksheet["B2"] = "2026-07-23"
worksheet.merge_cells("A4:A5")
worksheet.merge_cells("B4:C4")
worksheet["A4"] = "地区"
worksheet["B4"] = "销售"
worksheet["B5"] = "Q1"
worksheet["C5"] = "Q2"
worksheet.append(["华南", 88, 92])
output = io.BytesIO()
workbook.save(output)
workbook.close()
parsed = parse_text_content(output.getvalue(), filename="two-level.xlsx")
assert parsed.records == (
{"地区": "华南", "销售.Q1": 88, "销售.Q2": 92},
)
def test_xlsx_header_inference_supports_title_and_three_header_levels() -> None:
workbook = Workbook()
worksheet = workbook.active
worksheet.merge_cells("A1:D1")
worksheet["A1"] = "年度销售分析报告"
worksheet["A2"] = "统计日期"
worksheet["B2"] = "2026-07-23"
worksheet.merge_cells("A4:A6")
worksheet.merge_cells("B4:D4")
worksheet.merge_cells("B5:C5")
worksheet.merge_cells("D5:D6")
worksheet["A4"] = "地区"
worksheet["B4"] = "销售"
worksheet["B5"] = "国内"
worksheet["D5"] = "海外"
worksheet["B6"] = "Q1"
worksheet["C6"] = "Q2"
worksheet.append(["华东", 100, 120, 80])
output = io.BytesIO()
workbook.save(output)
workbook.close()
parsed = parse_text_content(output.getvalue(), filename="three-level.xlsx")
assert parsed.records == (
{
"地区": "华东",
"销售.国内.Q1": 100,
"销售.国内.Q2": 120,
"销售.海外": 80,
},
)
def test_xlsx_header_inference_keeps_an_ordinary_single_header_row() -> None:
parsed = parse_text_content(_xlsx_bytes(), filename="ordinary.xlsx")
assert tuple(parsed.records[0]) == ("name", "score", "created_at")
assert len(parsed.records) == 2
@pytest.mark.parametrize(
"target",
[
"./worksheets/../worksheets/sheet1.xml",
"./worksheets/%2e%2e/worksheets/sheet1.xml",
"../xl/worksheets/sheet1.xml",
"/xl/worksheets/./sheet1.xml",
],
)
def test_xlsx_worksheet_relationship_allows_safe_dot_segments(target: str) -> None:
raw = _xlsx_with_worksheet_relationship(_xlsx_bytes(), target)
parsed = parse_text_content(raw, filename="records.xlsx")
assert parsed.records[0]["name"] == "Alice"
@pytest.mark.parametrize(
"target",
[
"../../outside.xml",
"worksheets\\sheet1.xml",
"%2e%2e/%2e%2e/outside.xml",
"%252e%252e/%252e%252e/outside.xml",
"https://example.com/sheet1.xml",
],
)
def test_xlsx_worksheet_relationship_rejects_path_traversal(target: str) -> None:
raw = _xlsx_with_worksheet_relationship(_xlsx_bytes(), target)
with pytest.raises(ValueError, match="unsafe worksheet path"):
parse_text_content(raw, filename="unsafe.xlsx")
def test_xlsx_worksheet_relationship_rejects_external_and_missing_targets() -> None:
external = _xlsx_with_worksheet_relationship(
_xlsx_bytes(),
"https://example.com/sheet1.xml",
target_mode="External",
)
with pytest.raises(ValueError, match="external relationship"):
parse_text_content(external, filename="external.xlsx")
missing = _xlsx_with_worksheet_relationship(
_xlsx_bytes(),
"worksheets/missing.xml",
)
with pytest.raises(ValueError, match="target does not exist"):
parse_text_content(missing, filename="missing.xlsx")
@pytest.mark.parametrize(
("filename", "replacement"),
[
("legacy.doc", ".docx"),
("legacy.xls", ".xlsx"),
("legacy.ppt", ".pptx"),
],
)
def test_legacy_office_formats_require_conversion(filename: str, replacement: str) -> None:
with pytest.raises(ValueError, match=rf"convert the file to \{replacement}"):
parse_text_content(b"legacy", filename=filename)
def test_office_zip_bomb_and_invalid_pdf_are_rejected_before_parsing() -> None:
archive = io.BytesIO()
with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as package:
package.writestr("[Content_Types].xml", "<Types/>")
package.writestr("word/document.xml", b"A" * (2 * 1024 * 1024))
with pytest.raises(ValueError, match="unsafe compression ratio"):
parse_text_content(archive.getvalue(), filename="unsafe.docx")
active_xml = io.BytesIO()
with zipfile.ZipFile(active_xml, "w") as package:
package.writestr("[Content_Types].xml", "<Types/>")
package.writestr(
"word/document.xml",
'<!DOCTYPE document [<!ENTITY xxe SYSTEM "file:///etc/passwd">]><document/>',
)
with pytest.raises(ValueError, match="unsupported active XML"):
parse_text_content(active_xml.getvalue(), filename="active.docx")
with pytest.raises(ValueError, match="missing PDF header"):
parse_text_content(b"not a pdf", filename="broken.pdf")
blank_pdf = io.BytesIO()
blank_writer = PdfWriter()
blank_writer.add_blank_page(width=612, height=792)
blank_writer.write(blank_pdf)
with pytest.raises(ValueError, match="scanned PDF requires OCR"):
parse_text_content(blank_pdf.getvalue(), filename="scanned.pdf")
aes_pdf_without_open_password = parse_text_content(
_aes_encrypted_pdf(user_password=""),
filename="aes-no-password.pdf",
)
assert "Hello PDF" in aes_pdf_without_open_password.text
with pytest.raises(ValueError, match="password-protected PDF files are not supported"):
parse_text_content(
_aes_encrypted_pdf(user_password="secret"),
filename="aes-password.pdf",
)
def test_invalid_utf8_and_malformed_structured_content_fail_loudly() -> None:
with pytest.raises(ValueError, match="not valid UTF-8"):
parse_text_content(b"\xff\xfe", filename="broken.txt")
@@ -80,7 +471,89 @@ def test_desensitize_pii_returns_masked_text_and_counts() -> None:
assert counts == {"email": 1, "phone": 1, "id_card": 1, "total": 3}
@pytest.mark.parametrize("method", ["semantic", "heading", "fixed", "custom"])
def test_every_structured_preprocess_option_has_independent_behavior() -> None:
clean_source = [
{"id": "1", "name": "有效", "empty_column": ""},
{"id": "", "name": "缺少关键字段", "empty_column": ""},
{"id": "2", "name": "有效", "empty_column": ""},
]
assert preprocess_structured_records(clean_source, []) == clean_source
assert preprocess_structured_records(clean_source, ["clean_invalid"]) == [
{"id": "1", "name": "有效"},
{"id": "2", "name": "有效"},
]
nested = [{"id": 1, "profile": {"name": "张三", "level": 2}}]
assert "profile" in preprocess_structured_records(nested, [])[0]
assert preprocess_structured_records(nested, ["detect_structure"])[0] == {
"id": 1,
"profile.name": "张三",
"profile.level": 2,
}
duplicates = [
{"customer_id": "C-1", "value": "first"},
{"customer_id": "C-1", "value": "updated"},
{"customer_id": "", "value": "blank-one"},
{"customer_id": "", "value": "blank-two"},
]
assert len(preprocess_structured_records(duplicates, [])) == 4
deduplicated = preprocess_structured_records(duplicates, ["deduplicate"])
assert [record["value"] for record in deduplicated] == [
"first",
"blank-one",
"blank-two",
]
unnormalized = [{" User Name ": "ＡＢＣ\r\n第二行"}]
assert preprocess_structured_records(unnormalized, []) == unnormalized
assert preprocess_structured_records(unnormalized, ["normalize_format"]) == [
{"user_name": "ABC\n第二行"}
]
anomaly_source = [
{"id": 10_000 + index, "amount": amount, "text": "正常内容"}
for index, amount in enumerate((10, 10, 11, 11, 12, 12, 13, 1000))
]
assert len(preprocess_structured_records(anomaly_source, [])) == 8
filtered = preprocess_structured_records(anomaly_source, ["filter_anomaly"])
assert len(filtered) == 7
assert all(record["amount"] != 1000 for record in filtered)
assert max(record["id"] for record in filtered) > 10_000
sensitive = [{"姓名": "张三", "phone": "13800138000", "email": "a@b.com"}]
assert preprocess_structured_records(sensitive, []) == sensitive
masked = preprocess_structured_records(sensitive, ["desensitize"])[0]
assert masked == {"姓名": "[NAME]", "phone": "[PHONE]", "email": "[EMAIL]"}
def test_structured_desensitization_counts_and_document_helpers() -> None:
masked, counts = desensitize_structured_record(
{"联系人姓名": "李四", "说明": "邮箱 user@example.com，手机 13900139000"}
)
assert masked == {
"联系人姓名": "[NAME]",
"说明": "邮箱 [EMAIL]，手机 [PHONE]",
}
assert counts == {"email": 1, "phone": 1, "id_card": 0, "name": 1, "total": 3}
structure = detect_document_structure(
"# 第一章\n正文\n\n## 细节\n- 项目一\n- 项目二\n\n```python\nprint(1)\n```"
)
assert [heading.title for heading in structure.headings] == ["第一章", "细节"]
assert structure.list_block_count == 1
assert structure.code_block_count == 1
assert merge_short_blocks(["短一", "短二", "这是一段足够长的正文内容"], min_token_count=4)
assert "mojibake" in content_quality_flags("正常文字锟斤拷内容", min_chars=0, min_tokens=0)
assert is_near_duplicate(
"alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron",
"alpha beta gamma, delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron",
similarity_threshold=0.92,
max_hamming_distance=2,
)
@pytest.mark.parametrize("method", ["structure", "fixed", "custom"])
def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
text = "# 第一章\n" + "甲。" * 18 + "\n# 第二章\n" + "乙。" * 18
kwargs = {"custom_delimiter": "\\n"} if method == "custom" else {}
@@ -99,8 +572,38 @@ def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
assert all(chunk.start_line <= chunk.end_line for chunk in chunks)
def test_fixed_chunk_overlap_is_exact_when_chunks_are_large_enough() -> None:
text = " ".join(f"token{i}" for i in range(30))
def test_default_and_structure_chunking_split_headings_without_cross_section_overlap() -> None:
text = (
"# 第一章\n"
+ " ".join(f"alpha{i}" for i in range(18))
+ "\n# 第二章\n"
+ " ".join(f"beta{i}" for i in range(18))
)
normalized = normalize_text(text)
second_chapter_start = normalized.index("# 第二章")
kwargs = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4}
default_chunks = chunk_unstructured(text, **kwargs)
structure_chunks = chunk_unstructured(text, method="structure", **kwargs)
assert default_chunks == structure_chunks
assert len(structure_chunks) > 2
assert all(
chunk.content == normalized[chunk.start : chunk.end] for chunk in structure_chunks
)
assert all(
not (chunk.start < second_chapter_start < chunk.end) for chunk in structure_chunks
)
second_chapter_chunks = [
chunk for chunk in structure_chunks if chunk.start >= second_chapter_start
]
assert second_chapter_chunks[0].start == second_chapter_start
assert second_chapter_chunks[0].content.startswith("# 第二章")
def test_fixed_chunk_offsets_and_actual_token_overlap_are_exact() -> None:
text = " ".join(f"token{i}" for i in range(30))
normalized = normalize_text(text)
chunks = chunk_unstructured(
text,
method="fixed",
@@ -108,10 +611,16 @@ def test_fixed_chunk_overlap_is_exact_when_chunks_are_large_enough() -> None:
chunk_overlap=3,
min_chunk_size=4,
)
first_tokens = chunks[0].content.split()
second_tokens = chunks[1].content.split()
assert first_tokens[-3:] == second_tokens[:3]
assert chunks[0].token_count == 10
assert len(chunks) > 2
assert all(chunk.content == normalized[chunk.start : chunk.end] for chunk in chunks)
assert all(chunk.token_count == estimate_token_count(chunk.content) for chunk in chunks)
assert all(chunk.token_count == 10 for chunk in chunks[:-1])
for left, right in zip(chunks, chunks[1:]):
overlap_text = normalized[right.start : left.end]
assert right.start < left.end
assert estimate_token_count(overlap_text) == 3
assert left.content.endswith(overlap_text)
assert right.content.startswith(overlap_text)
def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
@@ -129,18 +638,7 @@ def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
assert chunks[1].start_line == 2
def test_heading_and_custom_boundaries_are_respected() -> None:
heading_text = "前言 " * 8 + "\n# 第二章\n" + "正文 " * 12
heading_chunks = chunk_unstructured(
heading_text,
method="heading",
chunk_size=20,
chunk_overlap=0,
min_chunk_size=4,
)
assert "# 第二章" not in heading_chunks[0].content
assert heading_chunks[1].content.startswith("#")
def test_custom_delimiter_is_preserved_as_the_chunk_boundary() -> None:
custom_chunks = chunk_unstructured(
"a b c d <CUT> e f g h i j",
method="custom",
@@ -172,6 +670,13 @@ def test_heading_and_custom_boundaries_are_respected() -> None:
)
def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None:
text = "前言。" * 15 + "\n" + block + "\n" + "结尾。" * 40
unprotected = chunk_unstructured(
text,
method="fixed",
chunk_size=40,
chunk_overlap=0,
min_chunk_size=10,
)
chunks = chunk_unstructured(
text,
method="fixed",
@@ -180,6 +685,7 @@ def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None
min_chunk_size=10,
**{field: True},
)
assert all(block not in chunk.content for chunk in unprotected)
assert any(block in chunk.content for chunk in chunks)
@@ -194,6 +700,8 @@ def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None
"cannot exceed",
),
({"method": "custom", "custom_delimiter": ""}, "custom_delimiter"),
({"method": "semantic"}, "unsupported chunk method"),
({"method": "heading"}, "unsupported chunk method"),
],
)
def test_chunk_configuration_validation(kwargs: dict[str, object], message: str) -> None: