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

1
.gitignore vendored
View File

@@ -44,6 +44,7 @@ pip-delete-this-directory.txt
# Runtime data and logs
runtime/
backend/runtime/
backend/storage/
logs/
backend/logs/
*.db

697
backend/app/api/v1/endpoints/data_process.py
View File

@@ -3,13 +3,15 @@ from __future__ import annotations
import hashlib
import ipaddress
import json
import logging
import os
import re
import socket
from contextlib import contextmanager
from dataclasses import asdict
from pathlib import Path
from typing import Any, Iterator, Literal
from urllib.parse import urlsplit
from urllib.parse import quote, urlsplit
import psycopg
from fastapi import (
@@ -18,22 +20,38 @@ from fastapi import (
Body,
Depends,
File,
Header,
HTTPException,
Query,
UploadFile,
)
from fastapi.responses import StreamingResponse
from psycopg.rows import dict_row
from app.modules.data_process.algorithms import (
ParsedText,
TextChunk,
canonical_record_json,
chunk_unstructured,
decode_utf8,
content_quality_flags,
desensitize_pii,
desensitize_structured_record,
detect_document_structure,
estimate_token_count,
extract_pdf_page_texts,
generate_standard_records,
is_near_duplicate,
near_duplicate_fingerprint,
parse_text_content,
preprocess_structured_records,
score_quality,
)
from app.modules.data_process.generation import generate_model_records
from app.modules.data_process.storage import (
LocalDataProcessStorage,
StagedSourceObject,
get_data_process_storage,
)
from app.modules.data_process.store import (
ConflictError,
DataProcessStore,
@@ -41,11 +59,12 @@ from app.modules.data_process.store import (
InvalidStateError,
NotFoundError,
get_data_process_store,
new_id,
)
from app.schemas.data_process import (
DataProcessStatus,
DataProcessTaskCreate,
DataProcessTaskUpdate,
DataProcessStatus,
ExternalPullRequest,
ExternalSourceRequest,
GenerateRequest,
@@ -57,12 +76,34 @@ from app.schemas.data_process import (
ResultUpdate,
)
router = APIRouter(prefix="/data-process")
logger = logging.getLogger(__name__)
MAX_SOURCE_FILE_BYTES = 200 * 1024 * 1024
MAX_SOURCE_FILE_COUNT = 20
MAX_SOURCE_BATCH_BYTES = 500 * 1024 * 1024
MAX_EXTERNAL_PULL_BYTES = 50 * 1024 * 1024
STRUCTURED_SOURCE_SUFFIXES = frozenset(
{".json", ".jsonl", ".ndjson", ".csv", ".tsv", ".xlsx"}
)
UNSTRUCTURED_SOURCE_SUFFIXES = frozenset(
{
".txt",
".md",
".markdown",
".pdf",
".docx",
".pptx",
".json",
".jsonl",
".ndjson",
}
)
SUPPORTED_SOURCE_SUFFIXES = STRUCTURED_SOURCE_SUFFIXES | UNSTRUCTURED_SOURCE_SUFFIXES
LEGACY_OFFICE_CONVERSIONS = {
".doc": ".docx",
".xls": ".xlsx",
".ppt": ".pptx",
}
def ok(data: Any = None, message: str = "ok") -> dict[str, Any]:
@@ -99,6 +140,57 @@ def _safe_file_name(value: str | None, fallback: str) -> str:
return name if name not in {"", ".", ".."} else fallback
def _range_not_satisfiable(size: int) -> HTTPException:
return HTTPException(
status_code=416,
detail={"code": 416, "message": "invalid source byte range", "data": None},
headers={"Content-Range": f"bytes */{size}"},
)
def _source_byte_range(value: str | None, size: int) -> tuple[int, int] | None:
if value is None:
return None
match = re.fullmatch(r"bytes=(\d*)-(\d*)", value.strip())
if match is None or size <= 0:
raise _range_not_satisfiable(size)
start_text, end_text = match.groups()
if not start_text and not end_text:
raise _range_not_satisfiable(size)
if start_text:
start = int(start_text)
end = int(end_text) if end_text else size - 1
if start >= size or end < start:
raise _range_not_satisfiable(size)
else:
suffix_length = int(end_text)
if suffix_length <= 0:
raise _range_not_satisfiable(size)
start = max(0, size - suffix_length)
end = size - 1
return start, min(end, size - 1)
def _commit_source_batch(
store: DataProcessStore,
storage: LocalDataProcessStorage,
task_id: str,
prepared: list[dict[str, Any]],
staged: list[StagedSourceObject],
) -> list[dict[str, Any]]:
storage.publish(staged)
try:
return store.add_source_files(task_id, prepared)
except Exception:
for item in staged:
try:
storage.delete(item.reference)
except Exception:
# 文件系统回滚失败不能覆盖数据库抛出的根因，并继续清理其余对象。
pass
raise
def _value(config: dict[str, Any], snake_name: str, camel_name: str, default: Any) -> Any:
if snake_name in config:
return config[snake_name]
@@ -126,6 +218,204 @@ def _preview_quality(content: str, config: dict[str, Any]) -> dict[str, Any]:
)
def _parse_stored_source(source: dict[str, Any]) -> ParsedText:
"""重新读取已入库的规范化正文，避免把二进制格式当二进制重复解析。"""
content = str(source.get("content") or "")
file_format = str(source.get("file_format") or "").lower()
if file_format == "xlsx":
# XLSX 上传阶段已安全解析为 JSONL 后入库。
return parse_text_content(content, file_format="jsonl")
if file_format in {"pdf", "docx", "pptx"}:
# 文档上传阶段已抽取文本，预览阶段只需要对正文切片。
return ParsedText(format=file_format, text=content, records=())
return parse_text_content(
content,
filename=source.get("name"),
file_format=file_format or None,
)
def _shift_chunk(chunk: TextChunk, offset: int, source_text: str) -> TextChunk:
start = offset + chunk.start
end = offset + chunk.end
return TextChunk(
content=chunk.content,
start=start,
end=end,
start_line=source_text.count("\n", 0, start) + 1,
end_line=source_text.count("\n", 0, max(start, end - 1)) + 1,
token_count=chunk.token_count,
)
def _merge_short_chunks(
chunks: list[TextChunk],
source_text: str,
*,
min_token_count: int,
chunk_size: int,
) -> list[TextChunk]:
"""在同一章节内合并短块，同时严格遵守切片大小上限。"""
merged: list[TextChunk] = []
index = 0
while index < len(chunks):
current = chunks[index]
if current.token_count >= min_token_count:
merged.append(current)
index += 1
continue
candidates: list[tuple[int, int, int]] = []
if merged:
candidates.append((merged[-1].start, current.end, -1))
if index + 1 < len(chunks):
candidates.append((current.start, chunks[index + 1].end, 1))
selected = next(
(
(start, end, direction)
for start, end, direction in candidates
if estimate_token_count(source_text[start:end]) <= chunk_size
),
None,
)
if selected is None:
merged.append(current)
index += 1
continue
start, end, direction = selected
combined = TextChunk(
content=source_text[start:end],
start=start,
end=end,
start_line=source_text.count("\n", 0, start) + 1,
end_line=source_text.count("\n", 0, max(start, end - 1)) + 1,
token_count=estimate_token_count(source_text[start:end]),
)
if direction < 0:
merged[-1] = combined
index += 1
else:
merged.append(combined)
index += 2
return merged
def _chunk_source_text(
text: str,
config: dict[str, Any],
preprocess_options: set[str],
) -> list[tuple[TextChunk, tuple[str, ...]]]:
"""按可选文档结构分段后切片，章节之间绝不共享 overlap。"""
method = str(_value(config, "chunk_method", "chunkMethod", "structure"))
detect_structure = (
method == "structure" or "detect_document_structure" in preprocess_options
)
preserve_context = "preserve_context" in preprocess_options
merge_short = "merge_short_content" in preprocess_options
chunk_size = int(_value(config, "chunk_size", "chunkSize", 800))
configured_minimum = int(_value(config, "min_chunk_size", "minChunkSize", 100))
minimum = configured_minimum if merge_short else 1
overlap = (
int(_value(config, "chunk_overlap", "chunkOverlap", 100))
if preserve_context
else 0
)
common = {
"method": method,
"chunk_size": chunk_size,
"chunk_overlap": overlap,
"min_chunk_size": minimum,
"custom_delimiter": str(
_value(config, "custom_delimiter", "customDelimiter", "") or ""
),
"preserve_code_blocks": bool(
_value(config, "preserve_code_blocks", "preserveCodeBlocks", False)
),
"preserve_tables": bool(
_value(config, "preserve_tables", "preserveTables", False)
),
"preserve_lists": bool(
_value(config, "preserve_lists", "preserveLists", False)
),
}
sections: list[tuple[int, int, tuple[str, ...]]] = [(0, len(text), ())]
if detect_structure:
structure = detect_document_structure(text)
if structure.headings:
sections = []
first_start = structure.headings[0].start
if first_start > 0 and text[:first_start].strip():
sections.append((0, first_start, ()))
stack: list[tuple[int, str]] = []
for index, heading in enumerate(structure.headings):
while stack and stack[-1][0] >= heading.level:
stack.pop()
stack.append((heading.level, heading.title))
end = (
structure.headings[index + 1].start
if index + 1 < len(structure.headings)
else len(text)
)
sections.append((heading.start, end, tuple(title for _, title in stack)))
result: list[tuple[TextChunk, tuple[str, ...]]] = []
for start, end, heading_path in sections:
section_text = text[start:end]
local_chunks = chunk_unstructured(section_text, **common)
shifted = [_shift_chunk(chunk, start, text) for chunk in local_chunks]
if merge_short:
shifted = _merge_short_chunks(
shifted,
text,
min_token_count=configured_minimum,
chunk_size=chunk_size,
)
result.extend((chunk, heading_path) for chunk in shifted)
return result
_NEGATION_MARKERS = frozenset({"不", "无", "未", "否", "没有", "并非", "not", "no", "never"})
def _safe_near_duplicate(left: str, right: str) -> bool:
"""保守判断近重复，数字或否定含义变化时始终保留。"""
if min(estimate_token_count(left), estimate_token_count(right)) < 20:
return False
if re.findall(r"\d+(?:\.\d+)?", left) != re.findall(r"\d+(?:\.\d+)?", right):
return False
left_lower = left.casefold()
right_lower = right.casefold()
left_negations = {marker for marker in _NEGATION_MARKERS if marker in left_lower}
right_negations = {marker for marker in _NEGATION_MARKERS if marker in right_lower}
if left_negations != right_negations:
return False
return is_near_duplicate(
left,
right,
similarity_threshold=0.92,
max_hamming_distance=2,
)
def _near_duplicate_band_keys(content: str) -> tuple[tuple[int, int], ...]:
"""将 64 位 SimHash 分为三段，汉明距离不超过 2 时至少命中一段。"""
fingerprint = int(near_duplicate_fingerprint(content), 16)
widths = (22, 21, 21)
shift = 0
keys: list[tuple[int, int]] = []
for index, width in enumerate(widths):
keys.append((index, (fingerprint >> shift) & ((1 << width) - 1)))
shift += width
return tuple(keys)
def _build_preview_items(
task: dict[str, Any], source_files: list[dict[str, Any]]
) -> list[dict[str, Any]]:
@@ -140,6 +430,7 @@ def _build_preview_items(
preprocess_options & {"deduplicate", "deduplicate_content"}
)
seen_content_hashes: set[str] = set()
seen_near_duplicate_bands: dict[tuple[int, int], list[str]] = {}
items: list[dict[str, Any]] = []
def append_item(item: dict[str, Any]) -> None:
@@ -155,38 +446,58 @@ def _build_preview_items(
items.append(item)
for source in source_files:
parsed = parse_text_content(
source.get("content") or "",
filename=source.get("name"),
file_format=source.get("file_format"),
)
parsed = _parse_stored_source(source)
if process_type == "unstructured":
chunks = chunk_unstructured(
parsed.text,
method=_value(config, "chunk_method", "chunkMethod", "semantic"),
chunk_size=int(_value(config, "chunk_size", "chunkSize", 800)),
chunk_overlap=int(_value(config, "chunk_overlap", "chunkOverlap", 100)),
min_chunk_size=int(_value(config, "min_chunk_size", "minChunkSize", 100)),
custom_delimiter=str(
_value(config, "custom_delimiter", "customDelimiter", "") or ""
),
preserve_code_blocks=bool(
_value(config, "preserve_code_blocks", "preserveCodeBlocks", False)
),
preserve_tables=bool(
_value(config, "preserve_tables", "preserveTables", False)
),
preserve_lists=bool(
_value(config, "preserve_lists", "preserveLists", False)
),
chunks = _chunk_source_text(parsed.text, config, preprocess_options)
for chunk, heading_path in chunks:
preprocess_flags = content_quality_flags(
chunk.content,
min_chars=0,
min_tokens=0,
)
for chunk in chunks:
flag_set = set(preprocess_flags)
if "clean_invalid_content" in preprocess_options and flag_set & {
"empty_content",
"low_printable_ratio",
"repetitive_content",
}:
continue
if "filter_low_quality" in preprocess_options and flag_set & {
"content_too_long",
"mojibake",
"low_printable_ratio",
"repetitive_content",
}:
continue
if "deduplicate_content" in preprocess_options:
band_keys = _near_duplicate_band_keys(chunk.content)
candidates = {
previous
for key in band_keys
for previous in seen_near_duplicate_bands.get(key, ())
}
if any(
_safe_near_duplicate(chunk.content, previous)
for previous in candidates
):
continue
for key in band_keys:
seen_near_duplicate_bands.setdefault(key, []).append(chunk.content)
content = chunk.content
pii_counts: dict[str, int] = {}
if should_desensitize:
content, pii_counts = desensitize_pii(content)
quality = _preview_quality(content, config)
quality["pii_replacements"] = pii_counts
quality["preprocess_flags"] = list(preprocess_flags)
chunk_method = str(
_value(config, "chunk_method", "chunkMethod", "structure")
)
if (
chunk_method == "structure"
or "detect_document_structure" in preprocess_options
):
quality["heading_path"] = list(heading_path)
append_item(
{
"source_file_id": source["id"],
@@ -203,19 +514,41 @@ def _build_preview_items(
)
continue
record_contents = [
json.dumps(record, ensure_ascii=False, separators=(",", ":"))
for record in parsed.records
if not should_clean_invalid
or any(value not in (None, "", [], {}) for value in record.values())
]
if not record_contents and parsed.text:
record_contents = [parsed.text]
for content in record_contents:
original_content = content
pii_counts = {}
structured_options = preprocess_options & {
"clean_invalid",
"detect_structure",
"deduplicate",
"normalize_format",
"filter_anomaly",
}
source_records = list(parsed.records)
processed_records = preprocess_structured_records(
source_records,
structured_options,
)
if not processed_records and parsed.text and not source_records:
processed_records = [{"value": parsed.text}]
same_cardinality = len(processed_records) == len(source_records)
for index, record in enumerate(processed_records):
original_record = source_records[index] if same_cardinality else record
original_content = json.dumps(
original_record,
ensure_ascii=False,
separators=(",", ":"),
)
pii_counts: dict[str, int] = {}
edited_record = record
if should_desensitize:
content, pii_counts = desensitize_pii(content)
edited_record, pii_counts = desensitize_structured_record(record)
content = (
canonical_record_json(edited_record)
if "normalize_format" in preprocess_options
else json.dumps(
edited_record,
ensure_ascii=False,
separators=(",", ":"),
)
)
quality = _preview_quality(content, config)
quality["pii_replacements"] = pii_counts
append_item(
@@ -488,47 +821,75 @@ async def upload_source_files(
task_id: str,
files: list[UploadFile] = File(...),
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
if not files:
raise fail(400, "at least one source file is required")
if len(files) > MAX_SOURCE_FILE_COUNT:
raise fail(413, f"a source batch may contain at most {MAX_SOURCE_FILE_COUNT} files")
prepared: list[dict[str, Any]] = []
staged: list[StagedSourceObject] = []
batch_id = storage.new_batch_id()
batch_size = 0
commit_attempted = False
try:
with api_errors():
store.get_task(task_id)
task = store.get_task(task_id)
process_type = str(task["process_type"])
if process_type == "external":
raise InvalidStateError(
"external tasks must import data through the external source endpoint"
)
allowed_suffixes = (
UNSTRUCTURED_SOURCE_SUFFIXES
if process_type == "unstructured"
else STRUCTURED_SOURCE_SUFFIXES
)
for upload in files:
raw = await upload.read(MAX_SOURCE_FILE_BYTES + 1)
if len(raw) > MAX_SOURCE_FILE_BYTES:
raise fail(413, f"source file exceeds {MAX_SOURCE_FILE_BYTES} bytes")
content = decode_utf8(raw)
name = _safe_file_name(upload.filename, "source.txt")
suffix = Path(name).suffix.lower()
if suffix not in {
".txt",
".md",
".markdown",
".csv",
".tsv",
".json",
".jsonl",
".ndjson",
}:
if suffix in LEGACY_OFFICE_CONVERSIONS:
replacement = LEGACY_OFFICE_CONVERSIONS[suffix]
raise fail(
415,
f"legacy {suffix} format is not supported; "
f"convert the file to {replacement} and upload again",
)
if suffix not in SUPPORTED_SOURCE_SUFFIXES:
raise fail(415, f"unsupported source file format: {suffix or 'none'}")
parsed = parse_text_content(content, filename=name)
if suffix not in allowed_suffixes:
raise fail(
415,
f"{suffix} is not supported for {process_type} data processing",
)
parsed = parse_text_content(raw, filename=name)
if not parsed.text:
raise fail(400, f"source file is empty: {name}")
normalized_raw = parsed.text.encode("utf-8")
batch_size += len(normalized_raw)
batch_size += len(raw)
if batch_size > MAX_SOURCE_BATCH_BYTES:
raise fail(413, f"source batch exceeds {MAX_SOURCE_BATCH_BYTES} bytes")
source_file_id = new_id("dpsf")
staged_object = storage.stage_bytes(
batch_id=batch_id,
task_id=task_id,
source_file_id=source_file_id,
version=1,
name=name,
content=raw,
)
staged.append(staged_object)
record_count = len(parsed.records) or (1 if parsed.text else 0)
prepared.append(
{
"id": source_file_id,
"storage_object_id": staged_object.reference,
"name": name,
"content": parsed.text,
"raw_size": len(normalized_raw),
"checksum_sha256": hashlib.sha256(normalized_raw).hexdigest(),
"raw_size": len(raw),
"checksum_sha256": hashlib.sha256(raw).hexdigest(),
"file_format": parsed.format,
"record_count": record_count,
"metadata": {
@@ -539,7 +900,12 @@ async def upload_source_files(
"created_by": None,
}
)
created = store.add_source_files(task_id, prepared)
# publish 无论成功或失败都会消费并清理暂存对象，外层不能再次 discard。
commit_attempted = True
created = _commit_source_batch(store, storage, task_id, prepared, staged)
finally:
if not commit_attempted:
storage.discard(staged)
return ok({"files": created}, "source files uploaded")
@@ -559,15 +925,145 @@ def source_file_content(
return ok(store.source_content_window(task_id, file_id, offset, limit))
@router.get("/{task_id}/source-files/{file_id}/raw")
def source_file_raw(
task_id: str,
file_id: str,
range_header: str | None = Header(default=None, alias="Range"),
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> StreamingResponse:
with api_errors():
source = store.get_source_file(task_id, file_id, include_content=False)
if str(source.get("file_format") or "").lower() != "pdf":
raise fail(415, "raw inline preview is only available for PDF source files")
storage_object_id = str(source.get("storage_object_id") or "")
actual_size = storage.file_size(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
if actual_size is None:
raise fail(410, "the original PDF is unavailable for this legacy source file")
expected_size = int(source.get("size_bytes") or 0)
if actual_size != expected_size:
raise ValueError("source object size does not match metadata")
selected_range = _source_byte_range(range_header, actual_size)
start, end = selected_range or (0, actual_size - 1)
length = end - start + 1
name = _safe_file_name(str(source.get("name") or "source.pdf"), "source.pdf")
headers = {
"Accept-Ranges": "bytes",
"Cache-Control": "private, no-store",
"Content-Disposition": f"inline; filename*=UTF-8''{quote(name, safe='')}",
"Content-Length": str(length),
"X-Accel-Buffering": "no",
"X-Content-Type-Options": "nosniff",
}
checksum = str(source.get("checksum_sha256") or "")
if checksum:
headers["ETag"] = f'"{checksum}"'
if selected_range is not None:
headers["Content-Range"] = f"bytes {start}-{end}/{actual_size}"
body = storage.iter_bytes(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
expected_size=actual_size,
start=start,
length=length,
)
return StreamingResponse(
body,
status_code=206 if selected_range is not None else 200,
media_type="application/pdf",
headers=headers,
)
@router.get("/{task_id}/source-files/{file_id}/pdf-pages")
def source_file_pdf_pages(
task_id: str,
file_id: str,
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
"""返回切片全文偏移对应的 PDF 物理页范围。"""
with api_errors():
source = store.get_source_file(task_id, file_id, include_content=False)
if str(source.get("file_format") or "").lower() != "pdf":
raise fail(415, "PDF page mapping is only available for PDF source files")
storage_object_id = str(source.get("storage_object_id") or "")
actual_size = storage.file_size(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
if actual_size is None:
raise fail(410, "the original PDF is unavailable for this legacy source file")
expected_size = int(source.get("size_bytes") or 0)
if actual_size != expected_size:
raise ValueError("source object size does not match metadata")
raw = b"".join(
storage.iter_bytes(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
expected_size=actual_size,
)
)
pages = extract_pdf_page_texts(raw)
return ok(
{
"page_count": len(pages),
"pages": [
{
"page_number": page.page_number,
"source_start": page.source_start,
"source_end": page.source_end,
}
for page in pages
],
}
)
@router.delete("/{task_id}/source-files/{file_id}")
def delete_source_file(
task_id: str,
file_id: str,
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
with api_errors():
source = store.get_source_file(task_id, file_id, include_content=False)
storage_object_id = str(source.get("storage_object_id") or "")
storage.validate_owner(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
store.delete_source_file(task_id, file_id)
return ok({"deleted": file_id}, "source file removed")
cleanup_pending = False
try:
storage.delete(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
except Exception:
# 数据库软删除已经提交，不能再向客户端返回可重试的失败；保留逻辑引用，
# 由后续存储清理任务重试物理删除。
cleanup_pending = True
logger.exception(
"failed to remove data process source object after soft deletion",
extra={"task_id": task_id, "source_file_id": file_id},
)
return ok(
{"deleted": file_id, "storage_cleanup_pending": cleanup_pending},
"source file removed",
)
def _external_postgres_connection(payload: ExternalSourceRequest) -> psycopg.Connection[Any]:
@@ -635,7 +1131,11 @@ def test_external_source(
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
store.get_task(task_id)
task = store.get_task(task_id)
if str(task.get("process_type")) != "external":
raise InvalidStateError(
"external source access requires an external data processing task"
)
try:
with _external_postgres_connection(payload) as conn:
conn.execute("SELECT 1 AS ok").fetchone()
@@ -649,6 +1149,7 @@ def pull_external_source(
task_id: str,
payload: ExternalPullRequest,
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
query = (payload.query or "").strip()
if query.endswith(";"):
@@ -659,7 +1160,9 @@ def pull_external_source(
if first_token not in {"select", "with"}:
raise fail(400, "a read-only SELECT or WITH query is required for external pull")
with api_errors():
store.get_task(task_id)
task = store.get_task(task_id)
if str(task.get("process_type")) != "external":
raise InvalidStateError("external pull requires an external data processing task")
try:
with _external_postgres_connection(payload) as conn:
conn.execute("SET TRANSACTION READ ONLY")
@@ -686,20 +1189,40 @@ def pull_external_source(
raise fail(400, "external query returned no rows")
content = "".join(content_parts)
raw = content.encode("utf-8")
source = store.add_source_file(
name = _safe_file_name(payload.file_name, "external-data.jsonl")
source_file_id = new_id("dpsf")
staged = storage.stage_bytes(
batch_id=storage.new_batch_id(),
task_id=task_id,
source_file_id=source_file_id,
version=1,
name=name,
content=raw,
)
sources = _commit_source_batch(
store,
storage,
task_id,
name=_safe_file_name(payload.file_name, "external-data.jsonl"),
content=content,
raw_size=len(raw),
checksum_sha256=hashlib.sha256(raw).hexdigest(),
file_format="jsonl",
record_count=len(rows),
metadata={
[
{
"id": source_file_id,
"storage_object_id": staged.reference,
"name": name,
"content": content,
"raw_size": len(raw),
"checksum_sha256": hashlib.sha256(raw).hexdigest(),
"file_format": "jsonl",
"record_count": len(rows),
"metadata": {
"external_type": payload.type,
"external_host": urlsplit(payload.url).hostname,
"external_limit": payload.limit,
},
}
],
[staged],
)
source = sources[0]
return ok({"files": [source]}, "external source pulled")
@@ -724,7 +1247,7 @@ def _prepare_preview_items(
if not sources:
raise InvalidStateError("at least one source file is required")
items = _build_preview_items(task, sources)
if not items:
if not items and source_file_ids is None:
raise InvalidStateError("source files did not produce preview items")
return items
@@ -736,10 +1259,38 @@ def build_preview(
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
items = _prepare_preview_items(task_id, store, payload.source_file_ids)
created = store.replace_preview_items(task_id, items)
selected_ids = payload.source_file_ids
items = _prepare_preview_items(task_id, store, selected_ids)
created = store.replace_preview_items(
task_id,
items,
source_file_ids=selected_ids,
)
target_ids = selected_ids or [
str(source["id"]) for source in store.list_source_files(task_id)
]
file_counts = dict.fromkeys(target_ids, 0)
for item in created:
source_file_id = str(item.get("source_file_id") or "")
if source_file_id in file_counts:
file_counts[source_file_id] += 1
files = [
{
"source_file_id": source_file_id,
"preview_count": count,
"status": "completed" if count else "empty",
}
for source_file_id, count in file_counts.items()
]
return ok(
{"items": created, "total": len(created), "page": 1, "page_size": len(created)},
{
"items": created,
"total": len(created),
"page": 1,
"page_size": len(created),
"file_counts": file_counts,
"files": files,
},
"preview built",
)

1976
backend/app/modules/data_process/algorithms.py
View File

File diff suppressed because it is too large Load Diff

499
backend/app/modules/data_process/storage.py Normal file
View File

@@ -0,0 +1,499 @@
"""数据处理原始源文件的受控本地对象存储。"""
from __future__ import annotations
import os
import re
import stat
import unicodedata
import uuid
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path, PurePosixPath
from typing import Iterable, Iterator
from urllib.parse import quote, unquote, urlsplit
class DataProcessStorageError(ValueError):
"""本地对象引用或文件系统状态不安全。"""
@dataclass(frozen=True, slots=True)
class StagedSourceObject:
"""尚未发布的原始文件；绝对路径仅在存储模块内部流转。"""
reference: str
_temporary_path: Path
_relative_path: PurePosixPath
def _default_storage_root() -> Path:
return Path(__file__).resolve().parents[3] / "storage" / "data-process"
def _configured_storage_root() -> Path:
configured = os.getenv("DATA_PROCESS_STORAGE_DIR", "").strip()
if not configured:
return _default_storage_root()
path = Path(configured).expanduser()
# 相对配置固定以 backend 目录为基准，
# 避免从不同 cwd 启动时写入不同位置。
return path if path.is_absolute() else Path(__file__).resolve().parents[3] / path
def _safe_component(value: str, label: str) -> str:
if not value or value in {".", ".."} or len(value) > 128:
raise DataProcessStorageError(f"invalid {label}")
if not value[0].isalnum() or any(
not (character.isalnum() or character in {"-", "_", "."})
for character in value
):
raise DataProcessStorageError(f"invalid {label}")
return value
def _safe_basename(value: str) -> str:
if not value or len(value.encode("utf-8")) > 255:
raise DataProcessStorageError("invalid source file name")
if value != Path(value).name or "/" in value or "\\" in value or "\x00" in value:
raise DataProcessStorageError("invalid source file name")
if value in {".", ".."} or any(
unicodedata.category(character).startswith("C") for character in value
):
raise DataProcessStorageError("invalid source file name")
return value
class LocalDataProcessStorage:
"""只允许访问配置根目录下的版本化原始文件。"""
def __init__(self, root: str | os.PathLike[str] | Path | None = None) -> None:
configured = Path(root) if root is not None else _configured_storage_root()
configured = configured.expanduser()
if configured.exists() and configured.is_symlink():
raise DataProcessStorageError("data process storage root must not be a symlink")
configured.mkdir(parents=True, exist_ok=True, mode=0o700)
self._root = configured.resolve(strict=True)
# StagedSourceObject 本身是普通 dataclass，不能只依赖其中的路径字段判断
# 来源；只接受由当前存储实例实际签发的对象，
# 避免调用方伪造暂存路径。
self._issued_staged_objects: dict[Path, StagedSourceObject] = {}
self._ensure_directory(self._root / ".staging")
@property
def root(self) -> Path:
"""仅供运维和测试检查；API 响应不得序列化该属性。"""
return self._root
def new_batch_id(self) -> str:
return f"batch-{uuid.uuid4().hex}"
def stage_bytes(
self,
*,
batch_id: str,
task_id: str,
source_file_id: str,
version: int,
name: str,
content: bytes,
) -> StagedSourceObject:
batch_id = _safe_component(batch_id, "batch id")
task_id = _safe_component(task_id, "task id")
source_file_id = _safe_component(source_file_id, "source file id")
if isinstance(version, bool) or not isinstance(version, int) or version < 1:
raise DataProcessStorageError("invalid source file version")
basename = _safe_basename(name)
if not isinstance(content, bytes):
raise TypeError("content must be bytes")
batch_directory = self._ensure_directory(self._root / ".staging" / batch_id)
temporary_path = batch_directory / f"{source_file_id}-{uuid.uuid4().hex}.tmp"
flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY
if hasattr(os, "O_NOFOLLOW"):
flags |= os.O_NOFOLLOW
descriptor = os.open(temporary_path, flags, 0o600)
try:
with os.fdopen(descriptor, "wb", closefd=True) as stream:
stream.write(content)
stream.flush()
os.fsync(stream.fileno())
except Exception:
temporary_path.unlink(missing_ok=True)
raise
relative_path = PurePosixPath(
task_id,
source_file_id,
f"v{version}",
basename,
)
reference = (
"local://data-process/"
f"{task_id}/{source_file_id}/v{version}/{quote(basename, safe='')}"
)
staged = StagedSourceObject(reference, temporary_path, relative_path)
self._issued_staged_objects[temporary_path] = staged
return staged
def publish(self, objects: Iterable[StagedSourceObject]) -> None:
staged = list(objects)
published: list[StagedSourceObject] = []
try:
seen_temporary_paths: set[Path] = set()
for item in staged:
self._validate_staged_object(item, require_file=True)
if item._temporary_path in seen_temporary_paths:
raise DataProcessStorageError("duplicate staged source object")
seen_temporary_paths.add(item._temporary_path)
for item in staged:
final_path = self._path_for_relative(item._relative_path)
self._ensure_directory(final_path.parent)
if final_path.exists() or final_path.is_symlink():
raise DataProcessStorageError("source storage object already exists")
os.link(item._temporary_path, final_path, follow_symlinks=False)
published.append(item)
item._temporary_path.unlink()
self._fsync_directory(final_path.parent)
except Exception:
for item in reversed(published):
try:
self.delete(item.reference)
except Exception:
# 回滚必须尽量处理其余对象，并保留真正的发布异常。
pass
for item in staged:
try:
self.discard([item])
except Exception:
pass
raise
self.discard(staged)
def discard(self, objects: Iterable[StagedSourceObject]) -> None:
staged = list(objects)
for item in staged:
self._validate_staged_object(item, require_file=False)
batch_directories: set[Path] = set()
first_error: Exception | None = None
for item in staged:
temporary_path = item._temporary_path
try:
temporary_path.unlink(missing_ok=True)
except Exception as exc:
if first_error is None:
first_error = exc
else:
self._issued_staged_objects.pop(temporary_path, None)
batch_directories.add(temporary_path.parent)
for directory in batch_directories:
self._remove_empty_directory(directory)
if first_error is not None:
raise first_error
def read(self, reference: str) -> bytes | None:
"""读取 local 引用；旧 ``db://`` 对象返回 ``None`` 由数据库正文兜底。"""
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return None
descriptor, _ = self._open_read_descriptor(relative_path)
with os.fdopen(descriptor, "rb", closefd=True) as stream:
return stream.read()
def file_size(
self,
reference: str,
*,
expected_task_id: str,
expected_source_file_id: str,
) -> int | None:
"""返回受控 local 对象大小；旧 ``db://`` 对象没有原始文件。"""
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return None
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
descriptor, info = self._open_read_descriptor(relative_path)
os.close(descriptor)
return info.st_size
def iter_bytes(
self,
reference: str,
*,
expected_task_id: str,
expected_source_file_id: str,
expected_size: int,
start: int = 0,
length: int | None = None,
chunk_size: int = 256 * 1024,
) -> Iterator[bytes]:
"""按范围流式读取原始文件，避免 PDF 预览把大文件整体载入内存。"""
relative_path = self._relative_from_reference(reference)
if relative_path is None:
raise DataProcessStorageError("original source object is not available")
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
if start < 0 or expected_size < 0 or chunk_size < 1:
raise DataProcessStorageError("invalid source byte range")
descriptor, info = self._open_read_descriptor(relative_path)
if info.st_size != expected_size:
os.close(descriptor)
raise DataProcessStorageError("source object size does not match metadata")
remaining = expected_size - start if length is None else length
if remaining < 0 or start + remaining > expected_size:
os.close(descriptor)
raise DataProcessStorageError("invalid source byte range")
with os.fdopen(descriptor, "rb", closefd=True) as stream:
stream.seek(start)
while remaining:
chunk = stream.read(min(chunk_size, remaining))
if not chunk:
raise DataProcessStorageError("source object ended unexpectedly")
remaining -= len(chunk)
yield chunk
def validate_owner(
self,
reference: str,
*,
expected_task_id: str,
expected_source_file_id: str,
) -> bool:
"""校验 local 引用归属；旧 ``db://`` 引用无需文件系统处理。"""
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return False
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
return True
def _open_read_descriptor(
self,
relative_path: PurePosixPath,
) -> tuple[int, os.stat_result]:
path = self._path_for_relative(relative_path)
self._assert_controlled_parent(path)
try:
before_open = path.lstat()
except FileNotFoundError as exc:
raise DataProcessStorageError("source storage object does not exist") from exc
if stat.S_ISLNK(before_open.st_mode) or not stat.S_ISREG(before_open.st_mode):
raise DataProcessStorageError("source storage object is not a regular file")
flags = os.O_RDONLY
if hasattr(os, "O_NOFOLLOW"):
flags |= os.O_NOFOLLOW
descriptor = os.open(path, flags)
after_open = os.fstat(descriptor)
if (
not stat.S_ISREG(after_open.st_mode)
or before_open.st_dev != after_open.st_dev
or before_open.st_ino != after_open.st_ino
):
os.close(descriptor)
raise DataProcessStorageError("source storage object changed while opening")
return descriptor, after_open
def delete(
self,
reference: str,
*,
expected_task_id: str | None = None,
expected_source_file_id: str | None = None,
) -> bool:
"""删除受控 local 对象；旧 ``db://`` 引用保持不变。"""
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return False
if (expected_task_id is None) != (expected_source_file_id is None):
raise DataProcessStorageError("both expected storage owner fields are required")
if expected_task_id is not None and expected_source_file_id is not None:
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
path = self._path_for_relative(relative_path)
self._assert_controlled_parent(path)
try:
info = path.lstat()
except FileNotFoundError:
return False
if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode):
raise DataProcessStorageError("refusing to delete a non-regular storage object")
path.unlink()
self._fsync_directory(path.parent)
for directory in (path.parent, path.parent.parent, path.parent.parent.parent):
self._remove_empty_directory(directory)
return True
@staticmethod
def _assert_expected_owner(
relative_path: PurePosixPath,
*,
expected_task_id: str,
expected_source_file_id: str,
) -> None:
task_id = _safe_component(expected_task_id, "expected task id")
source_file_id = _safe_component(
expected_source_file_id,
"expected source file id",
)
if relative_path.parts[:2] != (task_id, source_file_id):
raise DataProcessStorageError("source storage object owner mismatch")
def _relative_from_reference(self, reference: str) -> PurePosixPath | None:
if reference.startswith("db://"):
return None
parsed = urlsplit(reference)
if parsed.scheme != "local" or parsed.netloc != "data-process":
raise DataProcessStorageError("unsupported source storage reference")
if parsed.query or parsed.fragment or "\\" in parsed.path:
raise DataProcessStorageError("unsafe source storage reference")
raw_parts = parsed.path.lstrip("/").split("/")
if len(raw_parts) != 4:
raise DataProcessStorageError("unsafe source storage reference")
if any(re.search(r"%(?![0-9A-Fa-f]{2})", part) for part in raw_parts):
raise DataProcessStorageError("unsafe source storage reference")
try:
decoded = [unquote(part, encoding="utf-8", errors="strict") for part in raw_parts]
except UnicodeDecodeError as exc:
raise DataProcessStorageError("unsafe source storage reference") from exc
if any("/" in part or "\\" in part for part in decoded):
raise DataProcessStorageError("unsafe source storage reference")
canonical_parts = [
quote(decoded[0], safe="-_."),
quote(decoded[1], safe="-_."),
quote(decoded[2], safe="-_."),
quote(decoded[3], safe=""),
]
if canonical_parts != raw_parts:
raise DataProcessStorageError("source storage reference is not canonical")
task_id = _safe_component(decoded[0], "task id")
source_file_id = _safe_component(decoded[1], "source file id")
version_component = decoded[2]
if not version_component.startswith("v") or not version_component[1:].isdigit():
raise DataProcessStorageError("invalid source file version")
version = int(version_component[1:])
if version < 1:
raise DataProcessStorageError("invalid source file version")
basename = _safe_basename(decoded[3])
return PurePosixPath(task_id, source_file_id, f"v{version}", basename)
def _path_for_relative(self, relative_path: PurePosixPath) -> Path:
if relative_path.is_absolute() or any(
part in {"", ".", ".."} for part in relative_path.parts
):
raise DataProcessStorageError("storage path escapes the configured root")
path = self._root.joinpath(*relative_path.parts)
self._assert_controlled_parent(path)
return path
def _validate_staged_object(
self,
item: StagedSourceObject,
*,
require_file: bool,
) -> None:
if not isinstance(item, StagedSourceObject):
raise DataProcessStorageError("invalid staged source object")
if self._issued_staged_objects.get(item._temporary_path) is not item:
raise DataProcessStorageError("staged source object was not issued by this storage")
expected_relative = self._relative_from_reference(item.reference)
if expected_relative is None or expected_relative != item._relative_path:
raise DataProcessStorageError("staged source object reference mismatch")
staging_root = self._root / ".staging"
try:
relative_temporary = item._temporary_path.relative_to(staging_root)
except ValueError as exc:
raise DataProcessStorageError("staged source object escapes staging") from exc
if len(relative_temporary.parts) != 2:
raise DataProcessStorageError("invalid staged source object path")
_safe_component(relative_temporary.parts[0], "batch id")
_safe_basename(relative_temporary.parts[1])
self._assert_controlled_parent(item._temporary_path)
try:
info = item._temporary_path.lstat()
except FileNotFoundError:
if require_file:
raise DataProcessStorageError("staged source object does not exist") from None
return
if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode):
raise DataProcessStorageError("staged source object is not a regular file")
def _ensure_directory(self, directory: Path) -> Path:
try:
relative = directory.relative_to(self._root)
except ValueError as exc:
raise DataProcessStorageError("storage path escapes the configured root") from exc
current = self._root
for component in relative.parts:
current = current / component
try:
current.mkdir(mode=0o700)
except FileExistsError:
pass
info = current.lstat()
if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode):
raise DataProcessStorageError("storage path contains a symlink or non-directory")
return directory
def _assert_controlled_parent(self, path: Path) -> None:
try:
relative_parent = path.parent.relative_to(self._root)
except ValueError as exc:
raise DataProcessStorageError("storage path escapes the configured root") from exc
current = self._root
for component in relative_parent.parts:
current = current / component
if not current.exists():
continue
info = current.lstat()
if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode):
raise DataProcessStorageError("storage path contains a symlink or non-directory")
@staticmethod
def _fsync_directory(directory: Path) -> None:
descriptor = os.open(directory, os.O_RDONLY)
try:
os.fsync(descriptor)
finally:
os.close(descriptor)
def _remove_empty_directory(self, directory: Path) -> None:
if directory in {self._root, self._root / ".staging"}:
return
self._assert_controlled_parent(directory / "placeholder")
try:
directory.rmdir()
except (FileNotFoundError, OSError):
return
@lru_cache
def get_data_process_storage() -> LocalDataProcessStorage:
return LocalDataProcessStorage()
__all__ = [
"DataProcessStorageError",
"LocalDataProcessStorage",
"StagedSourceObject",
"get_data_process_storage",
]

94
backend/app/modules/data_process/store.py
View File

@@ -15,7 +15,6 @@ from psycopg.rows import dict_row
from app.core.config import get_settings
from app.modules.data_process.algorithms import estimate_token_count, stable_split
TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"}
EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"}
@@ -69,6 +68,34 @@ def _serialize_value(value: Any) -> Any:
return value
def _source_storage_descriptor(
payload: dict[str, Any],
task_id: str,
file_id: str,
) -> tuple[str, dict[str, Any]]:
storage_object_id = str(
payload.get("storage_object_id")
or f"db://data-process/{task_id}/{file_id}/v1"
)
expected_local_prefix = f"local://data-process/{task_id}/{file_id}/v1/"
expected_database_reference = f"db://data-process/{task_id}/{file_id}/v1"
if storage_object_id.startswith(expected_local_prefix) and len(storage_object_id) > len(
expected_local_prefix
):
storage_backend = "local"
elif storage_object_id == expected_database_reference:
storage_backend = "database"
elif storage_object_id.startswith(("local://data-process/", "db://data-process/")):
raise DataProcessStoreError("source storage object owner mismatch")
else:
raise DataProcessStoreError("unsupported source storage object reference")
metadata = {
**(payload.get("metadata") or {}),
"storage_backend": storage_backend,
}
return storage_object_id, metadata
def _decode_row(row: dict[str, Any] | None) -> dict[str, Any] | None:
if row is None:
return None
@@ -343,6 +370,8 @@ class DataProcessStore:
record_count: int,
metadata: dict[str, Any] | None = None,
created_by: str | None = None,
source_file_id: str | None = None,
storage_object_id: str | None = None,
) -> dict[str, Any]:
return self.add_source_files(
task_id,
@@ -356,6 +385,8 @@ class DataProcessStore:
"record_count": record_count,
"metadata": metadata or {},
"created_by": created_by,
"id": source_file_id,
"storage_object_id": storage_object_id,
}
],
)[0]
@@ -376,12 +407,12 @@ class DataProcessStore:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
for payload in files:
file_id = new_id("dpsf")
storage_object_id = f"db://data-process/{task_id}/{file_id}/v1"
metadata_payload = {
"storage_backend": "database",
**(payload.get("metadata") or {}),
}
file_id = str(payload.get("id") or new_id("dpsf"))
storage_object_id, metadata_payload = _source_storage_descriptor(
payload,
task_id,
file_id,
)
row = conn.execute(
"""
INSERT INTO data_process_source_files
@@ -530,14 +561,59 @@ class DataProcessStore:
)
def replace_preview_items(
self, task_id: str, items: Sequence[dict[str, Any]]
self,
task_id: str,
items: Sequence[dict[str, Any]],
*,
source_file_ids: Sequence[str] | None = None,
) -> list[dict[str, Any]]:
selected_ids = (
list(dict.fromkeys(str(file_id) for file_id in source_file_ids))
if source_file_ids is not None
else None
)
if selected_ids is not None:
if not selected_ids or any(not file_id for file_id in selected_ids):
raise ValueError("source_file_ids must contain non-empty ids")
selected_set = set(selected_ids)
unexpected = {
str(item.get("source_file_id") or "")
for item in items
if str(item.get("source_file_id") or "") not in selected_set
}
if unexpected:
raise ValueError("preview items contain an unselected source file")
now = utcnow()
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
conn.execute("DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,))
if selected_ids is None:
conn.execute(
"DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
)
else:
rows = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL AND id=ANY(%s)
""",
(task_id, selected_ids),
).fetchall()
found = {str(row["id"]) for row in rows}
missing = set(selected_ids) - found
if missing:
raise NotFoundError(
f"source files not found: {', '.join(sorted(missing))}"
)
conn.execute(
"""
DELETE FROM data_process_preview_items
WHERE task_id=%s AND source_file_id=ANY(%s)
""",
(task_id, selected_ids),
)
created: list[dict[str, Any]] = []
for item in items:
row = conn.execute(

35
backend/app/schemas/data_process.py
View File

@@ -13,6 +13,24 @@ def _config_value(config: dict[str, Any], snake_name: str, camel_name: str, defa
def _validate_process_config(config: dict[str, Any]) -> None:
chunk_method = _config_value(config, "chunk_method", "chunkMethod", "structure")
if not isinstance(chunk_method, str) or chunk_method not in {
"structure",
"fixed",
"custom",
}:
raise ValueError("chunk_method must be one of: structure, fixed, custom")
custom_delimiter = _config_value(
config,
"custom_delimiter",
"customDelimiter",
"",
)
if chunk_method == "custom" and (
not isinstance(custom_delimiter, str) or not custom_delimiter
):
raise ValueError("custom_delimiter is required for custom chunking")
split = _config_value(config, "dataset_split", "datasetSplit", None)
if split is not None:
if not isinstance(split, dict) or set(split) != {"train", "validation", "test"}:
@@ -140,6 +158,23 @@ class PreviewBuildRequest(BaseModel):
replace_existing: Literal[True] = True
source_file_ids: list[str] | None = None
source_file_id: str | None = None
@model_validator(mode="after")
def validate_source_file_selection(self) -> "PreviewBuildRequest":
if self.source_file_ids is not None and self.source_file_id is not None:
raise ValueError("source_file_id and source_file_ids cannot be used together")
values = self.source_file_ids
if values is None and self.source_file_id is not None:
values = [self.source_file_id]
if values is None:
return self
normalized = list(dict.fromkeys(str(value).strip() for value in values))
if not normalized or any(not value for value in normalized):
raise ValueError("at least one non-empty source file id is required")
self.source_file_ids = normalized
self.source_file_id = None
return self
class PreviewItemCreate(BaseModel):

5
backend/pyproject.toml
View File

@@ -16,6 +16,11 @@ dependencies = [
"PyJWT>=2.8.0",
"passlib[bcrypt]>=1.7.4",
"python-dotenv>=1.0.1",
"pypdf[crypto]>=5.0.0",
"python-docx>=1.1.2",
"openpyxl>=3.1.5",
"python-pptx>=1.0.2",
"llama-index-core==0.14.23",
]
[project.optional-dependencies]

5
backend/requirements.txt
View File

@@ -10,3 +10,8 @@ httpx>=0.27.0
PyJWT>=2.8.0
passlib[bcrypt]>=1.7.4
python-dotenv>=1.0.1
pypdf[crypto]>=5.0.0
python-docx>=1.1.2
openpyxl>=3.1.5
python-pptx>=1.0.2
llama-index-core==0.14.23

544
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:
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:

808
backend/tests/test_data_process_api.py
View File

@@ -1,13 +1,23 @@
from __future__ import annotations
from copy import deepcopy
from io import BytesIO
from pathlib import Path
from typing import Any
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from openpyxl import Workbook
from app.api.v1.endpoints import data_process as data_process_endpoint
from app.api.v1.endpoints.data_process import router
from app.modules.data_process.algorithms import normalize_text
from app.modules.data_process.storage import (
DataProcessStorageError,
LocalDataProcessStorage,
get_data_process_storage,
)
from app.modules.data_process.store import InvalidStateError, NotFoundError, get_data_process_store
@@ -87,11 +97,26 @@ class FakeDataProcessStore:
def add_source_file(self, task_id: str, **payload: Any) -> dict[str, Any]:
self.get_task(task_id)
values = deepcopy(payload)
source_id = str(values.pop("id", None) or self._id("dpsf"))
storage_object_id = str(
values.pop("storage_object_id", None)
or f"db://data-process/{task_id}/{source_id}/v1"
)
raw_size = int(values.pop("raw_size"))
metadata = deepcopy(values.pop("metadata", {}))
metadata.setdefault(
"storage_backend",
"local" if storage_object_id.startswith("local://data-process/") else "database",
)
source = {
"id": self._id("dpsf"),
"id": source_id,
"task_id": task_id,
"version_no": 1,
**deepcopy(payload),
"storage_object_id": storage_object_id,
"size_bytes": raw_size,
"metadata": metadata,
**values,
}
self.sources[task_id].append(source)
self.tasks[task_id]["input_count"] += payload["record_count"]
@@ -163,14 +188,27 @@ class FakeDataProcessStore:
self.results[task_id] = []
def replace_preview_items(
self, task_id: str, items: list[dict[str, Any]]
self,
task_id: str,
items: list[dict[str, Any]],
*,
source_file_ids: list[str] | None = None,
) -> list[dict[str, Any]]:
self.previews[task_id] = [
created = [
{"id": self._id("dpp"), "task_id": task_id, **deepcopy(item)} for item in items
]
if source_file_ids is None:
self.previews[task_id] = created
else:
selected = set(source_file_ids)
self.previews[task_id] = [
item
for item in self.previews[task_id]
if item["source_file_id"] not in selected
] + created
self.results[task_id] = []
self.tasks[task_id]["progress"] = 20
return deepcopy(self.previews[task_id])
return deepcopy(created)
def list_preview_items(
self,
@@ -389,16 +427,59 @@ class FakeDataProcessStore:
return {"dataset": deepcopy(dataset), "created": True}
def make_client() -> tuple[TestClient, FakeDataProcessStore]:
def make_client(
tmp_path: Path,
) -> tuple[TestClient, FakeDataProcessStore, LocalDataProcessStorage]:
store = FakeDataProcessStore()
storage = LocalDataProcessStorage(tmp_path / "data-process")
app = FastAPI()
app.include_router(router, prefix="/modelTF")
app.dependency_overrides[get_data_process_store] = lambda: store
return TestClient(app), store
app.dependency_overrides[get_data_process_storage] = lambda: storage
return TestClient(app), store, storage
def test_data_process_full_contract_without_database() -> None:
client, store = make_client()
def _stored_files(storage: LocalDataProcessStorage) -> list[Path]:
return [path for path in storage.root.rglob("*") if path.is_file() or path.is_symlink()]
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() + 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())
result.extend(value)
result.extend(b"\nendobj\n")
xref_offset = len(result)
result.extend(f"xref\n0 {len(objects) + 1}\n".encode())
result.extend(b"0000000000 65535 f \n")
for offset in offsets[1:]:
result.extend(f"{offset:010d} 00000 n \n".encode())
result.extend(
(
f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n"
f"startxref\n{xref_offset}\n%%EOF\n"
).encode()
)
return bytes(result)
def test_data_process_full_contract_without_database(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
created = client.post(
"/modelTF/data-process",
json={
@@ -506,8 +587,143 @@ def test_data_process_full_contract_without_database() -> None:
)
def test_external_source_never_returns_fake_success() -> None:
client, _ = make_client()
def test_preview_build_replaces_only_selected_files_and_reports_file_counts(
tmp_path: Path,
) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "逐文件预览", "process_type": "structured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files=[
("files", ("first.jsonl", b'{"id":1}\n', "application/jsonl")),
(
"files",
("second.jsonl", b'{"id":2}\n{"id":3}\n', "application/jsonl"),
),
],
)
assert uploaded.status_code == 200
first_source, second_source = uploaded.json()["data"]["files"]
first_build = client.post(
f"/modelTF/data-process/{task_id}/preview/build",
json={"source_file_ids": [first_source["id"]]},
)
assert first_build.status_code == 200
first_data = first_build.json()["data"]
assert first_data["file_counts"] == {first_source["id"]: 1}
assert first_data["files"] == [
{
"source_file_id": first_source["id"],
"preview_count": 1,
"status": "completed",
}
]
first_item = first_data["items"][0]
edited = client.put(
f"/modelTF/data-process/{task_id}/preview/{first_item['id']}",
json={"edited_content": "人工确认后的第一文件预览"},
)
assert edited.status_code == 200
second_build = client.post(
f"/modelTF/data-process/{task_id}/preview/build",
json={"source_file_id": second_source["id"]},
)
assert second_build.status_code == 200
second_data = second_build.json()["data"]
assert second_data["file_counts"] == {second_source["id"]: 2}
assert second_data["files"] == [
{
"source_file_id": second_source["id"],
"preview_count": 2,
"status": "completed",
}
]
assert {item["source_file_id"] for item in store.previews[task_id]} == {
first_source["id"],
second_source["id"],
}
preserved_first = next(
item
for item in store.previews[task_id]
if item["source_file_id"] == first_source["id"]
)
assert preserved_first["id"] == first_item["id"]
assert preserved_first["edited_content"] == "人工确认后的第一文件预览"
previous_second_ids = {
item["id"]
for item in store.previews[task_id]
if item["source_file_id"] == second_source["id"]
}
next(
source
for source in store.sources[task_id]
if source["id"] == second_source["id"]
)["content"] = '{"id":4}\n'
rebuilt = client.post(
f"/modelTF/data-process/{task_id}/preview/build",
json={"source_file_ids": [second_source["id"]]},
)
assert rebuilt.status_code == 200
assert rebuilt.json()["data"]["file_counts"] == {second_source["id"]: 1}
current_second_ids = {
item["id"]
for item in store.previews[task_id]
if item["source_file_id"] == second_source["id"]
}
assert current_second_ids.isdisjoint(previous_second_ids)
assert len(current_second_ids) == 1
assert next(
item
for item in store.previews[task_id]
if item["source_file_id"] == first_source["id"]
)["id"] == first_item["id"]
def test_preview_build_rejects_unknown_and_cross_task_source_file_ids(
tmp_path: Path,
) -> None:
client, _, _ = make_client(tmp_path)
first_task_id = client.post(
"/modelTF/data-process",
json={"name": "归属任务一", "process_type": "structured", "config": {}},
).json()["data"]["id"]
second_task_id = client.post(
"/modelTF/data-process",
json={"name": "归属任务二", "process_type": "structured", "config": {}},
).json()["data"]["id"]
foreign_source = client.post(
f"/modelTF/data-process/{second_task_id}/source-files",
files={"files": ("foreign.jsonl", b'{"id":2}\n', "application/jsonl")},
).json()["data"]["files"][0]
unknown = client.post(
f"/modelTF/data-process/{first_task_id}/preview/build",
json={"source_file_ids": ["dpsf_not_found"]},
)
assert unknown.status_code == 404
foreign = client.post(
f"/modelTF/data-process/{first_task_id}/preview/build",
json={"source_file_id": foreign_source["id"]},
)
assert foreign.status_code == 404
ambiguous = client.post(
f"/modelTF/data-process/{first_task_id}/preview/build",
json={
"source_file_id": foreign_source["id"],
"source_file_ids": [foreign_source["id"]],
},
)
assert ambiguous.status_code == 422
def test_external_source_never_returns_fake_success(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "外部数据", "process_type": "external", "config": {}},
@@ -520,8 +736,8 @@ def test_external_source_never_returns_fake_success() -> None:
assert response.json()["detail"]["code"] == 501
def test_config_validation_and_stop_state() -> None:
client, store = make_client()
def test_config_validation_and_stop_state(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
invalid = client.post(
"/modelTF/data-process",
json={
@@ -537,6 +753,28 @@ def test_config_validation_and_stop_state() -> None:
)
assert invalid.status_code == 422
legacy_semantic = client.post(
"/modelTF/data-process",
json={
"name": "旧切分策略",
"process_type": "unstructured",
"config": {"chunk_method": "semantic"},
},
)
assert legacy_semantic.status_code == 422
assert "chunk_method" in legacy_semantic.text
missing_custom_delimiter = client.post(
"/modelTF/data-process",
json={
"name": "缺少自定义分隔符",
"process_type": "unstructured",
"config": {"chunk_method": "custom"},
},
)
assert missing_custom_delimiter.status_code == 422
assert "custom_delimiter" in missing_custom_delimiter.text
task_id = client.post(
"/modelTF/data-process",
json={"name": "可停止任务", "process_type": "structured", "config": {}},
@@ -547,11 +785,11 @@ def test_config_validation_and_stop_state() -> None:
assert stopped.json()["data"]["status"] == "stopped"
def test_upload_batch_is_atomic_and_empty_files_are_rejected() -> None:
client, store = make_client()
def test_upload_batch_is_atomic_and_empty_files_are_rejected(tmp_path: Path) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "批量上传", "process_type": "structured", "config": {}},
json={"name": "批量上传", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
duplicate_batch = client.post(
@@ -563,6 +801,18 @@ def test_upload_batch_is_atomic_and_empty_files_are_rejected() -> None:
)
assert duplicate_batch.status_code == 400
assert store.sources[task_id] == []
assert _stored_files(storage) == []
parse_failure = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files=[
("files", ("valid.txt", "先暂存的内容".encode(), "text/plain")),
("files", ("broken.txt", b"\xff", "text/plain")),
],
)
assert parse_failure.status_code == 400
assert store.sources[task_id] == []
assert _stored_files(storage) == []
empty = client.post(
f"/modelTF/data-process/{task_id}/source-files",
@@ -570,10 +820,45 @@ def test_upload_batch_is_atomic_and_empty_files_are_rejected() -> None:
)
assert empty.status_code == 400
assert store.sources[task_id] == []
assert _stored_files(storage) == []
def test_preprocess_deduplicates_and_quality_filter_removes_short_results() -> None:
client, _ = make_client()
def test_upload_preserves_store_error_when_storage_rollback_fails(
tmp_path: Path,
monkeypatch: Any,
) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "回滚异常", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
cleanup_attempts: list[str] = []
def fail_store(*_: Any, **__: Any) -> list[dict[str, Any]]:
raise ValueError("simulated database transaction failure")
def fail_cleanup(reference: str, **_: Any) -> bool:
cleanup_attempts.append(reference)
raise OSError("simulated storage cleanup failure")
monkeypatch.setattr(store, "add_source_files", fail_store)
monkeypatch.setattr(storage, "delete", fail_cleanup)
response = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("rollback.txt", b"rollback payload", "text/plain")},
)
assert response.status_code == 400
assert response.json()["detail"]["message"] == "simulated database transaction failure"
assert len(cleanup_attempts) == 1
assert store.sources[task_id] == []
def test_preprocess_deduplicates_and_quality_filter_removes_short_results(
tmp_path: Path,
) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={
@@ -651,8 +936,8 @@ def test_stale_generation_worker_cannot_overwrite_new_run(monkeypatch: Any) -> N
assert store.tasks[task_id]["status"] == "running"
def test_result_status_cannot_be_forged_by_client() -> None:
client, _ = make_client()
def test_result_status_cannot_be_forged_by_client(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "状态保护", "process_type": "structured", "config": {}},
@@ -664,8 +949,8 @@ def test_result_status_cannot_be_forged_by_client() -> None:
assert response.status_code == 422
def test_start_rebuilds_preview_and_generates_in_one_request() -> None:
client, _ = make_client()
def test_start_rebuilds_preview_and_generates_in_one_request(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "一键处理", "process_type": "structured", "config": {}},
@@ -691,14 +976,487 @@ def test_start_rebuilds_preview_and_generates_in_one_request() -> None:
assert client.get(f"/modelTF/data-process/{task_id}/results").json()["data"]["total"] == 1
def test_unsupported_upload_format_returns_415() -> None:
client, _ = make_client()
def test_unsupported_upload_format_returns_415(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "格式限制", "process_type": "structured", "config": {}},
).json()["data"]["id"]
response = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("document.pdf", b"not a pdf", "application/pdf")},
files={"files": ("payload.exe", b"not supported", "application/octet-stream")},
)
assert response.status_code == 415
legacy = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("document.doc", b"legacy", "application/msword")},
)
assert legacy.status_code == 415
assert "convert the file to .docx" in legacy.json()["detail"]["message"]
def test_xlsx_upload_is_accepted_as_structured_records(tmp_path: Path) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "XLSX 上传", "process_type": "structured", "config": {}},
).json()["data"]["id"]
workbook = Workbook()
worksheet = workbook.active
worksheet.append(["question", "answer"])
worksheet.append(["问题一", "答案一"])
worksheet.append(["问题二", "答案二"])
output = BytesIO()
workbook.save(output)
workbook.close()
original_bytes = output.getvalue()
response = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={
"files": (
"records.xlsx",
original_bytes,
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
)
},
)
assert response.status_code == 200
source = response.json()["data"]["files"][0]
assert source["file_format"] == "xlsx"
assert source["record_count"] == 2
assert source["size_bytes"] == len(original_bytes)
assert source["storage_object_id"].startswith("local://data-process/")
assert str(storage.root) not in response.text
assert storage.read(source["storage_object_id"]) == original_bytes
stored_source = store.get_source_file(task_id, source["id"])
assert stored_source["id"] == source["id"]
assert stored_source["storage_object_id"] == source["storage_object_id"]
assert stored_source["metadata"]["storage_backend"] == "local"
assert stored_source["metadata"]["original_size_bytes"] == len(original_bytes)
assert '"question":"问题一"' in stored_source["content"]
content = client.get(
f"/modelTF/data-process/{task_id}/source-files/{source['id']}/content"
)
assert content.status_code == 200
assert '"answer":"答案二"' in content.json()["data"]["content"]
preview = client.post(f"/modelTF/data-process/{task_id}/preview/build")
assert preview.status_code == 200
assert preview.json()["data"]["total"] == 2
def test_pdf_raw_preview_streams_original_file_and_supports_ranges(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "PDF 原件预览", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
original_pdf = _minimal_pdf()
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("说明 文档.pdf", original_pdf, "application/pdf")},
).json()["data"]["files"][0]
raw_url = f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/raw"
full = client.get(raw_url)
assert full.status_code == 200
assert full.content == original_pdf
assert full.headers["content-type"] == "application/pdf"
assert full.headers["accept-ranges"] == "bytes"
assert full.headers["cache-control"] == "private, no-store"
assert full.headers["content-length"] == str(len(original_pdf))
assert full.headers["content-disposition"].startswith("inline;")
assert "%E8%AF%B4%E6%98%8E%20%E6%96%87%E6%A1%A3.pdf" in full.headers[
"content-disposition"
]
assert full.headers["etag"] == f'"{uploaded["checksum_sha256"]}"'
partial = client.get(raw_url, headers={"Range": "bytes=5-14"})
assert partial.status_code == 206
assert partial.content == original_pdf[5:15]
assert partial.headers["content-range"] == f"bytes 5-14/{len(original_pdf)}"
assert partial.headers["content-length"] == "10"
suffix = client.get(raw_url, headers={"Range": "bytes=-8"})
assert suffix.status_code == 206
assert suffix.content == original_pdf[-8:]
invalid = client.get(raw_url, headers={"Range": "bytes=0-1,4-5"})
assert invalid.status_code == 416
assert invalid.headers["content-range"] == f"bytes */{len(original_pdf)}"
pages_url = (
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/pdf-pages"
)
pages = client.get(pages_url)
assert pages.status_code == 200
assert pages.json()["data"] == {
"page_count": 1,
"pages": [
{
"page_number": 1,
"source_start": 0,
"source_end": len("Hello PDF"),
}
],
}
legacy_id = "dpsf_legacy_pdf"
store.add_source_file(
task_id,
id=legacy_id,
storage_object_id=f"db://data-process/{task_id}/{legacy_id}/v1",
name="legacy.pdf",
content="legacy extracted PDF text",
raw_size=len(original_pdf),
checksum_sha256="a" * 64,
file_format="pdf",
record_count=1,
metadata={"legacy": True},
)
legacy = client.get(
f"/modelTF/data-process/{task_id}/source-files/{legacy_id}/raw"
)
assert legacy.status_code == 410
legacy_pages = client.get(
f"/modelTF/data-process/{task_id}/source-files/{legacy_id}/pdf-pages"
)
assert legacy_pages.status_code == 410
def test_raw_inline_preview_rejects_non_pdf_source(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "非 PDF 原件", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("notes.txt", b"plain source text", "text/plain")},
).json()["data"]["files"][0]
response = client.get(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/raw"
)
assert response.status_code == 415
pages_response = client.get(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/pdf-pages"
)
assert pages_response.status_code == 415
def test_delete_source_removes_owned_local_object_and_accepts_legacy_db_reference(
tmp_path: Path,
) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "删除原件", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("原件.txt", "本地原始内容".encode(), "text/plain")},
).json()["data"]["files"][0]
reference = uploaded["storage_object_id"]
assert storage.read(reference) == "本地原始内容".encode()
deleted = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}"
)
assert deleted.status_code == 200
assert deleted.json()["data"]["storage_cleanup_pending"] is False
with pytest.raises(DataProcessStorageError, match="does not exist"):
storage.read(reference)
legacy_id = "dpsf_legacy"
store.add_source_file(
task_id,
id=legacy_id,
storage_object_id=f"db://data-process/{task_id}/{legacy_id}/v1",
name="legacy.txt",
content="旧记录正文",
raw_size=len("旧记录正文".encode()),
checksum_sha256="a" * 64,
file_format="txt",
record_count=1,
metadata={"legacy": True},
)
legacy_deleted = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{legacy_id}"
)
assert legacy_deleted.status_code == 200
assert legacy_deleted.json()["data"]["storage_cleanup_pending"] is False
def test_delete_reports_pending_cleanup_after_database_soft_delete(
tmp_path: Path,
monkeypatch: Any,
) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "待清理原件", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("pending.txt", b"pending cleanup", "text/plain")},
).json()["data"]["files"][0]
def fail_cleanup(*_: Any, **__: Any) -> bool:
raise OSError("simulated storage failure")
monkeypatch.setattr(storage, "delete", fail_cleanup)
response = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}"
)
assert response.status_code == 200
assert response.json()["data"]["storage_cleanup_pending"] is True
with pytest.raises(NotFoundError):
store.get_source_file(task_id, uploaded["id"])
assert storage.read(uploaded["storage_object_id"]) == b"pending cleanup"
def test_delete_rejects_polluted_reference_owned_by_another_source(tmp_path: Path) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "归属校验", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("safe.txt", b"owned content", "text/plain")},
).json()["data"]["files"][0]
target_reference = uploaded["storage_object_id"]
polluted_id = "dpsf_polluted"
store.add_source_file(
task_id,
id=polluted_id,
storage_object_id=target_reference,
name="polluted.txt",
content="polluted",
raw_size=8,
checksum_sha256="b" * 64,
file_format="txt",
record_count=1,
metadata={},
)
rejected = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{polluted_id}"
)
assert rejected.status_code == 400
assert storage.read(target_reference) == b"owned content"
assert store.get_source_file(task_id, polluted_id)["id"] == polluted_id
def test_upload_format_must_match_process_type(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
structured_id = client.post(
"/modelTF/data-process",
json={"name": "结构化格式约束", "process_type": "structured", "config": {}},
).json()["data"]["id"]
structured_pdf = client.post(
f"/modelTF/data-process/{structured_id}/source-files",
files={"files": ("manual.pdf", b"not parsed", "application/pdf")},
)
assert structured_pdf.status_code == 415
unstructured_id = client.post(
"/modelTF/data-process",
json={"name": "非结构化格式约束", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
unstructured_xlsx = client.post(
f"/modelTF/data-process/{unstructured_id}/source-files",
files={"files": ("records.xlsx", b"not parsed", "application/octet-stream")},
)
assert unstructured_xlsx.status_code == 415
external_id = client.post(
"/modelTF/data-process",
json={"name": "外部数据格式约束", "process_type": "external", "config": {}},
).json()["data"]["id"]
external_upload = client.post(
f"/modelTF/data-process/{external_id}/source-files",
files={"files": ("records.jsonl", b'{"id":1}', "application/jsonl")},
)
assert external_upload.status_code == 409
def _preview_task(
content: str,
*,
options: list[str],
config: dict[str, Any] | None = None,
source_id: str = "source-1",
file_format: str = "txt",
) -> list[dict[str, Any]]:
task_config = {
"preprocess_options": options,
"chunk_size": 200,
"chunk_overlap": 20,
"min_chunk_size": 20,
**(config or {}),
}
return data_process_endpoint._build_preview_items(
{"process_type": "unstructured", "config": task_config},
[
{
"id": source_id,
"name": f"{source_id}.{file_format}",
"file_format": file_format,
"content": content,
}
],
)
def test_default_and_structure_preview_split_headings_without_cross_section_overlap() -> None:
content = (
"# 第一章\n"
+ " ".join(f"alpha{index}" for index in range(18))
+ "\n# 第二章\n"
+ " ".join(f"beta{index}" for index in range(18))
)
normalized = normalize_text(content)
second_chapter_start = normalized.index("# 第二章")
common_config = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4}
default_items = _preview_task(
content,
options=["preserve_context"],
config=common_config,
)
structure_items = _preview_task(
content,
options=["preserve_context"],
config={**common_config, "chunk_method": "structure"},
)
def snapshot(items: list[dict[str, Any]]) -> list[tuple[Any, ...]]:
return [
(
item["original_content"],
item["source_start"],
item["source_end"],
item["source_start_line"],
item["source_end_line"],
)
for item in items
]
assert snapshot(default_items) == snapshot(structure_items)
assert all(
item["original_content"]
== normalized[item["source_start"] : item["source_end"]]
for item in structure_items
)
assert all(
not (item["source_start"] < second_chapter_start < item["source_end"])
for item in structure_items
)
second_chapter_items = [
item for item in structure_items if item["source_start"] >= second_chapter_start
]
assert second_chapter_items[0]["source_start"] == second_chapter_start
assert second_chapter_items[0]["original_content"].startswith("# 第二章")
def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None:
repeated = "@" * 120
assert len(_preview_task(repeated, options=[])) == 1
assert _preview_task(repeated, options=["clean_invalid_content"]) == []
structured_text = "# 第一章\n" + "甲。" * 30 + "\n# 第二章\n" + "乙。" * 30
detected = _preview_task(
structured_text,
options=["detect_document_structure"],
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
)
undetected = _preview_task(
structured_text,
options=[],
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
)
assert all("heading_path" in item["quality_score"] for item in detected)
assert {tuple(item["quality_score"]["heading_path"]) for item in detected} == {
("第一章",),
("第二章",),
}
assert all("heading_path" not in item["quality_score"] for item in undetected)
assert all(not ("第一章" in item["edited_content"] and "第二章" in item["edited_content"]) for item in detected)
short_lead = "a b. c d e f g h i j k l m n o p q r s t u v w x y z"
without_merge = _preview_task(
short_lead,
options=[],
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
)
with_merge = _preview_task(
short_lead,
options=["merge_short_content"],
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
)
assert without_merge[0]["token_count"] < 5
assert with_merge[0]["token_count"] >= 5
mojibake = "这是无法可靠读取的内容，锟斤拷锟斤拷锟斤拷，需要预先过滤。"
assert len(_preview_task(mojibake, options=[])) == 1
assert _preview_task(mojibake, options=["filter_low_quality"]) == []
first = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega"
second = "alpha beta gamma, delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega"
sources = [
{"id": "near-1", "name": "one.txt", "file_format": "txt", "content": first},
{"id": "near-2", "name": "two.txt", "file_format": "txt", "content": second},
]
base_task = {
"process_type": "unstructured",
"config": {
"chunk_method": "fixed",
"chunk_size": 200,
"chunk_overlap": 0,
"min_chunk_size": 1,
"preprocess_options": [],
},
}
assert len(data_process_endpoint._build_preview_items(base_task, sources)) == 2
deduplicated_task = deepcopy(base_task)
deduplicated_task["config"]["preprocess_options"] = ["deduplicate_content"]
assert len(data_process_endpoint._build_preview_items(deduplicated_task, sources)) == 1
context_text = " ".join(f"token{index}" for index in range(45))
no_context = _preview_task(
context_text,
options=[],
config={"chunk_method": "fixed", "chunk_size": 20, "chunk_overlap": 5},
)
with_context = _preview_task(
context_text,
options=["preserve_context"],
config={"chunk_method": "fixed", "chunk_size": 20, "chunk_overlap": 5},
)
assert no_context[1]["source_start"] >= no_context[0]["source_end"]
assert with_context[1]["source_start"] < with_context[0]["source_end"]
sensitive = "联系人：张三，手机 13800138000，邮箱 user@example.com。"
plain = _preview_task(sensitive, options=[])[0]
masked = _preview_task(sensitive, options=["desensitize"])[0]
assert "张三" in plain["edited_content"]
assert "联系人:[NAME]" in masked["edited_content"]
assert "[PHONE]" in masked["edited_content"]
assert "[EMAIL]" in masked["edited_content"]
def test_stored_binary_document_text_is_not_reparsed_as_binary() -> None:
for file_format in ("pdf", "docx", "pptx"):
items = _preview_task(
f"{file_format.upper()} 已抽取正文，可直接进入切片处理。",
options=[],
file_format=file_format,
)
assert len(items) == 1
assert "已抽取正文" in items[0]["edited_content"]

242
backend/tests/test_data_process_storage.py Normal file
View File

@@ -0,0 +1,242 @@
from __future__ import annotations
from pathlib import Path, PurePosixPath
import pytest
from app.modules.data_process import storage as storage_module
from app.modules.data_process.storage import (
DataProcessStorageError,
LocalDataProcessStorage,
StagedSourceObject,
)
def _stage(
storage: LocalDataProcessStorage,
*,
batch_id: str = "batch-main",
task_id: str = "task-1",
source_file_id: str = "source-1",
version: int = 1,
name: str = "source.txt",
content: bytes = b"payload",
) -> StagedSourceObject:
return storage.stage_bytes(
batch_id=batch_id,
task_id=task_id,
source_file_id=source_file_id,
version=version,
name=name,
content=content,
)
def _create_symlink(link: Path, target: Path, *, target_is_directory: bool = False) -> None:
try:
link.symlink_to(target, target_is_directory=target_is_directory)
except (NotImplementedError, OSError) as exc:
pytest.skip(f"当前平台不支持创建测试所需的符号链接: {exc}")
def _assert_staging_empty(storage: LocalDataProcessStorage) -> None:
assert list((storage.root / ".staging").iterdir()) == []
def test_stage_publish_read_delete_roundtrip_with_unicode_filename(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
content = "第一行\n第二行，100% 完成".encode()
staged = _stage(
storage,
name="中文 数据 100%.csv",
content=content,
)
assert "%20" in staged.reference
assert "%25" in staged.reference
storage.publish([staged])
assert storage.read(staged.reference) == content
assert storage.delete(staged.reference) is True
assert storage.delete(staged.reference) is False
_assert_staging_empty(storage)
def test_db_reference_is_left_to_database_storage(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
assert storage.read("db://source-files/source-1") is None
assert storage.delete("db://source-files/source-1") is False
def test_owned_source_can_be_streamed_by_byte_range(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
content = b"0123456789abcdef"
staged = _stage(storage, content=content)
storage.publish([staged])
assert storage.file_size(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
) == len(content)
assert b"".join(storage.iter_bytes(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
expected_size=len(content),
start=4,
length=6,
chunk_size=2,
)) == b"456789"
with pytest.raises(DataProcessStorageError, match="owner mismatch"):
storage.file_size(
staged.reference,
expected_task_id="another-task",
expected_source_file_id="source-1",
)
with pytest.raises(DataProcessStorageError, match="does not match metadata"):
b"".join(storage.iter_bytes(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
expected_size=len(content) + 1,
))
@pytest.mark.parametrize(
"reference",
[
"local://data-process/../source-1/v1/file.txt",
"local://data-process/task-1/source-1/v1/file%2Fname.txt",
"local://data-process/task-1/source-1/v1/file.txt?download=1",
"local://data-process/task-1/source-1/v1/file.txt#fragment",
"https://data-process/task-1/source-1/v1/file.txt",
],
ids=[
"parent-traversal",
"percent-encoded-slash",
"query",
"fragment",
"wrong-scheme",
],
)
def test_unsafe_references_are_rejected(tmp_path: Path, reference: str) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
with pytest.raises(DataProcessStorageError):
storage.read(reference)
with pytest.raises(DataProcessStorageError):
storage.delete(reference)
def test_publish_rejects_intermediate_directory_symlink(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
outside = tmp_path / "outside"
outside.mkdir()
staged = _stage(storage, task_id="linked-task")
_create_symlink(
storage.root / "linked-task",
outside,
target_is_directory=True,
)
with pytest.raises(DataProcessStorageError, match="symlink|non-directory"):
storage.publish([staged])
assert list(outside.iterdir()) == []
_assert_staging_empty(storage)
def test_target_symlink_is_never_followed_or_deleted(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
staged = _stage(storage, task_id="task-link", source_file_id="source-link")
outside_file = tmp_path / "outside.txt"
outside_file.write_bytes(b"outside sentinel")
final_path = storage.root.joinpath(*staged._relative_path.parts)
final_path.parent.mkdir(parents=True)
_create_symlink(final_path, outside_file)
with pytest.raises(DataProcessStorageError, match="already exists"):
storage.publish([staged])
with pytest.raises(DataProcessStorageError, match="regular file"):
storage.read(staged.reference)
with pytest.raises(DataProcessStorageError, match="non-regular"):
storage.delete(staged.reference)
assert final_path.is_symlink()
assert outside_file.read_bytes() == b"outside sentinel"
_assert_staging_empty(storage)
def test_publish_rolls_back_first_object_when_second_target_collides(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
existing = _stage(
storage,
batch_id="batch-existing",
source_file_id="source-existing",
content=b"existing content",
)
storage.publish([existing])
first = _stage(
storage,
batch_id="batch-new",
source_file_id="source-new",
content=b"must be rolled back",
)
colliding_second = _stage(
storage,
batch_id="batch-new",
source_file_id="source-existing",
content=b"must not replace existing content",
)
with pytest.raises(DataProcessStorageError, match="already exists"):
storage.publish([first, colliding_second])
with pytest.raises(DataProcessStorageError, match="does not exist"):
storage.read(first.reference)
assert storage.read(existing.reference) == b"existing content"
_assert_staging_empty(storage)
def test_publish_rejects_manually_forged_staged_object(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
temporary_path = storage.root / ".staging" / "batch-forged" / "forged.tmp"
temporary_path.parent.mkdir()
temporary_path.write_bytes(b"forged content")
relative_path = PurePosixPath("task-forged", "source-forged", "v1", "forged.txt")
forged = StagedSourceObject(
reference="local://data-process/task-forged/source-forged/v1/forged.txt",
_temporary_path=temporary_path,
_relative_path=relative_path,
)
with pytest.raises(DataProcessStorageError, match="was not issued"):
storage.publish([forged])
with pytest.raises(DataProcessStorageError, match="was not issued"):
storage.discard([forged])
assert temporary_path.read_bytes() == b"forged content"
assert not storage.root.joinpath(*relative_path.parts).exists()
def test_relative_storage_configuration_is_anchored_to_backend_root(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
relative_configuration = Path("relative-storage") / tmp_path.name
backend_root = Path(storage_module.__file__).resolve().parents[3]
monkeypatch.chdir(tmp_path)
monkeypatch.setenv("DATA_PROCESS_STORAGE_DIR", str(relative_configuration))
storage_module.get_data_process_storage.cache_clear()
try:
configured_root = storage_module._configured_storage_root()
assert configured_root == backend_root / relative_configuration
assert not configured_root.exists()
finally:
storage_module.get_data_process_storage.cache_clear()

54
backend/tests/test_data_process_store.py Normal file
View File

@@ -0,0 +1,54 @@
from __future__ import annotations
import pytest
from app.modules.data_process.store import (
DataProcessStoreError,
_source_storage_descriptor,
)
def test_source_storage_descriptor_accepts_owned_local_and_legacy_db_references() -> None:
task_id = "dpt_task"
source_file_id = "dpsf_source"
local_reference = (
f"local://data-process/{task_id}/{source_file_id}/v1/source%20100%25.csv"
)
reference, metadata = _source_storage_descriptor(
{
"storage_object_id": local_reference,
"metadata": {"storage_backend": "spoofed", "content_type": "text/csv"},
},
task_id,
source_file_id,
)
assert reference == local_reference
assert metadata == {"storage_backend": "local", "content_type": "text/csv"}
legacy_reference, legacy_metadata = _source_storage_descriptor(
{"metadata": {"legacy": True}},
task_id,
source_file_id,
)
assert legacy_reference == f"db://data-process/{task_id}/{source_file_id}/v1"
assert legacy_metadata == {"storage_backend": "database", "legacy": True}
@pytest.mark.parametrize(
"reference",
[
"local://data-process/dpt_other/dpsf_source/v1/source.txt",
"db://data-process/dpt_task/dpsf_other/v1",
"/var/tmp/source.txt",
],
)
def test_source_storage_descriptor_rejects_unowned_or_unsupported_references(
reference: str,
) -> None:
with pytest.raises(DataProcessStoreError):
_source_storage_descriptor(
{"storage_object_id": reference},
"dpt_task",
"dpsf_source",
)