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:
@@ -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
145
backend/app/modules/data_process/algorithms/__init__.py
Normal file
145
backend/app/modules/data_process/algorithms/__init__.py
Normal 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",
|
||||
]
|
||||
119
backend/app/modules/data_process/algorithms/format_detection.py
Normal file
119
backend/app/modules/data_process/algorithms/format_detection.py
Normal 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),
|
||||
)
|
||||
@@ -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',
|
||||
]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
700
backend/app/modules/data_process/algorithms/parsers/office.py
Normal file
700
backend/app/modules/data_process/algorithms/parsers/office.py
Normal 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
|
||||
297
backend/app/modules/data_process/algorithms/parsers/pdf.py
Normal file
297
backend/app/modules/data_process/algorithms/parsers/pdf.py
Normal 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)
|
||||
343
backend/app/modules/data_process/algorithms/quality.py
Normal file
343
backend/app/modules/data_process/algorithms/quality.py
Normal 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,
|
||||
)
|
||||
@@ -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
|
||||
328
backend/app/modules/data_process/algorithms/text_utils.py
Normal file
328
backend/app/modules/data_process/algorithms/text_utils.py
Normal 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_utils(normalize_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)
|
||||
81
backend/app/modules/data_process/algorithms/transforms.py
Normal file
81
backend/app/modules/data_process/algorithms/transforms.py
Normal 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
|
||||
266
backend/app/modules/data_process/algorithms/types.py
Normal file
266
backend/app/modules/data_process/algorithms/types.py
Normal 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
69
backend/app/modules/data_process/store/__init__.py
Normal file
69
backend/app/modules/data_process/store/__init__.py
Normal 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",
|
||||
]
|
||||
305
backend/app/modules/data_process/store/base.py
Normal file
305
backend/app/modules/data_process/store/base.py
Normal 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)
|
||||
504
backend/app/modules/data_process/store/datasets.py
Normal file
504
backend/app/modules/data_process/store/datasets.py
Normal 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()
|
||||
275
backend/app/modules/data_process/store/generation.py
Normal file
275
backend/app/modules/data_process/store/generation.py
Normal 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"),
|
||||
}
|
||||
539
backend/app/modules/data_process/store/preview.py
Normal file
539
backend/app/modules/data_process/store/preview.py
Normal 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())
|
||||
324
backend/app/modules/data_process/store/results.py
Normal file
324
backend/app/modules/data_process/store/results.py
Normal 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 {}
|
||||
321
backend/app/modules/data_process/store/source_files.py
Normal file
321
backend/app/modules/data_process/store/source_files.py
Normal 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),
|
||||
)
|
||||
871
backend/app/modules/data_process/store/tasks.py
Normal file
871
backend/app/modules/data_process/store/tasks.py
Normal 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),
|
||||
)
|
||||
Reference in New Issue
Block a user