refactor: 完整重构 data_process 模块并修复拆分遗留缺陷

将 algorithms.py / store.py 拆分为 algorithms/ 与 store/ 子包,并修复
机械拆分造成的导入与辅助函数缺失:

- algorithms/: 补全各子模块依赖与 17 个私有辅助函数、8 个常量;重写
  __init__.py 移除坏的 importlib 兜底,分层导入并以局部 import 断开
  text_utils<->parsers、quality<->structured_processing 循环依赖。
- store/: 补回 DataProcessStoreError / hashlib / _serialize_value /
  estimate_token_count 等缺失导入,包入口导出测试与调用方依赖的私有
  辅助函数。
- 删除旧单文件 algorithms.py / store.py 及重构残留(_algorithms_old、
  backups、refactor 脚本、REFACTORING 文档)。

algorithms 与 store 测试套件 91 项全部通过。
This commit is contained in:
caoxiaozhu
2026-08-12 14:25:37 +08:00
parent ae39c45a73
commit ca012893f7
23 changed files with 6815 additions and 6200 deletions

View File

@@ -1 +1,16 @@
"""Data processing module.""" """数据处理模块。
本模块已重构为多个子模块以提高可维护性:
- store/: 数据持久层,拆分自原 store.py (2869行 → 8个文件)
- algorithms/: 算法和解析器,拆分自原 algorithms.py (3330行 → 11个文件)
使用方式:
from app.modules.data_process.store import DataProcessStore
from app.modules.data_process.algorithms import estimate_token_count
"""
# 注意:为了避免循环导入,不在此处导入所有内容
# 请直接从子模块导入所需功能
__all__ = ["store", "algorithms"]

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,145 @@
"""数据处理算法模块。
重构为多个按职责拆分的子模块:
- types.py: 类型定义、常量与 dataclass
- text_utils.py: 文本解码、归一化与格式检测
- format_detection.py: 文档结构检测
- parsers/: PDF / Office / JSON / CSV 解析器
- quality.py: 质量评分与去重
- transforms.py: 数据集分割
- structured_processing.py: 结构化数据预处理
使用方式:
from app.modules.data_process.algorithms import estimate_token_count
子模块之间存在导入分层,顶层按依赖顺序导入以避免循环导入。
"""
from __future__ import annotations
# Layer 0: 类型与常量(叶子,不依赖内部模块)
from .types import (
DatasetSplit,
DocumentHeading,
DocumentNoiseSpan,
DocumentStructure,
MAX_QA_PAIRS_PER_ITEM,
ParsedText,
PdfPageText,
ProcessedStructuredRecord,
QualityScore,
SUPPORTED_TEXT_FORMATS,
StructuredPreprocessOption,
TextFormat,
_MAX_WORKBOOK_COLUMNS,
_MAX_WORKBOOK_HEADER_SCAN_ROWS,
)
# Layer 2: 文本工具(仅依赖 types对 parsers 的依赖在函数体内延迟导入)
from .text_utils import (
_normalize_spreadsheet_value,
decode_utf8,
detect_text_format,
normalize_text,
parse_text_content,
parse_utf8_text,
structured_json_dumps,
)
# Layer 1: 解析器(依赖 types 与 text_utils
from .parsers import (
_infer_xlsx_header_region,
_rewrite_xlsx_workbook_relationships,
_validate_office_archive,
_xlsx_sheet_merge_ranges,
detect_pdf_document_noise,
extract_pdf_page_texts,
remove_document_noise,
)
# Layer 3: 数据转换与质量评分
from .transforms import stable_split, stable_split_assignments
from .quality import (
content_quality_flags,
deduplicate_structured_records,
estimate_token_count,
fingerprints_are_near_duplicate,
is_low_quality_content,
is_near_duplicate,
near_duplicate_fingerprint,
record_fingerprint,
score_quality,
)
# Layer 4: 结构化数据处理(依赖 text_utils / quality / transforms / parsers
from .structured_processing import (
canonical_record_json,
desensitize_pii,
desensitize_structured_record,
expand_to_context_boundaries,
extract_structured_records,
filter_anomalous_structured_records,
flatten_structured_record,
generate_standard_records,
merge_short_blocks,
normalize_structured_record,
preprocess_structured_records,
preprocess_structured_records_with_lineage,
protected_context_ranges,
)
# Layer 5: 文档结构检测(依赖 structured_processing
from .format_detection import detect_document_structure
__all__ = [
"DatasetSplit",
"DocumentHeading",
"DocumentNoiseSpan",
"DocumentStructure",
"MAX_QA_PAIRS_PER_ITEM",
"ParsedText",
"PdfPageText",
"ProcessedStructuredRecord",
"QualityScore",
"SUPPORTED_TEXT_FORMATS",
"StructuredPreprocessOption",
"TextFormat",
"_MAX_WORKBOOK_COLUMNS",
"_MAX_WORKBOOK_HEADER_SCAN_ROWS",
"_normalize_spreadsheet_value",
"canonical_record_json",
"content_quality_flags",
"decode_utf8",
"deduplicate_structured_records",
"desensitize_pii",
"desensitize_structured_record",
"detect_document_structure",
"detect_pdf_document_noise",
"detect_text_format",
"estimate_token_count",
"expand_to_context_boundaries",
"extract_pdf_page_texts",
"extract_structured_records",
"filter_anomalous_structured_records",
"fingerprints_are_near_duplicate",
"flatten_structured_record",
"generate_standard_records",
"is_low_quality_content",
"is_near_duplicate",
"merge_short_blocks",
"near_duplicate_fingerprint",
"normalize_structured_record",
"normalize_text",
"parse_text_content",
"parse_utf8_text",
"preprocess_structured_records",
"preprocess_structured_records_with_lineage",
"protected_context_ranges",
"record_fingerprint",
"remove_document_noise",
"score_quality",
"stable_split",
"stable_split_assignments",
"structured_json_dumps",
]

View File

@@ -0,0 +1,119 @@
"""数据处理算法 - 格式检测。"""
from __future__ import annotations
import re
from .structured_processing import protected_context_ranges
from .text_utils import normalize_text
from .types import DocumentHeading, DocumentStructure
def detect_document_structure(text: str) -> DocumentStructure:
"""识别 Markdown、中文章节和数字编号标题及不可拆分块。"""
normalized = normalize_text(text)
if not normalized:
return DocumentStructure(
line_count=0,
paragraph_count=0,
headings=(),
code_block_count=0,
table_block_count=0,
list_block_count=0,
)
code_ranges = protected_context_ranges(
normalized,
preserve_code_blocks=True,
preserve_tables=False,
preserve_lists=False,
)
table_candidates = protected_context_ranges(
normalized,
preserve_code_blocks=False,
preserve_tables=True,
preserve_lists=False,
)
list_candidates = protected_context_ranges(
normalized,
preserve_code_blocks=False,
preserve_tables=False,
preserve_lists=True,
)
table_ranges = tuple(
item
for item in table_candidates
if not any(
item[0] < code_end and item[1] > code_start
for code_start, code_end in code_ranges
)
)
list_ranges = tuple(
item
for item in list_candidates
if not any(
item[0] < code_end and item[1] > code_start
for code_start, code_end in code_ranges
)
)
markdown_heading = re.compile(r"^\s*(?P<marks>#{1,6})\s+(?P<title>.+?)\s*#*\s*$")
chinese_heading = re.compile(
r"^\s*(?P<title>第[一二三四五六七八九十百千万0-9]+[章节篇部分].*)$"
)
numbered_heading = re.compile(
r"^\s*(?P<number>\d+(?:\.\d+)*)(?:[、.]|\s+)\s*(?P<title>\S.*)$"
)
headings: list[DocumentHeading] = []
cursor = 0
for line_number, raw_line in enumerate(normalized.splitlines(keepends=True), start=1):
line = raw_line.rstrip("\n")
line_end = cursor + len(line)
if not any(range_start <= cursor < range_end for range_start, range_end in code_ranges):
match = markdown_heading.match(line)
if match:
headings.append(
DocumentHeading(
level=len(match.group("marks")),
title=normalize_text(match.group("title")),
line_number=line_number,
start=cursor,
end=line_end,
)
)
else:
match = chinese_heading.match(line)
if match:
headings.append(
DocumentHeading(
level=1,
title=normalize_text(match.group("title")),
line_number=line_number,
start=cursor,
end=line_end,
)
)
else:
match = numbered_heading.match(line)
if match:
headings.append(
DocumentHeading(
level=min(6, match.group("number").count(".") + 1),
title=normalize_text(match.group("title")),
line_number=line_number,
start=cursor,
end=line_end,
)
)
cursor += len(raw_line)
paragraphs = [part for part in re.split(r"\n\s*\n", normalized) if part.strip()]
return DocumentStructure(
line_count=len(normalized.splitlines()),
paragraph_count=len(paragraphs),
headings=tuple(headings),
code_block_count=len(code_ranges),
table_block_count=len(table_ranges),
list_block_count=len(list_ranges),
)

View File

@@ -0,0 +1,19 @@
"""文档解析器模块。"""
from .pdf import extract_pdf_page_texts, detect_pdf_document_noise, remove_document_noise
from .office import (
_validate_office_archive,
_rewrite_xlsx_workbook_relationships,
_xlsx_sheet_merge_ranges,
_infer_xlsx_header_region,
)
__all__ = [
'extract_pdf_page_texts',
'detect_pdf_document_noise',
'remove_document_noise',
'_validate_office_archive',
'_rewrite_xlsx_workbook_relationships',
'_xlsx_sheet_merge_ranges',
'_infer_xlsx_header_region',
]

View File

@@ -0,0 +1,89 @@
"""数据处理算法 - CSV 解析。"""
from __future__ import annotations
import csv
import io
import re
from typing import Any
from ..text_utils import normalize_text
def _source_line_offsets(
text: str,
start_line: int,
end_line: int,
) -> tuple[int, int]:
"""把 1-based 物理行范围转换为左闭右开的字符范围。"""
line_starts = [0]
line_starts.extend(match.end() for match in re.finditer("\n", text))
if start_line < 1 or end_line < start_line or end_line > len(line_starts):
raise ValueError("source line range is outside normalized text")
source_start = line_starts[start_line - 1]
source_end = (
line_starts[end_line] - 1
if end_line < len(line_starts)
else len(text)
)
return source_start, source_end
def _extract_csv_records_with_locators(
text: str,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""从 CSV 文本中提取记录及其行级定位信息。"""
normalized_text = normalize_text(text)
if not normalized_text:
return [], []
try:
dialect = csv.Sniffer().sniff(normalized_text[:8192], delimiters=",\t;")
except csv.Error:
dialect = csv.excel
reader = csv.DictReader(io.StringIO(normalized_text), dialect=dialect)
if not reader.fieldnames:
raise ValueError("CSV header is required")
headers = [normalize_text(header or "") for header in reader.fieldnames]
if any(not header for header in headers):
raise ValueError("CSV header cannot be empty")
if len(set(headers)) != len(headers):
raise ValueError("CSV headers must be unique")
reader.fieldnames = headers
records: list[dict[str, Any]] = []
locators: list[dict[str, Any]] = []
source_lines = normalized_text.splitlines()
previous_end_line = reader.line_num
for row in reader:
end_line = reader.line_num
start_line = previous_end_line + 1
previous_end_line = end_line
while start_line < end_line and not source_lines[start_line - 1].strip():
start_line += 1
if None in row:
raise ValueError("CSV row has more fields than the header")
normalized_row = {
key: normalize_text(value or "")
for key, value in row.items()
}
if any(value for value in normalized_row.values()):
records.append(normalized_row)
source_start, source_end = _source_line_offsets(
normalized_text,
start_line,
end_line,
)
locators.append(
{
"kind": "csv",
"record_index": len(records),
"start_line": start_line,
"end_line": end_line,
"source_start": source_start,
"source_end": source_end,
}
)
return records, locators

View File

@@ -0,0 +1,395 @@
"""数据处理算法 - JSON/JSONL 解析。"""
from __future__ import annotations
import json
import re
from collections.abc import Mapping, Sequence
from decimal import Decimal
from typing import Any
from ..text_utils import _normalize_format, _normalize_value, normalize_text
from ..types import (
_DuplicateJsonKeyError,
_JSON_ENVELOPE_KEYS,
_JSON_RECORD_ARRAY_KEYS,
_JSON_RESPONSE_METADATA_KEYS,
_JSON_WRAPPER_METADATA_KEYS,
_MAX_JSON_DEPTH,
)
def _record_from_value(value: Any, *, normalize: bool = True) -> dict[str, Any]:
if isinstance(value, Mapping):
return dict(_normalize_value(value)) if normalize else dict(value)
return {"value": _normalize_value(value) if normalize else value}
def _json_pointer_segment(value: Any) -> str:
return str(value).replace("~", "~0").replace("/", "~1")
def _reject_duplicate_json_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
result: dict[str, Any] = {}
for key, value in pairs:
if key in result:
raise _DuplicateJsonKeyError(f"duplicate JSON object key: {key!r}")
result[key] = value
return result
def _reject_json_constant(value: str) -> Any:
raise ValueError(f"non-finite JSON number is not allowed: {value}")
def _skip_json_whitespace(text: str, offset: int) -> int:
while offset < len(text) and text[offset] in " \t\r\n":
offset += 1
return offset
def _validate_json_nesting(text: str) -> None:
"""在构造 Python 对象前限制容器深度,避免依赖解释器递归阈值。"""
depth = 0
in_string = False
escaped = False
for char in text:
if in_string:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
in_string = False
continue
if char == '"':
in_string = True
elif char in "[{":
depth += 1
if depth > _MAX_JSON_DEPTH:
raise ValueError(
f"JSON nesting exceeds the supported depth of {_MAX_JSON_DEPTH}"
)
elif char in "]}":
depth = max(0, depth - 1)
def _strict_json_loads(text: str) -> tuple[Any, int, int]:
"""严格解析单个 JSON 值并返回其左闭右开源码区间。"""
start = _skip_json_whitespace(text, 0)
if start >= len(text):
raise ValueError("JSON content is empty")
_validate_json_nesting(text)
decoder = json.JSONDecoder(
object_pairs_hook=_reject_duplicate_json_keys,
parse_float=Decimal,
parse_int=int,
parse_constant=_reject_json_constant,
strict=True,
)
try:
payload, end = decoder.raw_decode(text, start)
except json.JSONDecodeError as exc:
raise ValueError(
f"invalid JSON at line {exc.lineno}, column {exc.colno}: {exc.msg}"
) from exc
except RecursionError as exc:
raise ValueError("JSON nesting exceeds the supported depth") from exc
except _DuplicateJsonKeyError as exc:
raise ValueError(str(exc)) from exc
except ValueError as exc:
# parse_int/parse_float/parse_constant 的异常也必须稳定映射为客户端错误。
raise ValueError(f"invalid JSON number: {exc}") from exc
trailing = _skip_json_whitespace(text, end)
if trailing != len(text):
line = text.count("\n", 0, trailing) + 1
line_start = text.rfind("\n", 0, trailing) + 1
column = trailing - line_start + 1
raise ValueError(
f"invalid JSON at line {line}, column {column}: extra data"
)
return payload, start, end
def _json_value_end(text: str, start: int) -> int:
"""在已验证 JSON 中定位一个值的结束偏移,不对数值做二次解析。"""
if start >= len(text):
raise ValueError("invalid JSON source span")
first = text[start]
if first == '"':
escaped = False
for offset in range(start + 1, len(text)):
char = text[offset]
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
return offset + 1
raise ValueError("invalid JSON source span")
if first in "[{":
stack = [first]
in_string = False
escaped = False
for offset in range(start + 1, len(text)):
char = text[offset]
if in_string:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
in_string = False
continue
if char == '"':
in_string = True
elif char in "[{":
stack.append(char)
elif char in "]}":
expected = "[" if char == "]" else "{"
if not stack or stack[-1] != expected:
raise ValueError("invalid JSON source span")
stack.pop()
if not stack:
return offset + 1
raise ValueError("invalid JSON source span")
end = start
while end < len(text) and text[end] not in " \t\r\n,]}":
end += 1
if end == start:
raise ValueError("invalid JSON source span")
return end
def _json_object_value_spans(
text: str,
start: int,
end: int,
) -> dict[str, tuple[int, int]]:
"""返回已验证 JSON 对象直接子字段的值区间。"""
if text[start] != "{" or text[end - 1] != "}":
raise ValueError("JSON source value is not an object")
result: dict[str, tuple[int, int]] = {}
offset = _skip_json_whitespace(text, start + 1)
key_decoder = json.JSONDecoder()
while offset < end - 1:
key, key_end = key_decoder.raw_decode(text, offset)
if not isinstance(key, str):
raise ValueError("invalid JSON object key")
offset = _skip_json_whitespace(text, key_end)
if offset >= end or text[offset] != ":":
raise ValueError("invalid JSON object member")
value_start = _skip_json_whitespace(text, offset + 1)
value_end = _json_value_end(text, value_start)
result[key] = (value_start, value_end)
offset = _skip_json_whitespace(text, value_end)
if offset >= end - 1:
break
if text[offset] != ",":
raise ValueError("invalid JSON object member")
offset = _skip_json_whitespace(text, offset + 1)
return result
def _json_array_item_spans(
text: str,
start: int,
end: int,
) -> list[tuple[int, int]]:
"""返回已验证 JSON 数组中每个直接元素的源码区间。"""
if text[start] != "[" or text[end - 1] != "]":
raise ValueError("JSON source value is not an array")
result: list[tuple[int, int]] = []
offset = _skip_json_whitespace(text, start + 1)
while offset < end - 1:
item_end = _json_value_end(text, offset)
result.append((offset, item_end))
offset = _skip_json_whitespace(text, item_end)
if offset >= end - 1:
break
if text[offset] != ",":
raise ValueError("invalid JSON array item")
offset = _skip_json_whitespace(text, offset + 1)
return result
def _json_span_at_path(
text: str,
root_span: tuple[int, int],
path: Sequence[str],
) -> tuple[int, int]:
span = root_span
for key in path:
try:
span = _json_object_value_spans(text, *span)[key]
except KeyError as exc:
raise ValueError(f"JSON source path cannot be located: {key}") from exc
return span
def _pure_json_record_wrapper(
payload: Any,
) -> tuple[list[Any], tuple[str, ...]] | None:
"""识别不会与业务字段冲突的纯记录包装对象。"""
if not isinstance(payload, Mapping):
return None
def direct_wrapper(
value: Mapping[str, Any],
metadata_keys: frozenset[str] | set[str] = _JSON_WRAPPER_METADATA_KEYS,
) -> tuple[list[Any], tuple[str, ...]] | None:
candidates = [
key
for key in _JSON_RECORD_ARRAY_KEYS
if isinstance(value.get(key), list)
]
if len(candidates) != 1:
return None
record_key = candidates[0]
records = value[record_key]
if any(not isinstance(record, Mapping) for record in records):
return None
if any(
key != record_key and key not in metadata_keys
for key in value
):
return None
return records, (record_key,)
direct = direct_wrapper(payload)
if direct is not None:
return direct
envelope_keys = [
key
for key in _JSON_ENVELOPE_KEYS
if isinstance(payload.get(key), Mapping)
]
if len(envelope_keys) != 1:
return None
envelope_key = envelope_keys[0]
if any(
key != envelope_key and key not in _JSON_RESPONSE_METADATA_KEYS
for key in payload
):
return None
nested = direct_wrapper(
payload[envelope_key],
_JSON_RESPONSE_METADATA_KEYS,
)
if nested is None:
return None
records, nested_path = nested
return records, (envelope_key, *nested_path)
def _json_record_locator(
text: str,
*,
record_index: int,
json_pointer: str,
span: tuple[int, int],
) -> dict[str, Any]:
source_start, source_end = span
start_line = text.count("\n", 0, source_start) + 1
last_character = max(source_start, source_end - 1)
end_line = text.count("\n", 0, last_character) + 1
return {
"kind": "json",
"record_index": record_index,
"json_pointer": json_pointer,
"source_start": source_start,
"source_end": source_end,
"start_line": start_line,
"end_line": end_line,
}
def _extract_structured_records_with_locators(
text: str,
file_format: str,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""从 JSON、JSONL 或 CSV 中提取记录。
JSON 根数组始终表示多条记录;对象仅在满足纯包装契约时展开,其他
对象均视为一条业务记录。JSON 字段和值在解析阶段保持原样,只有用户
明确选择 ``normalize_format`` 后才会规范化。
"""
normalized_format = _normalize_format(file_format)
if normalized_format not in {"json", "jsonl", "csv"}:
raise ValueError("structured record extraction only supports JSON, JSONL and CSV")
if normalized_format == "json":
if _skip_json_whitespace(text, 0) == len(text):
return [], []
payload, root_start, root_end = _strict_json_loads(text)
values: Sequence[Any]
pointer_path: tuple[str, ...] = ()
record_spans: list[tuple[int, int]]
if isinstance(payload, list):
values = payload
record_spans = _json_array_item_spans(text, root_start, root_end)
else:
wrapper = _pure_json_record_wrapper(payload)
if wrapper is None:
values = [payload]
record_spans = [(root_start, root_end)]
else:
values, pointer_path = wrapper
array_span = _json_span_at_path(
text,
(root_start, root_end),
pointer_path,
)
record_spans = _json_array_item_spans(text, *array_span)
if len(record_spans) != len(values):
raise ValueError("JSON record source spans do not match parsed records")
records = [_record_from_value(value, normalize=False) for value in values]
pointer_prefix = "".join(
f"/{_json_pointer_segment(segment)}" for segment in pointer_path
)
locators = [
_json_record_locator(
text,
record_index=index + 1,
json_pointer=(
f"{pointer_prefix}/{index}"
if pointer_path or isinstance(payload, list)
else ""
),
span=record_spans[index],
)
for index in range(len(records))
]
return records, locators
if normalized_format == "jsonl":
records: list[dict[str, Any]] = []
locators: list[dict[str, Any]] = []
source_offset = 0
for line_number, line in enumerate(text.split("\n"), start=1):
line_content_end = len(line)
if _skip_json_whitespace(line, 0) == line_content_end:
source_offset += len(line) + 1
continue
try:
value, value_start, value_end = _strict_json_loads(line)
except ValueError as exc:
raise ValueError(
f"invalid JSONL at line {line_number}: {exc}"
) from exc
records.append(_record_from_value(value, normalize=False))
locators.append(
{
"kind": "jsonl",
"record_index": len(records),
"start_line": line_number,
"end_line": line_number,
"source_start": source_offset + value_start,
"source_end": source_offset + value_end,
}
)
source_offset += len(line) + 1
return records, locators
# CSV 分支独立放在 csv_parser 中,避免与 JSON 机制耦合。
from .csv_parser import _extract_csv_records_with_locators
return _extract_csv_records_with_locators(text)

View File

@@ -0,0 +1,700 @@
"""数据处理算法 - Office 文档解析。"""
from __future__ import annotations
import io
import re
import unicodedata
import zipfile
import xml.etree.ElementTree as ET
from collections.abc import Mapping, Sequence
from pathlib import PurePosixPath
from typing import Any
from urllib.parse import unquote, urlsplit
from docx import Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from openpyxl import load_workbook
from openpyxl.utils.cell import range_boundaries
from pptx import Presentation
from ..text_utils import _append_bounded_text, _normalize_spreadsheet_value, normalize_text
from ..types import (
_MAX_ARCHIVE_COMPRESSION_RATIO,
_MAX_ARCHIVE_ENTRIES,
_MAX_ARCHIVE_ENTRY_BYTES,
_MAX_ARCHIVE_UNCOMPRESSED_BYTES,
_MAX_PRESENTATION_SLIDES,
_MAX_WORKBOOK_CELLS,
_MAX_WORKBOOK_COLUMNS,
_MAX_WORKBOOK_HEADER_ROWS,
_MAX_WORKBOOK_HEADER_SCAN_ROWS,
_MAX_WORKBOOK_MERGED_RANGES,
_MAX_WORKBOOK_ROWS,
_MAX_WORKBOOK_SCANNED_ROWS,
_MAX_WORKBOOK_SHEETS,
TextFormat,
)
_XLSX_REPORT_METADATA_PATTERN = re.compile(
r"^(?:报表|报告|标题|说明|备注|制表|统计|日期|时间|期间|"
r"report|title|note|remark|date|time|period)\b",
flags=re.IGNORECASE,
)
def _validate_office_archive(raw: bytes, file_format: TextFormat) -> None:
"""在交给 Office 解析库前限制 ZIP 包规模并拒绝活动 XML。"""
required_members = {
"docx": {"[Content_Types].xml", "word/document.xml"},
"xlsx": {"[Content_Types].xml", "xl/workbook.xml"},
"pptx": {"[Content_Types].xml", "ppt/presentation.xml"},
}
try:
with zipfile.ZipFile(io.BytesIO(raw)) as archive:
members = archive.infolist()
if len(members) > _MAX_ARCHIVE_ENTRIES:
raise ValueError(
f"{file_format.upper()} archive contains too many entries "
f"(limit {_MAX_ARCHIVE_ENTRIES})"
)
names: set[str] = set()
total_size = 0
for member in members:
path = PurePosixPath(member.filename.replace("\\", "/"))
if path.is_absolute() or ".." in path.parts:
raise ValueError(
f"{file_format.upper()} archive contains an unsafe member path"
)
if member.filename in names:
raise ValueError(
f"{file_format.upper()} archive contains duplicate member names"
)
names.add(member.filename)
if member.flag_bits & 0x1:
raise ValueError(f"encrypted {file_format.upper()} files are not supported")
if member.is_dir():
continue
if member.file_size > _MAX_ARCHIVE_ENTRY_BYTES:
raise ValueError(
f"{file_format.upper()} archive entry exceeds "
f"{_MAX_ARCHIVE_ENTRY_BYTES} bytes"
)
total_size += member.file_size
if total_size > _MAX_ARCHIVE_UNCOMPRESSED_BYTES:
raise ValueError(
f"{file_format.upper()} archive expands beyond "
f"{_MAX_ARCHIVE_UNCOMPRESSED_BYTES} bytes"
)
if member.file_size >= 1024 * 1024:
if member.compress_size <= 0:
raise ValueError(
f"{file_format.upper()} archive has an unsafe compression ratio"
)
ratio = member.file_size / member.compress_size
if ratio > _MAX_ARCHIVE_COMPRESSION_RATIO:
raise ValueError(
f"{file_format.upper()} archive has an unsafe compression ratio"
)
missing = required_members[file_format] - names
if missing:
raise ValueError(
f"invalid {file_format.upper()} package: missing "
f"{', '.join(sorted(missing))}"
)
for member in members:
if member.is_dir() or not member.filename.lower().endswith((".xml", ".rels")):
continue
with archive.open(member) as stream:
prefix = stream.read(min(member.file_size, 1024 * 1024)).upper()
if b"<!DOCTYPE" in prefix or b"<!ENTITY" in prefix:
raise ValueError(
f"{file_format.upper()} archive contains unsupported active XML"
)
except zipfile.BadZipFile as exc:
raise ValueError(f"invalid {file_format.upper()} file: not an Office ZIP package") from exc
def _extract_docx_text(raw: bytes) -> str:
_validate_office_archive(raw, "docx")
try:
document = Document(io.BytesIO(raw))
except Exception as exc:
raise ValueError(f"invalid DOCX file: {exc}") from exc
parts: list[str] = []
total = 0
for child in document.element.body.iterchildren():
if isinstance(child, CT_P):
total = _append_bounded_text(parts, Paragraph(child, document).text, total)
continue
if isinstance(child, CT_Tbl):
table = Table(child, document)
for row in table.rows:
cells = [normalize_text(cell.text) for cell in row.cells]
total = _append_bounded_text(parts, "\t".join(cells), total)
return normalize_text("\n\n".join(parts))
def _presentation_shape_text(shape: Any) -> list[str]:
if getattr(shape, "has_table", False):
return [
"\t".join(normalize_text(cell.text) for cell in row.cells)
for row in shape.table.rows
]
if getattr(shape, "has_text_frame", False):
return [shape.text]
child_shapes = getattr(shape, "shapes", None)
if child_shapes is not None:
values: list[str] = []
for child in child_shapes:
values.extend(_presentation_shape_text(child))
return values
return []
def _extract_pptx_text(raw: bytes) -> str:
_validate_office_archive(raw, "pptx")
try:
presentation = Presentation(io.BytesIO(raw))
except Exception as exc:
raise ValueError(f"invalid PPTX file: {exc}") from exc
if len(presentation.slides) > _MAX_PRESENTATION_SLIDES:
raise ValueError(
f"PPTX contains too many slides (limit {_MAX_PRESENTATION_SLIDES})"
)
parts: list[str] = []
total = 0
for slide in presentation.slides:
slide_parts: list[str] = []
for shape in slide.shapes:
slide_parts.extend(_presentation_shape_text(shape))
total = _append_bounded_text(parts, "\n".join(slide_parts), total)
return normalize_text("\n\n".join(parts))
def _xml_local_name(tag: str) -> str:
return tag.rsplit("}", 1)[-1]
def _resolve_xlsx_relationship_target(
archive: zipfile.ZipFile,
target: str,
*,
source_part: str = "xl/workbook.xml",
) -> str:
"""按 OPC URI 规则解析内部关系目标,并保证结果仍位于 ZIP 根内。"""
raw_target = target.strip()
if (
raw_target != target
or not raw_target
or "\\" in raw_target
or any(unicodedata.category(char).startswith("C") for char in raw_target)
or re.search(r"%(?![0-9A-Fa-f]{2})", raw_target)
):
raise ValueError("XLSX workbook contains an unsafe worksheet path")
try:
parsed = urlsplit(raw_target)
except ValueError as exc:
raise ValueError("XLSX workbook contains an unsafe worksheet path") from exc
if parsed.scheme or parsed.netloc or parsed.query or parsed.fragment:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
if re.search(r"%(?:2[fF]|5[cC]|0{2})", parsed.path):
raise ValueError("XLSX workbook contains an unsafe worksheet path")
try:
decoded_path = unquote(parsed.path, encoding="utf-8", errors="strict")
except UnicodeDecodeError as exc:
raise ValueError("XLSX workbook contains an unsafe worksheet path") from exc
if not decoded_path or "\\" in decoded_path or "\x00" in decoded_path or "%" in decoded_path:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
parts = [] if decoded_path.startswith("/") else list(PurePosixPath(source_part).parent.parts)
for part in decoded_path.lstrip("/").split("/"):
if part in {"", "."}:
continue
if part == "..":
if not parts:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
parts.pop()
continue
if unicodedata.category(part[0]).startswith("C") or any(
unicodedata.category(char).startswith("C") for char in part
):
raise ValueError("XLSX workbook contains an unsafe worksheet path")
parts.append(part)
if not parts:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
member_name = "/".join(parts)
try:
member = archive.getinfo(member_name)
except KeyError as exc:
raise ValueError(
f"XLSX worksheet relationship target does not exist: {member_name}"
) from exc
if member.is_dir():
raise ValueError("XLSX worksheet relationship target must be a file")
return member_name
def _rewrite_xlsx_workbook_relationships(
raw: bytes,
replacements: Mapping[str, str],
) -> bytes:
"""把已验证的 worksheet Target 改为解析库稳定支持的包内绝对路径。"""
relationships_member = "xl/_rels/workbook.xml.rels"
output = io.BytesIO()
with zipfile.ZipFile(io.BytesIO(raw)) as source, zipfile.ZipFile(output, "w") as target:
target.comment = source.comment
for member in source.infolist():
if member.filename == relationships_member:
root = ET.fromstring(source.read(member))
pending = dict(replacements)
for element in root:
if _xml_local_name(element.tag) != "Relationship":
continue
relationship_id = element.attrib.get("Id")
if relationship_id in pending:
element.set("Target", pending.pop(relationship_id))
if pending:
raise ValueError(
"XLSX workbook relationship changed during normalization"
)
content = ET.tostring(root, encoding="utf-8", xml_declaration=True)
target.writestr(member, content)
continue
if member.is_dir():
target.writestr(member, b"")
continue
with source.open(member) as source_stream, target.open(
member,
"w",
force_zip64=True,
) as target_stream:
while chunk := source_stream.read(1024 * 1024):
target_stream.write(chunk)
return output.getvalue()
def _xlsx_sheet_merge_ranges(
raw: bytes,
) -> tuple[
dict[str, tuple[tuple[int, int, int, int], ...]],
dict[str, str],
]:
"""流式读取 XLSX 合并单元格,不把工作表 XML 整体载入内存。"""
relationship_namespace = (
"http://schemas.openxmlformats.org/officeDocument/2006/relationships"
)
try:
with zipfile.ZipFile(io.BytesIO(raw)) as archive:
relationships: dict[str, tuple[str, str, str]] = {}
with archive.open("xl/_rels/workbook.xml.rels") as stream:
for _, element in ET.iterparse(stream, events=("end",)):
if _xml_local_name(element.tag) != "Relationship":
element.clear()
continue
relationship_id = element.attrib.get("Id")
target = element.attrib.get("Target")
target_mode = element.attrib.get("TargetMode", "Internal")
relationship_type = element.attrib.get("Type", "")
if relationship_id:
if relationship_id in relationships:
raise ValueError(
"XLSX workbook contains duplicate relationship identifiers"
)
relationships[relationship_id] = (
target or "",
target_mode,
relationship_type,
)
element.clear()
sheet_paths: dict[str, str] = {}
normalized_targets: dict[str, str] = {}
with archive.open("xl/workbook.xml") as stream:
for _, element in ET.iterparse(stream, events=("end",)):
if _xml_local_name(element.tag) != "sheet":
element.clear()
continue
title = element.attrib.get("name")
relationship_id = element.attrib.get(
f"{{{relationship_namespace}}}id"
)
if title and relationship_id:
relationship = relationships.get(relationship_id)
if relationship is None:
raise ValueError(
f"XLSX worksheet {title!r} references a missing relationship"
)
target, target_mode, relationship_type = relationship
if target_mode.strip().lower() != "internal":
raise ValueError(
f"XLSX worksheet {title!r} uses an external relationship"
)
if not relationship_type.endswith("/worksheet"):
raise ValueError(
f"XLSX worksheet {title!r} has an invalid relationship type"
)
member_name = _resolve_xlsx_relationship_target(
archive,
target,
)
sheet_paths[title] = member_name
canonical_target = f"/{member_name}"
if target != canonical_target:
normalized_targets[relationship_id] = canonical_target
element.clear()
result: dict[str, tuple[tuple[int, int, int, int], ...]] = {}
total_ranges = 0
for title, member_name in sheet_paths.items():
ranges: list[tuple[int, int, int, int]] = []
with archive.open(member_name) as stream:
for _, element in ET.iterparse(stream, events=("end",)):
if _xml_local_name(element.tag) != "mergeCell":
element.clear()
continue
reference = element.attrib.get("ref")
if reference:
try:
boundaries = range_boundaries(reference)
except ValueError as exc:
raise ValueError(
f"XLSX worksheet {title!r} has an invalid merged range"
) from exc
ranges.append(boundaries)
total_ranges += 1
if total_ranges > _MAX_WORKBOOK_MERGED_RANGES:
raise ValueError(
"XLSX workbook contains too many merged ranges "
f"(limit {_MAX_WORKBOOK_MERGED_RANGES})"
)
element.clear()
result[title] = tuple(ranges)
return result, normalized_targets
except (KeyError, ET.ParseError, zipfile.BadZipFile) as exc:
raise ValueError(f"invalid XLSX workbook structure: {exc}") from exc
def _xlsx_header_end_row(
first_row: int,
rows: Mapping[int, Sequence[Any]],
merged_ranges: Sequence[tuple[int, int, int, int]],
) -> int:
"""在固定深度内闭包表头合并关系,忽略越界或跨空行的可疑级联。"""
header_end = first_row
maximum_end = first_row + _MAX_WORKBOOK_HEADER_ROWS - 1
has_header_hierarchy = any(
max_column > min_column and min_row == first_row
for min_column, min_row, max_column, _ in merged_ranges
)
changed = True
while changed:
changed = False
for min_column, min_row, max_column, max_row in merged_ranges:
if min_row < first_row or min_row > header_end or max_row < first_row:
continue
if max_column == min_column and not has_header_hierarchy:
continue
candidate = max_row + 1 if max_column > min_column else max_row
if candidate <= header_end or candidate > maximum_end:
continue
if (
max_column > min_column
and len(_xlsx_nonempty_values(rows.get(min_row, ()))) < 2
and len(_xlsx_nonempty_values(rows.get(candidate, ()))) < 2
):
continue
if any(
not rows.get(row_number)
for row_number in range(header_end + 1, candidate + 1)
):
continue
header_end = candidate
changed = True
return header_end
def _xlsx_headers(
title: str,
rows: Mapping[int, Sequence[Any]],
first_row: int,
header_end: int,
merged_ranges: Sequence[tuple[int, int, int, int]],
) -> list[str]:
width = max((len(row) for row in rows.values()), default=0)
for min_column, min_row, max_column, max_row in merged_ranges:
if min_row <= header_end and max_row >= first_row:
width = max(width, max_column)
if width > _MAX_WORKBOOK_COLUMNS:
raise ValueError(
f"XLSX worksheet {title!r} exceeds {_MAX_WORKBOOK_COLUMNS} columns"
)
matrix = [
list(rows.get(row_number, ())) + [None] * (width - len(rows.get(row_number, ())))
for row_number in range(first_row, header_end + 1)
]
for min_column, min_row, max_column, max_row in merged_ranges:
if min_row > header_end or max_row < first_row:
continue
source_row = max(first_row, min_row) - first_row
source_column = min_column - 1
source = matrix[source_row][source_column]
for row_number in range(max(first_row, min_row), min(header_end, max_row) + 1):
for column_number in range(min_column, max_column + 1):
matrix[row_number - first_row][column_number - 1] = source
headers: list[str] = []
for column in range(width):
components: list[str] = []
for row in matrix:
component = normalize_text(str(row[column] or ""))
if component and (not components or component != components[-1]):
components.append(component)
header = ".".join(components)
if not header:
raise ValueError(f"XLSX worksheet {title!r} contains an empty header")
headers.append(header)
if len(set(headers)) != len(headers):
raise ValueError(f"XLSX worksheet {title!r} contains duplicate headers")
return headers
def _xlsx_nonempty_values(row: Sequence[Any]) -> list[Any]:
return [value for value in row if value not in {None, ""}]
def _infer_xlsx_header_region(
title: str,
rows: Mapping[int, Sequence[Any]],
merged_ranges: Sequence[tuple[int, int, int, int]],
) -> tuple[int, int, list[str]]:
"""从有界前缀中推断表头,区分报表说明、多级表头和正文合并。"""
ordered_rows = sorted(rows)
if not ordered_rows:
return 0, 0, []
first_nonempty_row = ordered_rows[0]
horizontal_merge_rows = {
min_row
for min_column, min_row, max_column, _ in merged_ranges
if max_column > min_column
}
candidates: list[tuple[float, int, int, list[str]]] = []
for first_row in ordered_rows:
raw_values = _xlsx_nonempty_values(rows[first_row])
if not raw_values:
continue
if len(raw_values) < 2 and first_row in horizontal_merge_rows:
continue
header_end = _xlsx_header_end_row(first_row, rows, merged_ranges)
header_rows = {
row_number: rows[row_number]
for row_number in range(first_row, header_end + 1)
if row_number in rows
}
try:
headers = _xlsx_headers(
title,
header_rows,
first_row,
header_end,
merged_ranges,
)
except ValueError:
if header_end == first_row:
continue
header_end = first_row
try:
headers = _xlsx_headers(
title,
{first_row: rows[first_row]},
first_row,
first_row,
merged_ranges,
)
except ValueError:
continue
data_rows = [rows[row_number] for row_number in ordered_rows if row_number > header_end]
if any(len(row) > len(headers) for row in data_rows):
continue
text_ratio = sum(isinstance(value, str) for value in raw_values) / len(raw_values)
score = text_ratio * 6 + min(len(raw_values), 4)
if data_rows:
first_data_values = _xlsx_nonempty_values(data_rows[0])
score += 4 * min(len(first_data_values), len(headers)) / len(headers)
if first_data_values:
score += (
2
* sum(
not isinstance(value, str)
for value in first_data_values
)
/ len(first_data_values)
)
first_label = normalize_text(str(raw_values[0]))
if len(raw_values) <= 2 and _XLSX_REPORT_METADATA_PATTERN.match(first_label):
score -= 8
candidates.append((score, first_row, header_end, headers))
if not candidates:
first_row = first_nonempty_row
headers = _xlsx_headers(
title,
{first_row: rows[first_row]},
first_row,
first_row,
(),
)
return first_row, first_row, headers
_, first_row, header_end, headers = max(
candidates,
key=lambda candidate: (candidate[0], -candidate[1]),
)
return first_row, header_end, headers
def _extract_xlsx_records(
raw: bytes,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
_validate_office_archive(raw, "xlsx")
merged_by_sheet, normalized_targets = _xlsx_sheet_merge_ranges(raw)
workbook_raw = (
_rewrite_xlsx_workbook_relationships(raw, normalized_targets)
if normalized_targets
else raw
)
try:
workbook = load_workbook(
io.BytesIO(workbook_raw),
read_only=True,
data_only=True,
keep_links=False,
)
except Exception as exc:
raise ValueError(f"invalid XLSX file: {exc}") from exc
records: list[dict[str, Any]] = []
locators: list[dict[str, Any]] = []
total_cells = 0
try:
if len(workbook.worksheets) > _MAX_WORKBOOK_SHEETS:
raise ValueError(
f"XLSX contains too many worksheets (limit {_MAX_WORKBOOK_SHEETS})"
)
for sheet_index, worksheet in enumerate(workbook.worksheets):
reset_dimensions = getattr(worksheet, "reset_dimensions", None)
if callable(reset_dimensions):
reset_dimensions()
merged_ranges = merged_by_sheet.get(worksheet.title, ())
sheet_rows = 0
scanned_rows = 0
row_iterator = enumerate(
worksheet.iter_rows(values_only=True),
start=1,
)
buffered_rows: dict[int, Sequence[Any]] = {}
def normalized_row_values(
row: Sequence[Any],
sheet_title: str = worksheet.title,
) -> list[Any]:
values = list(row)
while values and values[-1] in {None, ""}:
values.pop()
if len(values) > _MAX_WORKBOOK_COLUMNS:
raise ValueError(
f"XLSX worksheet {sheet_title!r} exceeds "
f"{_MAX_WORKBOOK_COLUMNS} columns"
)
return values
for row_number, row in row_iterator:
scanned_rows += 1
if scanned_rows > _MAX_WORKBOOK_SCANNED_ROWS:
raise ValueError(
f"XLSX worksheet {worksheet.title!r} exceeds "
f"{_MAX_WORKBOOK_SCANNED_ROWS} scanned rows"
)
values = normalized_row_values(row)
if not values or all(value in {None, ""} for value in values):
continue
buffered_rows[row_number] = tuple(values)
if len(buffered_rows) >= _MAX_WORKBOOK_HEADER_SCAN_ROWS:
break
if not buffered_rows:
continue
_, header_end_row, headers = _infer_xlsx_header_region(
worksheet.title,
buffered_rows,
merged_ranges,
)
def append_record(
row_number: int,
values: Sequence[Any],
record_headers: Sequence[str] = tuple(headers),
locator_sheet_index: int = sheet_index,
sheet_title: str = worksheet.title,
) -> None:
nonlocal total_cells, sheet_rows
row_values = list(values)
if len(row_values) > len(record_headers):
raise ValueError(
f"XLSX worksheet {sheet_title!r} has a row wider than its header"
)
row_values.extend([None] * (len(record_headers) - len(row_values)))
record = {
header: _normalize_spreadsheet_value(value)
for header, value in zip(record_headers, row_values, strict=True)
}
if not any(value not in {"", None} for value in record.values()):
return
sheet_record_index = sheet_rows
sheet_rows += 1
total_cells += len(record_headers)
if sheet_rows > _MAX_WORKBOOK_ROWS:
raise ValueError(
f"XLSX worksheet {sheet_title!r} exceeds "
f"{_MAX_WORKBOOK_ROWS} data rows"
)
if total_cells > _MAX_WORKBOOK_CELLS:
raise ValueError(
f"XLSX workbook exceeds {_MAX_WORKBOOK_CELLS} populated cells"
)
records.append(record)
locators.append(
{
"kind": "xlsx",
"record_index": len(records),
"sheet_index": locator_sheet_index,
"sheet_name": sheet_title,
"row_number": row_number,
"sheet_record_index": sheet_record_index,
}
)
for row_number, values in buffered_rows.items():
if row_number > header_end_row:
append_record(row_number, values)
for row_number, row in row_iterator:
scanned_rows += 1
if scanned_rows > _MAX_WORKBOOK_SCANNED_ROWS:
raise ValueError(
f"XLSX worksheet {worksheet.title!r} exceeds "
f"{_MAX_WORKBOOK_SCANNED_ROWS} scanned rows"
)
values = normalized_row_values(row)
if not values or all(value in {None, ""} for value in values):
continue
append_record(row_number, values)
finally:
workbook.close()
return records, locators

View File

@@ -0,0 +1,297 @@
"""数据处理算法 - PDF 文档解析。"""
from __future__ import annotations
import io
import math
import re
import unicodedata
from collections.abc import Sequence
from typing import Literal
from pypdf import PdfReader
from ..text_utils import normalize_text
from ..types import (
DocumentNoiseSpan,
PdfPageText,
_MAX_EXTRACTED_TEXT_CHARS,
_MAX_PDF_PAGES,
_PdfLine,
)
_PDF_PAGE_NUMBER_LINE_PATTERN = re.compile(
r"^(?:页次\s*)?(?:第\s*)?(?P<page>\d+)\s*页\s*"
r"(?:(?:[/]\s*)?共\s*(?P<total>\d+)\s*页)?$"
)
_PDF_FRACTION_PAGE_LINE_PATTERN = re.compile(
r"^[—–-]?\s*(?P<page>\d+)\s*[/]\s*(?P<total>\d+)\s*[—–-]?$"
)
_PDF_CLASSIFICATION_LABEL_PATTERN = re.compile(
r"^(?:(?:秘密等级|密级)\s*)?(?:商密|秘密|机密|绝密)"
r"\s*(?:[【\[(][^】\])]{1,8}[】\])])?$"
)
_TOC_TITLE_PATTERN = re.compile(r"^(?:目\s*录|contents)$", re.IGNORECASE)
_TOC_LEADER_ENTRY_PATTERN = re.compile(
r"(?:[..…·•]\s*){3,}\s*\d{1,4}\s*$"
)
_TOC_NUMBERED_ENTRY_PATTERN = re.compile(
r"^(?:第[\u3400-\u4dbf\u4e00-\u9fff]{1,12}章|附表\s*\d+|\d+(?:\.\d+)+)"
r"\s+.+\s+\d{1,4}\s*$",
re.IGNORECASE,
)
_MARGIN_TEMPLATE_KEYWORDS = (
"",
"页次",
"版本",
"文件编码",
"秘密等级",
"密级",
"商密",
"confidential",
)
def extract_pdf_page_texts(raw: bytes) -> tuple[PdfPageText, ...]:
"""提取 PDF 各页文本,并保留与切片字符偏移一致的页范围。"""
if b"%PDF-" not in raw[:1024]:
raise ValueError("invalid PDF file: missing PDF header")
try:
reader = PdfReader(io.BytesIO(raw), strict=True)
except Exception as exc:
raise ValueError(f"invalid PDF file: {exc}") from exc
if reader.is_encrypted and not reader.decrypt(""):
raise ValueError("password-protected PDF files are not supported")
if len(reader.pages) > _MAX_PDF_PAGES:
raise ValueError(f"PDF contains too many pages (limit {_MAX_PDF_PAGES})")
pages: list[PdfPageText] = []
total = 0
has_text = False
for page_number, page in enumerate(reader.pages, start=1):
try:
text = normalize_text(page.extract_text() or "")
except Exception as exc:
raise ValueError(
f"failed to extract text from PDF page {page_number}: {exc}"
) from exc
if not text:
pages.append(
PdfPageText(
page_number=page_number,
text="",
source_start=total,
source_end=total,
)
)
continue
if has_text:
total += 2
start = total
total += len(text)
if total > _MAX_EXTRACTED_TEXT_CHARS:
raise ValueError(
f"extracted document text exceeds {_MAX_EXTRACTED_TEXT_CHARS} characters"
)
pages.append(
PdfPageText(
page_number=page_number,
text=text,
source_start=start,
source_end=total,
)
)
has_text = True
if not has_text:
raise ValueError("PDF contains no extractable text; scanned PDF requires OCR")
return tuple(pages)
def _pdf_page_lines(page: PdfPageText) -> tuple[_PdfLine, ...]:
lines: list[_PdfLine] = []
local_offset = 0
for raw_line in page.text.splitlines(keepends=True):
content = raw_line.rstrip("\r\n")
leading = len(content) - len(content.lstrip())
trailing = len(content.rstrip())
text = content.strip()
if text:
lines.append(
_PdfLine(
text=text,
start=page.source_start + local_offset + leading,
end=page.source_start + local_offset + trailing,
)
)
local_offset += len(raw_line)
return tuple(lines)
def _is_standalone_page_number(
text: str,
*,
physical_page: int,
page_count: int,
) -> bool:
normalized = unicodedata.normalize("NFKC", text).strip()
match = _PDF_PAGE_NUMBER_LINE_PATTERN.fullmatch(normalized)
if match is None:
match = _PDF_FRACTION_PAGE_LINE_PATTERN.fullmatch(normalized)
if match is None or int(match.group("page")) != physical_page:
return False
total = match.groupdict().get("total")
return total is None or int(total) == page_count
def _margin_signature(text: str) -> str:
normalized = unicodedata.normalize("NFKC", text).casefold()
normalized = re.sub(r"\s+", " ", normalized).strip()
if any(keyword in normalized for keyword in _MARGIN_TEMPLATE_KEYWORDS):
normalized = re.sub(r"\d+", "#", normalized)
return normalized
def _has_margin_metadata_keyword(text: str) -> bool:
normalized = unicodedata.normalize("NFKC", text).casefold()
return any(keyword in normalized for keyword in _MARGIN_TEMPLATE_KEYWORDS)
def _has_meaningful_margin_signature(signature: str) -> bool:
return len(re.sub(r"[#\W_]+", "", signature, flags=re.UNICODE)) >= 2
def _is_toc_leader_entry(text: str) -> bool:
return bool(_TOC_LEADER_ENTRY_PATTERN.search(text))
def _is_toc_numbered_entry(text: str) -> bool:
return bool(_TOC_NUMBERED_ENTRY_PATTERN.fullmatch(text))
def detect_pdf_document_noise(
pages: Sequence[PdfPageText],
) -> tuple[DocumentNoiseSpan, ...]:
"""识别 PDF 中的独立页码、重复页边内容和高置信目录。
规则只查看每页顶部 5 行和底部 3 行来推断页眉页脚;目录必须有
明显的点引导线密度,避免仅因正文中出现“目录”或章节标题而误删。
"""
page_lines = tuple(_pdf_page_lines(page) for page in pages)
detected: dict[tuple[int, int], DocumentNoiseSpan] = {}
def mark(
line: _PdfLine,
kind: Literal["page_number", "repeated_margin", "table_of_contents"],
) -> None:
detected.setdefault(
(line.start, line.end),
DocumentNoiseSpan(
start=line.start,
end=line.end,
kind=kind,
),
)
for page, lines in zip(pages, page_lines, strict=True):
for line in lines:
if _is_standalone_page_number(
line.text,
physical_page=page.page_number,
page_count=len(pages),
):
mark(line, "page_number")
outer_margin_lines = (*lines[:2], *lines[-2:])
for line in outer_margin_lines:
if _PDF_CLASSIFICATION_LABEL_PATTERN.fullmatch(line.text):
mark(line, "repeated_margin")
# 只在三页及以上文档中推断通用页眉页脚,避免短文档误删。
if len(pages) >= 3:
signature_pages: dict[str, set[int]] = {}
candidate_lines: list[tuple[int, _PdfLine, str]] = []
for page_index, lines in enumerate(page_lines):
boundary_lines = (
*((line, index < 2) for index, line in enumerate(lines[:5])),
*((line, index < 2) for index, line in enumerate(reversed(lines[-3:]))),
)
seen_ranges: set[tuple[int, int]] = set()
for line, is_outer_margin in boundary_lines:
line_range = (line.start, line.end)
if (
line_range in seen_ranges
or line_range in detected
or len(line.text) > 160
):
continue
seen_ranges.add(line_range)
if not is_outer_margin and not _has_margin_metadata_keyword(line.text):
continue
signature = _margin_signature(line.text)
if not _has_meaningful_margin_signature(signature):
continue
signature_pages.setdefault(signature, set()).add(page_index)
candidate_lines.append((page_index, line, signature))
minimum_pages = max(3, math.ceil(len(pages) * 0.3))
repeated_signatures = {
signature
for signature, matching_pages in signature_pages.items()
if len(matching_pages) >= minimum_pages
}
for _, line, signature in candidate_lines:
if signature in repeated_signatures:
mark(line, "repeated_margin")
# 先依据强证据判定目录页,再补充删除少量不带点引导线的编号目录项。
toc_active = False
for lines in page_lines:
content_lines = [
line for line in lines if (line.start, line.end) not in detected
]
leader_entries = [line for line in content_lines if _is_toc_leader_entry(line.text)]
titles = [line for line in content_lines if _TOC_TITLE_PATTERN.fullmatch(line.text)]
starts_toc = bool(titles and len(leader_entries) >= 2)
is_toc_dense = bool(
len(leader_entries) >= 3
and len(leader_entries) / max(1, len(content_lines)) >= 0.5
)
if not (starts_toc or (toc_active and is_toc_dense)):
toc_active = False
continue
toc_active = True
for line in content_lines:
if (
line in titles
or _is_toc_leader_entry(line.text)
or _is_toc_numbered_entry(line.text)
):
mark(line, "table_of_contents")
return tuple(sorted(detected.values(), key=lambda span: (span.start, span.end)))
def remove_document_noise(
text: str,
spans: Sequence[DocumentNoiseSpan],
*,
source_offset: int = 0,
) -> str:
"""按原文绝对偏移移除噪声,不改动调用方保留的原文及偏移。"""
text_end = source_offset + len(text)
intersections = sorted(
(
max(0, span.start - source_offset),
min(len(text), span.end - source_offset),
)
for span in spans
if span.start < text_end and span.end > source_offset
)
if not intersections:
return text
parts: list[str] = []
cursor = 0
for start, end in intersections:
if end <= cursor:
continue
if start > cursor:
parts.append(text[cursor:start])
cursor = end
parts.append(text[cursor:])
cleaned = normalize_text("".join(parts))
return re.sub(r"\n{3,}", "\n\n", cleaned)
def _extract_pdf_text(raw: bytes) -> str:
return "\n\n".join(page.text for page in extract_pdf_page_texts(raw) if page.text)

View File

@@ -0,0 +1,343 @@
"""数据处理算法 - 质量评分和去重。"""
from __future__ import annotations
import hashlib
import json
import re
import unicodedata
from collections import Counter
from collections.abc import Iterable, Mapping, Sequence
from copy import deepcopy
from typing import Any
from .text_utils import normalize_text
from .types import (
_MAX_ANOMALY_TEXT_CHARS,
_MOJIBAKE_MARKERS,
_TOKEN_PATTERN,
ProcessedStructuredRecord,
QualityScore,
)
def estimate_token_count(text: str) -> int:
"""粗略估计文本的 token 数量。"""
return len(_TOKEN_PATTERN.findall(text))
def content_quality_flags(
text: str,
*,
min_chars: int = 20,
min_tokens: int = 5,
max_chars: int = _MAX_ANOMALY_TEXT_CHARS,
) -> tuple[str, ...]:
"""返回非结构化内容的确定性低质量原因。"""
if min_chars < 0 or min_tokens < 0 or max_chars <= 0:
raise ValueError("content quality limits must be non-negative")
normalized = normalize_text(text)
if not normalized:
return ("empty_content",)
flags: list[str] = []
if len(normalized) < min_chars or estimate_token_count(normalized) < min_tokens:
flags.append("content_too_short")
if len(normalized) > max_chars:
flags.append("content_too_long")
if any(marker in normalized for marker in _MOJIBAKE_MARKERS):
flags.append("mojibake")
nonspace = [char for char in normalized if not char.isspace()]
if nonspace:
readable_ratio = sum(
char.isprintable()
and unicodedata.category(char) not in {"Co", "Cs", "Cn"}
for char in nonspace
) / len(nonspace)
if readable_ratio < 0.85:
flags.append("low_printable_ratio")
if len(nonspace) >= 100:
most_common = Counter(nonspace).most_common(1)[0][1]
if most_common / len(nonspace) > 0.9:
flags.append("repetitive_content")
return tuple(dict.fromkeys(flags))
def is_low_quality_content(
text: str,
*,
min_chars: int = 20,
min_tokens: int = 5,
max_chars: int = _MAX_ANOMALY_TEXT_CHARS,
) -> bool:
"""判断内容是否命中任一低质量规则。"""
return bool(
content_quality_flags(
text,
min_chars=min_chars,
min_tokens=min_tokens,
max_chars=max_chars,
)
)
def _deduplicate_structured_entries(
entries: Sequence[ProcessedStructuredRecord],
) -> list[ProcessedStructuredRecord]:
"""仅按整条 canonical JSON 稳定去重,避免误删同 ID 的更新记录。"""
# canonical_record_json 位于 structured_processing延迟导入以断开循环依赖。
from .structured_processing import canonical_record_json
exact_seen: set[str] = set()
unique: list[ProcessedStructuredRecord] = []
for entry in entries:
record = entry.record
fingerprint = hashlib.sha256(canonical_record_json(record).encode("utf-8")).hexdigest()
if fingerprint in exact_seen:
continue
exact_seen.add(fingerprint)
unique.append(
ProcessedStructuredRecord(
entry.source_index,
deepcopy(dict(record)),
)
)
return unique
def deduplicate_structured_records(
records: Sequence[Mapping[str, Any]],
) -> list[dict[str, Any]]:
"""仅按整条 canonical JSON 稳定去重。"""
entries = [
ProcessedStructuredRecord(index, deepcopy(dict(record)))
for index, record in enumerate(records)
]
return [entry.record for entry in _deduplicate_structured_entries(entries)]
def _near_duplicate_features(text: str, shingle_size: int) -> tuple[str, ...]:
if isinstance(shingle_size, bool) or not isinstance(shingle_size, int):
raise TypeError("shingle_size must be an integer")
if shingle_size <= 0:
raise ValueError("shingle_size must be greater than 0")
tokens = re.findall(
r"[\u3400-\u4dbf\u4e00-\u9fff]|[A-Za-z0-9_]+",
normalize_text(text).casefold(),
)
if not tokens:
return ()
if len(tokens) < shingle_size:
return ("\x1f".join(tokens),)
return tuple(
"\x1f".join(tokens[index : index + shingle_size])
for index in range(len(tokens) - shingle_size + 1)
)
def near_duplicate_fingerprint(text: str, *, shingle_size: int = 3) -> str:
"""生成 64 位 SimHash 指纹,用于低成本近重复候选筛选。"""
if isinstance(shingle_size, bool) or not isinstance(shingle_size, int):
raise TypeError("shingle_size must be an integer")
if shingle_size <= 0:
raise ValueError("shingle_size must be greater than 0")
features = Counter(_near_duplicate_features(text, shingle_size))
if not features:
return "0" * 16
vector = [0] * 64
for feature, weight in features.items():
digest = int.from_bytes(hashlib.sha256(feature.encode("utf-8")).digest()[:8], "big")
for bit in range(64):
vector[bit] += weight if digest & (1 << bit) else -weight
fingerprint = sum(1 << bit for bit, value in enumerate(vector) if value >= 0)
return f"{fingerprint:016x}"
def fingerprints_are_near_duplicate(
left: str,
right: str,
*,
max_hamming_distance: int = 3,
) -> bool:
"""比较两个 64 位十六进制 SimHash 指纹。"""
if isinstance(max_hamming_distance, bool) or not isinstance(max_hamming_distance, int):
raise TypeError("max_hamming_distance must be an integer")
if not 0 <= max_hamming_distance <= 64:
raise ValueError("max_hamming_distance must be in [0, 64]")
if not re.fullmatch(r"[0-9a-fA-F]{16}", left) or not re.fullmatch(
r"[0-9a-fA-F]{16}", right
):
raise ValueError("fingerprints must be 16-character hexadecimal strings")
distance = (int(left, 16) ^ int(right, 16)).bit_count()
return distance <= max_hamming_distance
def is_near_duplicate(
left: str,
right: str,
*,
shingle_size: int = 3,
similarity_threshold: float = 0.9,
max_hamming_distance: int = 3,
) -> bool:
"""结合词片 Jaccard 和 SimHash 判断两段内容是否近重复。"""
if isinstance(similarity_threshold, bool) or not isinstance(
similarity_threshold, (int, float)
):
raise TypeError("similarity_threshold must be a number")
if not 0 <= similarity_threshold <= 1:
raise ValueError("similarity_threshold must be in [0, 1]")
if isinstance(max_hamming_distance, bool) or not isinstance(max_hamming_distance, int):
raise TypeError("max_hamming_distance must be an integer")
if not 0 <= max_hamming_distance <= 64:
raise ValueError("max_hamming_distance must be in [0, 64]")
left_normalized = normalize_text(left)
right_normalized = normalize_text(right)
if not left_normalized or not right_normalized:
return left_normalized == right_normalized
if left_normalized.casefold() == right_normalized.casefold():
return True
left_features = set(_near_duplicate_features(left_normalized, shingle_size))
right_features = set(_near_duplicate_features(right_normalized, shingle_size))
union = left_features | right_features
similarity = len(left_features & right_features) / len(union) if union else 1.0
if similarity >= similarity_threshold:
return True
return fingerprints_are_near_duplicate(
near_duplicate_fingerprint(left_normalized, shingle_size=shingle_size),
near_duplicate_fingerprint(right_normalized, shingle_size=shingle_size),
max_hamming_distance=max_hamming_distance,
)
def record_fingerprint(record: Mapping[str, Any]) -> str:
"""计算与字典键顺序无关的稳定记录指纹。"""
canonical = {
"instruction": normalize_text(str(record.get("instruction") or "")),
"input": normalize_text(str(record.get("input") or "")),
"output": normalize_text(str(record.get("output") or "")),
}
raw = json.dumps(canonical, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
def _readability_score(text: str) -> float:
if not text:
return 0.0
nonspace = [char for char in text if not char.isspace()]
if not nonspace:
return 0.0
printable_ratio = sum(char.isprintable() for char in nonspace) / len(nonspace)
useful_ratio = sum(
char.isalnum() or "\u3400" <= char <= "\u9fff" or unicodedata.category(char).startswith("P")
for char in nonspace
) / len(nonspace)
return round(100 * (0.65 * printable_ratio + 0.35 * useful_ratio), 2)
def _internal_duplicate_score(text: str) -> float:
units = [unit.strip().lower() for unit in re.split(r"[\n。!?;]+", text) if unit.strip()]
if len(units) <= 1:
return 100.0
return round(100 * len(set(units)) / len(units), 2)
def _source_relevance_score(record: Mapping[str, Any], source_content: str) -> float:
"""估算结果与来源文本的词元覆盖率。
这是无外部模型依赖、可重复的首版评分。没有来源文本(例如人工新增结果)
时不扣分;存在来源时,以结果中的有效词元被来源覆盖的比例计分。
"""
source = normalize_text(source_content)
if not source:
return 100.0
candidate = normalize_text(
"\n".join(
str(record.get(field) or "") for field in ("instruction", "input", "output")
)
)
def semantic_tokens(text: str) -> set[str]:
return {
token.lower()
for token in _TOKEN_PATTERN.findall(text)
if token.isalnum() or "\u3400" <= token <= "\u9fff"
}
source_tokens = semantic_tokens(source)
candidate_tokens = semantic_tokens(candidate)
if not candidate_tokens:
return 0.0
if not source_tokens:
return 0.0
return round(100 * len(candidate_tokens & source_tokens) / len(candidate_tokens), 2)
def score_quality(
record: Mapping[str, Any],
*,
min_output_length: int = 20,
source_content: str = "",
known_fingerprints: Iterable[str] = (),
threshold: float = 60.0,
) -> QualityScore:
"""按完整性、长度、可读性、来源相关性和重复度计算质量分。"""
if min_output_length <= 0:
raise ValueError("min_output_length must be greater than 0")
if not 0 <= threshold <= 100:
raise ValueError("threshold must be in [0, 100]")
instruction = normalize_text(str(record.get("instruction") or ""))
input_text = normalize_text(str(record.get("input") or ""))
output = normalize_text(str(record.get("output") or ""))
flags: list[str] = []
completeness = 100.0
if not instruction:
completeness -= 50
flags.append("missing_instruction")
if not output:
completeness -= 50
flags.append("missing_output")
output_length = len(output)
length_score = round(min(100.0, output_length / min_output_length * 100), 2)
if output_length < min_output_length:
flags.append("output_too_short")
readability = _readability_score("\n".join((instruction, input_text, output)))
if readability < 70:
flags.append("low_readability")
relevance = _source_relevance_score(record, source_content)
if source_content and relevance < 30:
flags.append("low_source_relevance")
fingerprint = record_fingerprint(record)
known = set(known_fingerprints)
duplicate = 0.0 if fingerprint in known else _internal_duplicate_score(output)
if duplicate == 0:
flags.append("duplicate_record")
elif duplicate < 70:
flags.append("repetitive_output")
overall = round(
completeness * 0.35
+ length_score * 0.20
+ readability * 0.20
+ relevance * 0.15
+ duplicate * 0.10,
2,
)
hard_valid = bool(instruction and output)
return QualityScore(
overall=overall,
completeness=completeness,
length=length_score,
readability=readability,
relevance=relevance,
duplicate=duplicate,
is_valid=hard_valid and overall >= threshold,
flags=tuple(flags),
fingerprint=fingerprint,
)

View File

@@ -0,0 +1,809 @@
"""数据处理算法 - 结构化数据处理。"""
from __future__ import annotations
import hashlib
import json
import math
import re
from collections.abc import Iterable, Mapping, Sequence
from copy import deepcopy
from datetime import date, datetime, time
from decimal import Decimal
from typing import Any
from .parsers.json_parser import _extract_structured_records_with_locators
from .quality import (
_deduplicate_structured_entries,
content_quality_flags,
estimate_token_count,
)
from .text_utils import _normalize_field_name, normalize_text, structured_json_dumps
from .transforms import stable_split_assignments
from .types import (
_CHINESE_NAME_CONTEXT_PATTERN,
_EMAIL_PATTERN,
_ENGLISH_NAME_CONTEXT_PATTERN,
_ID_CARD_PATTERN,
_IDENTITY_FIELD_PATTERN,
_MAX_STRUCTURED_DEPTH,
_MAX_STRUCTURED_FIELDS,
_NAME_FIELD_NAMES,
_PHONE_PATTERN,
_STRUCTURED_OPTIONS,
MAX_QA_PAIRS_PER_ITEM,
ProcessedStructuredRecord,
StructuredPreprocessOption,
)
def _canonical_value(value: Any) -> Any:
if value is None or isinstance(value, (bool, int)):
return value
if isinstance(value, Decimal):
if not value.is_finite():
raise ValueError("non-finite JSON number is not allowed")
return value
if isinstance(value, float):
if not math.isfinite(value):
raise ValueError("non-finite JSON number is not allowed")
return value
if isinstance(value, str):
return normalize_text(value)
if isinstance(value, (datetime, date, time)):
return value.isoformat()
if isinstance(value, Mapping):
normalized: dict[str, Any] = {}
for key, item in sorted(value.items(), key=lambda pair: str(pair[0])):
normalized_key = normalize_text(str(key))
if not normalized_key:
raise ValueError("structured record contains an empty field name")
if normalized_key in normalized:
raise ValueError(
f"structured record fields collide after normalization: {normalized_key}"
)
normalized[normalized_key] = _canonical_value(item)
return normalized
if isinstance(value, (list, tuple)):
return [_canonical_value(item) for item in value]
if isinstance(value, (set, frozenset)):
items = [_canonical_value(item) for item in value]
return sorted(
items,
key=lambda item: structured_json_dumps(item, sort_keys=True),
)
if isinstance(value, (bytes, bytearray, memoryview)):
return bytes(value).hex()
return normalize_text(str(value))
def _is_empty_value(value: Any) -> bool:
if value is None:
return True
if isinstance(value, str):
return not normalize_text(value)
if isinstance(value, Mapping):
return not value or all(_is_empty_value(item) for item in value.values())
if isinstance(value, (list, tuple, set, frozenset)):
return not value or all(_is_empty_value(item) for item in value)
return False
def _is_name_field(field: Any) -> bool:
raw_field = normalize_text(str(field))
if not raw_field:
return False
# 只匹配明确表示自然人姓名的字段,避免将 table_name、product_name、
# chinese_name 等业务名称或元数据字段误判为个人敏感信息。
if _normalize_field_name(raw_field, "snake_case") in _NAME_FIELD_NAMES:
return True
# detect_structure 会使用点号生成扁平化路径(例如 profile.name此时仅
# 判断最后一个路径段,不能退回到宽泛的 ``*_name`` 后缀匹配。
if "." not in raw_field:
return False
leaf_field = raw_field.rsplit(".", 1)[-1]
return _normalize_field_name(leaf_field, "snake_case") in _NAME_FIELD_NAMES
def _embedded_structure(value: Any) -> Any:
if not isinstance(value, str):
return value
candidate = value.strip()
if not candidate or candidate[0] not in "[{":
return value
try:
parsed = json.loads(candidate)
except json.JSONDecodeError:
return value
return parsed if isinstance(parsed, (Mapping, list)) else value
def _structured_options(options: Iterable[str] | Mapping[str, Any]) -> set[str]:
if isinstance(options, str):
raise TypeError("options must be an iterable or mapping of option names")
if isinstance(options, Mapping):
enabled = {str(key) for key, value in options.items() if bool(value)}
else:
enabled = {str(option) for option in options}
unknown = enabled - _STRUCTURED_OPTIONS
if unknown:
raise ValueError(f"unsupported structured preprocess options: {', '.join(sorted(unknown))}")
return enabled
def _clean_invalid_structured_entries(
entries: Sequence[ProcessedStructuredRecord],
) -> list[ProcessedStructuredRecord]:
if not entries:
return []
fields: list[str] = []
for entry in entries:
record = entry.record
for field in record:
if field not in fields:
fields.append(field)
active_fields = [
field
for field in fields
if any(not _is_empty_value(entry.record.get(field)) for entry in entries)
]
if not active_fields:
return []
cleaned: list[ProcessedStructuredRecord] = []
for entry in entries:
record = entry.record
values = {field: deepcopy(record.get(field)) for field in active_fields}
# 清洗只依据整行是否为空。外键、父级 ID 等字段天然允许为空,不能
# 因为字段名以 *_id 结尾就把它们全部提升为联合必填项。
if all(_is_empty_value(value) for value in values.values()):
continue
cleaned.append(ProcessedStructuredRecord(entry.source_index, values))
return cleaned
def _percentile(values: Sequence[float], fraction: float) -> float:
if not values:
raise ValueError("cannot calculate a percentile of an empty sequence")
ordered = sorted(values)
position = (len(ordered) - 1) * fraction
lower = math.floor(position)
upper = math.ceil(position)
if lower == upper:
return ordered[lower]
weight = position - lower
return ordered[lower] * (1 - weight) + ordered[upper] * weight
def _filter_anomalous_structured_entries(
entries: Sequence[ProcessedStructuredRecord],
*,
iqr_multiplier: float = 1.5,
) -> list[ProcessedStructuredRecord]:
"""按字段级数值 IQR、乱码和极端文本长度过滤异常记录。"""
if iqr_multiplier <= 0:
raise ValueError("iqr_multiplier must be greater than 0")
numeric_values: dict[str, list[float]] = {}
text_lengths: dict[str, list[float]] = {}
for entry in entries:
record = entry.record
for field, value in record.items():
if (
isinstance(value, (int, float))
and not isinstance(value, bool)
and not _IDENTITY_FIELD_PATTERN.search(
_normalize_field_name(field, "snake_case")
)
):
number = float(value)
if math.isfinite(number):
numeric_values.setdefault(field, []).append(number)
elif isinstance(value, str) and value:
text_lengths.setdefault(field, []).append(float(len(value)))
numeric_bounds: dict[str, tuple[float, float]] = {}
for field, values in numeric_values.items():
# 小样本不做统计异常判断,避免把合法长尾值误删。
if len(values) < 8:
continue
first_quartile = _percentile(values, 0.25)
third_quartile = _percentile(values, 0.75)
spread = third_quartile - first_quartile
numeric_bounds[field] = (
first_quartile - iqr_multiplier * spread,
third_quartile + iqr_multiplier * spread,
)
text_upper_bounds: dict[str, float] = {}
for field, lengths in text_lengths.items():
if len(lengths) < 8:
continue
first_quartile = _percentile(lengths, 0.25)
third_quartile = _percentile(lengths, 0.75)
spread = third_quartile - first_quartile
text_upper_bounds[field] = max(512.0, third_quartile + 3 * spread)
accepted: list[ProcessedStructuredRecord] = []
for entry in entries:
record = entry.record
anomalous = False
for field, value in record.items():
if (
isinstance(value, (int, float))
and not isinstance(value, bool)
and not _IDENTITY_FIELD_PATTERN.search(
_normalize_field_name(field, "snake_case")
)
):
number = float(value)
if not math.isfinite(number):
anomalous = True
break
bounds = numeric_bounds.get(field)
if bounds and not bounds[0] <= number <= bounds[1]:
anomalous = True
break
if isinstance(value, str):
flags = content_quality_flags(value, min_chars=0, min_tokens=0)
if {"content_too_long", "mojibake", "low_printable_ratio"} & set(flags):
anomalous = True
break
upper_bound = text_upper_bounds.get(field)
if upper_bound is not None and len(value) > upper_bound:
anomalous = True
break
if not anomalous:
accepted.append(
ProcessedStructuredRecord(
entry.source_index,
deepcopy(dict(record)),
)
)
return accepted
def _preview_content(item: Mapping[str, Any]) -> str:
for field in ("edited_content", "editedContent", "original_content", "originalContent", "content"):
value = item.get(field)
if value is not None:
return normalize_text(str(value))
return ""
def _standard_fields(content: str) -> tuple[str, str, str]:
if not content:
return "", "", ""
try:
payload = json.loads(content)
except json.JSONDecodeError:
payload = None
if isinstance(payload, Mapping):
instruction = next(
(
str(payload[key])
for key in ("instruction", "question", "prompt")
if payload.get(key) is not None
),
"",
)
input_text = next(
(str(payload[key]) for key in ("input", "context") if payload.get(key) is not None),
"",
)
output = next(
(str(payload[key]) for key in ("output", "answer", "response") if payload.get(key) is not None),
"",
)
if instruction or output:
return normalize_text(instruction), normalize_text(input_text), normalize_text(output)
question_answer = re.match(
r"^\s*(?:问|question)\s*[:]\s*(.+?)(?:\n|\r\n?)\s*(?:答|answer)\s*[:]\s*(.+)\s*$",
content,
flags=re.IGNORECASE | re.DOTALL,
)
if question_answer:
return normalize_text(question_answer.group(1)), "", normalize_text(question_answer.group(2))
lines = [line.strip() for line in content.splitlines() if line.strip()]
first_line = re.sub(r"^(?:问|question)\s*[:]\s*", "", lines[0], flags=re.IGNORECASE)
output = normalize_text("\n".join(lines[1:])) if len(lines) > 1 else normalize_text(content)
return normalize_text(first_line), "", output
def _protected_markdown_ranges(
text: str,
*,
preserve_code_blocks: bool,
preserve_tables: bool,
preserve_lists: bool,
) -> list[tuple[int, int]]:
"""找出不应从中间切开的 Markdown 代码块、表格和连续列表。"""
lines: list[tuple[int, int, str]] = []
cursor = 0
for raw_line in text.splitlines(keepends=True):
end = cursor + len(raw_line)
lines.append((cursor, end, raw_line.rstrip("\r\n")))
cursor = end
if cursor < len(text) or not lines:
lines.append((cursor, len(text), text[cursor:]))
ranges: list[tuple[int, int]] = []
code_line_indexes: set[int] = set()
if preserve_code_blocks:
open_block: tuple[int, str, int] | None = None
for index, (start, end, content) in enumerate(lines):
fence = re.match(r"^\s*(`{3,}|~{3,})", content)
if not fence:
continue
marker = fence.group(1)[0]
length = len(fence.group(1))
if open_block is None:
open_block = (index, marker, length)
continue
first_index, open_marker, open_length = open_block
if marker == open_marker and length >= open_length:
ranges.append((lines[first_index][0], end))
code_line_indexes.update(range(first_index, index + 1))
open_block = None
if open_block is not None:
first_index = open_block[0]
ranges.append((lines[first_index][0], len(text)))
code_line_indexes.update(range(first_index, len(lines)))
if preserve_tables:
index = 0
while index + 1 < len(lines):
if index in code_line_indexes:
index += 1
continue
header = lines[index][2].strip()
separator = lines[index + 1][2].strip().strip("|")
cells = [cell.strip() for cell in separator.split("|")]
if (
"|" not in header
or len(cells) < 2
or not all(re.fullmatch(r":?-{3,}:?", cell) for cell in cells)
):
index += 1
continue
end_index = index + 1
while (
end_index + 1 < len(lines)
and end_index + 1 not in code_line_indexes
and lines[end_index + 1][2].strip()
and "|" in lines[end_index + 1][2]
):
end_index += 1
ranges.append((lines[index][0], lines[end_index][1]))
index = end_index + 1
if preserve_lists:
list_pattern = re.compile(r"^\s*(?:[-+*]|\d+[.)])\s+\S")
continuation_pattern = re.compile(r"^\s{2,}\S")
index = 0
while index < len(lines):
if index in code_line_indexes or not list_pattern.match(lines[index][2]):
index += 1
continue
end_index = index
item_count = 1
while end_index + 1 < len(lines) and end_index + 1 not in code_line_indexes:
next_line = lines[end_index + 1][2]
if list_pattern.match(next_line):
item_count += 1
end_index += 1
elif continuation_pattern.match(next_line):
end_index += 1
else:
break
if item_count >= 2:
ranges.append((lines[index][0], lines[end_index][1]))
index = end_index + 1
merged: list[tuple[int, int]] = []
for start, end in sorted(ranges):
if merged and start < merged[-1][1]:
merged[-1] = (merged[-1][0], max(merged[-1][1], end))
else:
merged.append((start, end))
return merged
def extract_structured_records(text: str, file_format: str) -> list[dict[str, Any]]:
"""从 JSON、JSONL 或 CSV 中提取规范化记录。"""
records, _ = _extract_structured_records_with_locators(text, file_format)
return records
def canonical_record_json(record: Mapping[str, Any]) -> str:
"""生成与字段顺序无关、可用于比较和落库的 canonical JSON。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
return structured_json_dumps(_canonical_value(record), sort_keys=True)
def normalize_structured_record(
record: Mapping[str, Any],
*,
field_name_style: str = "snake_case",
) -> dict[str, Any]:
"""规范字段名、Unicode/空白、容器类型和不可 JSON 化的标量。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
normalized: dict[str, Any] = {}
for key, value in record.items():
normalized_key = _normalize_field_name(key, field_name_style)
if not normalized_key:
raise ValueError("structured record contains an empty field name")
if normalized_key in normalized:
raise ValueError(
f"structured record fields collide after normalization: {normalized_key}"
)
normalized[normalized_key] = _canonical_value(value)
return dict(sorted(normalized.items()))
def flatten_structured_record(
record: Mapping[str, Any],
*,
separator: str = ".",
) -> dict[str, Any]:
"""把嵌套对象展平;数组保留为 canonical JSON 兼容值。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
if not separator:
raise ValueError("separator cannot be empty")
flattened: dict[str, Any] = {}
def visit(value: Any, path: tuple[str, ...], depth: int) -> None:
if depth > _MAX_STRUCTURED_DEPTH:
raise ValueError(
f"structured record nesting exceeds {_MAX_STRUCTURED_DEPTH} levels"
)
value = _embedded_structure(value)
if isinstance(value, Mapping):
if not value and path:
key = separator.join(path)
flattened[key] = {}
return
for child_key, child_value in value.items():
normalized_key = normalize_text(str(child_key))
if not normalized_key:
raise ValueError("structured record contains an empty field name")
visit(child_value, (*path, normalized_key), depth + 1)
return
key = separator.join(path)
if key in flattened:
raise ValueError(f"structured record fields collide while flattening: {key}")
flattened[key] = _canonical_value(value)
if len(flattened) > _MAX_STRUCTURED_FIELDS:
raise ValueError(
f"structured record exceeds {_MAX_STRUCTURED_FIELDS} flattened fields"
)
for field, value in record.items():
field_name = normalize_text(str(field))
if not field_name:
raise ValueError("structured record contains an empty field name")
visit(value, (field_name,), 1)
return flattened
def filter_anomalous_structured_records(
records: Sequence[Mapping[str, Any]],
*,
iqr_multiplier: float = 1.5,
) -> list[dict[str, Any]]:
"""按字段级数值 IQR、乱码和极端文本长度过滤异常记录。"""
entries = [
ProcessedStructuredRecord(index, deepcopy(dict(record)))
for index, record in enumerate(records)
]
return [
entry.record
for entry in _filter_anomalous_structured_entries(
entries,
iqr_multiplier=iqr_multiplier,
)
]
def desensitize_pii(text: str) -> tuple[str, dict[str, int]]:
"""掩码邮箱、手机号、身份证号及有明确上下文的姓名。"""
if not isinstance(text, str):
raise TypeError("text must be str")
counts: dict[str, int] = {"email": 0, "phone": 0, "id_card": 0}
def replace(pattern: re.Pattern[str], replacement: str, kind: str, value: str) -> str:
def replacer(_: re.Match[str]) -> str:
counts[kind] += 1
return replacement
return pattern.sub(replacer, value)
masked = replace(_EMAIL_PATTERN, "[EMAIL]", "email", text)
masked = replace(_ID_CARD_PATTERN, "[ID_CARD]", "id_card", masked)
masked = replace(_PHONE_PATTERN, "[PHONE]", "phone", masked)
def replace_context_name(match: re.Match[str]) -> str:
counts["name"] = counts.get("name", 0) + 1
return f"{match.group('label')}{match.group('separator')}[NAME]"
masked = _CHINESE_NAME_CONTEXT_PATTERN.sub(replace_context_name, masked)
masked = _ENGLISH_NAME_CONTEXT_PATTERN.sub(replace_context_name, masked)
counts["total"] = sum(counts.values())
return masked, counts
def desensitize_structured_record(
record: Mapping[str, Any],
) -> tuple[dict[str, Any], dict[str, int]]:
"""递归脱敏结构化姓名字段及任意文本中的手机号、邮箱、身份证号。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
counts: dict[str, int] = {"email": 0, "phone": 0, "id_card": 0}
def add_counts(values: Mapping[str, int]) -> None:
for kind, count in values.items():
if kind != "total" and count:
counts[kind] = counts.get(kind, 0) + count
def visit(value: Any, field: Any = "") -> Any:
if isinstance(value, Mapping):
return {key: visit(item, key) for key, item in value.items()}
if isinstance(value, list):
return [visit(item, field) for item in value]
if isinstance(value, tuple):
return [visit(item, field) for item in value]
if _is_name_field(field) and not _is_empty_value(value):
counts["name"] = counts.get("name", 0) + 1
return "[NAME]"
if isinstance(value, str):
masked, found = desensitize_pii(value)
add_counts(found)
return masked
return deepcopy(value)
masked = {key: visit(value, key) for key, value in record.items()}
counts["total"] = sum(counts.values())
return masked, counts
def preprocess_structured_records_with_lineage(
records: Iterable[Mapping[str, Any]],
options: Iterable[str] | Mapping[str, Any],
) -> list[ProcessedStructuredRecord]:
"""执行结构化预处理,并保留每条结果在原始输入中的稳定索引。"""
enabled = _structured_options(options)
current: list[ProcessedStructuredRecord] = []
for source_index, record in enumerate(records):
if not isinstance(record, Mapping):
if "clean_invalid" in enabled:
continue
raise TypeError("structured records must contain mappings")
value = deepcopy(dict(record))
if "detect_structure" in enabled:
value = flatten_structured_record(value)
if "normalize_format" in enabled:
value = normalize_structured_record(value)
current.append(ProcessedStructuredRecord(source_index, value))
if "clean_invalid" in enabled:
current = _clean_invalid_structured_entries(current)
if "filter_anomaly" in enabled:
current = _filter_anomalous_structured_entries(current)
if "deduplicate" in enabled:
current = _deduplicate_structured_entries(current)
if "desensitize" in enabled:
current = [
ProcessedStructuredRecord(
entry.source_index,
desensitize_structured_record(entry.record)[0],
)
for entry in current
]
return current
def preprocess_structured_records(
records: Iterable[Mapping[str, Any]],
options: Iterable[str] | Mapping[str, Any],
) -> list[dict[str, Any]]:
"""按界面选项执行确定性、无副作用的结构化数据预处理。"""
return [
entry.record
for entry in preprocess_structured_records_with_lineage(records, options)
]
def protected_context_ranges(
text: str,
*,
preserve_code_blocks: bool = True,
preserve_tables: bool = True,
preserve_lists: bool = True,
) -> tuple[tuple[int, int], ...]:
"""返回代码块、表格和列表的不可拆分区间。
返回的偏移量基于规范化后的文本;调用方应在同一份 ``normalize_text``
结果上使用这些区间。
"""
normalized = normalize_text(text)
return tuple(
_protected_markdown_ranges(
normalized,
preserve_code_blocks=preserve_code_blocks,
preserve_tables=preserve_tables,
preserve_lists=preserve_lists,
)
)
def expand_to_context_boundaries(
text: str,
start: int,
end: int,
*,
preserve_paragraph: bool = True,
preserve_code_blocks: bool = True,
preserve_tables: bool = True,
preserve_lists: bool = True,
) -> tuple[int, int]:
"""将一个文本区间扩展到段落及受保护 Markdown 结构边界。"""
normalized = normalize_text(text)
if isinstance(start, bool) or isinstance(end, bool):
raise TypeError("start and end must be integers")
if not isinstance(start, int) or not isinstance(end, int):
raise TypeError("start and end must be integers")
if not 0 <= start <= end <= len(normalized):
raise ValueError("start and end must define a valid normalized text range")
expanded_start = start
expanded_end = end
if preserve_paragraph and normalized:
paragraph_start = normalized.rfind("\n\n", 0, start)
expanded_start = 0 if paragraph_start < 0 else paragraph_start + 2
paragraph_end = normalized.find("\n\n", end)
expanded_end = len(normalized) if paragraph_end < 0 else paragraph_end
ranges = protected_context_ranges(
normalized,
preserve_code_blocks=preserve_code_blocks,
preserve_tables=preserve_tables,
preserve_lists=preserve_lists,
)
changed = True
while changed:
changed = False
for range_start, range_end in ranges:
overlaps = range_start < expanded_end and range_end > expanded_start
contains_boundary = (
range_start <= expanded_start < range_end
or range_start < expanded_end <= range_end
)
if not overlaps and not contains_boundary:
continue
next_start = min(expanded_start, range_start)
next_end = max(expanded_end, range_end)
if (next_start, next_end) != (expanded_start, expanded_end):
expanded_start, expanded_end = next_start, next_end
changed = True
return expanded_start, expanded_end
def merge_short_blocks(
blocks: Iterable[str],
*,
min_token_count: int = 100,
separator: str = "\n\n",
) -> list[str]:
"""按原顺序合并短内容块,并把末尾残块归入前一块。"""
if isinstance(min_token_count, bool) or not isinstance(min_token_count, int):
raise TypeError("min_token_count must be an integer")
if min_token_count <= 0:
raise ValueError("min_token_count must be greater than 0")
if not isinstance(separator, str):
raise TypeError("separator must be str")
merged: list[str] = []
pending: list[str] = []
pending_tokens = 0
for block in blocks:
if not isinstance(block, str):
raise TypeError("blocks must contain strings")
normalized = normalize_text(block)
if not normalized:
continue
token_count = estimate_token_count(normalized)
if not pending and token_count >= min_token_count:
merged.append(normalized)
continue
pending.append(normalized)
pending_tokens += token_count
if pending_tokens >= min_token_count:
merged.append(separator.join(pending))
pending = []
pending_tokens = 0
if pending:
tail = separator.join(pending)
if merged:
merged[-1] = separator.join((merged[-1], tail))
else:
merged.append(tail)
return merged
def generate_standard_records(
preview_items: Iterable[Mapping[str, Any]],
*,
qa_pairs_per_item: int = 1,
semantic_enrichment: bool = False,
split: Mapping[str, int] | None = None,
split_seed: str = "",
) -> list[dict[str, Any]]:
"""把预览内容确定性转换为标准 instruction/input/output 记录。
该函数只负责本地标准化,不冒充 LLM服务层可将其作为无模型模式或
LLM 响应解析后的统一落库步骤。
"""
if not 1 <= qa_pairs_per_item <= MAX_QA_PAIRS_PER_ITEM:
raise ValueError(
f"qa_pairs_per_item must be in [1, {MAX_QA_PAIRS_PER_ITEM}]"
)
prefixes = (
"请结合实际情况说明:",
"请用通俗易懂的方式说明:",
"请从实际应用角度说明:",
"请简洁自然地说明:",
"请详细解答:",
)
results: list[dict[str, Any]] = []
for item_index, item in enumerate(preview_items):
content = _preview_content(item)
instruction, input_text, output = _standard_fields(content)
preview_id = str(item.get("id") or f"preview-{item_index + 1}")
for variant_index in range(qa_pairs_per_item):
variant_instruction = instruction
if variant_index:
if semantic_enrichment:
prefix = prefixes[variant_index % len(prefixes)]
if variant_index >= len(prefixes):
prefix = (
f"{prefix.removesuffix('')}"
f"(问法 {variant_index + 1}"
)
variant_instruction = f"{prefix}{instruction}"
else:
variant_instruction = f"{instruction}(问法 {variant_index + 1})"
raw_id = f"{preview_id}:{variant_index + 1}"
result_id = f"result_{hashlib.sha256(raw_id.encode('utf-8')).hexdigest()[:16]}"
status = "valid" if variant_instruction and output else "invalid"
results.append(
{
"id": result_id,
"preview_item_id": preview_id,
"instruction": variant_instruction,
"input": input_text,
"output": output,
"original_instruction": variant_instruction,
"original_input": input_text,
"original_output": output,
"status": status,
"split": "train",
}
)
assignments = stable_split_assignments(
[str(result["id"]) for result in results],
split,
seed=split_seed,
)
for result, assignment in zip(results, assignments, strict=True):
result["split"] = assignment
return results

View File

@@ -0,0 +1,328 @@
"""数据处理算法 - 文本处理工具。"""
from __future__ import annotations
import csv
import io
import json
import math
import re
import unicodedata
from datetime import date, datetime, time
from decimal import Decimal
from pathlib import Path
from typing import Any, Literal
from collections.abc import Mapping
from .types import ParsedText, TextFormat
# 格式别名映射
_FORMAT_ALIASES: dict[str, TextFormat] = {
"json": "json",
"jsonl": "jsonl",
"ndjson": "jsonl",
"csv": "csv",
"tsv": "csv",
"md": "markdown",
"markdown": "markdown",
"txt": "txt",
"text": "txt",
"pdf": "pdf",
"docx": "docx",
"xlsx": "xlsx",
"pptx": "pptx",
}
# 旧版 Office 格式映射
_LEGACY_OFFICE_FORMATS: dict[str, str] = {
"doc": "docx",
"xls": "xlsx",
"ppt": "pptx",
}
# Office Open XML 格式集合
_OFFICE_OPEN_XML_FORMATS = {"docx", "xlsx", "pptx"}
# 文本提取限制
_MAX_EXTRACTED_TEXT_CHARS = 20_000_000
def decode_utf8(raw: bytes | bytearray | memoryview | str) -> str:
"""严格解码 UTF-8 文本,并移除可选 BOM。
不使用 ``errors='replace'``,避免上传内容损坏后仍被静默接收。
"""
if isinstance(raw, str):
return raw.removeprefix("\ufeff")
if not isinstance(raw, (bytes, bytearray, memoryview)):
raise TypeError("raw must be bytes-like or str")
try:
return bytes(raw).decode("utf-8-sig")
except UnicodeDecodeError as exc:
raise ValueError(f"content is not valid UTF-8 at byte {exc.start}") from exc
def parse_utf8_text(raw: bytes | bytearray | memoryview | str) -> str:
"""``decode_utf8`` 的语义化别名,供上传服务直接调用。"""
return decode_utf8(raw)
def normalize_text(text: str) -> str:
"""规范 Unicode、换行和行尾空白同时保留段落结构。"""
if not isinstance(text, str):
raise TypeError("text must be str")
normalized = unicodedata.normalize("NFKC", text.removeprefix("\ufeff"))
normalized = normalized.replace("\r\n", "\n").replace("\r", "\n")
normalized = "".join(
char
for char in normalized
if char in {"\n", "\t"} or not unicodedata.category(char).startswith("C")
)
lines = [re.sub(r"[\t \f\v]+$", "", line) for line in normalized.split("\n")]
return "\n".join(lines).strip()
def _normalize_format(value: str | None) -> TextFormat | None:
if value is None:
return None
normalized = value.strip().lower().removeprefix(".")
if normalized in _LEGACY_OFFICE_FORMATS:
replacement = _LEGACY_OFFICE_FORMATS[normalized]
raise ValueError(
f"legacy .{normalized} format is not supported; "
f"convert the file to .{replacement} and upload again"
)
try:
return _FORMAT_ALIASES[normalized]
except KeyError as exc:
raise ValueError(f"unsupported text format: {value}") from exc
def detect_text_format(
*,
filename: str | None = None,
text: str = "",
file_format: str | None = None,
) -> TextFormat:
"""按显式格式、扩展名和内容特征依次识别文本格式。"""
explicit = _normalize_format(file_format)
if explicit:
return explicit
if filename:
suffix = Path(filename).suffix.lower().removeprefix(".")
if suffix in _LEGACY_OFFICE_FORMATS:
replacement = _LEGACY_OFFICE_FORMATS[suffix]
raise ValueError(
f"legacy .{suffix} format is not supported; "
f"convert the file to .{replacement} and upload again"
)
detected = _FORMAT_ALIASES.get(suffix)
if detected:
return detected
stripped = text.strip()
if stripped:
if stripped[0] in "[{":
try:
json.loads(stripped)
except json.JSONDecodeError:
pass
else:
return "json"
nonempty_lines = [line for line in stripped.splitlines() if line.strip()]
if len(nonempty_lines) > 1:
try:
for line in nonempty_lines:
json.loads(line)
except json.JSONDecodeError:
pass
else:
return "jsonl"
if re.search(r"(?m)^(?:#{1,6}\s+|```|~~~)", stripped) or re.search(
r"(?m)^\s*\|.+\|\s*$", stripped
):
return "markdown"
sample = stripped[:8192]
try:
dialect = csv.Sniffer().sniff(sample, delimiters=",\t;")
rows = list(csv.reader(io.StringIO(sample), dialect))
if len(rows) >= 2 and len(rows[0]) >= 2:
return "csv"
except csv.Error:
pass
return "txt"
def _append_bounded_text(parts: list[str], value: Any, total: int) -> int:
text = normalize_text(str(value or ""))
if not text:
return total
total += len(text)
if total > _MAX_EXTRACTED_TEXT_CHARS:
raise ValueError(
f"extracted document text exceeds {_MAX_EXTRACTED_TEXT_CHARS} characters"
)
parts.append(text)
return total
def _normalize_spreadsheet_value(value: Any) -> Any:
if value is None:
return ""
if isinstance(value, str):
return normalize_text(value)
if isinstance(value, (datetime, date, time)):
return value.isoformat()
if isinstance(value, (bool, int, float)):
return value
return normalize_text(str(value))
def _normalize_value(value: Any) -> Any:
if isinstance(value, str):
return normalize_text(value)
if isinstance(value, Mapping):
return {normalize_text(str(key)): _normalize_value(item) for key, item in value.items()}
if isinstance(value, list):
return [_normalize_value(item) for item in value]
return value
def parse_text_content(
raw: bytes | bytearray | memoryview | str,
*,
filename: str | None = None,
file_format: str | None = None,
) -> ParsedText:
"""安全解析 UTF-8 文本、文本型 PDF 和现代 Office 文件。"""
# 解析器依赖 text_utilsnormalize_text 等),这里延迟导入以断开循环依赖。
from .parsers.json_parser import _extract_structured_records_with_locators
from .parsers.office import _extract_docx_text, _extract_pptx_text, _extract_xlsx_records
from .parsers.pdf import _extract_pdf_text
detected_format = detect_text_format(
filename=filename,
text="",
file_format=file_format,
)
if detected_format in _OFFICE_OPEN_XML_FORMATS or detected_format == "pdf":
binary = _binary_bytes(raw, detected_format)
if detected_format == "pdf":
text = _extract_pdf_text(binary)
return ParsedText(format=detected_format, text=text, records=())
if detected_format == "docx":
text = _extract_docx_text(binary)
return ParsedText(format=detected_format, text=text, records=())
if detected_format == "pptx":
text = _extract_pptx_text(binary)
return ParsedText(format=detected_format, text=text, records=())
records, record_locators = _extract_xlsx_records(binary)
text = "\n".join(
json.dumps(record, ensure_ascii=False, separators=(",", ":"))
for record in records
)
return ParsedText(
format=detected_format,
text=normalize_text(text),
records=tuple(records),
record_locators=tuple(record_locators),
)
decoded_text = decode_utf8(raw)
detected_format = detect_text_format(
filename=filename,
text=decoded_text,
file_format=file_format,
)
# JSON/JSONL 是有损规范化的禁区NFKC、控制字符删除或 trim 都可能改变字段值、
# 掩盖非法输入,甚至把原本合法的字符串变成语法错误。其他格式保持历史行为。
text = (
decoded_text
if detected_format in {"json", "jsonl"}
else normalize_text(decoded_text)
)
records: list[dict[str, Any]] = []
record_locators: list[dict[str, Any]] = []
if detected_format in {"json", "jsonl", "csv"}:
records, record_locators = _extract_structured_records_with_locators(
text,
detected_format,
)
return ParsedText(
format=detected_format,
text=text,
records=tuple(records),
record_locators=tuple(record_locators),
)
def _normalize_field_name(value: Any, style: str) -> str:
name = normalize_text(str(value))
if style == "preserve":
return name
if style == "lower":
return name.lower()
if style != "snake_case":
raise ValueError("field_name_style must be snake_case, lower or preserve")
name = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", name)
name = re.sub(r"[^\w\u3400-\u4dbf\u4e00-\u9fff]+", "_", name, flags=re.UNICODE)
return re.sub(r"_+", "_", name).strip("_").lower()
def structured_json_dumps(value: Any, *, sort_keys: bool = False) -> str:
"""序列化紧凑 JSON并把 ``Decimal`` 保持为原值对应的 JSON 数字。
标准库会要求先把 ``Decimal`` 转成 float 或字符串;前者可能静默舍入,
后者会改变 JSON 类型。这里直接输出有限 Decimal 的十进制表示。
"""
def serialize(item: Any) -> str:
if item is None:
return "null"
if item is True:
return "true"
if item is False:
return "false"
if isinstance(item, int):
return str(item)
if isinstance(item, Decimal):
if not item.is_finite():
raise ValueError("non-finite JSON number is not allowed")
return str(item)
if isinstance(item, float):
if not math.isfinite(item):
raise ValueError("non-finite JSON number is not allowed")
return json.dumps(item, allow_nan=False)
if isinstance(item, str):
return json.dumps(item, ensure_ascii=False)
if isinstance(item, Mapping):
pairs: list[tuple[str, Any]] = []
seen_keys: set[str] = set()
for key, child in item.items():
if not isinstance(key, str):
raise TypeError("JSON object keys must be strings")
if key in seen_keys:
raise ValueError(f"duplicate JSON object key: {key!r}")
seen_keys.add(key)
pairs.append((key, child))
if sort_keys:
pairs.sort(key=lambda pair: pair[0])
return "{" + ",".join(
f"{json.dumps(key, ensure_ascii=False)}:{serialize(child)}"
for key, child in pairs
) + "}"
if isinstance(item, (list, tuple)):
return "[" + ",".join(serialize(child) for child in item) + "]"
raise TypeError(f"value of type {type(item).__name__} is not JSON serializable")
return serialize(value)
def _binary_bytes(
raw: bytes | bytearray | memoryview | str,
file_format: TextFormat,
) -> bytes:
if isinstance(raw, str):
raise ValueError(f"{file_format.upper()} content must be uploaded as binary data")
if not isinstance(raw, (bytes, bytearray, memoryview)):
raise TypeError("raw must be bytes-like or str")
return bytes(raw)

View File

@@ -0,0 +1,81 @@
"""数据处理算法 - 数据集转换和分割。"""
from __future__ import annotations
import hashlib
import math
from collections.abc import Mapping, Sequence
from .types import DatasetSplit
def stable_split(
value: str | int,
split: Mapping[str, int] | None = None,
*,
seed: str = "",
) -> DatasetSplit:
"""按稳定哈希将记录划分到 train/validation/test。"""
ratios = dict(split or {"train": 80, "validation": 10, "test": 10})
required = {"train", "validation", "test"}
if set(ratios) != required:
raise ValueError("split must contain exactly train, validation and test")
if any(isinstance(value, bool) or not isinstance(value, int) or value < 0 for value in ratios.values()):
raise ValueError("split ratios must be non-negative integers")
if sum(ratios.values()) != 100:
raise ValueError("split ratios must sum to 100")
digest = hashlib.sha256(f"{seed}:{value}".encode("utf-8")).digest()
bucket = int.from_bytes(digest[:8], "big") % 10_000
train_boundary = ratios["train"] * 100
validation_boundary = train_boundary + ratios["validation"] * 100
if bucket < train_boundary:
return "train"
if bucket < validation_boundary:
return "validation"
return "test"
def stable_split_assignments(
values: Sequence[str | int],
split: Mapping[str, int] | None = None,
*,
seed: str = "",
) -> list[DatasetSplit]:
"""按稳定顺序和精确配额批量划分数据集。
单条哈希分桶只能在大样本下近似比例。这里先按哈希稳定排序,再用
最大余数法计算各切分配额,确保小数据集也严格遵循配置比例。
"""
ratios = dict(split or {"train": 80, "validation": 10, "test": 10})
# 复用单条划分的参数校验,避免两套规则逐渐漂移。
stable_split("validation", ratios, seed=seed)
if not values:
return []
split_order: tuple[DatasetSplit, ...] = ("train", "validation", "test")
exact = {name: len(values) * ratios[name] / 100 for name in split_order}
quotas = {name: math.floor(exact[name]) for name in split_order}
remaining = len(values) - sum(quotas.values())
remainder_order = sorted(
split_order,
key=lambda name: (-(exact[name] - quotas[name]), split_order.index(name)),
)
for name in remainder_order[:remaining]:
quotas[name] += 1
ranked_indices = sorted(
range(len(values)),
key=lambda index: (
hashlib.sha256(f"{seed}:{values[index]}".encode("utf-8")).digest(),
index,
),
)
assignments: list[DatasetSplit] = ["train"] * len(values)
cursor = 0
for name in split_order:
for index in ranked_indices[cursor : cursor + quotas[name]]:
assignments[index] = name
cursor += quotas[name]
return assignments

View File

@@ -0,0 +1,266 @@
"""数据处理模块使用的无副作用算法。
本模块不访问数据库、文件系统或网络,便于 API、后台任务和测试共同复用。
所有偏移量均为 Python 字符串偏移量。
"""
from __future__ import annotations
import csv
import hashlib
import io
import json
import math
import re
import unicodedata
import xml.etree.ElementTree as ET
import zipfile
from collections import Counter
from collections.abc import Iterable, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from datetime import date, datetime, time
from decimal import Decimal
from pathlib import Path, PurePosixPath
from typing import Any, Literal
from urllib.parse import unquote, urlsplit
from docx import Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from openpyxl import load_workbook
from openpyxl.utils.cell import range_boundaries
from pptx import Presentation
from pypdf import PdfReader
# 避免循环导入:直接定义常量而不是从 constants 导入
MAX_QA_PAIRS_PER_ITEM = 50
TextFormat = Literal[
"json",
"jsonl",
"csv",
"markdown",
"txt",
"pdf",
"docx",
"xlsx",
"pptx",
]
DatasetSplit = Literal["train", "validation", "test"]
StructuredPreprocessOption = Literal[
"clean_invalid",
"detect_structure",
"deduplicate",
"normalize_format",
"filter_anomaly",
"desensitize",
]
SUPPORTED_TEXT_FORMATS: tuple[TextFormat, ...] = (
"json",
"jsonl",
"csv",
"markdown",
"txt",
"pdf",
"docx",
"xlsx",
"pptx",
)
_FORMAT_ALIASES: dict[str, TextFormat] = {
"json": "json",
"jsonl": "jsonl",
"ndjson": "jsonl",
"csv": "csv",
"tsv": "csv",
"md": "markdown",
"markdown": "markdown",
"txt": "txt",
"text": "txt",
"pdf": "pdf",
"docx": "docx",
"xlsx": "xlsx",
"pptx": "pptx",
}
_LEGACY_OFFICE_FORMATS: dict[str, str] = {
"doc": "docx",
"xls": "xlsx",
"ppt": "pptx",
}
_OFFICE_OPEN_XML_FORMATS = {"docx", "xlsx", "pptx"}
_MAX_ARCHIVE_ENTRIES = 10_000
_MAX_ARCHIVE_UNCOMPRESSED_BYTES = 512 * 1024 * 1024
_MAX_ARCHIVE_ENTRY_BYTES = 128 * 1024 * 1024
_MAX_ARCHIVE_COMPRESSION_RATIO = 200
_MAX_EXTRACTED_TEXT_CHARS = 20_000_000
_MAX_PDF_PAGES = 2_000
_MAX_PRESENTATION_SLIDES = 2_000
_MAX_WORKBOOK_SHEETS = 100
_MAX_WORKBOOK_ROWS = 100_000
_MAX_WORKBOOK_SCANNED_ROWS = 200_000
_MAX_WORKBOOK_COLUMNS = 256
_MAX_WORKBOOK_CELLS = 2_000_000
_MAX_WORKBOOK_HEADER_ROWS = 8
_MAX_WORKBOOK_HEADER_SCAN_ROWS = 64
_MAX_WORKBOOK_MERGED_RANGES = 100_000
_MAX_STRUCTURED_FIELDS = 1_024
_MAX_STRUCTURED_DEPTH = 16
_MAX_JSON_DEPTH = 64
_MAX_ANOMALY_TEXT_CHARS = 1_000_000
_STRUCTURED_OPTIONS = {
"clean_invalid",
"detect_structure",
"deduplicate",
"normalize_format",
"filter_anomaly",
"desensitize",
}
_IDENTITY_FIELD_PATTERN = re.compile(r"(?:^|[._])(?:id|uuid|key|code)$|(?:^|[._]).+_id$")
_MOJIBAKE_MARKERS = ("\ufffd", "锟斤拷", "烫烫烫", "屯屯屯", "Ã", "Â", "â€")
_JSON_RECORD_ARRAY_KEYS = ("records", "data", "items", "rows")
_JSON_ENVELOPE_KEYS = ("response", "payload")
_JSON_WRAPPER_METADATA_KEYS = frozenset(
{
"page",
"page_size",
"pageSize",
"per_page",
"perPage",
"total",
"total_count",
"totalCount",
"count",
"offset",
"limit",
"cursor",
"next_cursor",
"nextCursor",
"has_more",
"hasMore",
}
)
_JSON_RESPONSE_METADATA_KEYS = _JSON_WRAPPER_METADATA_KEYS | {
"success",
"status",
"code",
"message",
"error",
}
_NAME_FIELD_NAMES = {
"name",
"full_name",
"fullname",
"real_name",
"contact_name",
"customer_name",
"recipient_name",
"姓名",
"中文姓名",
"真实姓名",
"联系人",
"联系人姓名",
"客户姓名",
"收件人",
"收件人姓名",
}
_EMAIL_PATTERN = re.compile(
r"(?<![\w.+-])[A-Za-z0-9.!#$%&'*+/=?^_`{|}~-]+"
r"@[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?"
r"(?:\.[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?)+(?![\w.-])"
)
_PHONE_PATTERN = re.compile(r"(?<!\d)(?:(?:\+|00)?86[-\s]?)?1[3-9]\d{9}(?!\d)")
_ID_CARD_PATTERN = re.compile(r"(?<!\d)(?:\d{17}[\dXx]|\d{15})(?!\d)")
_CHINESE_NAME_CONTEXT_PATTERN = re.compile(
r"(?P<label>姓名|真实姓名|联系人(?:姓名)?|收件人)"
r"(?P<separator>\s*(?:[:=]|为)\s*|\s+)"
r"(?P<name>[\u3400-\u4dbf\u4e00-\u9fff·]{2,8})"
)
_ENGLISH_NAME_CONTEXT_PATTERN = re.compile(
r"(?im)(?P<label>full\s+name|contact\s+name|name)"
r"(?P<separator>\s*[:=]\s*)"
r"(?P<name>[A-Za-z][A-Za-z'-]*(?:[ \t]+[A-Za-z][A-Za-z'-]*){0,3})"
)
_TOKEN_PATTERN = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]|[A-Za-z0-9_]+|[^\s]")
@dataclass(frozen=True, slots=True)
class ParsedText:
"""文本、文档或工作簿的统一解析结果。"""
format: TextFormat
text: str
records: tuple[dict[str, Any], ...]
record_locators: tuple[dict[str, Any], ...] = ()
@dataclass(frozen=True, slots=True)
class ProcessedStructuredRecord:
"""保留原始记录索引的结构化预处理结果。"""
source_index: int
record: dict[str, Any]
@dataclass(frozen=True, slots=True)
class PdfPageText:
"""PDF 物理页在统一提取文本中的字符范围。"""
page_number: int
text: str
source_start: int
source_end: int
@dataclass(frozen=True, slots=True)
class DocumentNoiseSpan:
"""PDF 中可安全从展示内容移除的文本范围。"""
start: int
end: int
kind: Literal["page_number", "repeated_margin", "table_of_contents"]
@dataclass(frozen=True, slots=True)
class QualityScore:
"""标准 instruction/input/output 记录的可解释质量分。"""
overall: float
completeness: float
length: float
readability: float
relevance: float
duplicate: float
is_valid: bool
flags: tuple[str, ...]
fingerprint: str
@dataclass(frozen=True, slots=True)
class DocumentHeading:
"""文档标题的位置和层级。"""
level: int
title: str
line_number: int
start: int
end: int
@dataclass(frozen=True, slots=True)
class DocumentStructure:
"""无需外部模型即可复现的文档结构摘要。"""
line_count: int
paragraph_count: int
headings: tuple[DocumentHeading, ...]
code_block_count: int
table_block_count: int
list_block_count: int
@dataclass(frozen=True, slots=True)
class _PdfLine:
text: str
start: int
end: int
class _DuplicateJsonKeyError(ValueError):
"""严格 JSON 解析时发现同一对象内的重复键。"""

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,69 @@
"""数据处理存储层。"""
from .base import (
StoreBase,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
utcnow,
new_id,
repeat_task_id,
json_dumps,
TASK_STATUSES,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_decode_row,
_preview_config_changed,
_reasoning_output_is_valid,
_source_storage_descriptor,
)
from .tasks import TasksMixin
from .source_files import SourceFilesMixin
from .preview import PreviewMixin
from .generation import GenerationMixin
from .results import ResultsMixin
from .datasets import DatasetsMixin
class DataProcessStore(
StoreBase,
TasksMixin,
SourceFilesMixin,
PreviewMixin,
GenerationMixin,
ResultsMixin,
DatasetsMixin,
):
"""数据处理持久层。
构造函数不会连接数据库或执行迁移。部署方必须显式执行 002 SQL
或在受控的管理命令中调用 :meth:`ensure_schema`,避免应用启动时
修改远程数据库。
"""
pass
def get_data_process_store() -> DataProcessStore:
"""获取数据处理存储实例。"""
return DataProcessStore()
__all__ = [
"DataProcessStore",
"DataProcessStoreError",
"NotFoundError",
"ConflictError",
"InvalidStateError",
"get_data_process_store",
"utcnow",
"new_id",
"repeat_task_id",
"TASK_STATUSES",
"EDITABLE_STATUSES",
"_decode_row",
"_preview_config_changed",
"_reasoning_output_is_valid",
"_source_storage_descriptor",
]

View File

@@ -0,0 +1,305 @@
"""数据处理存储层 - 基础设施。
包含:异常类、工具函数、常量定义、基类。
"""
from __future__ import annotations
import hashlib
import json
import re
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import UTC, date, datetime
from decimal import Decimal
from pathlib import Path
from typing import Any
import psycopg
from psycopg.rows import dict_row
from app.core.config import get_settings
# 常量定义
TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"}
EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"}
ACTIVE_PREVIEW_STATUSES = {"queued", "running"}
WORKFLOW_STEPS = {"create", "model", "upload", "preview", "generate", "results"}
_PREVIEW_CONFIG_ALIASES = {
"preprocess_options": "preprocessOptions",
"chunk_method": "chunkMethod",
"chunk_size": "chunkSize",
"chunk_overlap": "chunkOverlap",
"min_chunk_size": "minChunkSize",
"semantic_breakpoint_percentile": "semanticBreakpointPercentile",
"preserve_tables": "preserveTables",
"preserve_code_blocks": "preserveCodeBlocks",
"preserve_lists": "preserveLists",
}
_UNSTRUCTURED_PREVIEW_DEFAULTS: dict[str, Any] = {
"chunk_method": "layout_hybrid",
"chunk_size": 800,
"chunk_overlap": 100,
"min_chunk_size": 100,
"semantic_breakpoint_percentile": 95,
"preserve_tables": True,
"preserve_code_blocks": True,
"preserve_lists": True,
}
_REGENERATION_MARKER_KEY = "_regeneration_prepared"
_REPEAT_SOURCE_TASK_KEY = "_repeat_source_task_id"
_REPEAT_REQUEST_KEY = "_repeat_request_id"
_INTERNAL_CONFIG_KEYS = {
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
}
# 异常类
class DataProcessStoreError(RuntimeError):
pass
class NotFoundError(DataProcessStoreError):
pass
class ConflictError(DataProcessStoreError):
pass
class InvalidStateError(DataProcessStoreError):
pass
# 工具函数
def utcnow() -> str:
return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")
def new_id(prefix: str) -> str:
return f"{prefix}_{uuid.uuid4().hex[:20]}"
def repeat_task_id(source_task_id: str, request_id: str) -> str:
"""按源任务和请求幂等键生成稳定的新任务 ID。"""
digest = hashlib.sha256(f"{source_task_id}:{request_id}".encode()).hexdigest()
return f"dpt_{digest[:20]}"
def json_dumps(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
def _database_url(value: str) -> str:
return value.replace("postgresql+psycopg://", "postgresql://")
def _json_value(value: Any, default: Any) -> Any:
if value is None or value == "":
return default
if isinstance(value, (dict, list)):
return value
try:
return json.loads(value)
except (TypeError, json.JSONDecodeError):
return default
def _task_output_type(task: dict[str, Any]) -> str:
config = _json_value(task.get("config"), {})
if not isinstance(config, dict):
return "standard"
return str(config.get("output_type") or config.get("outputType") or "standard")
def _task_reasoning_detail(task: dict[str, Any]) -> str:
config = _json_value(task.get("config"), {})
if not isinstance(config, dict):
return "normal"
return str(
config.get("reasoning_detail")
or config.get("reasoningDetail")
or "normal"
)
def _reasoning_output_is_valid(value: Any) -> bool:
match = re.fullmatch(
r"\s*<think>\s*(?P<reasoning>[\s\S]*?)\s*</think>\s*(?P<answer>[\s\S]+?)\s*",
str(value or ""),
flags=re.IGNORECASE,
)
return bool(
match
and match.group("reasoning").strip()
and match.group("answer").strip()
and all(
tag not in part.lower()
for tag in ("<think", "</think")
for part in (match.group("reasoning"), match.group("answer"))
)
)
def _dpo_fields_are_valid(row: dict[str, Any]) -> bool:
chosen = str(row.get("chosen") or "").strip()
rejected = str(row.get("rejected") or "").strip()
return bool(chosen and rejected and chosen != rejected)
def _preview_config_value(config: dict[str, Any], key: str, default: Any) -> Any:
if key in config:
return config[key]
return config.get(_PREVIEW_CONFIG_ALIASES[key], default)
def _normalized_preprocess_options(config: dict[str, Any]) -> Any:
value = _preview_config_value(config, "preprocess_options", [])
if isinstance(value, (list, tuple, set)):
return tuple(sorted({str(item) for item in value}))
return value
def _preview_config_projection(
process_type: str,
config: dict[str, Any],
) -> dict[str, Any]:
"""只投影会改变预览切片的配置。
生成模型、提示词、温度等参数不影响源文切片,因此不应该
破坏用户已经校对过的预览内容。
"""
projection: dict[str, Any] = {
"preprocess_options": _normalized_preprocess_options(config),
}
if process_type != "unstructured":
return projection
for key, default in _UNSTRUCTURED_PREVIEW_DEFAULTS.items():
projection[key] = _preview_config_value(config, key, default)
return projection
def _preview_config_changed(
process_type: str,
current_config: dict[str, Any],
next_config: dict[str, Any],
) -> bool:
return _preview_config_projection(process_type, current_config) != _preview_config_projection(
process_type, next_config
)
def _regeneration_marker(task: dict[str, Any]) -> dict[str, Any] | None:
config = task.get("config")
if not isinstance(config, dict):
return None
marker = config.get(_REGENERATION_MARKER_KEY)
if not isinstance(marker, dict) or marker.get("prepared") is not True:
return None
return marker
def _is_regeneration_prepared(task: dict[str, Any]) -> bool:
return _regeneration_marker(task) is not None
def _business_config(config: dict[str, Any] | None) -> dict[str, Any]:
"""过滤只供服务端维护的工作流标记。"""
return {
key: value
for key, value in (config or {}).items()
if key not in _INTERNAL_CONFIG_KEYS
}
def _public_task(item: dict[str, Any] | None) -> dict[str, Any] | None:
"""从 API 任务快照中移除服务端内部工作流标记。"""
if item is None:
return None
public = dict(item)
config = public.get("config")
if isinstance(config, dict):
public["config"] = _business_config(config)
return public
def _serialize_value(value: Any) -> Any:
if isinstance(value, (datetime, date)):
return value.isoformat().replace("+00:00", "Z")
if isinstance(value, Decimal):
return float(value)
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
item = {key: _serialize_value(value) for key, value in row.items()}
for key, default in {
"config": {},
"metadata": {},
"quality_score": {},
"versions": [],
"output_datasets": [],
}.items():
if key in item:
item[key] = _json_value(item[key], default)
return item
class StoreBase:
"""数据处理存储基类。"""
def __init__(self, database_url: str | None = None) -> None:
self.database_url = _database_url(database_url or get_settings().database_url)
@contextmanager
def connect(self) -> Iterator[psycopg.Connection[dict[str, Any]]]:
with psycopg.connect(self.database_url, row_factory=dict_row) as conn:
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
def ensure_schema(self) -> None:
"""显式安装数据处理表API 路由和应用启动流程不会调用此方法。"""
schema_path = Path(__file__).resolve().parents[3] / "db" / "sql" / "002_data_process.sql"
sql = schema_path.read_text(encoding="utf-8")
with self.connect() as conn, conn.cursor() as cursor:
cursor.execute(sql)

View File

@@ -0,0 +1,504 @@
"""数据处理存储层 - 数据集发布。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import hashlib
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
from ..algorithms import stable_split_assignments
class DatasetsMixin:
"""数据集发布 Mixin。"""
def get_generation_model(self, model_id: str) -> dict[str, Any]:
with self.connect() as conn:
row = conn.execute(
"""
SELECT id, name, type, purpose, model_source, description, path,
api_url, api_key, online_model_name, create_time
FROM models WHERE id=%s
""",
(model_id,),
).fetchone()
if not row:
raise NotFoundError("generation model not found")
return _decode_row(row) or {}
def save_generation_model_snapshot(
self,
task_id: str,
model_snapshot: dict[str, Any],
*,
generation_run_id: str,
) -> dict[str, Any]:
# API 密钥仅用于本次调用,绝不能进入任务配置、详情响应或审计快照。
safe_snapshot = {
key: value for key, value in model_snapshot.items() if key != "api_key"
}
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task["status"] != "running"
or task.get("generation_run_id") != generation_run_id
):
raise InvalidStateError("generation run is no longer active")
config = dict(task.get("config") or {})
config["generation_model_snapshot"] = safe_snapshot
row = conn.execute(
"""
UPDATE data_process_tasks SET config=%s, updated_at=%s
WHERE id=%s AND generation_run_id=%s RETURNING *
""",
(json_dumps(config), utcnow(), task_id, generation_run_id),
).fetchone()
return _decode_row(row) or {}
def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
"""按精确配额发布训练、验证、测试三个独立数据集。"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if _is_regeneration_prepared(task):
raise InvalidStateError(
"regeneration must start and complete before publishing"
)
if task["status"] != "completed":
raise InvalidStateError("only a completed task can be published")
if not task.get("results_confirmed"):
raise InvalidStateError("results must be confirmed before publishing")
rows = conn.execute(
"""
SELECT * FROM data_process_results
WHERE task_id=%s ORDER BY created_at, id
""",
(task_id,),
).fetchall()
if not rows:
raise InvalidStateError("task has no results to publish")
invalid_count = sum(
1
for row in rows
if row["status"] == "invalid"
or not str(row.get("instruction") or "").strip()
or not str(row.get("output") or "").strip()
or (
_task_output_type(task) == "reasoning"
and not _reasoning_output_is_valid(row.get("output"))
)
or (
_task_output_type(task) == "dpo"
and not _dpo_fields_are_valid(row)
)
)
if invalid_count:
raise InvalidStateError(f"task contains {invalid_count} invalid results")
now = utcnow()
requested_split = payload.get("split") or {
"train": 80,
"validation": 10,
"test": 10,
}
assignments = stable_split_assignments(
[str(row["id"]) for row in rows],
requested_split,
seed=task_id,
)
if _task_output_type(task) == "dpo":
records = [
{
"instruction": row["instruction"],
"input": row["input"],
"chosen": row["chosen"],
"rejected": row["rejected"],
"split": assignment,
}
for row, assignment in zip(rows, assignments, strict=True)
]
else:
records = [
{
"instruction": row["instruction"],
"input": row["input"],
"output": row["output"],
"split": assignment,
}
for row, assignment in zip(rows, assignments, strict=True)
]
split_order = ("train", "validation", "test")
split_counts = {
split_name: assignments.count(split_name) for split_name in split_order
}
split_specs: list[dict[str, Any]] = []
for split_name in split_order:
split_records = [
(source_row, record)
for source_row, record in zip(rows, records, strict=True)
if record["split"] == split_name
]
file_id = new_id("dfile")
version_id = new_id("dfv")
content = "".join(
json_dumps(record) + "\n" for _, record in split_records
)
raw = content.encode("utf-8")
split_specs.append(
{
"split": split_name,
"records": split_records,
"file_id": file_id,
"version_id": version_id,
"content": content,
"raw": raw,
"checksum": hashlib.sha256(raw).hexdigest(),
"storage_object_id": (
f"db://data-process/{task_id}/{file_id}/v1"
),
}
)
source_result_ids = [row["id"] for row in rows]
common_metadata = {
"source": "data_process",
"storage_backend": "database",
"source_task_id": task_id,
"output_type": _task_output_type(task),
"reasoning_detail": _task_reasoning_detail(task),
"source_file_ids": [item["id"] for item in self._source_ids(conn, task_id)],
"source_result_ids": source_result_ids,
"format": (
"dpo"
if _task_output_type(task) == "dpo"
else payload.get("format") or "alpaca_jsonl"
),
"split": requested_split,
}
existing_datasets = conn.execute(
"""
SELECT * FROM datasets
WHERE source='task' AND deleted_at IS NULL
AND (
source_task_id=%s
OR (source_task_id IS NULL AND task_id=%s)
)
ORDER BY created_at, id
""",
(task_id, task_id),
).fetchall()
existing_by_split: dict[str, dict[str, Any]] = {}
primary_existing = None
for existing in existing_datasets:
existing_metadata = _json_value(existing.get("metadata"), {})
existing_split = str(existing_metadata.get("dataset_split") or "")
if existing_split in split_order:
existing_by_split[existing_split] = existing
if str(existing["id"]) == str(task.get("output_dataset_id") or ""):
primary_existing = existing
if primary_existing and "train" not in existing_by_split:
# 兼容旧版“一个数据集包含三个文件”的发布物,原数据集复用为训练集。
existing_by_split["train"] = primary_existing
existing_group_metadata = _json_value(
(primary_existing or {}).get("metadata"), {}
)
base_dataset_name = str(
existing_group_metadata.get("base_dataset_name")
or payload["dataset_name"]
).strip()
for suffix in ("-训练集", "-验证集", "-测试集"):
if base_dataset_name.endswith(suffix):
base_dataset_name = base_dataset_name[: -len(suffix)].rstrip()
break
split_group_id = str(
existing_group_metadata.get("split_group_id")
or f"dsg_{hashlib.sha256(task_id.encode()).hexdigest()[:20]}"
)
dataset_ids = {
spec["split"]: str(existing_by_split[spec["split"]]["id"])
if spec["split"] in existing_by_split
else new_id("dataset")
for spec in split_specs
}
created_any = any(
spec["split"] not in existing_by_split for spec in split_specs
)
split_labels = {
"train": "训练集",
"validation": "验证集",
"test": "测试集",
}
dataset_types = {"train": "train", "validation": "val", "test": "test"}
published_datasets: list[dict[str, Any]] = []
try:
for spec in split_specs:
split_name = str(spec["split"])
dataset_id = dataset_ids[split_name]
existing_dataset = existing_by_split.get(split_name)
dataset_metadata = {
**common_metadata,
"base_dataset_name": base_dataset_name,
"dataset_split": split_name,
"split_group_id": split_group_id,
"split_dataset_ids": dataset_ids,
"split_counts": {
name: split_counts[name] if name == split_name else 0
for name in split_order
},
}
dataset_name = f"{base_dataset_name}-{split_labels[split_name]}"
if existing_dataset:
conn.execute(
"DELETE FROM dataset_records WHERE dataset_id=%s", (dataset_id,)
)
conn.execute(
"""DELETE FROM dataset_file_versions
WHERE dataset_file_id IN
(SELECT id FROM dataset_files WHERE dataset_id=%s)""",
(dataset_id,),
)
conn.execute(
"DELETE FROM dataset_files WHERE dataset_id=%s", (dataset_id,)
)
dataset = conn.execute(
"""
UPDATE datasets
SET name=%s, type=%s, storage_type=%s, size=%s, size_bytes=%s,
count=%s, record_count=%s, description=%s, metadata=%s,
updated_at=%s
WHERE id=%s RETURNING *
""",
(
dataset_name,
dataset_types[split_name],
payload.get("storage_type") or "local",
f"{len(spec['raw'])} B",
len(spec["raw"]),
len(spec["records"]),
len(spec["records"]),
payload.get("description") or task.get("description") or "",
json_dumps(dataset_metadata),
now,
dataset_id,
),
).fetchone()
else:
dataset = conn.execute(
"""
INSERT INTO datasets
(id, name, type, storage_type, source, task_id, source_task_id,
size, size_bytes, count, record_count, description, metadata,
tenant_id, project_id, owner_id, created_by, create_time,
created_at, updated_at)
VALUES (
%s, %s, %s, %s, 'task', %s, %s,
%s, %s, %s, %s, %s, %s,
%s, %s, %s, %s, %s, %s, %s
)
RETURNING *
""",
(
dataset_id,
dataset_name,
dataset_types[split_name],
payload.get("storage_type") or "local",
task_id,
task_id,
f"{len(spec['raw'])} B",
len(spec["raw"]),
len(spec["records"]),
len(spec["records"]),
payload.get("description") or task.get("description") or "",
json_dumps(dataset_metadata),
task.get("tenant_id"),
task.get("project_id"),
task.get("owner_id"),
payload.get("created_by") or task.get("created_by"),
now,
now,
now,
),
).fetchone()
file_metadata = {**dataset_metadata, "file_split": split_name}
version = {
"id": spec["version_id"],
"version_no": 1,
"version": 1,
"description": f"data process {split_name} publish",
"checksum_sha256": spec["checksum"],
"size_bytes": len(spec["raw"]),
"record_count": len(spec["records"]),
"created_at": now,
"create_time": now,
"source_task_id": task_id,
"storage_object_id": spec["storage_object_id"],
}
conn.execute(
"""
INSERT INTO dataset_files
(id, dataset_id, name, storage_object_id, size, content,
active_version_id, versions, create_time, current_version_id,
size_bytes, record_count, file_format, checksum_sha256, version_no,
source_task_id, tenant_id, project_id, created_by, metadata,
created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, 1, %s, %s, %s, %s, %s, %s, %s)
""",
(
spec["file_id"],
dataset_id,
f"{base_dataset_name}.{split_name}.jsonl",
spec["storage_object_id"],
f"{len(spec['raw'])} B",
spec["content"],
spec["version_id"],
json_dumps([version]),
now,
spec["version_id"],
len(spec["raw"]),
len(spec["records"]),
"jsonl",
spec["checksum"],
task_id,
task.get("tenant_id"),
task.get("project_id"),
payload.get("created_by") or task.get("created_by"),
json_dumps(file_metadata),
now,
now,
),
)
conn.execute(
"""
INSERT INTO dataset_file_versions
(id, dataset_file_id, version_no, storage_object_id, content_preview,
description, size_bytes, record_count, checksum_sha256,
source_task_id, metadata, created_by, created_at)
VALUES (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""",
(
spec["version_id"],
spec["file_id"],
spec["storage_object_id"],
spec["content"][:2000],
f"data process {split_name} publish",
len(spec["raw"]),
len(spec["records"]),
spec["checksum"],
task_id,
json_dumps(file_metadata),
payload.get("created_by") or task.get("created_by"),
now,
),
)
for line_number, (source_row, record) in enumerate(
spec["records"], start=1
):
conn.execute(
"""
INSERT INTO dataset_records
(id, dataset_id, dataset_file_id, version_id, line_no, split,
instruction, input, output, raw, status, source_task_id,
source_result_id, preview_item_id, created_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, %s, %s)
""",
(
new_id("drec"),
dataset_id,
spec["file_id"],
spec["version_id"],
line_number,
record["split"],
record["instruction"],
record["input"],
record.get("output") or record.get("chosen") or "",
json_dumps(
{
**record,
"source_task_id": task_id,
"source_result_id": source_row["id"],
"preview_item_id": source_row.get("preview_item_id"),
}
),
source_row["status"],
task_id,
source_row["id"],
source_row.get("preview_item_id"),
now,
),
)
published_datasets.append(_decode_row(dataset) or {})
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("dataset name already exists") from exc
train_dataset_id = dataset_ids.get("train")
if not train_dataset_id:
raise InvalidStateError("published split does not contain training data")
conn.execute(
"""
UPDATE data_process_tasks
SET output_dataset_id=%s, updated_at=%s, updated_by=%s
WHERE id=%s
""",
(train_dataset_id, now, payload.get("created_by"), task_id),
)
train_dataset = next(
item
for item in published_datasets
if _json_value(item.get("metadata"), {}).get("dataset_split") == "train"
)
return {
"dataset": train_dataset,
"datasets": published_datasets,
"output_datasets": published_datasets,
"created": created_any,
"split_counts": split_counts,
}
@staticmethod
def _source_ids(
conn: psycopg.Connection[dict[str, Any]], task_id: str
) -> list[dict[str, Any]]:
return conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id
""",
(task_id,),
).fetchall()

View File

@@ -0,0 +1,275 @@
"""数据处理存储层 - 生成管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
class GenerationMixin:
"""生成管理 Mixin。"""
def _invalidate_results(
self,
conn: psycopg.Connection[dict[str, Any]],
task: dict[str, Any],
task_id: str,
now: str,
) -> None:
if _is_regeneration_prepared(task):
conn.execute(
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
(now, task_id),
)
return
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=20, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
generation_run_id=NULL, results_confirmed=FALSE, updated_at=%s
WHERE id=%s
""",
(now, task_id),
)
def start_generation(self, task_id: str, *, replace_existing: bool = True) -> dict[str, Any]:
if not replace_existing:
raise DataProcessStoreError("incremental generation is not supported")
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
regeneration_prepared = _is_regeneration_prepared(task)
if task.get("output_dataset_id") and not regeneration_prepared:
raise InvalidStateError("published task cannot be regenerated")
if task["status"] == "running":
raise ConflictError("data process task is already running")
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise ConflictError("preview is still running")
preview_count = conn.execute(
"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE task_id=%s",
(task_id,),
).fetchone()["count"]
if not preview_count:
raise InvalidStateError("preview must be built before generation")
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
now = utcnow()
generation_run_id = new_id("dprun")
next_config = dict(task.get("config") or {})
next_config.pop(_REGENERATION_MARKER_KEY, None)
row = conn.execute(
"""
UPDATE data_process_tasks
SET config=%s, status='running', progress=30, failure_reason=NULL,
started_at=%s, completed_at=NULL, output_dataset_id=NULL,
output_count=0, filtered_count=0, duplicate_count=0, error_count=0,
generation_run_id=%s, results_confirmed=FALSE,
workflow_step='generate', updated_at=%s
WHERE id=%s
RETURNING *
""",
(json_dumps(next_config), now, generation_run_id, now, task_id),
).fetchone()
return _decode_row(row) or {}
def stop_task(self, task_id: str) -> dict[str, Any]:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] != "running":
raise InvalidStateError("only a running task can be stopped")
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_tasks
SET status='stopped', failure_reason=NULL, generation_run_id=NULL,
updated_at=%s
WHERE id=%s RETURNING *
""",
(now, task_id),
).fetchone()
return _decode_row(row) or {}
def generation_is_running(self, task_id: str, generation_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
SELECT status, generation_run_id
FROM data_process_tasks
WHERE id=%s AND deleted_at IS NULL
""",
(task_id,),
).fetchone()
return bool(
row
and row.get("status") == "running"
and row.get("generation_run_id") == generation_run_id
)
def update_generation_progress(
self,
task_id: str,
generation_run_id: str,
processed_count: int,
total_count: int,
) -> bool:
ratio = processed_count / max(1, total_count)
progress = min(95.0, 30.0 + ratio * 65.0)
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET progress=%s, updated_at=%s
WHERE id=%s AND status='running' AND generation_run_id=%s
RETURNING id
""",
(progress, utcnow(), task_id, generation_run_id),
).fetchone()
return row is not None
def complete_generation(
self,
task_id: str,
results: Sequence[dict[str, Any]],
*,
generation_run_id: str,
filtered_count: int = 0,
duplicate_count: int = 0,
error_count: int = 0,
) -> dict[str, Any]:
now = utcnow()
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task["status"] != "running"
or task.get("generation_run_id") != generation_run_id
):
return task
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
for result in results:
conn.execute(
"""
INSERT INTO data_process_results
(id, task_id, preview_item_id, instruction, input, output,
chosen, rejected, original_instruction, original_input,
original_output, original_chosen, original_rejected, status, error,
split, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, %s, %s, %s, %s, %s)
""",
(
result.get("id") or new_id("dpr"),
task_id,
result.get("preview_item_id"),
result.get("instruction") or "",
result.get("input") or "",
result.get("output") or "",
result.get("chosen") or "",
result.get("rejected") or "",
result.get("original_instruction", result.get("instruction") or ""),
result.get("original_input", result.get("input") or ""),
result.get("original_output", result.get("output") or ""),
result.get("original_chosen", result.get("chosen") or ""),
result.get("original_rejected", result.get("rejected") or ""),
result.get("status") or "valid",
result.get("error"),
result.get("split"),
json_dumps(result.get("quality_score") or {}),
now,
now,
),
)
row = conn.execute(
"""
UPDATE data_process_tasks
SET status='completed', progress=100, output_count=%s, filtered_count=%s,
duplicate_count=%s, error_count=%s, failure_reason=NULL,
completed_at=%s, generation_run_id=NULL, results_confirmed=FALSE,
updated_at=%s
WHERE id=%s AND generation_run_id=%s RETURNING *
""",
(
len(results),
filtered_count,
duplicate_count,
error_count,
now,
now,
task_id,
generation_run_id,
),
).fetchone()
return _decode_row(row) or {}
def mark_failed(
self, task_id: str, reason: str, *, generation_run_id: str
) -> dict[str, Any]:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task["status"] != "running"
or task.get("generation_run_id") != generation_run_id
):
return task
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_tasks
SET status='failed', failure_reason=%s, completed_at=%s,
generation_run_id=NULL, updated_at=%s
WHERE id=%s AND generation_run_id=%s RETURNING *
""",
(reason[:4000], now, now, task_id, generation_run_id),
).fetchone()
return _decode_row(row) or {}
def progress(self, task_id: str) -> dict[str, Any]:
task = self.get_task(task_id)
return {
"task_id": task["id"],
"status": task["status"],
"progress": float(task.get("progress") or 0),
"input_count": int(task.get("input_count") or 0),
"output_count": int(task.get("output_count") or 0),
"filtered_count": int(task.get("filtered_count") or 0),
"duplicate_count": int(task.get("duplicate_count") or 0),
"error_count": int(task.get("error_count") or 0),
"failure_reason": task.get("failure_reason"),
"results_confirmed": bool(task.get("results_confirmed")),
"started_at": task.get("started_at"),
"completed_at": task.get("completed_at"),
}

View File

@@ -0,0 +1,539 @@
"""数据处理存储层 - 预览管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
_serialize_value,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
from ..algorithms import estimate_token_count # noqa: E402
class PreviewMixin:
"""预览管理 Mixin。"""
def replace_preview_items(
self,
task_id: str,
items: Sequence[dict[str, Any]],
*,
source_file_ids: Sequence[str] | None = None,
preview_run_id: 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")
preview_file_count = len(selected_ids) if selected_ids is not None else len(
{str(item.get("source_file_id") or "") for item in items}
)
is_direct_build = preview_run_id is None
now = utcnow()
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if is_direct_build:
self._ensure_editable(task)
elif (
task.get("preview_run_id") != preview_run_id
or task.get("preview_status") != "running"
):
raise InvalidStateError("preview run is no longer active")
regeneration_prepared = _is_regeneration_prepared(task)
if not regeneration_prepared:
conn.execute("DELETE FROM data_process_results 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(
"""
INSERT INTO data_process_preview_items
(id, task_id, source_file_id, original_content, edited_content,
source_start, source_end, source_start_line, source_end_line,
token_count, status, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
item.get("id") or new_id("dpp"),
task_id,
item.get("source_file_id"),
item.get("original_content") or "",
item.get("edited_content", item.get("original_content") or ""),
item.get("source_start"),
item.get("source_end"),
item.get("source_start_line"),
item.get("source_end_line"),
max(0, int(item.get("token_count") or 0)),
item.get("status") or "original",
json_dumps(item.get("quality_score") or {}),
now,
now,
),
).fetchone()
created.append(_decode_row(row) or {})
if regeneration_prepared:
if is_direct_build:
conn.execute(
"""
UPDATE data_process_tasks
SET workflow_step='preview', preview_status='completed',
preview_progress=100, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=%s,
preview_completed_files=%s, updated_at=%s
WHERE id=%s
""",
(preview_file_count, preview_file_count, now, task_id),
)
else:
conn.execute(
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
(now, task_id),
)
else:
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=20, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
results_confirmed=FALSE,
workflow_step=CASE WHEN %s THEN 'preview' ELSE workflow_step END,
preview_status=CASE WHEN %s THEN 'completed' ELSE preview_status END,
preview_progress=CASE WHEN %s THEN 100 ELSE preview_progress END,
preview_run_id=CASE WHEN %s THEN NULL ELSE preview_run_id END,
preview_failure_reason=CASE WHEN %s THEN NULL ELSE preview_failure_reason END,
preview_total_files=CASE WHEN %s THEN %s ELSE preview_total_files END,
preview_completed_files=CASE WHEN %s THEN %s ELSE preview_completed_files END,
updated_at=%s
WHERE id=%s
""",
(
is_direct_build,
is_direct_build,
is_direct_build,
is_direct_build,
is_direct_build,
is_direct_build,
preview_file_count,
is_direct_build,
preview_file_count,
now,
task_id,
),
)
return created
def start_preview(
self,
task_id: str,
*,
source_file_ids: Sequence[str] | None = None,
) -> tuple[dict[str, Any], list[str]]:
"""创建一轮持久化切分任务,并返回本轮固定的源文件集合。"""
requested_ids = (
list(dict.fromkeys(str(file_id) for file_id in source_file_ids))
if source_file_ids is not None
else None
)
if requested_ids is not None and (
not requested_ids or any(not file_id for file_id in requested_ids)
):
raise ValueError("source_file_ids must contain non-empty ids")
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
if requested_ids is None:
rows = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
ORDER BY created_at, id
""",
(task_id,),
).fetchall()
else:
rows = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL AND id=ANY(%s)
ORDER BY created_at, id
""",
(task_id, requested_ids),
).fetchall()
selected_ids = [str(row["id"]) for row in rows]
if not selected_ids:
raise InvalidStateError("at least one source file is required")
if requested_ids is not None:
missing = set(requested_ids) - set(selected_ids)
if missing:
raise NotFoundError(
f"source files not found: {', '.join(sorted(missing))}"
)
regeneration_prepared = _is_regeneration_prepared(task)
if not regeneration_prepared:
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
preview_run_id = new_id("dpprun")
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_tasks
SET status=CASE WHEN %s THEN status ELSE 'pending' END,
progress=CASE WHEN %s THEN progress ELSE 0 END,
output_count=CASE WHEN %s THEN output_count ELSE 0 END,
filtered_count=CASE WHEN %s THEN filtered_count ELSE 0 END,
duplicate_count=CASE WHEN %s THEN duplicate_count ELSE 0 END,
error_count=CASE WHEN %s THEN error_count ELSE 0 END,
failure_reason=CASE WHEN %s THEN failure_reason ELSE NULL END,
results_confirmed=CASE WHEN %s THEN results_confirmed ELSE FALSE END,
workflow_step='upload', preview_status='queued', preview_progress=0,
preview_run_id=%s, preview_failure_reason=NULL,
preview_total_files=%s, preview_completed_files=0,
updated_at=%s
WHERE id=%s AND deleted_at IS NULL
RETURNING *
""",
(
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
preview_run_id,
len(selected_ids),
now,
task_id,
),
).fetchone()
return _public_task(_decode_row(row)) or {}, selected_ids
def mark_preview_running(self, task_id: str, preview_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET preview_status='running', updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status='queued' AND preview_run_id=%s
RETURNING id
""",
(utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def preview_is_running(self, task_id: str, preview_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
SELECT preview_status, preview_run_id
FROM data_process_tasks
WHERE id=%s AND deleted_at IS NULL
""",
(task_id,),
).fetchone()
return bool(
row
and row.get("preview_status") in ACTIVE_PREVIEW_STATUSES
and row.get("preview_run_id") == preview_run_id
)
def update_preview_progress(
self,
task_id: str,
preview_run_id: str,
completed_files: int,
total_files: int,
) -> bool:
total = max(1, total_files)
completed = min(max(0, completed_files), total)
progress = completed / total * 100
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET preview_progress=%s, preview_completed_files=%s, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status='running' AND preview_run_id=%s
RETURNING id
""",
(progress, completed, utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def complete_preview(self, task_id: str, preview_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET workflow_step='preview', preview_status='completed',
preview_progress=100, preview_run_id=NULL,
preview_failure_reason=NULL,
preview_completed_files=preview_total_files, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status='running' AND preview_run_id=%s
RETURNING id
""",
(utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def mark_preview_failed(
self,
task_id: str,
reason: str,
*,
preview_run_id: str,
) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET preview_status='failed', preview_run_id=NULL,
preview_failure_reason=%s, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status IN ('queued', 'running') AND preview_run_id=%s
RETURNING id
""",
(reason[:4000], utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def preview_progress(self, task_id: str) -> dict[str, Any]:
task = self.get_task(task_id)
return {
"task_id": task["id"],
"workflow_step": task.get("workflow_step") or "create",
"preview_status": task.get("preview_status") or "idle",
"preview_progress": float(task.get("preview_progress") or 0),
"preview_run_id": task.get("preview_run_id"),
"preview_failure_reason": task.get("preview_failure_reason"),
"preview_total_files": int(task.get("preview_total_files") or 0),
"preview_completed_files": int(task.get("preview_completed_files") or 0),
}
def list_preview_items(
self,
task_id: str,
*,
source_file_id: str | None = None,
page: int = 1,
page_size: int = 200,
keyword: str | None = None,
) -> dict[str, Any]:
self.get_task(task_id)
clauses = ["task_id=%s"]
params: list[Any] = [task_id]
if source_file_id:
clauses.append("source_file_id=%s")
params.append(source_file_id)
if keyword:
clauses.append("(original_content ILIKE %s OR edited_content ILIKE %s)")
pattern = f"%{keyword.strip()}%"
params.extend([pattern, pattern])
where = " AND ".join(clauses)
with self.connect() as conn:
total = conn.execute(
f"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE {where}", params
).fetchone()["count"]
rows = conn.execute(
f"""
SELECT * FROM data_process_preview_items
WHERE {where}
ORDER BY source_file_id NULLS LAST, source_start NULLS LAST, created_at, id
LIMIT %s OFFSET %s
""",
[*params, page_size, (page - 1) * page_size],
).fetchall()
return {
"items": [_decode_row(row) for row in rows],
"total": int(total),
"page": page,
"page_size": page_size,
}
def get_preview_item(self, task_id: str, preview_id: str) -> dict[str, Any]:
with self.connect() as conn:
row = conn.execute(
"SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
(preview_id, task_id),
).fetchone()
if not row:
raise NotFoundError("preview item not found")
return _decode_row(row) or {}
def create_preview_item(self, task_id: str, item: dict[str, Any]) -> dict[str, Any]:
now = utcnow()
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
if item.get("source_file_id"):
source = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
""",
(item["source_file_id"], task_id),
).fetchone()
if not source:
raise NotFoundError("source file not found")
row = conn.execute(
"""
INSERT INTO data_process_preview_items
(id, task_id, source_file_id, original_content, edited_content,
source_start, source_end, source_start_line, source_end_line,
token_count, status, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
new_id("dpp"),
task_id,
item.get("source_file_id"),
item.get("original_content") or "",
item.get("edited_content") or "",
item.get("source_start"),
item.get("source_end"),
item.get("source_start_line"),
item.get("source_end_line"),
max(0, int(item.get("token_count") or 0)),
item.get("status") or "manual",
json_dumps(item.get("quality_score") or {}),
now,
now,
),
).fetchone()
self._invalidate_results(conn, task, task_id, now)
return _decode_row(row) or {}
def update_preview_item(
self, task_id: str, preview_id: str, payload: dict[str, Any]
) -> dict[str, Any]:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
existing = conn.execute(
"SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
(preview_id, task_id),
).fetchone()
if not existing:
raise NotFoundError("preview item not found")
expected_updated_at = payload.get("expected_updated_at")
current_updated_at = _serialize_value(existing.get("updated_at"))
if expected_updated_at and expected_updated_at != current_updated_at:
raise ConflictError("preview item was modified by another request")
edited = payload["edited_content"]
status = payload.get("status")
if not status:
if not edited.strip():
status = "invalid"
elif edited == existing["original_content"]:
status = "original"
else:
status = "modified"
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_preview_items
SET edited_content=%s, token_count=%s, status=%s, quality_score=%s,
updated_at=%s
WHERE id=%s AND task_id=%s
RETURNING *
""",
(
edited,
estimate_token_count(edited),
status,
json_dumps(payload.get("quality_score") or {}),
now,
preview_id,
task_id,
),
).fetchone()
self._invalidate_results(conn, task, task_id, now)
return _decode_row(row) or {}
def delete_preview_item(self, task_id: str, preview_id: str) -> None:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
row = conn.execute(
"DELETE FROM data_process_preview_items WHERE id=%s AND task_id=%s RETURNING id",
(preview_id, task_id),
).fetchone()
if not row:
raise NotFoundError("preview item not found")
self._invalidate_results(conn, task, task_id, utcnow())

View File

@@ -0,0 +1,324 @@
"""数据处理存储层 - 结果管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
_serialize_value,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
class ResultsMixin:
"""结果管理 Mixin。"""
def confirm_results(self, task_id: str) -> dict[str, Any]:
"""确认第六步结果,确认前再次校验所有生成记录。"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] != "completed":
raise InvalidStateError("only a completed task can confirm results")
if task.get("workflow_step") != "results":
raise InvalidStateError("workflow must be on results before confirmation")
if task.get("results_confirmed"):
return task
rows = conn.execute(
"""
SELECT status, instruction, output, chosen, rejected
FROM data_process_results
WHERE task_id=%s
""",
(task_id,),
).fetchall()
if not rows:
raise InvalidStateError("task has no results to confirm")
invalid_count = sum(
1
for row in rows
if row["status"] == "invalid"
or not str(row.get("instruction") or "").strip()
or not str(row.get("output") or "").strip()
or (
_task_output_type(task) == "reasoning"
and not _reasoning_output_is_valid(row.get("output"))
)
or (
_task_output_type(task) == "dpo"
and not _dpo_fields_are_valid(row)
)
)
if invalid_count:
raise InvalidStateError(
f"task contains {invalid_count} invalid results"
)
row = conn.execute(
"""
UPDATE data_process_tasks
SET results_confirmed=TRUE, updated_at=%s
WHERE id=%s RETURNING *
""",
(utcnow(), task_id),
).fetchone()
return _decode_row(row) or {}
def list_results(
self,
task_id: str,
*,
page: int = 1,
page_size: int = 100,
status: str | None = None,
split: str | None = None,
keyword: str | None = None,
) -> dict[str, Any]:
self.get_task(task_id)
clauses = ["task_id=%s"]
params: list[Any] = [task_id]
if status:
clauses.append("status=%s")
params.append(status)
if split:
clauses.append("split=%s")
params.append(split)
if keyword:
clauses.append("(instruction ILIKE %s OR input ILIKE %s OR output ILIKE %s)")
pattern = f"%{keyword.strip()}%"
params.extend([pattern, pattern, pattern])
where = " AND ".join(clauses)
with self.connect() as conn:
total = conn.execute(
f"SELECT COUNT(*) AS count FROM data_process_results WHERE {where}", params
).fetchone()["count"]
rows = conn.execute(
f"""
SELECT * FROM data_process_results WHERE {where}
ORDER BY created_at, id LIMIT %s OFFSET %s
""",
[*params, page_size, (page - 1) * page_size],
).fetchall()
return {
"items": [_decode_row(row) for row in rows],
"total": int(total),
"page": page,
"page_size": page_size,
}
def get_result(self, task_id: str, result_id: str) -> dict[str, Any]:
with self.connect() as conn:
row = conn.execute(
"SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
(result_id, task_id),
).fetchone()
if not row:
raise NotFoundError("data process result not found")
return _decode_row(row) or {}
def update_result(
self, task_id: str, result_id: str, payload: dict[str, Any]
) -> dict[str, Any]:
allowed = {
"instruction", "input", "output", "chosen", "rejected", "quality_score"
}
values = {key: value for key, value in payload.items() if key in allowed}
if "quality_score" in values:
values["quality_score"] = json_dumps(values["quality_score"])
if not values:
raise DataProcessStoreError("no result fields supplied")
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] == "running":
raise InvalidStateError("results cannot be edited while generation is running")
if task.get("output_dataset_id"):
raise InvalidStateError("published results cannot be edited")
current = conn.execute(
"SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
(result_id, task_id),
).fetchone()
if not current:
raise NotFoundError("data process result not found")
expected_updated_at = payload.get("expected_updated_at")
current_updated_at = _serialize_value(current.get("updated_at"))
if expected_updated_at and expected_updated_at != current_updated_at:
raise ConflictError("data process result was modified by another request")
output_type = _task_output_type(task)
if output_type == "dpo" and "chosen" in values:
values["output"] = values["chosen"]
merged = {**current, **values}
quality = payload.get("quality_score") or {}
instruction_valid = bool(str(merged.get("instruction") or "").strip())
output_valid = bool(str(merged.get("output") or "").strip())
reasoning_valid = (
output_type != "reasoning"
or _reasoning_output_is_valid(merged.get("output"))
)
dpo_valid = output_type != "dpo" or _dpo_fields_are_valid(merged)
hard_valid = instruction_valid and output_valid and reasoning_valid and dpo_valid
quality_valid = bool(quality.get("is_valid", hard_valid))
changed = any(
str(merged.get(field) or "")
!= str(merged.get(f"original_{field}") or "")
for field in (
("instruction", "input", "chosen", "rejected")
if output_type == "dpo"
else ("instruction", "input", "output")
)
)
status = "invalid" if not hard_valid or not quality_valid else (
"modified" if changed else "valid"
)
values["status"] = status
flags = quality.get("flags") if isinstance(quality, dict) else None
format_error = (
"思维链输出必须包含非空的 <think>...</think> 推理过程和最终答案"
if instruction_valid and output_valid and not reasoning_valid
else "DPO 输出必须包含不同的非空 Chosen 和 Rejected 回答"
if instruction_valid and not dpo_valid
else "Instruction 和 Output 不能为空"
if not instruction_valid or not output_valid
else None
)
values["error"] = ", ".join(str(flag) for flag in flags or []) or (
format_error
or ("quality validation failed" if status == "invalid" else None)
)
values["updated_at"] = utcnow()
assignments = ", ".join(f"{key}=%s" for key in values)
row = conn.execute(
f"""UPDATE data_process_results SET {assignments}
WHERE id=%s AND task_id=%s RETURNING *""",
[*values.values(), result_id, task_id],
).fetchone()
conn.execute(
"""
UPDATE data_process_tasks
SET error_count=(
SELECT COUNT(*) FROM data_process_results
WHERE task_id=%s AND status='invalid'
), updated_at=%s
WHERE id=%s
""",
(task_id, utcnow(), task_id),
)
return _decode_row(row) or {}
def replace_generated_result(
self,
task_id: str,
result_id: str,
replacement: dict[str, Any],
*,
expected_updated_at: str,
) -> dict[str, Any]:
"""用新模型结果原位替换失败项,并将新内容设为恢复基线。"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] != "completed" or task.get("workflow_step") != "results":
raise InvalidStateError("task is not editing generation results")
if task.get("results_confirmed"):
raise InvalidStateError("confirmed results cannot be regenerated")
if task.get("output_dataset_id"):
raise InvalidStateError("published results cannot be regenerated")
current = conn.execute(
"""SELECT * FROM data_process_results
WHERE id=%s AND task_id=%s FOR UPDATE""",
(result_id, task_id),
).fetchone()
if not current:
raise NotFoundError("data process result not found")
if current.get("status") != "invalid":
raise InvalidStateError("only an invalid result can be regenerated")
current_updated_at = _serialize_value(current.get("updated_at"))
if expected_updated_at != current_updated_at:
raise ConflictError("data process result was modified by another request")
instruction = str(replacement.get("instruction") or "").strip()
input_text = str(replacement.get("input") or "").strip()
output = str(replacement.get("output") or "").strip()
chosen = str(replacement.get("chosen") or "").strip()
rejected = str(replacement.get("rejected") or "").strip()
quality_score = replacement.get("quality_score") or {}
if not instruction or not output or not bool(quality_score.get("is_valid")):
raise InvalidStateError("regenerated result did not pass quality validation")
if _task_output_type(task) == "reasoning" and not _reasoning_output_is_valid(output):
raise InvalidStateError("regenerated reasoning result has an invalid output format")
if _task_output_type(task) == "dpo" and not _dpo_fields_are_valid(replacement):
raise InvalidStateError("regenerated DPO result has invalid preference fields")
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_results
SET instruction=%s, input=%s, output=%s, chosen=%s, rejected=%s,
original_instruction=%s, original_input=%s, original_output=%s,
original_chosen=%s, original_rejected=%s,
status='valid', error=NULL, quality_score=%s, updated_at=%s
WHERE id=%s AND task_id=%s
RETURNING *
""",
(
instruction,
input_text,
output,
chosen,
rejected,
instruction,
input_text,
output,
chosen,
rejected,
json_dumps(quality_score),
now,
result_id,
task_id,
),
).fetchone()
conn.execute(
"""
UPDATE data_process_tasks
SET error_count=(
SELECT COUNT(*) FROM data_process_results
WHERE task_id=%s AND status='invalid'
), updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)
return _decode_row(row) or {}

View File

@@ -0,0 +1,321 @@
"""数据处理存储层 - 源文件管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
class SourceFilesMixin:
"""源文件管理 Mixin。"""
def list_source_files(self, task_id: str) -> list[dict[str, Any]]:
self.get_task(task_id)
with self.connect() as conn:
rows = conn.execute(
"""
SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content_preview, metadata,
tenant_id, project_id, created_by, created_at, updated_at
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
ORDER BY created_at, id
""",
(task_id,),
).fetchall()
return [_decode_row(row) or {} for row in rows]
def add_source_file(
self,
task_id: str,
*,
name: str,
content: str,
raw_size: int,
checksum_sha256: str,
file_format: str,
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,
[
{
"name": name,
"content": content,
"raw_size": raw_size,
"checksum_sha256": checksum_sha256,
"file_format": file_format,
"record_count": record_count,
"metadata": metadata or {},
"created_by": created_by,
"id": source_file_id,
"storage_object_id": storage_object_id,
}
],
)[0]
def add_source_files(
self,
task_id: str,
files: Sequence[dict[str, Any]],
) -> list[dict[str, Any]]:
"""在同一事务中登记一个上传批次,任一文件失败则全部回滚。"""
if not files:
raise DataProcessStoreError("at least one source file is required")
now = utcnow()
created: list[dict[str, Any]] = []
try:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
for payload in files:
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
(id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content, content_preview,
metadata, tenant_id, project_id,
created_by, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s, %s,
%s, %s, %s, %s)
RETURNING id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content_preview, metadata,
tenant_id, project_id, created_by, created_at, updated_at
""",
(
file_id,
task_id,
storage_object_id,
payload["name"],
payload["raw_size"],
payload["record_count"],
payload["file_format"],
payload["checksum_sha256"],
payload["content"],
str(payload["content"])[:2000],
json_dumps(metadata_payload),
task.get("tenant_id"),
task.get("project_id"),
payload.get("created_by") or task.get("created_by"),
now,
now,
),
).fetchone()
created.append(_decode_row(row) or {})
preview_row = conn.execute(
"""
SELECT COUNT(*) AS count FROM data_process_preview_items
WHERE task_id=%s
""",
(task_id,),
).fetchone()
preview_count = int((preview_row or {}).get("count") or 0)
if _is_regeneration_prepared(task):
conn.execute(
"""
UPDATE data_process_tasks
SET input_count=(
SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
), workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0, updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)
else:
# 前端会只对本次新增的源文件构建预览,因此保留旧文件切片,
# 但普通未发布任务的旧生成结果已经不再有效。
conn.execute(
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
)
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=%s, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
generation_run_id=NULL, results_confirmed=FALSE,
workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0,
started_at=NULL, completed_at=NULL,
input_count=(
SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
), updated_at=%s
WHERE id=%s
""",
(20 if preview_count else 0, task_id, now, task_id),
)
except psycopg.errors.UniqueViolation as exc:
raise ConflictError(
"the same source file content is already attached to this task"
) from exc
return created
def get_source_file(
self, task_id: str, file_id: str, *, include_content: bool = True
) -> dict[str, Any]:
# 先验证父任务仍然可见,避免软删除任务后通过已知文件 ID 读取正文。
self.get_task(task_id)
content_column = ", content" if include_content else ""
with self.connect() as conn:
row = conn.execute(
f"""
SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content_preview, metadata,
tenant_id, project_id, created_by, created_at, updated_at{content_column}
FROM data_process_source_files
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
""",
(file_id, task_id),
).fetchone()
if not row:
raise NotFoundError("source file not found")
return _decode_row(row) or {}
def source_content_window(
self, task_id: str, file_id: str, offset: int, limit: int
) -> dict[str, Any]:
source_file = self.get_source_file(task_id, file_id, include_content=True)
content = str(source_file.pop("content", ""))
window = content[offset : offset + limit]
return {
"file": source_file,
"content": window,
"offset": offset,
"limit": limit,
"total_chars": len(content),
"has_more": offset + len(window) < len(content),
}
def source_content_lines(
self,
task_id: str,
file_id: str,
start_line: int,
line_count: int,
) -> dict[str, Any]:
source_file = self.get_source_file(task_id, file_id, include_content=True)
content = str(source_file.pop("content", ""))
lines = content.splitlines(keepends=True)
start_index = min(len(lines), start_line - 1)
selected = lines[start_index : start_index + line_count]
end_line = start_index + len(selected)
return {
"file": source_file,
"content": "".join(selected),
"start_line": start_line,
"end_line": end_line,
"line_count": len(selected),
"total_lines": len(lines),
"has_more": end_line < len(lines),
}
def delete_source_file(self, task_id: str, file_id: str) -> None:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
row = conn.execute(
"""
UPDATE data_process_source_files
SET deleted_at=%s, updated_at=%s
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
RETURNING id
""",
(utcnow(), utcnow(), file_id, task_id),
).fetchone()
if not row:
raise NotFoundError("source file not found")
conn.execute(
"DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_id,)
)
now = utcnow()
if _is_regeneration_prepared(task):
conn.execute(
"""
UPDATE data_process_tasks
SET input_count=(SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL),
workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0, updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)
else:
conn.execute(
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
)
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=0, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
results_confirmed=FALSE,
workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0,
input_count=(SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL),
updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)

View File

@@ -0,0 +1,871 @@
"""数据处理存储层 - 任务管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
_serialize_value,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
_INTERNAL_CONFIG_KEYS,
)
from ..algorithms import estimate_token_count
class TasksMixin:
"""任务管理 Mixin。"""
def list_tasks(
self,
*,
page: int = 1,
page_size: int = 20,
keyword: str | None = None,
status: str | None = None,
process_type: str | None = None,
tenant_id: str | None = None,
project_id: str | None = None,
) -> dict[str, Any]:
clauses = ["task.deleted_at IS NULL"]
params: list[Any] = []
if keyword:
clauses.append("(task.name ILIKE %s OR COALESCE(task.description, '') ILIKE %s)")
pattern = f"%{keyword.strip()}%"
params.extend([pattern, pattern])
if status:
clauses.append("task.status = %s")
params.append(status)
if process_type:
clauses.append("task.process_type = %s")
params.append(process_type)
if tenant_id:
clauses.append("task.tenant_id = %s")
params.append(tenant_id)
if project_id:
clauses.append("task.project_id = %s")
params.append(project_id)
where = " AND ".join(clauses)
with self.connect() as conn:
total = conn.execute(
f"SELECT COUNT(*) AS count FROM data_process_tasks task WHERE {where}",
params,
).fetchone()["count"]
rows = conn.execute(
f"""
SELECT task.*,
(SELECT COUNT(*) FROM data_process_source_files source_file
WHERE source_file.task_id=task.id
AND source_file.deleted_at IS NULL) AS source_file_count
FROM data_process_tasks task
WHERE {where}
ORDER BY task.created_at DESC, task.id DESC
LIMIT %s OFFSET %s
""",
[*params, page_size, (page - 1) * page_size],
).fetchall()
return {
"items": [_public_task(_decode_row(row)) for row in rows],
"total": int(total),
"page": page,
"page_size": page_size,
}
def create_task(self, payload: dict[str, Any]) -> dict[str, Any]:
task_id = new_id("dpt")
now = utcnow()
try:
with self.connect() as conn:
row = conn.execute(
"""
INSERT INTO data_process_tasks
(id, name, description, status, process_type, source_dataset_id, config,
progress, results_confirmed, tenant_id, project_id, owner_id, created_by, updated_by,
created_at, updated_at)
VALUES (%s, %s, %s, 'pending', %s, %s, %s, 0, FALSE,
%s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
task_id,
payload["name"],
payload.get("description") or "",
payload["process_type"],
payload.get("source_dataset_id"),
json_dumps(_business_config(payload.get("config"))),
payload.get("tenant_id"),
payload.get("project_id"),
payload.get("owner_id"),
payload.get("created_by"),
payload.get("created_by"),
now,
now,
),
).fetchone()
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc
return _public_task(_decode_row(row)) or {}
@staticmethod
def _repeat_response(
conn: psycopg.Connection[dict[str, Any]],
row: dict[str, Any],
*,
source_task_id: str,
created: bool,
) -> dict[str, Any]:
task_id = str(row["id"])
counts = conn.execute(
"""
SELECT
(SELECT COUNT(*) FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL) AS source_file_count,
(SELECT COUNT(*) FROM data_process_preview_items
WHERE task_id=%s) AS preview_count
""",
(task_id, task_id),
).fetchone() or {}
task = _public_task(_decode_row(row)) or {}
task["source_file_count"] = int(counts.get("source_file_count") or 0)
task["preview_count"] = int(counts.get("preview_count") or 0)
return {
"task": task,
"source_task_id": source_task_id,
"created": created,
"copied_source_file_count": task["source_file_count"],
"copied_preview_count": task["preview_count"],
}
def find_repeated_task(
self,
source_task_id: str,
request_id: str,
) -> dict[str, Any] | None:
"""查找同一幂等请求已创建的新任务。"""
task_id = repeat_task_id(source_task_id, request_id)
with self.connect() as conn:
row = conn.execute(
"SELECT * FROM data_process_tasks WHERE id=%s",
(task_id,),
).fetchone()
if row is None:
return None
decoded = _decode_row(row) or {}
config = decoded.get("config") or {}
if (
config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id
or config.get(_REPEAT_REQUEST_KEY) != request_id
):
raise ConflictError("再次生成请求与现有任务冲突")
if decoded.get("deleted_at"):
raise ConflictError("此次再次生成创建的任务已被删除,请重新发起")
return self._repeat_response(
conn,
row,
source_task_id=source_task_id,
created=False,
)
def repeat_task(
self,
source_task_id: str,
*,
expected_updated_at: str,
request_id: str,
file_copies: dict[str, dict[str, str]],
) -> dict[str, Any]:
"""复制已确认任务的配置、源文件和预览,结果与发布数据保持独立。"""
task_id = repeat_task_id(source_task_id, request_id)
now = utcnow()
try:
with self.connect() as conn:
existing = conn.execute(
"SELECT * FROM data_process_tasks WHERE id=%s FOR UPDATE",
(task_id,),
).fetchone()
if existing is not None:
decoded = _decode_row(existing) or {}
config = decoded.get("config") or {}
if (
config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id
or config.get(_REPEAT_REQUEST_KEY) != request_id
):
raise ConflictError("再次生成请求与现有任务冲突")
if decoded.get("deleted_at"):
raise ConflictError("此次再次生成创建的任务已被删除,请重新发起")
return self._repeat_response(
conn,
existing,
source_task_id=source_task_id,
created=False,
)
source_task = self._task_in_connection(
conn,
source_task_id,
for_update=True,
)
if (
source_task.get("status") != "completed"
or source_task.get("results_confirmed") is False
):
raise InvalidStateError("只有已完成并确认结果的任务可以再次生成")
if source_task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise ConflictError("源任务仍在处理切分,暂时不能再次生成")
if expected_updated_at != _serialize_value(source_task.get("updated_at")):
raise ConflictError("源任务已被其他操作修改,请刷新后重试")
source_files = conn.execute(
"""
SELECT * FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
ORDER BY created_at, id
""",
(source_task_id,),
).fetchall()
source_file_ids = {str(row["id"]) for row in source_files}
if source_file_ids != set(file_copies):
raise ConflictError("源文件快照已变化,请刷新后重试")
previews = conn.execute(
"""
SELECT * FROM data_process_preview_items
WHERE task_id=%s
ORDER BY source_file_id NULLS LAST, source_start NULLS LAST,
created_at, id
""",
(source_task_id,),
).fetchall()
if not previews:
raise InvalidStateError("源任务没有可用于再次生成的切分结果")
suffix = f"(再次生成-{task_id[-6:]}"
base_name = str(source_task.get("name") or "数据处理任务")
repeated_name = f"{base_name[: max(1, 150 - len(suffix))]}{suffix}"
repeated_config = _business_config(source_task.get("config") or {})
repeated_config[_REPEAT_SOURCE_TASK_KEY] = source_task_id
repeated_config[_REPEAT_REQUEST_KEY] = request_id
input_count = sum(int(row.get("record_count") or 0) for row in source_files)
task_row = conn.execute(
"""
INSERT INTO data_process_tasks
(id, name, description, status, process_type, source_dataset_id,
output_dataset_id, config, progress, input_count, output_count,
filtered_count, duplicate_count, error_count, failure_reason,
generation_run_id, results_confirmed, workflow_step,
preview_status, preview_progress, preview_run_id,
preview_failure_reason, preview_total_files,
preview_completed_files, tenant_id, project_id, owner_id,
approval_status, created_by, updated_by, created_at, updated_at)
VALUES
(%s, %s, %s, 'pending', %s, %s, NULL, %s, 20, %s, 0,
0, 0, 0, NULL, NULL, FALSE, 'preview', 'completed', 100,
NULL, NULL, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
task_id,
repeated_name,
source_task.get("description") or "",
source_task["process_type"],
source_task.get("source_dataset_id"),
json_dumps(repeated_config),
input_count,
len(source_files),
len(source_files),
source_task.get("tenant_id"),
source_task.get("project_id"),
source_task.get("owner_id"),
source_task.get("approval_status") or "not_required",
source_task.get("created_by"),
source_task.get("created_by"),
now,
now,
),
).fetchone()
file_id_map: dict[str, str] = {}
for source in source_files:
old_file_id = str(source["id"])
copy = file_copies[old_file_id]
new_file_id = str(copy["id"])
storage_object_id, metadata = _source_storage_descriptor(
{
"storage_object_id": copy["storage_object_id"],
"metadata": _json_value(source.get("metadata"), {}),
},
task_id,
new_file_id,
)
file_id_map[old_file_id] = new_file_id
conn.execute(
"""
INSERT INTO data_process_source_files
(id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content,
content_preview, metadata, tenant_id, project_id, created_by,
created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s,
%s, %s, %s, %s, %s)
""",
(
new_file_id,
task_id,
storage_object_id,
source["name"],
source.get("size_bytes") or 0,
source.get("record_count") or 0,
source.get("file_format"),
source["checksum_sha256"],
source.get("content") or "",
source.get("content_preview"),
json_dumps(metadata),
source_task.get("tenant_id"),
source_task.get("project_id"),
source.get("created_by") or source_task.get("created_by"),
now,
now,
),
)
for preview in previews:
old_source_file_id = preview.get("source_file_id")
conn.execute(
"""
INSERT INTO data_process_preview_items
(id, task_id, source_file_id, original_content, edited_content,
source_start, source_end, source_start_line, source_end_line,
token_count, status, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s)
""",
(
new_id("dpp"),
task_id,
file_id_map.get(str(old_source_file_id))
if old_source_file_id
else None,
preview.get("original_content") or "",
preview.get("edited_content") or "",
preview.get("source_start"),
preview.get("source_end"),
preview.get("source_start_line"),
preview.get("source_end_line"),
max(0, int(preview.get("token_count") or 0)),
preview.get("status") or "original",
json_dumps(_json_value(preview.get("quality_score"), {})),
now,
now,
),
)
return self._repeat_response(
conn,
task_row or {},
source_task_id=source_task_id,
created=True,
)
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("再次生成任务名称或请求发生冲突,请重试") from exc
def get_task(self, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
lock = " FOR UPDATE" if for_update else ""
with self.connect() as conn:
if for_update:
row = conn.execute(
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
(task_id,),
).fetchone()
else:
row = conn.execute(
"""
SELECT task.*,
(SELECT COUNT(*) FROM data_process_source_files source
WHERE source.task_id=task.id AND source.deleted_at IS NULL)
AS source_file_count,
(SELECT COUNT(*) FROM data_process_preview_items preview
WHERE preview.task_id=task.id) AS preview_count,
(SELECT COALESCE(json_agg(json_build_object(
'id', dataset.id,
'name', dataset.name,
'type', dataset.type,
'count', dataset.count,
'dataset_split', CASE dataset.type
WHEN 'train' THEN 'train'
WHEN 'val' THEN 'validation'
WHEN 'test' THEN 'test'
ELSE NULL
END
) ORDER BY CASE dataset.type
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4 END), '[]'::json)
FROM datasets dataset
WHERE dataset.source='task'
AND dataset.deleted_at IS NULL
AND (
dataset.source_task_id=task.id
OR (dataset.source_task_id IS NULL AND dataset.task_id=task.id)
))
AS output_datasets,
CASE
WHEN task.started_at IS NOT NULL AND task.completed_at IS NOT NULL
THEN EXTRACT(EPOCH FROM (task.completed_at - task.started_at))
ELSE NULL
END AS duration_seconds
FROM data_process_tasks task
WHERE task.id=%s AND task.deleted_at IS NULL
""",
(task_id,),
).fetchone()
if not row:
raise NotFoundError("data process task not found")
return _public_task(_decode_row(row)) or {}
def _task_in_connection(
self,
conn: psycopg.Connection[dict[str, Any]],
task_id: str,
*,
for_update: bool = False,
) -> dict[str, Any]:
lock = " FOR UPDATE" if for_update else ""
row = conn.execute(
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
(task_id,),
).fetchone()
if not row:
raise NotFoundError("data process task not found")
return _decode_row(row) or {}
@staticmethod
def _ensure_editable(task: dict[str, Any]) -> None:
if task["status"] not in EDITABLE_STATUSES:
raise InvalidStateError(f"task cannot be edited while status is {task['status']}")
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise InvalidStateError("task cannot be edited while preview is running")
if task.get("output_dataset_id") and not _is_regeneration_prepared(task):
raise InvalidStateError("published task cannot be edited")
def update_workflow_step(self, task_id: str, workflow_step: str) -> dict[str, Any]:
"""独立保存向导位置,不触发配置或结果失效逻辑。"""
if workflow_step not in WORKFLOW_STEPS:
raise ValueError("invalid data process workflow step")
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET workflow_step=%s, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
RETURNING *
""",
(workflow_step, utcnow(), task_id),
).fetchone()
if not row:
raise NotFoundError("data process task not found")
return _public_task(_decode_row(row)) or {}
def recover_legacy_aborted_regeneration(self, task_id: str) -> dict[str, Any]:
"""恢复旧版在真正开始生成前误删的上一轮结果。
旧实现会在 ``POST /regenerate`` 时立即把已发布任务置为 pending、
清空结果并解除输出指针。三个已发布数据集仍是独立完整产物,因此只在
这个特征完全匹配时,使用其记录恢复结果和任务状态。该操作幂等,不会
触碰正常的新建待生成任务或已经开始的新一轮生成。
"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task.get("status") != "pending"
or task.get("generation_run_id")
or task.get("output_dataset_id")
or int(task.get("output_count") or 0) != 0
):
return {"recovered": False, "result_count": 0}
result_count = int(
(
conn.execute(
"SELECT COUNT(*) AS count FROM data_process_results WHERE task_id=%s",
(task_id,),
).fetchone()
or {}
).get("count")
or 0
)
if result_count:
return {"recovered": False, "result_count": result_count}
datasets = conn.execute(
"""
SELECT id, type, count, created_at
FROM datasets
WHERE source='task' AND deleted_at IS NULL
AND (
source_task_id=%s
OR (source_task_id IS NULL AND task_id=%s)
)
ORDER BY CASE type
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4
END, created_at, id
""",
(task_id, task_id),
).fetchall()
train_dataset = next(
(dataset for dataset in datasets if dataset.get("type") == "train"),
None,
)
if not train_dataset:
return {"recovered": False, "result_count": 0}
dataset_ids = [str(dataset["id"]) for dataset in datasets]
records = conn.execute(
"""
SELECT id, dataset_id, line_no, split, instruction, input, output,
raw, status, source_result_id, preview_item_id, created_at
FROM dataset_records
WHERE dataset_id = ANY(%s)
ORDER BY created_at, dataset_id, line_no NULLS LAST, id
""",
(dataset_ids,),
).fetchall()
if not records:
return {"recovered": False, "result_count": 0}
preview_rows = conn.execute(
"SELECT id FROM data_process_preview_items WHERE task_id=%s",
(task_id,),
).fetchall()
preview_ids = {str(row["id"]) for row in preview_rows}
used_result_ids: set[str] = set()
recovered_count = 0
for record in records:
raw = _json_value(record.get("raw"), {})
raw = raw if isinstance(raw, dict) else {}
candidate_id = str(
record.get("source_result_id")
or raw.get("source_result_id")
or ""
)
result_id = (
candidate_id
if candidate_id and candidate_id not in used_result_ids
else new_id("dpr")
)
used_result_ids.add(result_id)
candidate_preview_id = str(
record.get("preview_item_id")
or raw.get("preview_item_id")
or ""
)
preview_item_id = (
candidate_preview_id if candidate_preview_id in preview_ids else None
)
instruction = str(record.get("instruction") or raw.get("instruction") or "")
input_text = str(record.get("input") or raw.get("input") or "")
chosen = str(raw.get("chosen") or "")
rejected = str(raw.get("rejected") or "")
output = str(
record.get("output") or raw.get("output") or chosen or ""
)
split = str(record.get("split") or raw.get("split") or "") or None
status = str(record.get("status") or "valid")
if status not in {"valid", "modified", "invalid"}:
status = "valid"
created_at = record.get("created_at") or utcnow()
conn.execute(
"""
INSERT INTO data_process_results
(id, task_id, preview_item_id, instruction, input, output,
chosen, rejected, original_instruction, original_input,
original_output, original_chosen, original_rejected, status,
error, split, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, NULL, %s, '{}', %s, %s)
""",
(
result_id,
task_id,
preview_item_id,
instruction,
input_text,
output,
chosen,
rejected,
instruction,
input_text,
output,
chosen,
rejected,
status,
split,
created_at,
created_at,
),
)
conn.execute(
"""
UPDATE dataset_records
SET source_result_id=%s, preview_item_id=%s
WHERE id=%s
""",
(result_id, preview_item_id, record["id"]),
)
recovered_count += 1
now = utcnow()
conn.execute(
"""
UPDATE data_process_tasks
SET status='completed', progress=100, output_dataset_id=%s,
output_count=%s, filtered_count=0, duplicate_count=0,
error_count=(SELECT COUNT(*) FROM data_process_results
WHERE task_id=%s AND status='invalid'),
failure_reason=NULL, results_confirmed=TRUE,
workflow_step='results', updated_at=%s
WHERE id=%s
""",
(
train_dataset["id"],
recovered_count,
task_id,
now,
task_id,
),
)
return {"recovered": True, "result_count": recovered_count}
def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
allowed = {
"name",
"description",
"process_type",
"source_dataset_id",
}
values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed}
if not values and payload.get("config") is None:
return self.get_task(task_id)
try:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
regeneration_prepared = _is_regeneration_prepared(task)
if regeneration_prepared and any(
key in payload and payload.get(key) != task.get(key)
for key in ("process_type", "source_dataset_id")
):
raise InvalidStateError(
"process type and source dataset cannot change during regeneration"
)
if payload.get("config") is not None:
next_config = _business_config(payload["config"])
current_config = dict(task.get("config") or {})
for key in _INTERNAL_CONFIG_KEYS:
if key in current_config:
next_config[key] = current_config[key]
values["config"] = json_dumps(next_config)
invalidates_results = (
("config" in payload and payload.get("config") != task.get("config"))
or (
"process_type" in payload
and payload.get("process_type") != task.get("process_type")
)
or (
"source_dataset_id" in payload
and payload.get("source_dataset_id") != task.get("source_dataset_id")
)
)
if invalidates_results and not regeneration_prepared:
values.update(
{
"status": "pending",
"progress": 0,
"output_count": 0,
"filtered_count": 0,
"duplicate_count": 0,
"error_count": 0,
"failure_reason": None,
"generation_run_id": None,
"results_confirmed": False,
"preview_status": "idle",
"preview_progress": 0,
"preview_run_id": None,
"preview_failure_reason": None,
"preview_total_files": 0,
"preview_completed_files": 0,
"started_at": None,
"completed_at": None,
}
)
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 (
"process_type" in payload
and payload.get("process_type") != task.get("process_type")
):
conn.execute(
"""
UPDATE data_process_source_files
SET deleted_at=%s, updated_at=%s
WHERE task_id=%s AND deleted_at IS NULL
""",
(utcnow(), utcnow(), task_id),
)
values["input_count"] = 0
values["updated_at"] = utcnow()
assignments = ", ".join(f"{key}=%s" for key in values)
row = conn.execute(
f"UPDATE data_process_tasks SET {assignments} WHERE id=%s RETURNING *",
[*values.values(), task_id],
).fetchone()
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc
return _public_task(_decode_row(row)) or {}
def prepare_regeneration(
self,
task_id: str,
payload: dict[str, Any],
) -> dict[str, Any]:
"""非破坏性地保存重新生成配置。
准备阶段保留任务当前状态、结果、切片及已发布数据集。真正开始
生成时,才在同一事务内切换运行状态并清理上一轮结果。
"""
# 先修复曾被旧版 prepare 提前清空的任务,再建立新的非破坏性草稿标记。
self.recover_legacy_aborted_regeneration(task_id)
try:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] == "running":
raise ConflictError("running task cannot be prepared for regeneration")
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise ConflictError("running preview cannot be prepared for regeneration")
current_updated_at = _serialize_value(task.get("updated_at"))
if payload["expected_updated_at"] != current_updated_at:
raise ConflictError("data process task was modified by another request")
process_type = str(payload["process_type"])
if process_type != str(task["process_type"]):
raise InvalidStateError("process_type cannot be changed during regeneration")
current_config = dict(task.get("config") or {})
next_config = _business_config(payload.get("config"))
for key in (_REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY):
if key in current_config:
next_config[key] = current_config[key]
preview_invalidated = _preview_config_changed(
process_type,
current_config,
next_config,
)
now = utcnow()
next_config[_REGENERATION_MARKER_KEY] = {
"prepared": True,
"preview_invalidated": preview_invalidated,
"prepared_at": now,
}
# 002 迁移前发布的数据集只有 task_id。先补齐新关联字段保证
# 解除任务输出指针后,详情和后续重新发布仍能定位原来的三份数据集。
conn.execute(
"""
UPDATE datasets
SET source_task_id=%s, updated_at=%s
WHERE source='task' AND source_task_id IS NULL AND task_id=%s
AND deleted_at IS NULL
""",
(task_id, now, task_id),
)
published_row = conn.execute(
"""
SELECT EXISTS(
SELECT 1 FROM datasets
WHERE source='task' AND deleted_at IS NULL
AND (
source_task_id=%s
OR (source_task_id IS NULL AND task_id=%s)
)
) AS exists
""",
(task_id, task_id),
).fetchone()
published_outputs_preserved = bool(task.get("output_dataset_id")) or bool(
published_row and published_row.get("exists")
)
row = conn.execute(
"""
UPDATE data_process_tasks
SET name=%s, description=%s, config=%s, updated_at=%s
WHERE id=%s
RETURNING *
""",
(
payload["name"],
payload.get("description") or "",
json_dumps(next_config),
now,
task_id,
),
).fetchone()
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc
return {
"task": _public_task(_decode_row(row)) or {},
"preview_invalidated": preview_invalidated,
"published_outputs_preserved": published_outputs_preserved,
}
def delete_task(self, task_id: str, *, deleted_by: str | None = None) -> None:
with self.connect() as conn:
self._task_in_connection(conn, task_id, for_update=True)
now = utcnow()
conn.execute(
"""
UPDATE data_process_tasks
SET status=CASE WHEN status='running' THEN 'stopped' ELSE status END,
generation_run_id=NULL,
preview_status=CASE
WHEN preview_status IN ('queued', 'running') THEN 'cancelled'
ELSE preview_status
END,
preview_run_id=NULL,
deleted_at=%s, deleted_by=%s, updated_at=%s
WHERE id=%s
""",
(now, deleted_by, now, task_id),
)